Make data optional in all ParetoNBD methods

pymc-labs · Mar 15, 2024 · 72238c5 · 72238c5
1 parent d9c19c4
commit 72238c5
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/pymc_marketing/clv/models/pareto_nbd.py b/pymc_marketing/clv/models/pareto_nbd.py
@@ -915,7 +915,6 @@ def distribution_new_customer_dropout(
         data: pd.DataFrame, optional
             DataFrame containing the following columns:
                 * `customer_id`: unique customer identifier
-                * `T`: time between the first purchase and the end of the observation period.
                 * covariates: Purchase and dropout covariate columns if original model had any.
             If not provided, the method will use the fit dataset.
         random_seed: RandomState, optional
@@ -934,7 +933,7 @@ def distribution_new_customer_dropout(
 
     def distribution_new_customer_purchase_rate(
         self,
-        data,
+        data: Optional[pd.DataFrame] = None,
         *,
         random_seed: Optional[RandomState] = None,
     ) -> xarray.Dataset:
@@ -945,6 +944,11 @@ def distribution_new_customer_purchase_rate(
 
         Parameters
         ----------
+        data: pd.DataFrame, optional
+            DataFrame containing the following columns:
+                * `customer_id`: unique customer identifier
+                * covariates: Purchase and dropout covariate columns if original model had any.
+            If not provided, the method will use the fit dataset.
         random_seed : RandomState, optional
             Random state to use for sampling.
 
@@ -961,7 +965,7 @@ def distribution_new_customer_purchase_rate(
 
     def distribution_new_customer_recency_frequency(
         self,
-        data,
+        data: Optional[pd.DataFrame] = None,
         *,
         T: Optional[Union[int, np.ndarray, pd.Series]] = None,
         random_seed: Optional[RandomState] = None,