{"spec_id":"heatmap-cohort-retention","library":"plotnine","language":"python","code":"\"\"\" anyplot.ai\nheatmap-cohort-retention: Cohort Retention Heatmap\nLibrary: plotnine 0.15.8 | Python 3.13.15\nQuality: 95/100 | Updated: 2026-08-17\n\"\"\"\n\nimport os\n\nimport numpy as np\nimport pandas as pd\nfrom matplotlib.patches import FancyBboxPatch\nfrom plotnine import (\n    aes,\n    element_blank,\n    element_rect,\n    element_text,\n    geom_text,\n    geom_tile,\n    ggplot,\n    labs,\n    scale_fill_gradient,\n    scale_x_continuous,\n    scale_y_discrete,\n    theme,\n    theme_minimal,\n)\n\n\n# Theme-adaptive chrome\nTHEME = os.getenv(\"ANYPLOT_THEME\", \"light\")\nPAGE_BG = \"#FAF8F1\" if THEME == \"light\" else \"#1A1A17\"\nELEVATED_BG = \"#FFFDF6\" if THEME == \"light\" else \"#242420\"\nINK = \"#1A1A17\" if THEME == \"light\" else \"#F0EFE8\"\nINK_SOFT = \"#4A4A44\" if THEME == \"light\" else \"#B8B7B0\"\nRULE = (26 / 255, 26 / 255, 23 / 255, 0.15) if THEME == \"light\" else (240 / 255, 239 / 255, 232 / 255, 0.15)\n\n# Imprint sequential colormap (brand green -> blue) for single-polarity continuous data\n# Green (brand) reads as \"good\" -> high retention; blue anchors low retention\nSEQ_HIGH_RETENTION = \"#009E73\"\nSEQ_LOW_RETENTION = \"#4467A3\"\n\n# Data\nnp.random.seed(42)\ncohorts = [\n    \"Jan 2024\",\n    \"Feb 2024\",\n    \"Mar 2024\",\n    \"Apr 2024\",\n    \"May 2024\",\n    \"Jun 2024\",\n    \"Jul 2024\",\n    \"Aug 2024\",\n    \"Sep 2024\",\n    \"Oct 2024\",\n]\nn_cohorts = len(cohorts)\ncohort_sizes = [1200, 1350, 980, 1100, 1450, 1280, 1050, 1380, 1150, 1020]\n# Mar 2024 (index 2) suffered a pricing-change churn spike -> visibly worse retention\nchurn_event_idx = 2\n\nrows = []\nfor i, cohort in enumerate(cohorts):\n    max_periods = n_cohorts - i\n    for period in range(max_periods):\n        if period == 0:\n            retention = 100.0\n        else:\n            base_decay = 100 * np.exp(-0.22 * period)\n            noise = np.random.uniform(-3, 3)\n            trend_bonus = i * 2.2  # onboarding steadily improves for later cohorts\n            churn_penalty = 14 if i == churn_event_idx else 0\n            retention = np.clip(base_decay + noise + trend_bonus - churn_penalty, 5, 100)\n        rows.append(\n            {\"cohort\": cohort, \"period\": period, \"retention_rate\": round(retention, 1), \"cohort_size\": cohort_sizes[i]}\n        )\n\ndf = pd.DataFrame(rows)\n\n# Y-axis labels carry cohort size; reversed order puts Jan 2024 at the top, Oct 2024 at the bottom\ndf[\"cohort_label\"] = df.apply(lambda r: f\"{r['cohort']} (n={r['cohort_size']:,})\", axis=1)\ncohort_labels = [f\"{c} (n={s:,})\" for c, s in zip(cohorts, cohort_sizes, strict=True)]\ndf[\"cohort_label\"] = pd.Categorical(df[\"cohort_label\"], categories=cohort_labels[::-1], ordered=True)\n\ndf[\"label\"] = df[\"retention_rate\"].apply(lambda v: f\"{v:.0f}%\")\n\n# Compare an early vs. a later cohort at the same period for storytelling\ncompare_period = 4\nearly_val = df[(df[\"cohort\"] == \"Jan 2024\") & (df[\"period\"] == compare_period)][\"retention_rate\"].values[0]\nlater_val = df[(df[\"cohort\"] == \"Jun 2024\") & (df[\"period\"] == compare_period)][\"retention_rate\"].values[0]\nimprovement = later_val - early_val\ncohort_trend_pp = 2.2  # per-cohort onboarding bonus baked into the synthetic retention formula above\n\n# Plot\nplot = (\n    ggplot(df, aes(x=\"period\", y=\"cohort_label\", fill=\"retention_rate\"))\n    + geom_tile(color=PAGE_BG, size=0.8)\n    + geom_text(aes(label=\"label\"), size=3.1, color=\"#FFFFFF\", fontweight=\"bold\")\n    + scale_fill_gradient(low=SEQ_LOW_RETENTION, high=SEQ_HIGH_RETENTION, limits=(0, 100), name=\"Retention %\")\n    + scale_x_continuous(breaks=range(n_cohorts), labels=[f\"Month {i}\" for i in range(n_cohorts)])\n    + scale_y_discrete(expand=(0.06, 0))\n    + labs(\n        x=\"Months Since Signup\",\n        y=\"\",\n        title=\"heatmap-cohort-retention · python · plotnine · anyplot.ai\",\n        subtitle=\"Monthly cohort retention — newer cohorts retain better; Mar 2024 shows a pricing-change churn spike\",\n    )\n    + theme_minimal()\n    + theme(\n        figure_size=(6, 6),\n        plot_title=element_text(size=12, ha=\"center\", weight=\"bold\", color=INK),\n        plot_subtitle=element_text(size=8, ha=\"center\", color=INK_SOFT, style=\"italic\"),\n        axis_title_x=element_text(size=10, color=INK),\n        axis_text_x=element_text(size=8, color=INK_SOFT, angle=45, ha=\"right\"),\n        axis_text_y=element_text(size=8, color=INK_SOFT),\n        legend_title=element_text(size=9, weight=\"bold\", color=INK),\n        legend_text=element_text(size=8, color=INK_SOFT),\n        legend_background=element_rect(fill=ELEVATED_BG, color=None),\n        panel_grid_major=element_blank(),\n        panel_grid_minor=element_blank(),\n        panel_border=element_rect(color=RULE, fill=None, size=0.5),\n        plot_background=element_rect(fill=PAGE_BG, color=PAGE_BG),\n        panel_background=element_rect(fill=PAGE_BG, color=PAGE_BG),\n    )\n)\n\n# Render, then drop into the underlying matplotlib Figure/Axes (a capability\n# unique to plotnine's matplotlib backend, unlike R ggplot2's grid graphics) to\n# draw a rounded-corner callout that fills the empty triangle beneath the data\n# and carries two data-backed insights instead of leaving that panel space bare.\nfig = plot.draw()\nax = fig.axes[0]\ncallout_box = FancyBboxPatch(\n    (3.6, 0.6),\n    9.6 - 3.6,\n    4.4 - 0.6,\n    transform=ax.transData,\n    boxstyle=\"round,pad=0,rounding_size=0.25\",\n    facecolor=ELEVATED_BG,\n    edgecolor=RULE,\n    linewidth=1.0,\n    zorder=5,\n)\nax.add_patch(callout_box)\nax.text(\n    6.6,\n    3.3,\n    f\"Month {compare_period} retention improved\\n+{improvement:.0f}pp from Jan → Jun 2024\",\n    transform=ax.transData,\n    ha=\"center\",\n    va=\"center\",\n    fontsize=9,\n    color=INK,\n    fontweight=\"bold\",\n    zorder=6,\n)\nax.text(\n    6.6,\n    1.7,\n    f\"Each newer cohort trends ~+{cohort_trend_pp:.1f}pp per\\nMonth vs. the prior cohort (onboarding gains)\",\n    transform=ax.transData,\n    ha=\"center\",\n    va=\"center\",\n    fontsize=7.5,\n    color=INK_SOFT,\n    zorder=6,\n)\n\n# Save\nfig.savefig(f\"plot-{THEME}.png\", dpi=400)\n"}