This video introduces the concept of cohort analysis using a real-world scenario: SpaceStream's user retention challenge. You will learn what a cohort is and why tracking user activity over time is essential for subscription-based services. The video explains three key metrics: **retention rate** (the proportion of the original cohort active in each month), **churn rate** (the fraction that becomes inactive), and **survival rate** (the proportion of users who have stayed active every month up to the current point). The walkthrough demonstrates how to structure your dataset in pandas, with columns for each month and rows for each user, and shows step-by-step calculations for each metric using pandas operations. By the end, you will see how to interpret the results and understand user behavior trends over multiple months.

You are now tasked with calculating advanced retention metrics for SpaceStream, an intergalactic holovision service. As the Lead Data Analyst, you will analyze a cohort of 5 users over three months, tracking who stays loyal and who drifts away. Your goal is to compute three critical metrics for each month: **Retention Rate**, **Churn Rate**, and **Survival Rate**.

Begin by examining the provided dataset, where each user is marked as active (`1`) or inactive (`0`) for each month. The columns `month_0`, `month_1`, and `month_2` represent activity across three consecutive months. Your solution will require you to use `pandas` to process this dataset and extract the necessary metrics for each month.

import pandas as pd

data = {
    "user_id": [1, 2, 3, 4, 5],
    "month_0": [1, 1, 1, 1, 1],  # Everyone starts active
    "month_1": [1, 0, 1, 0, 1],  # 3 users active
    "month_2": [1, 0, 0, 0, 0],  # 1 user active
}
df = pd.DataFrame(data)

# Calculating retention rate: fraction of original cohort active in each month
cohort_size = len(df)
retention_rate = [df[f"month_{i}"].sum() / cohort_size for i in range(3)]

# Calculating churn rate: 1 - retention rate
churn_rate = [1 - r for r in retention_rate]

# Calculating survival rate: fraction of users still active in ALL months up to i
survival_rate = []
for i in range(3):
    still_active = df[[f"month_{j}" for j in range(i + 1)]].all(axis=1).sum()
    survival_rate.append(still_active / cohort_size)

print("retention_rate:", retention_rate)
print("churn_rate:", churn_rate)
print("survival_rate:", survival_rate)

This code calculates the required metrics for each month. The **retention rate** measures what fraction of the original cohort is active in a given month. The **churn rate** is simply one minus the retention rate, indicating the proportion that is no longer active. The **survival rate** checks for users who have remained continuously active from the beginning up to the current month - requiring a user to have a `1` in every month so far.


import unittest
import user_code
import ast
import re   
import importlib
import csv
import unittest
import pandas as pd
import importlib

class TestTask(unittest.TestCase):
    def test_basic_case(self):
        import user_code
        importlib.reload(user_code)
        data = {
            "user_id": [1, 2, 3, 4, 5],
            "month_0": [1, 1, 1, 1, 1],
            "month_1": [1, 0, 1, 0, 1],
            "month_2": [1, 0, 0, 0, 0],
        }
        df = pd.DataFrame(data)
        func = getattr(user_code, 'calculate_cohort_metrics', None)
        _dynamic_test(self, callable(func), "Function is defined", "Function calculate_cohort_metrics is not defined")
        if not callable(func):
            return
        result = func(df)
        _dynamic_test(self, result is not None and isinstance(result, (tuple, list)) and len(result) == 3,
                      "Function returns three lists",
                      f"Function must return three lists, got: {result}")
        if result is None or not isinstance(result, (tuple, list)) or len(result) != 3:
            return
        retention_rate, churn_rate, survival_rate = result
        _dynamic_test(self, all(isinstance(x, list) for x in [retention_rate, churn_rate, survival_rate]),
                      "All return values are lists",
                      f"Returned values are not all lists: {result}")
        _dynamic_test(self, retention_rate == [1.0, 0.6, 0.2],
                      "Correct retention_rate for given data",
                      f"Expected retention_rate [1.0, 0.6, 0.2], got {retention_rate}")
        _dynamic_test(self, churn_rate == [0.0, 0.4, 0.8],
                      "Correct churn_rate for given data",
                      f"Expected churn_rate [0.0, 0.4, 0.8], got {churn_rate}")
        _dynamic_test(self, survival_rate == [1.0, 0.6, 0.2],
                      "Correct survival_rate for given data",
                      f"Expected survival_rate [1.0, 0.6, 0.2], got {survival_rate}")

    def test_survival_differs_from_retention(self):
        import user_code
        importlib.reload(user_code)
        # User 2 is inactive in month_1 but returns in month_2
        data = {
            "user_id": [1, 2, 3, 4],
            "month_0": [1, 1, 1, 1],
            "month_1": [1, 0, 1, 1],
            "month_2": [1, 1, 0, 1],
        }
        df = pd.DataFrame(data)
        func = getattr(user_code, 'calculate_cohort_metrics', None)
        result = func(df)
        retention_rate, churn_rate, survival_rate = result
        # retention: [1.0, 0.75, 0.75]
        # churn:    [0.0, 0.25, 0.25]
        # survival: [1.0, 0.75, 0.5] (only users 1 and 4 are active in all months)
        _dynamic_test(self, retention_rate == [1.0, 0.75, 0.75],
                      "Retention rate correct when users return after inactivity",
                      f"Expected [1.0, 0.75, 0.75], got {retention_rate}")
        _dynamic_test(self, churn_rate == [0.0, 0.25, 0.25],
                      "Churn rate correct when users return after inactivity",
                      f"Expected [0.0, 0.25, 0.25], got {churn_rate}")
        _dynamic_test(self, survival_rate == [1.0, 0.75, 0.5],
                      "Survival rate correct when users return after inactivity",
                      f"Expected [1.0, 0.75, 0.5], got {survival_rate}")

    def test_all_churned(self):
        import user_code
        importlib.reload(user_code)
        data = {
            "user_id": [1, 2, 3],
            "month_0": [1, 1, 1],
            "month_1": [0, 0, 0],
            "month_2": [0, 0, 0],
        }
        df = pd.DataFrame(data)
        func = getattr(user_code, 'calculate_cohort_metrics', None)
        result = func(df)
        retention_rate, churn_rate, survival_rate = result
        _dynamic_test(self, retention_rate == [1.0, 0.0, 0.0],
                      "Retention rate is zero after all churn",
                      f"Expected [1.0, 0.0, 0.0], got {retention_rate}")
        _dynamic_test(self, churn_rate == [0.0, 1.0, 1.0],
                      "Churn rate is one after all churn",
                      f"Expected [0.0, 1.0, 1.0], got {churn_rate}")
        _dynamic_test(self, survival_rate == [1.0, 0.0, 0.0],
                      "Survival rate is zero after all churn",
                      f"Expected [1.0, 0.0, 0.0], got {survival_rate}")

    def test_single_user(self):
        import user_code
        importlib.reload(user_code)
        data = {
            "user_id": [42],
            "month_0": [1],
            "month_1": [1],
            "month_2": [0],
        }
        df = pd.DataFrame(data)
        func = getattr(user_code, 'calculate_cohort_metrics', None)
        result = func(df)
        retention_rate, churn_rate, survival_rate = result
        _dynamic_test(self, retention_rate == [1.0, 1.0, 0.0],
                      "Retention rate correct for single user",
                      f"Expected [1.0, 1.0, 0.0], got {retention_rate}")
        _dynamic_test(self, churn_rate == [0.0, 0.0, 1.0],
                      "Churn rate correct for single user",
                      f"Expected [0.0, 0.0, 1.0], got {churn_rate}")
        _dynamic_test(self, survival_rate == [1.0, 1.0, 0.0],
                      "Survival rate correct for single user",
                      f"Expected [1.0, 1.0, 0.0], got {survival_rate}")

    def test_month_column_order(self):
        import user_code
        importlib.reload(user_code)
        data = {
            "user_id": [1,2,3],
            "month_2": [1,0,1],
            "month_0": [1,1,1],
            "month_1": [1,0,1],
        }
        df = pd.DataFrame(data)[["user_id", "month_2", "month_0", "month_1"]]  # out of order
        func = getattr(user_code, 'calculate_cohort_metrics', None)
        result = func(df)
        # Should sort columns by month number
        # month_0: [1,1,1] sum=3
        # month_1: [1,0,1] sum=2
        # month_2: [1,0,1] sum=2
        retention_rate, churn_rate, survival_rate = result
        _dynamic_test(self, retention_rate == [1.0, 0.6666666666666666, 0.6666666666666666],
                      "Handles month columns out of order",
                      f"Expected [1.0, 0.666..., 0.666...], got {retention_rate}")
        _dynamic_test(self, churn_rate == [0.0, 0.33333333333333337, 0.33333333333333337],
                      "Handles churn for out of order columns",
                      f"Expected [0.0, 0.333..., 0.333...], got {churn_rate}")
        _dynamic_test(self, survival_rate == [1.0, 0.6666666666666666, 0.6666666666666666],
                      "Handles survival for out of order columns",
                      f"Expected [1.0, 0.666..., 0.666...], got {survival_rate}")

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\\s{2,}", " ", text)
    text = re.sub(r"\\s*([,:?])\\s*", r"\\1 ", text)
    return text.strip()

def change_var(code: str, var_name: str, value: str) -> str:
    tree = ast.parse(code)
    lines = code.splitlines()
    changed = False
    # Collect all assignment nodes to modify
    assign_nodes = [
        (i, node)
        for i, node in enumerate(tree.body)
        if isinstance(node, ast.Assign)
        and any(isinstance(target, ast.Name) and target.id == var_name for target in node.targets)
    ]

    # If nothing to change, return unmodified code
    if not assign_nodes:
        return code

    # Perform replacements for all matching assignments (from last to first to not break line offsets)
    for i, node in reversed(assign_nodes):
        start_line = node.lineno - 1
        line = lines[start_line]
        indent = ' ' * (len(line) - len(line.lstrip()))
        lines[start_line] = f"{indent}{var_name} = {value}"
        next_line = len(lines)
        for next_node in tree.body[i+1:]:
            if hasattr(next_node, 'lineno'):
                next_line = next_node.lineno - 1
                break
        if next_line > start_line + 1:
            lines[start_line+1:next_line] = []
        changed = True

    return '\\n'.join(lines) if changed else code

if __name__ == "__main__":
    unittest.main()


Challenge: Mapping the Lifecycle of a SpaceStream Explorer

Løsning