{"spec_id":"heatmap-cohort-retention","library":"altair","language":"python","code":"\"\"\" anyplot.ai\nheatmap-cohort-retention: Cohort Retention Heatmap\nLibrary: altair 6.2.1 | Python 3.13.14\nQuality: 94/100 | Updated: 2026-06-20\n\"\"\"\n\nimport os\n\nimport altair as alt\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\n\n\n# Theme tokens\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\"\nINK_MUTED = \"#6B6A63\" if THEME == \"light\" else \"#A8A79F\"\n\n# Imprint sequential colormap: brand green → blue (single-polarity continuous data)\nSEQ_LOW = \"#009E73\"  # brand green — low retention\nSEQ_HIGH = \"#4467A3\"  # blue — high retention\nANYPLOT_AMBER = \"#DDCC77\"  # warning / caution — callout annotations\n\n# Data\nnp.random.seed(42)\n\ncohort_labels = [\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(cohort_labels)\nn_periods = 10\ncohort_sizes = np.random.randint(800, 2500, size=n_cohorts)\n\n# Build retention data with realistic decay patterns\nrows = []\nfor i, cohort in enumerate(cohort_labels):\n    max_periods = n_cohorts - i\n    for period in range(max_periods):\n        if period == 0:\n            retention = 100.0\n        elif period == 1:\n            retention = np.random.uniform(55, 72)\n        else:\n            decay = np.random.uniform(0.85, 0.95)\n            retention = rows[-1][\"retention_rate\"] * decay\n            retention += np.random.uniform(-2, 2)\n            retention = max(5, min(retention, 100))\n        rows.append(\n            {\n                \"cohort\": cohort,\n                \"cohort_label\": f\"{cohort} (n={cohort_sizes[i]:,})\",\n                \"period\": period,\n                \"period_label\": f\"Month {period}\",\n                \"retention_rate\": round(retention, 1),\n            }\n        )\n\ndf = pd.DataFrame(rows)\n\ncohort_order = [f\"{c} (n={s:,})\" for c, s in zip(cohort_labels, cohort_sizes, strict=True)]\nperiod_order = [f\"Month {p}\" for p in range(n_periods)]\n\n# Average Month 0→1 retention drop — used for cliff callout annotation\navg_month1_retention = df[df[\"period\"] == 1][\"retention_rate\"].mean()\navg_drop = round(100 - avg_month1_retention)\n\n# Title — len=55 < 67, no font-size shrink needed; default 16px applies\ntitle = \"heatmap-cohort-retention · python · altair · anyplot.ai\"\ntitle_fontsize = 16\n\n# Heatmap rectangles — Imprint sequential palette\nheatmap = (\n    alt.Chart(df)\n    .mark_rect(stroke=INK_SOFT, strokeWidth=0.5, cornerRadius=3)\n    .encode(\n        x=alt.X(\n            \"period_label:O\",\n            title=\"Months Since Signup\",\n            sort=period_order,\n            axis=alt.Axis(\n                labelFontSize=10,\n                titleFontSize=12,\n                titleFontWeight=\"bold\",\n                labelAngle=-40,\n                domainWidth=0,\n                tickWidth=0,\n                titlePadding=12,\n                labelPadding=8,\n            ),\n        ),\n        y=alt.Y(\n            \"cohort_label:O\",\n            title=\"Signup Cohort\",\n            sort=cohort_order,\n            axis=alt.Axis(\n                labelFontSize=10,\n                titleFontSize=12,\n                titleFontWeight=\"bold\",\n                domainWidth=0,\n                tickWidth=0,\n                titlePadding=12,\n                labelPadding=6,\n            ),\n        ),\n        color=alt.Color(\n            \"retention_rate:Q\",\n            scale=alt.Scale(domain=[0, 100], range=[SEQ_LOW, SEQ_HIGH]),\n            legend=alt.Legend(\n                title=\"Retention %\",\n                titleFontSize=10,\n                titleFontWeight=\"bold\",\n                labelFontSize=10,\n                gradientLength=200,\n                gradientThickness=14,\n                orient=\"right\",\n                offset=10,\n            ),\n        ),\n        tooltip=[\n            alt.Tooltip(\"cohort:N\", title=\"Cohort\"),\n            alt.Tooltip(\"period_label:O\", title=\"Period\"),\n            alt.Tooltip(\"retention_rate:Q\", title=\"Retention %\", format=\".1f\"),\n        ],\n    )\n)\n\n# Retention rate value inside each cell\ntext = (\n    alt.Chart(df)\n    .mark_text(fontSize=11, fontWeight=\"bold\")\n    .encode(\n        x=alt.X(\"period_label:O\", sort=period_order),\n        y=alt.Y(\"cohort_label:O\", sort=cohort_order),\n        text=alt.Text(\"retention_rate:Q\", format=\".0f\"),\n        color=alt.value(\"#F0EFE8\"),\n    )\n)\n\n# Percent symbol — smaller, offset right for typographic polish\npct = (\n    alt.Chart(df)\n    .mark_text(fontSize=8, fontWeight=\"normal\", dx=14)\n    .encode(\n        x=alt.X(\"period_label:O\", sort=period_order),\n        y=alt.Y(\"cohort_label:O\", sort=cohort_order),\n        text=alt.value(\"%\"),\n        color=alt.value(\"rgba(240,239,232,0.7)\"),\n    )\n)\n\n# Amber border on Jan 2024 row confirms the subtitle insight (earliest = strongest)\njan_df = df[df[\"cohort\"] == \"Jan 2024\"].copy()\njan_highlight = (\n    alt.Chart(jan_df)\n    .mark_rect(fill=None, stroke=ANYPLOT_AMBER, strokeWidth=2, cornerRadius=3)\n    .encode(x=alt.X(\"period_label:O\", sort=period_order), y=alt.Y(\"cohort_label:O\", sort=cohort_order))\n)\n\n# Callout annotation at top of Month 1 column for the ~Month 0→1 retention cliff\ncliff_df = pd.DataFrame(\n    [{\"period_label\": \"Month 1\", \"cohort_label\": cohort_order[0], \"annotation\": f\"avg −{avg_drop}%\"}]\n)\ncliff_label = (\n    alt.Chart(cliff_df)\n    .mark_text(fontSize=9, fontWeight=\"bold\", dy=-26, color=ANYPLOT_AMBER, clip=False)\n    .encode(\n        x=alt.X(\"period_label:O\", sort=period_order), y=alt.Y(\"cohort_label:O\", sort=cohort_order), text=\"annotation:N\"\n    )\n)\n\n# Combine layers — square canvas (2400×2400) for symmetric heatmap grid\nchart = (\n    alt.layer(heatmap, jan_highlight, text, pct, cliff_label)\n    .properties(\n        width=370,\n        height=440,\n        background=PAGE_BG,\n        title=alt.Title(\n            title,\n            fontSize=title_fontsize,\n            fontWeight=\"bold\",\n            anchor=\"middle\",\n            color=INK,\n            subtitle=\"Monthly SaaS user retention — earliest cohorts show strongest long-term engagement\",\n            subtitleFontSize=11,\n            subtitleColor=INK_MUTED,\n            subtitlePadding=6,\n        ),\n    )\n    .configure_view(fill=PAGE_BG, stroke=INK_SOFT, strokeWidth=0.5)\n    .configure_axis(domainColor=INK_SOFT, tickColor=INK_SOFT, gridOpacity=0.0, labelColor=INK_SOFT, titleColor=INK)\n    .configure_legend(fillColor=ELEVATED_BG, strokeColor=INK_SOFT, labelColor=INK_SOFT, titleColor=INK)\n)\n\n# Save PNG — then pad-only to exactly 2400×2400\nchart.save(f\"plot-{THEME}.png\", scale_factor=4.0)\n\nTW, TH = 2400, 2400\n_img = Image.open(f\"plot-{THEME}.png\").convert(\"RGB\")\n_w, _h = _img.size\nif _w > TW or _h > TH:\n    raise SystemExit(\n        f\"altair vl-convert produced {_w}×{_h}, exceeds target {TW}×{TH}. \"\n        f\"Shrink chart .properties(width=, height=) values and re-render.\"\n    )\nif _w < TW or _h < TH:\n    _canvas = Image.new(\"RGB\", (TW, TH), PAGE_BG)\n    _canvas.paste(_img, ((TW - _w) // 2, (TH - _h) // 2))\n    _canvas.save(f\"plot-{THEME}.png\")\n\n# Interactive HTML (no padding applied — only PNGs are gated)\nchart.save(f\"plot-{THEME}.html\")\n"}