How to Build a Cohort Retention Matrix in R and Python

A build-along companion to the concept piece. Point R or Python at a raw transaction log and produce the same retention grid, cohort by cohort.

Knowing what a cohort retention matrix reveals is one thing. I talked about this in my first post on cohorts

Producing one from your own sales data is the step that makes it useful. This walkthrough builds the matrix twice, once in R and once in Python, from a raw customer transaction log.

A companion piece covers what the finished grid tells you about customer quality and budget decisions. Here the focus is the pipeline, the handful of steps that turn a list of orders into rows of cohorts and columns of months.

The input is a customer transaction export, the kind that carries a customer ID, an order date, and a quantity. The public UCI Online Retail workbook stands in for it, and the same code runs against any table shaped that way.

Setting Up the Data in R

The first step is calling the packages that will be used. The core libraries are ones for accessing data (the readxl package in this case), cleaning and rearranging data (dplyr and tidyr)

The first step is adding yout data. in this instance, the data is being downloaded from a URL — an Online Retail workbook from UCLA Machine Learning data repository.

The read_excel function loads the workbook into a data frame. The function is applied to an object you name — in this example, tmp, for a temporary file.

The filter function drops rows with no customer ID, which the raw file leaves blank on thousands of lines, and it removes returns that carry a non-positive quantity. floor_date collapses each timestamp to the first day of its month so orders fall into clean monthly buckets. The distinct function keeps one row per customer per month, since a cohort cares only whether a customer was active, not how many times.

Assigning Cohorts and Measuring Retention

The next step is assigning cohorts. This is where an application of filtering functions are applied to your dataset. The major difference is that the filters are based on time periods that you want your cohort matrix to analyze.

The group_by function paired with min sets each customer's cohort to their earliest active month. The interval call divided by months(1) converts the gap between that cohort month and every later order month into a whole-number period offset. The count function tallies active customers for each cohort and period, the left__join_ function attaches each cohort's starting size, and the mutate function turns raw counts into a retention share.

Reshaping to the Matrix

You can create a retention matrix to allow for further adjustments. The pivot_wider function moves the period values into columns, so each row becomes a cohort and each column a month since acquisition. The empty cells in the upper right are the recent cohorts that have not aged enough to fill later periods, the triangular shape the concept piece described. arrange orders the rows from oldest cohort to newest.

Visualizing the Retention Matrix

The ubiquitous ggplot package is used to created the actual matrix, leveraging the typical functions and elements to craft the visualization. The geom_tile function from the package renders each cell of the retention matrix, one tile per cohort row and period column. The geom_text function with percent prints the retention share inside each cell, placing the values from the pivoted matrix directly on the grid.

The scale_fill_gradient function encodes that share as color across the light-to-teal Zimana palette, so a horizontal read traces a cohort's retention curve and a vertical read compares cohorts at the same age. The blank upper right is the triangular gap, the recent cohorts still short on history.

The syntax will create the following cohort matrix. One note — ggplot tendsn to prefer long data, so in this instance the retention object was used for simplicity in the syntax.

Cohort Analysis in Python

Now let's recreate the same cohort analysis in Python. To do this, we need a few modules — numpy, pandas, seaborn and matplot, which are the usual Python choices for data modeling and visualization.

The pandas flow mirrors the R steps with familiar names. For examople, the to_period("M") call is the pandas answer to floor_date, collapsing each timestamp into a monthly bucket. The groupby call with transform("min") stamps every row with its customer's first month, and subtracting the two period columns yields the integer offset through .n. The pivot call reshapes the long table into the matrix, and seaborn's heatmap with fmt=".0%" prints the same percentage cells the ggplot version shows. Polars users can swap the pandas block for a lazy group_by and pivot with no change to the logic.

Here is whatr the matrix looks like once the script is run.

What Cohort Analysis Means for Your Work

The cohort matrix, be it R or Python, brings a transformation to transactional data, especially for your data that has not been explored deeply. It is enough insight for you to build a matrix out of curiousity. Once the data pipeline for your matrix runs, pointing it at a fresh export refreshes every cohort without touching the code.

Wire the retention share into the KPIs you already track. When you do, you will realixe that the your cohort grid stops being a one-time chart. It becomes a standing read on opportunities to improve customer quality, updated as often as you pull the data.