Implement ParetoNBD with covariates

pymc_marketing/clv/distributions.py

@@ -320,7 +320,8 @@
             dropout_time = rng.exponential(scale=1 / mu)
             wait = rng.exponential(scale=1 / lam)
 
-            while t + wait < min(dropout_time, T):
+            final_t = min(dropout_time, T)
+            while (t + wait) < final_t:
                 t += wait
                 n += 1
                 wait = rng.exponential(scale=1 / lam)

pymc_marketing/clv/models/beta_geo.py

@@ -352,17 +352,18 @@
             alpha = pm.HalfFlat("alpha")
             r = pm.HalfFlat("r")
 
-            # This is the shape with fit_method="map"
-            if self.fit_result.dims == {"chain": 1, "draw": 1}:
-                shape_kwargs = {"shape": 1000}
-            else:
-                shape_kwargs = {}
+            fit_result = self.fit_result
+            if fit_result.sizes["chain"] == 1 and fit_result.sizes["draw"] == 1:
+                # For map fit add a dummy draw dimension
+                fit_result = self.fit_result.squeeze("draw").expand_dims(
+                    draw=range(1000)
+                )
 
-            pm.Beta("population_dropout", alpha=a, beta=b, **shape_kwargs)
-            pm.Gamma("population_purchase_rate", alpha=r, beta=alpha, **shape_kwargs)
+            pm.Beta("population_dropout", alpha=a, beta=b)
+            pm.Gamma("population_purchase_rate", alpha=r, beta=alpha)
 
             return pm.sample_posterior_predictive(
-                self.fit_result,
+                fit_result,
                 var_names=var_names,
                 random_seed=random_seed,
             ).posterior_predictive

tests/clv/models/test_basic.py

@@ -7,6 +7,7 @@
 from arviz import InferenceData, from_dict
 
 from pymc_marketing.clv.models.basic import CLVModel
+from tests.clv.utils import set_model_fit
 
 
 class CLVModelTest(CLVModel):
@@ -54,10 +55,10 @@ def test_create_distribution_from_wrong_prior(self):
         model = CLVModelTest()
         with pytest.raises(
             ValueError,
-            match="Distribution definately_not_PyMC_dist does not exist in PyMC",
+            match="Distribution definitely_not_PyMC_dist does not exist in PyMC",
         ):
             model._create_distribution(
-                {"dist": "definately_not_PyMC_dist", "kwargs": {"alpha": 1, "beta": 1}}
+                {"dist": "definitely_not_PyMC_dist", "kwargs": {"alpha": 1, "beta": 1}}
             )
 
     def test_fit_mcmc(self):
@@ -165,9 +166,7 @@ def test_thin_fit_result(self):
         model = CLVModelTest(data=data)
         model.build_model()
         fake_idata = from_dict(dict(x=np.random.normal(size=(4, 1000))))
-        fake_idata.add_groups(dict(fit_data=data.to_xarray()))
-        model.set_idata_attrs(fake_idata)
-        model.idata = fake_idata
+        set_model_fit(model, fake_idata)
 
         thin_model = model.thin_fit_result(keep_every=20)
         assert thin_model is not model

tests/clv/models/test_gamma_gamma.py

@@ -10,6 +10,7 @@
     GammaGammaModel,
     GammaGammaModelIndividual,
 )
+from tests.clv.utils import set_model_fit
 
 
 class BaseTestGammaGammaModel:
@@ -232,8 +233,7 @@ def test_new_customer_spend(self, distribution):
         fake_fit = pm.sample_prior_predictive(
             samples=1000, model=model.model, random_seed=self.rng
         )
-        fake_fit.add_groups(dict(posterior=fake_fit.prior))
-        model.idata = fake_fit
+        set_model_fit(model, fake_fit.prior)
         # Closed formula solution for the mean and var of the population spend (eqs 3, 4 from [1])  # noqa: E501
         expected_preds_mean = p_mean * v_mean / (q_mean - 1)
         expected_preds_std = np.sqrt(