Skip to content

Cohort Analysis

Cohort retention matrix from transaction data.

Index = cohort start period; integer columns = periods since start. For churn per purchase-day ordinal, use customer.TransactionChurn; render with plots.cohort.plot.

CohortAnalysis

Computes a cohort retention matrix from transaction data.

Customers are grouped into cohorts by the period of their first purchase. .df (pandas) has one row per period from the first cohort to the last (index min_period_shopped, the period's start) and integer columns period_since from 0 to the longest gap observed, so it need not be square. Each cell is agg_func of aggregation_column over that cohort's transactions in that period. Cells with no transactions, including cohorts with no customers and periods past the end of the data, are 0, the same as a genuine 0. The query runs at construction; there is no .table.

Source code in openretailscience/analysis/cohort.py 
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
class CohortAnalysis:
    """Computes a cohort retention matrix from transaction data.

    Customers are grouped into cohorts by the period of their first purchase. ``.df``
    (pandas) has one row per period from the first cohort to the last (index
    ``min_period_shopped``, the period's start) and integer columns ``period_since`` from 0
    to the longest gap observed, so it need not be square. Each cell is ``agg_func`` of
    ``aggregation_column`` over that cohort's transactions in that period. Cells with no
    transactions, including cohorts with no customers and periods past the end of the data,
    are 0, the same as a genuine 0. The query runs at construction; there is no ``.table``.
    """

    VALID_PERIODS: ClassVar[tuple[str, ...]] = ("year", "quarter", "month", "week", "day")

    def __init__(
        self,
        df: pd.DataFrame | ibis.Table,
        aggregation_column: str,
        agg_func: str = "nunique",
        period: str = "month",
        percentage: bool = False,
    ) -> None:
        """Initialize the CohortAnalysis object.

        Args:
            df (pd.DataFrame | ibis.Table): Transaction data. ``transaction_date`` must be a
                date or datetime column; strings are not parsed.
            aggregation_column (str): The column to aggregate. For customer retention use
                ``customer_id`` with ``agg_func="nunique"``.
            agg_func (str, optional): Ibis column aggregation name, case-sensitive (e.g.
                "nunique", "count", "sum", "mean", "min", "max"). "nunique" on ``unit_spend``
                counts distinct prices, not customers. Defaults to "nunique".
            period (str): Period for cohort analysis: "year", "quarter", "month", "week", or "day"
                (case-insensitive; short forms like "m"/"d" accepted).
            percentage (bool): If True, each row is divided by its period-0 value, giving
                fractions (1.0 = 100%) rounded to 2 decimals. A cohort with 0 in period 0 is all 0.

        Raises:
            TypeError: If `df` is not a pandas DataFrame or an Ibis Table.
            TypeError: If `period` is not a string.
            ValueError: If `period` is not one of the allowed values.
            ValueError: If `df` is missing required columns
                (`customer_id`, `transaction_date`, or `aggregation_column`).
        """
        cols = ColumnHelper()

        period = ensure_period(period, self.VALID_PERIODS, "period")

        required_cols = [
            cols.customer_id,
            cols.transaction_date,
            aggregation_column,
        ]
        ensure_data_has_columns(df, required_cols)

        self.df = self._calculate_cohorts(
            df=df,
            agg_func=agg_func,
            period=period,
            aggregation_column=aggregation_column,
            percentage=percentage,
        )

    def _fill_cohort_gaps(
        self,
        cohort_analysis_table: pd.DataFrame,
        period: str,
    ) -> pd.DataFrame:
        """Fills gaps in the cohort analysis table for missing periods.

        Args:
            cohort_analysis_table (pd.DataFrame): The cohort analysis table to fill gaps in.
            period (str): The period of analysis (year, quarter, month, week, or day).

        Returns:
            pd.DataFrame: Cohort table with missing periods filled.
        """
        cohort_analysis_table.index = pd.to_datetime(cohort_analysis_table.index)

        min_period = cohort_analysis_table.index.min()
        max_period = cohort_analysis_table.index.max()

        # Week steps by 7 days from the (truncated) min period rather than freq="W", which
        # anchors on Sunday and would miss the Monday week-starts that truncate("week") yields.
        period_lookup = {"year": "YS", "quarter": "QS", "month": "MS", "week": "7D", "day": "D"}
        full_range = pd.date_range(start=min_period, end=max_period, freq=period_lookup[period])
        cohort_analysis_table = cohort_analysis_table.reindex(full_range, fill_value=0)
        if cohort_analysis_table.shape[1] > 0:
            max_period_since = cohort_analysis_table.columns.max()
            all_periods = range(max_period_since + 1)
            cohort_analysis_table = cohort_analysis_table.reindex(columns=all_periods, fill_value=0)
        return cohort_analysis_table

    def _calculate_cohorts(
        self,
        df: pd.DataFrame | ibis.Table,
        aggregation_column: str,
        agg_func: str = "nunique",
        period: str = "month",
        percentage: bool = False,
    ) -> pd.DataFrame:
        """Computes the cohort matrix described on the class.

        Args:
            df (pd.DataFrame | ibis.Table): The dataset containing transaction data.
            aggregation_column (str): The column to aggregate.
            agg_func (str, optional): Ibis aggregation name. Defaults to "nunique".
            period (str): Period for cohort analysis: "year", "quarter", "month", "week", or "day".
            percentage (bool): If True, divide each row by its period-0 value.

        Returns:
            pd.DataFrame: The cohort matrix.
        """
        cols = ColumnHelper()

        ibis_table = ensure_ibis_table(df)

        filtered_table = ibis_table.mutate(
            period_shopped=ibis_table[cols.transaction_date].truncate(period),
            period_value=ibis_table[aggregation_column],
        )

        customer_cohort = filtered_table.group_by(cols.customer_id).aggregate(
            min_period_shopped=filtered_table.period_shopped.min(),
        )

        cohort_table = (
            filtered_table.join(customer_cohort, [cols.customer_id])
            .group_by("min_period_shopped", "period_shopped")
            .aggregate(period_value=getattr(filtered_table.period_value, agg_func)())
        )

        cohort_df = cohort_table.execute()
        # period_since is computed in pandas (see _periods_between) rather than via Ibis
        # .delta(unit=period); the truncate and group-by above still execute in the database.
        cohort_df["period_since"] = _periods_between(
            cohort_df["min_period_shopped"],
            cohort_df["period_shopped"],
            period,
        )
        cohort_df = cohort_df.drop_duplicates(subset=["min_period_shopped", "period_since"])

        cohort_analysis_table = cohort_df.pivot(
            index="min_period_shopped",
            columns="period_since",
            values="period_value",
        )

        if percentage:
            first_period = cohort_analysis_table[0].replace(0, np.nan)
            cohort_analysis_table = cohort_analysis_table.div(first_period, axis=0).round(2)

        cohort_analysis_table = cohort_analysis_table.fillna(0)
        cohort_analysis_table = self._fill_cohort_gaps(cohort_analysis_table, period)
        cohort_analysis_table.index.name = "min_period_shopped"

        return cohort_analysis_table

__init__(df, aggregation_column, agg_func='nunique', period='month', percentage=False)

Initialize the CohortAnalysis object.

Parameters:

Name Type Description Default
df DataFrame | Table

Transaction data. transaction_date must be a date or datetime column; strings are not parsed.

 required
aggregation_column str

The column to aggregate. For customer retention use customer_id with agg_func="nunique".

 required
agg_func str

Ibis column aggregation name, case-sensitive (e.g. "nunique", "count", "sum", "mean", "min", "max"). "nunique" on unit_spend counts distinct prices, not customers. Defaults to "nunique".

 'nunique'
period str

Period for cohort analysis: "year", "quarter", "month", "week", or "day" (case-insensitive; short forms like "m"/"d" accepted).

 'month'
percentage bool

If True, each row is divided by its period-0 value, giving fractions (1.0 = 100%) rounded to 2 decimals. A cohort with 0 in period 0 is all 0.

 False

Raises:

Type Description
TypeError

If df is not a pandas DataFrame or an Ibis Table.
TypeError

If period is not a string.
ValueError

If period is not one of the allowed values.
ValueError

If df is missing required columns (customer_id, transaction_date, or aggregation_column).
Source code in openretailscience/analysis/cohort.py 
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def __init__(
    self,
    df: pd.DataFrame | ibis.Table,
    aggregation_column: str,
    agg_func: str = "nunique",
    period: str = "month",
    percentage: bool = False,
) -> None:
    """Initialize the CohortAnalysis object.

    Args:
        df (pd.DataFrame | ibis.Table): Transaction data. ``transaction_date`` must be a
            date or datetime column; strings are not parsed.
        aggregation_column (str): The column to aggregate. For customer retention use
            ``customer_id`` with ``agg_func="nunique"``.
        agg_func (str, optional): Ibis column aggregation name, case-sensitive (e.g.
            "nunique", "count", "sum", "mean", "min", "max"). "nunique" on ``unit_spend``
            counts distinct prices, not customers. Defaults to "nunique".
        period (str): Period for cohort analysis: "year", "quarter", "month", "week", or "day"
            (case-insensitive; short forms like "m"/"d" accepted).
        percentage (bool): If True, each row is divided by its period-0 value, giving
            fractions (1.0 = 100%) rounded to 2 decimals. A cohort with 0 in period 0 is all 0.

    Raises:
        TypeError: If `df` is not a pandas DataFrame or an Ibis Table.
        TypeError: If `period` is not a string.
        ValueError: If `period` is not one of the allowed values.
        ValueError: If `df` is missing required columns
            (`customer_id`, `transaction_date`, or `aggregation_column`).
    """
    cols = ColumnHelper()

    period = ensure_period(period, self.VALID_PERIODS, "period")

    required_cols = [
        cols.customer_id,
        cols.transaction_date,
        aggregation_column,
    ]
    ensure_data_has_columns(df, required_cols)

    self.df = self._calculate_cohorts(
        df=df,
        agg_func=agg_func,
        period=period,
        aggregation_column=aggregation_column,
        percentage=percentage,
    )