{"spec_id":"heatmap-cohort-retention","library":"pygal","language":"python","code":"\"\"\" anyplot.ai\nheatmap-cohort-retention: Cohort Retention Heatmap\nLibrary: pygal 3.1.3 | Python 3.13.14\nQuality: 89/100 | Updated: 2026-06-20\n\"\"\"\n\nimport os\nimport sys\n\nimport numpy as np\n\n\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# Remove script dir so \"pygal\" resolves to the installed package, not this file\n_cwd = sys.path.pop(0)\nfrom pygal.graph.graph import Graph\nfrom pygal.style import Style\n\n\nsys.path.insert(0, _cwd)\n\n\nclass CohortRetentionHeatmap(Graph):\n    _series_margin = 0\n\n    def __init__(self, *args, **kwargs):\n        self.matrix_data = kwargs.pop(\"matrix_data\", [])\n        self.row_labels = kwargs.pop(\"row_labels\", [])\n        self.col_labels = kwargs.pop(\"col_labels\", [])\n        self.cohort_sizes = kwargs.pop(\"cohort_sizes\", [])\n        self.colormap = kwargs.pop(\"colormap\", [])\n        super().__init__(*args, **kwargs)\n\n    def _lerp_color(self, c0, c1, t):\n        r = int(round(int(c0[1:3], 16) + (int(c1[1:3], 16) - int(c0[1:3], 16)) * t))\n        g = int(round(int(c0[3:5], 16) + (int(c1[3:5], 16) - int(c0[3:5], 16)) * t))\n        b = int(round(int(c0[5:7], 16) + (int(c1[5:7], 16) - int(c0[5:7], 16)) * t))\n        return f\"#{r:02X}{g:02X}{b:02X}\"\n\n    def _cell_color(self, value, min_val, max_val):\n        if max_val == min_val:\n            return self.colormap[len(self.colormap) // 2]\n        t = max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))\n        pos = t * (len(self.colormap) - 1)\n        lo = int(pos)\n        hi = min(lo + 1, len(self.colormap) - 1)\n        return self._lerp_color(self.colormap[lo], self.colormap[hi], pos - lo)\n\n    def _text_on(self, bg):\n        lum = (int(bg[1:3], 16) * 299 + int(bg[3:5], 16) * 587 + int(bg[5:7], 16) * 114) / 1000\n        return \"#F0EFE8\" if lum < 140 else \"#1A1A17\"\n\n    def _plot(self):\n        if not self.matrix_data:\n            return\n\n        n_rows = len(self.matrix_data)\n        n_cols = max(len(row) for row in self.matrix_data)\n        non_null = [v for row in self.matrix_data for v in row if v is not None]\n        min_val, max_val = min(non_null), max(non_null)\n\n        pw, ph = self.view.width, self.view.height\n        lm_l, lm_r, lm_t, lm_b = 460, 285, 130, 15\n        aw, ah = pw - lm_l - lm_r, ph - lm_t - lm_b\n        cw = aw / n_cols\n        ch = ah / (n_rows + 0.2)\n        gap = 4\n\n        gw = n_cols * (cw + gap) - gap\n        gh = n_rows * (ch + gap) - gap\n\n        x0 = self.view.x(0) + lm_l + (aw - gw) / 2\n        y0 = self.view.y(n_rows) + lm_t + (ah - gh - ch * 0.2) / 2\n\n        pn = self.nodes[\"plot\"]\n        col_fs = min(28, int(cw * 0.34))\n        row_fs = min(38, int(ch * 0.48))\n        size_fs = int(row_fs * 0.92)\n        val_fs = min(40, int(min(cw, ch) * 0.44))\n\n        # Column section header\n        nd = self.svg.node(pn, \"text\", x=x0 + gw / 2, y=y0 - 78)\n        nd.set(\"text-anchor\", \"middle\")\n        nd.set(\"fill\", INK_SOFT)\n        nd.set(\"style\", f\"font-size:{col_fs + 4}px;font-weight:600;font-family:'Segoe UI',Roboto,sans-serif\")\n        nd.text = \"Months Since Signup\"\n\n        # Column headers\n        for j, lbl in enumerate(self.col_labels):\n            cx = x0 + j * (cw + gap) + cw / 2\n            nd = self.svg.node(pn, \"text\", x=cx, y=y0 - 18)\n            nd.set(\"text-anchor\", \"middle\")\n            nd.set(\"fill\", INK)\n            nd.set(\"style\", f\"font-size:{col_fs}px;font-weight:700;font-family:'Segoe UI',Roboto,sans-serif\")\n            nd.text = str(lbl)\n\n        # Row labels with cohort sizes\n        for i, lbl in enumerate(self.row_labels):\n            ry = y0 + i * (ch + gap) + ch / 2\n            rx = x0 - 22\n            nd = self.svg.node(pn, \"text\", x=rx, y=ry + row_fs * 0.12)\n            nd.set(\"text-anchor\", \"end\")\n            nd.set(\"fill\", INK)\n            nd.set(\"style\", f\"font-size:{row_fs}px;font-weight:600;font-family:'Segoe UI',Roboto,sans-serif\")\n            nd.text = str(lbl)\n            if i < len(self.cohort_sizes):\n                nd2 = self.svg.node(pn, \"text\", x=rx, y=ry + row_fs * 0.12 + size_fs + 5)\n                nd2.set(\"text-anchor\", \"end\")\n                nd2.set(\"fill\", INK_MUTED)\n                nd2.set(\"style\", f\"font-size:{size_fs}px;font-style:italic;font-family:'Segoe UI',Roboto,sans-serif\")\n                nd2.text = f\"n={self.cohort_sizes[i]:,}\"\n\n        # Y-axis title (rotated)\n        ytx, yty = x0 - 355, y0 + gh / 2\n        nd = self.svg.node(pn, \"text\", x=ytx, y=yty)\n        nd.set(\"text-anchor\", \"middle\")\n        nd.set(\"fill\", INK_SOFT)\n        nd.set(\"style\", f\"font-size:{col_fs + 4}px;font-weight:600;font-family:'Segoe UI',Roboto,sans-serif\")\n        nd.set(\"transform\", f\"rotate(-90, {ytx}, {yty})\")\n        nd.text = \"Signup Cohort\"\n\n        # Cells\n        for i in range(n_rows):\n            for j in range(len(self.matrix_data[i])):\n                v = self.matrix_data[i][j]\n                if v is None:\n                    continue\n                color = self._cell_color(v, min_val, max_val)\n                tc = self._text_on(color)\n                cx = x0 + j * (cw + gap)\n                cy = y0 + i * (ch + gap)\n\n                grp = self.svg.node(pn, \"g\")\n                rect = self.svg.node(grp, \"rect\", x=cx, y=cy, width=cw, height=ch, rx=5, ry=5)\n                rect.set(\"fill\", color)\n                rect.set(\"stroke\", ELEVATED_BG)\n                rect.set(\"stroke-width\", \"1.5\")\n\n                co_lbl = self.row_labels[i] if i < len(self.row_labels) else \"\"\n                self._tooltip_data(grp, f\"{v:.1f}%\", cx + cw / 2, cy + ch / 2, xlabel=f\"{co_lbl} – Month {j}\")\n\n                vt = self.svg.node(grp, \"text\", x=cx + cw / 2, y=cy + ch / 2 + val_fs * 0.35)\n                vt.set(\"text-anchor\", \"middle\")\n                vt.set(\"fill\", tc)\n                vt.set(\"style\", f\"font-size:{val_fs}px;font-weight:600;font-family:'Segoe UI',Roboto,sans-serif\")\n                vt.text = f\"{v:.0f}%\"\n\n        # Colorbar\n        cb_w, cb_h = 48, gh * 0.80\n        cb_x, cb_y = x0 + gw + 55, y0 + (gh - cb_h) / 2\n        cb_ls = 28\n\n        defs = self.svg.node(pn, \"defs\")\n        grad = self.svg.node(defs, \"linearGradient\", id=\"cb-gradient\", x1=\"0\", y1=\"0\", x2=\"0\", y2=\"1\")\n        for fi in range(21):\n            f = fi / 20.0\n            val = max_val - (max_val - min_val) * f\n            color = self._cell_color(val, min_val, max_val)\n            stop = self.svg.node(grad, \"stop\", offset=f\"{f * 100}%\")\n            stop.set(\"stop-color\", color)\n\n        cbr = self.svg.node(pn, \"rect\", x=cb_x, y=cb_y, width=cb_w, height=cb_h, rx=4, ry=4)\n        cbr.set(\"fill\", \"url(#cb-gradient)\")\n        cbr.set(\"stroke\", INK_MUTED)\n        cbr.set(\"stroke-width\", \"1\")\n\n        for frac, val in [\n            (0.0, max_val),\n            (0.25, max_val * 0.75 + min_val * 0.25),\n            (0.5, (min_val + max_val) / 2),\n            (0.75, max_val * 0.25 + min_val * 0.75),\n            (1.0, min_val),\n        ]:\n            ty = cb_y + cb_h * frac\n            tk = self.svg.node(pn, \"line\", x1=cb_x + cb_w, y1=ty, x2=cb_x + cb_w + 10, y2=ty)\n            tk.set(\"stroke\", INK_SOFT)\n            tk.set(\"stroke-width\", \"1.5\")\n            tt = self.svg.node(pn, \"text\", x=cb_x + cb_w + 16, y=ty + cb_ls * 0.35)\n            tt.set(\"fill\", INK)\n            tt.set(\"style\", f\"font-size:{cb_ls}px;font-family:'Segoe UI',Roboto,sans-serif\")\n            tt.text = f\"{val:.0f}%\"\n\n        cbt = self.svg.node(pn, \"text\", x=cb_x + cb_w / 2, y=cb_y - 55)\n        cbt.set(\"text-anchor\", \"middle\")\n        cbt.set(\"fill\", INK_SOFT)\n        cbt.set(\"style\", f\"font-size:{cb_ls + 2}px;font-weight:600;font-family:'Segoe UI',Roboto,sans-serif\")\n        cbt.text = \"Retention %\"\n\n    def _compute(self):\n        n_rows = len(self.matrix_data) if self.matrix_data else 1\n        n_cols = max(len(row) for row in self.matrix_data) if self.matrix_data else 1\n        self._box.xmin = 0\n        self._box.xmax = n_cols\n        self._box.ymin = 0\n        self._box.ymax = n_rows\n\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_max_periods = 10\ncohort_sizes = [1200, 1350, 980, 1520, 1100, 1430, 1280, 1050, 1380, 1150]\nbase_retention = np.array([100.0, 65.0, 48.0, 40.0, 34.0, 30.0, 27.0, 25.0, 23.5, 22.0])\n\nmatrix = []\nfor i in range(n_cohorts):\n    n_periods = n_max_periods - i\n    row = []\n    for j in range(n_periods):\n        if j == 0:\n            row.append(100.0)\n        else:\n            improvement = i * 1.8\n            if i == 3:\n                improvement = -4.0\n            noise = np.random.uniform(-2.0, 2.0)\n            val = max(5.0, min(100.0, base_retention[j] + improvement + noise))\n            row.append(round(val, 1))\n    matrix.append(row)\n\nperiod_labels = [f\"Month {i}\" for i in range(n_max_periods)]\n\n# Imprint sequential colormap: #4467A3 (blue, low retention) → #009E73 (green, high retention)\n_n_stops = 12\nimprint_seq = []\nfor _i in range(_n_stops):\n    _t = _i / (_n_stops - 1)\n    _r = int(round(0x44 + (0x00 - 0x44) * _t))\n    _g = int(round(0x67 + (0x9E - 0x67) * _t))\n    _b = int(round(0xA3 + (0x73 - 0xA3) * _t))\n    imprint_seq.append(f\"#{_r:02X}{_g:02X}{_b:02X}\")\n\nIMPRINT_PALETTE = (\"#009E73\", \"#C475FD\", \"#4467A3\", \"#BD8233\", \"#AE3030\", \"#2ABCCD\", \"#954477\", \"#99B314\")\n\ncustom_style = Style(\n    background=PAGE_BG,\n    plot_background=PAGE_BG,\n    foreground=INK,\n    foreground_strong=INK,\n    foreground_subtle=INK_MUTED,\n    colors=IMPRINT_PALETTE,\n    title_font_size=66,\n    label_font_size=56,\n    major_label_font_size=44,\n    legend_font_size=44,\n    value_font_size=36,\n    stroke_width=2.5,\n)\n\ntitle = \"heatmap-cohort-retention · python · pygal · anyplot.ai\"\n\nchart = CohortRetentionHeatmap(\n    width=2400,\n    height=2400,\n    style=custom_style,\n    title=title,\n    matrix_data=matrix,\n    row_labels=cohort_labels,\n    col_labels=period_labels,\n    cohort_sizes=cohort_sizes,\n    colormap=imprint_seq,\n    show_legend=False,\n    margin=100,\n    margin_top=200,\n    margin_bottom=30,\n    margin_left=120,\n    margin_right=120,\n    show_x_labels=False,\n    show_y_labels=False,\n)\n\nchart.add(\"data\", [0])\n\n# Save\nchart.render_to_png(f\"plot-{THEME}.png\")\nchart.render_to_file(f\"plot-{THEME}.svg\")\n\nhtml_content = f\"\"\"<!DOCTYPE html>\n<html>\n<head>\n    <meta charset=\"utf-8\">\n    <title>heatmap-cohort-retention · python · pygal · anyplot.ai</title>\n    <style>\n        body {{ margin: 0; background: {PAGE_BG}; display: flex; justify-content: center; align-items: center; min-height: 100vh; }}\n    </style>\n</head>\n<body>\n    {chart.render(is_unicode=True)}\n</body>\n</html>\"\"\"\n\nwith open(f\"plot-{THEME}.html\", \"w\", encoding=\"utf-8\") as f:\n    f.write(html_content)\n"}