See retention at a glance with a cohort heatmap

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

cohorts = ["2024-01", "2024-02", "2024-03", "2024-04", "2024-05", "2024-06"]
months = [f"Month {m}" for m in range(1, 7)]
rng = np.random.default_rng(21)
base = np.linspace(0.7, 0.4, num=6)
matrix = np.vstack(
    [
        np.clip(base - idx * 0.03 + rng.normal(0, 0.01, size=base.size), 0.1, 0.9)
        for idx in range(len(cohorts))
    ]
)

fig, ax = plt.subplots(figsize=(6.4, 3.8))
im = ax.imshow(matrix, cmap="YlGnBu", vmin=0, vmax=1)

ax.set_xticks(range(len(months)), labels=months)
ax.set_yticks(range(len(cohorts)), labels=cohorts)
ax.set_title("Subscription retention cohort heatmap")

for i in range(matrix.shape[0]):
    for j in range(matrix.shape[1]):
        ax.text(j, i, f"{matrix[i, j]*100:.0f}%", ha="center", va="center", fontsize=9)

cbar = fig.colorbar(im, ax=ax, fraction=0.045, pad=0.02)
cbar.set_label("Retention rate")
ax.set_xlabel("Months since acquisition")
ax.set_ylabel("Acquisition cohort")

fig.tight_layout()

plt.show()