{"spec_id":"heatmap-cohort-retention","library":"matplotlib","language":"python","code":"\"\"\" anyplot.ai\nheatmap-cohort-retention: Cohort Retention Heatmap\nLibrary: matplotlib 3.11.0 | Python 3.13.14\nQuality: 94/100 | Updated: 2026-06-20\n\"\"\"\n\nimport os\nimport sys\n\n\nsys.path.pop(0)  # prevent this file from shadowing the installed matplotlib package\n\nimport matplotlib.colors as mcolors\nimport matplotlib.patches as mpatches\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom matplotlib.colors import LinearSegmentedColormap\n\n\n# Theme tokens — Imprint palette 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\"\nINK_MUTED = \"#6B6A63\" if THEME == \"light\" else \"#A8A79F\"\nANYPLOT_AMBER = \"#DDCC77\"  # warning / caution semantic anchor\n\n# Imprint sequential colormap — canonical green→blue (low retention → green, high retention → blue)\nimprint_seq = LinearSegmentedColormap.from_list(\"imprint_seq\", [\"#009E73\", \"#4467A3\"])\n\n# Data\nnp.random.seed(42)\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]\ncohort_sizes = [1200, 1350, 980, 1100, 1450, 1280, 1050, 1320, 1180, 1400]\nn_cohorts = len(cohort_labels)\nn_periods = n_cohorts\n\n# Generate realistic retention data with meaningful variation across cohorts\nretention = np.full((n_cohorts, n_periods), np.nan)\ndecay_profiles = [1.0, 1.15, 1.3, 1.1, 0.55, 0.65, 1.2, 0.85, 1.05, 0.75]\n\nfor i in range(n_cohorts):\n    max_periods = n_periods - i\n    retention[i, 0] = 100.0\n    for j in range(1, max_periods):\n        base_drop = (15 * np.exp(-0.25 * j) + 1.5) * decay_profiles[i]\n        noise = np.random.uniform(-2, 2)\n        retention[i, j] = max(retention[i, j - 1] - base_drop - noise, 5)\n\n# Find best-performing cohort (highest average retention across >= 4 periods)\navg_retention = [np.nanmean(retention[i, 1 : n_periods - i]) if n_periods - i >= 4 else 0.0 for i in range(n_cohorts)]\nbest_cohort = int(np.argmax(avg_retention))\n\n# Plot — square canvas for symmetric heatmap\nfig, ax = plt.subplots(figsize=(6, 6), dpi=400, facecolor=PAGE_BG)\nax.set_facecolor(PAGE_BG)\n\nnorm = mcolors.Normalize(vmin=0, vmax=100)\n\n# Draw heatmap cells using FancyBboxPatch for rounded corners\nfor i in range(n_cohorts):\n    for j in range(n_periods):\n        if np.isnan(retention[i, j]):\n            continue\n        val = retention[i, j]\n        color = imprint_seq(norm(val))\n        rect = mpatches.FancyBboxPatch(\n            (j - 0.47, i - 0.47),\n            0.94,\n            0.94,\n            boxstyle=mpatches.BoxStyle.Round(pad=0, rounding_size=0.08),\n            facecolor=color,\n            edgecolor=PAGE_BG,\n            linewidth=2.0,\n        )\n        ax.add_patch(rect)\n        # Adaptive text: light on dark cells, dark on light cells\n        luminance = 0.299 * color[0] + 0.587 * color[1] + 0.114 * color[2]\n        text_color = \"#FAF8F1\" if luminance < 0.45 else \"#1A1A17\"\n        ax.text(\n            j,\n            i,\n            f\"{val:.0f}%\",\n            ha=\"center\",\n            va=\"center\",\n            fontsize=9,\n            fontweight=\"bold\" if i == best_cohort else \"medium\",\n            color=text_color,\n        )\n\n# Highlight best cohort row with amber dashed border\nhighlight_rect = mpatches.FancyBboxPatch(\n    (-0.55, best_cohort - 0.55),\n    n_periods - best_cohort + 0.1,\n    1.1,\n    boxstyle=mpatches.BoxStyle.Round(pad=0, rounding_size=0.12),\n    facecolor=\"none\",\n    edgecolor=ANYPLOT_AMBER,\n    linewidth=2.5,\n    linestyle=\"--\",\n    zorder=5,\n)\nax.add_patch(highlight_rect)\n\n# Style\nax.set_xlim(-0.5, n_periods - 0.5)\nax.set_ylim(n_cohorts - 0.5, -0.5)\nax.set_xticks(range(n_periods))\nax.set_xticklabels([f\"Month {p}\" for p in range(n_periods)], fontsize=8, color=INK_SOFT, rotation=45, ha=\"right\")\nax.set_yticks(range(n_cohorts))\nytick_labels = []\nfor idx, (label, size) in enumerate(zip(cohort_labels, cohort_sizes, strict=True)):\n    text = f\"{label}  (n={size:,})\"\n    if idx == best_cohort:\n        text = f\"★ {text}\"\n    ytick_labels.append(text)\nax.set_yticklabels(ytick_labels, fontsize=8, color=INK_SOFT)\nax.set_xlabel(\"Months Since Signup\", fontsize=10, color=INK)\nax.set_ylabel(\"Signup Cohort\", fontsize=10, color=INK)\n\ntitle = \"heatmap-cohort-retention · python · matplotlib · anyplot.ai\"\ntitle_fontsize = max(8, round(12 * 67 / len(title))) if len(title) > 67 else 12\nax.set_title(title, fontsize=title_fontsize, fontweight=\"medium\", color=INK, pad=12)\n\nfor spine in ax.spines.values():\n    spine.set_visible(False)\nax.tick_params(axis=\"both\", length=0, labelcolor=INK_SOFT)\n\n# Colorbar with theme-adaptive chrome\nsm = plt.cm.ScalarMappable(cmap=imprint_seq, norm=norm)\nsm.set_array([])\ncbar = fig.colorbar(sm, ax=ax, shrink=0.6, aspect=25, pad=0.02)\ncbar.set_label(\"Retention Rate (%)\", fontsize=8, color=INK)\ncbar.ax.tick_params(labelsize=8, labelcolor=INK_SOFT, colors=INK_SOFT)\ncbar.outline.set_visible(False)\n\n# Save\nplt.tight_layout(pad=1.5)\nplt.savefig(f\"plot-{THEME}.png\", dpi=400, facecolor=PAGE_BG)\n"}