DEV Community

Malek Damiri
Malek Damiri

Posted on

Cohort Retention in pandas: Three Mistakes That Make Your Curve Lie

#datascience #analytics #python #pandas

Cohort retention tables look simple: group customers by the month they started, count how many are still around each month after, divide. Then the curve on the slide says retention falls off a cliff at month three, and it's not true.

Most retention charts that lie do it in one of three ways. Here is a small pandas version that avoids all three, with the numbers each mistake would have produced instead.

The data

One row per customer per month in which they paid. Synthetic, small enough to check by hand:

import pandas as pd

payments = pd.DataFrame({
    "customer_id": ["a", "a", "a", "a", "b", "b", "c", "c", "c",
                    "d", "e", "e", "e", "f", "f", "g"],
    "paid_month": pd.to_datetime([
        "2025-01-01", "2025-02-01", "2025-03-01", "2025-04-01",
        "2025-01-01", "2025-03-01",
        "2025-01-01", "2025-02-01", "2025-04-01",
        "2025-02-01",
        "2025-02-01", "2025-03-01", "2025-04-01",
        "2025-03-01", "2025-04-01",
        "2025-04-01",
    ]),
})
last_complete_month = pd.Timestamp("2025-04-01")
Enter fullscreen mode Exit fullscreen mode

last_complete_month matters more than it looks. It's the last month you have full data for, and it decides which cells in the table are real.

Build the table 

def months_between(later, earlier):
    return (later.year - earlier.year) * 12 + (later.month - earlier.month)


# 1. Cohort = first paid month.
payments["cohort"] = payments.groupby("customer_id")["paid_month"].transform("min")

# 2. Period 0 is the cohort month itself.
payments["period"] = months_between(payments["paid_month"].dt, payments["cohort"].dt)

# 3. Distinct paying customers per cohort and period.
counts = (
    payments.drop_duplicates(["customer_id", "paid_month"])
    .pivot_table(index="cohort", columns="period", values="customer_id", aggfunc="nunique")
    .fillna(0)
)

# 4. Divide by cohort size.
cohort_size = counts[0]
retention = counts.div(cohort_size, axis=0)

# 5. A cohort that hasn't reached a period yet is unknown, not zero.
months_observed = months_between(last_complete_month, retention.index)
observable = pd.DataFrame(
    {p: months_observed >= p for p in retention.columns}, index=retention.index
)
retention = retention.where(observable)

# 6. Weighted curve: retained customers over the cohorts that can be observed.
curve = (retention.mul(cohort_size, axis=0).sum()
         / observable.mul(cohort_size, axis=0).sum())
Enter fullscreen mode Exit fullscreen mode

The retention table:

period        0     1     2     3
cohort
2025-01-01  1.0  0.67  0.67  0.67
2025-02-01  1.0  0.50  0.50   NaN
2025-03-01  1.0  1.00   NaN   NaN
2025-04-01  1.0   NaN   NaN   NaN
Enter fullscreen mode Exit fullscreen mode

And the curve:

period
0    1.00
1    0.67
2    0.60
3    0.67
Enter fullscreen mode Exit fullscreen mode

Mistake 1: starting the count at 1

If the cohort month is "month 1", every column shifts by one, and the first column is always 100% by construction. Readers compare your month 1 to someone else's month 1 and draw the wrong conclusion. Pick period 0 for the starting month, and say so in the chart title.

Mistake 2: treating cells you can't see yet as zero

The April cohort started in the last month of data. It has no month 1 yet, not a month 1 of zero. If those cells are zeros, the curve collapses:

zeros for unreached cells: [1.0, 0.54, 0.29, 0.17]
correct weighted curve:    [1.0, 0.67, 0.60, 0.67]
Enter fullscreen mode Exit fullscreen mode

That's the "retention falls off a cliff" chart. It's an artifact of when you ran the query. Step 5 fixes it by marking those cells NaN and leaving them out of the curve.

The same bug has a quieter cousin: including the current, partial month as if it were complete. That's why the cutoff is last_complete_month and not "today".

Mistake 3: averaging the percentages

A plain mean gives every cohort the same weight, so a one-customer cohort counts as much as a three-customer cohort:

plain mean of percentages: [1.0, 0.72, 0.58, 0.67]
correct weighted curve:    [1.0, 0.67, 0.60, 0.67]
Enter fullscreen mode Exit fullscreen mode

Here the difference is small. With real data, a tiny early cohort of friendly design partners can pull the whole curve up for a year. Step 6 sums retained customers and divides by the customers who could have been retained, which is what "retention at month n" means.

Things to decide before you trust the chart

  • What counts as retained. Paid in the month, logged in, or still under contract? Each gives a different, defensible curve. Write the choice under the chart.
  • Reactivations. Customer b paid in January and March but not February. In this table they count in month 2 and not month 1, which is usually what you want for activity. For revenue retention you may want different rules.
  • Customer identity. If one company can have two customer IDs, merge them before step 1 or one customer becomes a churn and a new signup.

The table itself is ten lines of pandas. The part that takes judgment is deciding which cells are real.

This post was written with AI assistance.

Top comments (1)

Collapse
 
omyvnss profile image
Om Yaduvanshi •

all three are the same bug in different costumes, really. dividing by the wrong denominator: cohort size, unobservable cells, unweighted means. the framing i'd keep: every retention number is a ratio, so the first question is always 'what's the denominator here'.