Cohort Retention Analysis

import matplotlib.pyplot as plt
import numpy as np

# === STYLE CONFIG ===
COHORT_COLORS = ['#6366F1', '#8B5CF6', '#A855F7', '#D946EF']
COLORS = {
    'background': '#FFFFFF',
    'text': '#1E293B',
    'text_muted': '#64748B',
    'grid': '#F1F5F9',
}

# === DATA ===
weeks = np.arange(0, 9)
cohorts = {
    'Q1 2024': [100, 68, 52, 41, 35, 31, 28, 26, 25],
    'Q2 2024': [100, 72, 58, 48, 42, 38, 35, 33, None],
    'Q3 2024': [100, 75, 62, 54, 48, 44, None, None, None],
    'Q4 2024': [100, 78, 65, 58, None, None, None, None, None],
}

# === FIGURE ===
fig, ax = plt.subplots(figsize=(10, 6), dpi=100)
ax.set_facecolor(COLORS['background'])
fig.patch.set_facecolor(COLORS['background'])

# === PLOT ===
for (name, data), color in zip(cohorts.items(), COHORT_COLORS):
    valid_data = [(w, d) for w, d in zip(weeks, data) if d is not None]
    x_vals = [d[0] for d in valid_data]
    y_vals = [d[1] for d in valid_data]
    
    ax.plot(x_vals, y_vals, color=color, linewidth=2.5, label=name, zorder=2)
    ax.scatter(x_vals, y_vals, color=color, s=40, 
               edgecolors='white', linewidths=1.5, zorder=3)
    
    # End label
    ax.annotate(f'{y_vals[-1]}%', xy=(x_vals[-1], y_vals[-1]),
                xytext=(x_vals[-1] + 0.2, y_vals[-1]),
                fontsize=9, color=color, fontweight='500')

# === AXES ===
ax.set_xlim(-0.5, 9)
ax.set_ylim(0, 110)
ax.set_xticks(weeks)
ax.set_xticklabels([f'W{w}' for w in weeks])
ax.set_xlabel('Week', fontsize=10, color=COLORS['text'], labelpad=10)
ax.set_ylabel('Retention Rate (%)', fontsize=10, color=COLORS['text'], labelpad=10)

# === STYLING ===
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_color(COLORS['grid'])
ax.spines['bottom'].set_color(COLORS['grid'])

ax.yaxis.grid(True, color=COLORS['grid'], linewidth=1)
ax.set_axisbelow(True)
ax.tick_params(axis='both', colors=COLORS['text_muted'], labelsize=9, length=0, pad=8)

ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.15),
          ncol=4, frameon=False, fontsize=9, labelcolor=COLORS['text_muted'])

plt.tight_layout()
plt.show()