Fix ShiftedBetaGeoModel Cohort Bug (#2115)

* pred method bug comments

* fix sbg pred bug

---------

Co-authored-by: Juan Orduz <juanitorduz@gmail.com>

Author: ColtAllen

.gitignore

@@ -147,3 +147,6 @@ dmypy.json
 # Gallery images
 docs/source/gallery/images/
 docs/gettext/
+
+# Sandbox folder
+sandbox/

docs/source/notebooks/clv/dev/sBG_cohort.ipynb

@@ -928,9 +928,9 @@ class ShiftedBetaGeometric(Discrete):
     Hardie and Fader describe this distribution with the following PMF and survival functions in [1]_:
 
     .. math::
-        \mathbb{P}(T=t|\alpha,\beta) = (\frac{B(\alpha+1,\beta+t-1)}{B(\alpha,\beta}),t=1,2,...  \\
+        \mathbb{P}(T=t|\alpha,\beta) = \frac{B(\alpha+1,\beta+t-1)}{B(\alpha,\beta)},t=1,2,...  \\
         \begin{align}
-        \mathbb{S}(t|\alpha,\beta) = (\frac{B(\alpha,\beta+t)}{B(\alpha,\beta}),t=1,2,... \\
+        \mathbb{S}(t|\alpha,\beta) = \frac{B(\alpha,\beta+t)}{B(\alpha,\beta)},t=1,2,... \\
         \end{align}
 
     ========  ===============================================

pymc_marketing/clv/distributions.py

@@ -421,7 +421,7 @@ def _extract_predictive_variables(
 
             # Map cohort indices for each customer
             pred_cohort_idx = pd.Categorical(
-                customer_cohort_map.values, categories=self.cohorts
+                customer_cohort_map.values, categories=cohorts_present
             ).codes
 
             # Reconstruct customer-level parameters

pymc_marketing/clv/models/shifted_beta_geo.py

@@ -0,0 +1,106 @@
+import numpy as np
+import pandas as pd
+
+# Summary data from original research paper for ShiftedBetaGeoModel
+# From Table 1 in https://faculty.wharton.upenn.edu/wp-content/uploads/2012/04/Fader_hardie_jim_07.pdf
+sbg_research_data = pd.DataFrame(
+    {
+        "highend": [
+            100.0,
+            86.9,
+            74.3,
+            65.3,
+            59.3,
+            55.1,
+            51.7,
+            49.1,
+            46.8,
+            44.5,
+            42.7,
+            40.9,
+            39.4,
+        ],
+        "regular": [
+            100.0,
+            63.1,
+            46.8,
+            38.2,
+            32.6,
+            28.9,
+            26.2,
+            24.1,
+            22.3,
+            20.7,
+            19.4,
+            18.3,
+            17.3,
+        ],
+    }
+)
+
+
+def generate_sbg_data(
+    n_customers: int,
+    n_time_periods: int,
+    survival_data: pd.DataFrame,
+) -> pd.DataFrame:
+    """
+    Generate individual-level customer churn data from aggregate percentage alive data.
+
+    Parameters
+    ----------
+    n_customers : int
+        Number of customers to simulate for each cohort
+    n_time_periods : int
+        Number of time periods to include from the survival data
+    survival_data : pd.DataFrame, optional
+        DataFrame with columns representing cohorts and values as percentage alive.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame with columns: recency, T, cohort
+        Contains individual customer data for all cohorts in the survival data
+    """
+
+    def _individual_data_from_percentage_alive(percentage_alive, initial_customers):
+        """Convert percentage alive data to individual churn times."""
+        n_alive = np.asarray(percentage_alive / 100 * initial_customers, dtype=int)
+
+        died_at = np.zeros((initial_customers,), dtype=int)
+        counter = 0
+        for t, diff in enumerate((n_alive[:-1] - n_alive[1:]), start=1):
+            died_at[counter : counter + diff] = t
+            counter += diff
+
+        censoring_t = t + 1
+        died_at[counter:] = censoring_t
+
+        return died_at
+
+    # Truncate data to requested number of time periods
+    truncated_df = survival_data[:n_time_periods]
+
+    # Generate individual churn data for each cohort
+    datasets = []
+    for cohort_name in truncated_df.columns:
+        churn_data = _individual_data_from_percentage_alive(
+            truncated_df[cohort_name], n_customers
+        )
+
+        dataset = pd.DataFrame(
+            {
+                "recency": churn_data,
+                "T": n_time_periods,
+                "cohort": cohort_name,
+            }
+        )
+        datasets.append(dataset)
+
+    # Combine all cohorts into a single dataset
+    combined_dataset = pd.concat(datasets, ignore_index=True)
+
+    # Create customer_id column from index
+    combined_dataset["customer_id"] = combined_dataset.index + 1
+
+    return combined_dataset[["customer_id", "recency", "T", "cohort"]]

scripts/data_generators/clv_data_generation.py

@@ -928,6 +928,40 @@ def test_covariate_cols_only_in_config(self, covariate_test_data):
         model.build_model()
         assert "dropout_covariate" in model.model.coords
 
+    def test_predictions_with_covariates_subset_cohorts(self):
+        # Training data with 3 cohorts
+        train_data = pd.DataFrame(
+            {
+                "customer_id": range(30),
+                "recency": [3, 4, 5] * 10,
+                "T": [5] * 30,
+                "cohort": ["A"] * 10 + ["B"] * 10 + ["C"] * 10,
+                "channel": [0, 1, 0, 1, 0] * 6,
+            }
+        )
+
+        model = ShiftedBetaGeoModel(
+            data=train_data, model_config={"dropout_covariate_cols": ["channel"]}
+        )
+        model.fit(method="map", maxeval=10)
+
+        # Prediction data with subset of cohorts (NOT starting at index 0)
+        pred_data = pd.DataFrame(
+            {
+                "customer_id": [100, 101, 102],
+                "T": [3, 3, 3],
+                "cohort": ["B", "C", "C"],  # Missing cohort "A"
+                "channel": [1, 0, 1],
+            }
+        )
+
+        # Should not raise IndexError
+        prob_alive = model.expected_probability_alive(data=pred_data, future_t=1)
+        assert prob_alive.shape[-1] == 3  # 3 customers
+
+        retention = model.expected_retention_rate(data=pred_data, future_t=1)
+        assert retention.shape[-1] == 3
+
 
 class TestShiftedBetaGeoModelIndividual:
     @classmethod