Update BetaGeoModel API (#709)

* _extract_predictive_variables util * deprecation warnings * expected_purchases_new_customer * TestBetaGeoModel.setup_class * test cleanup and test_expected_purchases_new_customer * expected_probability_alive * expected_purchases * TODOs and docstrings * update runslow tests rtol for new test dataset * prob_alive_matrix plot fix and notebook testing * alive loop in bgnbd nb * quickstart nb fix * docstring and TODO revisions * docstring syntax * docstring edits * more docstring fixes * docstrings indent * docstring indent, clear codecov bug --------- Co-authored-by: Juan Orduz <[email protected]>
pymc-labs · Sep 10, 2024 · 7774e7c · 7774e7c
1 parent c74d369
commit 7774e7c
Show file tree

Hide file tree

Showing 31 changed files with 7,463 additions and 5,511 deletions.
diff --git a/docs/source/notebooks/clv/bg_nbd.ipynb b/docs/source/notebooks/clv/bg_nbd.ipynb
diff --git a/docs/source/notebooks/clv/clv_quickstart.ipynb b/docs/source/notebooks/clv/clv_quickstart.ipynb
diff --git a/docs/source/notebooks/mmm/mmm_budget_allocation_example.ipynb b/docs/source/notebooks/mmm/mmm_budget_allocation_example.ipynb
diff --git a/docs/source/notebooks/mmm/mmm_example.ipynb b/docs/source/notebooks/mmm/mmm_example.ipynb
diff --git a/docs/source/notebooks/mmm/mmm_lift_test.ipynb b/docs/source/notebooks/mmm/mmm_lift_test.ipynb
diff --git a/docs/source/notebooks/mmm/mmm_tvp_example.ipynb b/docs/source/notebooks/mmm/mmm_tvp_example.ipynb
diff --git a/data/budget_optimizer_model.nc → docs/source/notebooks/mmm/model.nc b/data/budget_optimizer_model.nc → docs/source/notebooks/mmm/model.nc
diff --git a/pymc_marketing/clv/models/beta_geo.py b/pymc_marketing/clv/models/beta_geo.py
diff --git a/pymc_marketing/clv/plotting.py b/pymc_marketing/clv/plotting.py
@@ -331,34 +331,27 @@ def plot_probability_alive_matrix(
         max_frequency=max_frequency,
         max_recency=max_recency,
     )
+    transaction_data = pd.DataFrame(
+        {
+            "customer_id": np.arange(mesh_recency.size),  # placeholder
+            "frequency": mesh_frequency.ravel(),
+            "recency": mesh_recency.ravel(),
+            "T": max_recency,
+        }
+    )
     # FIXME: This is a hotfix for ParetoNBDModel, as it has a different API from BetaGeoModel
-    #  We should harmonize them!
     if isinstance(model, ParetoNBDModel):
-        transaction_data = pd.DataFrame(
-            {
-                "customer_id": np.arange(mesh_recency.size),  # placeholder
-                "frequency": mesh_frequency.ravel(),
-                "recency": mesh_recency.ravel(),
-                "T": max_recency,
-            }
-        )
-
         Z = (
             model.expected_probability_alive(
                 data=transaction_data,
-                future_t=0,  # TODO: This can be a function parameter in the case of ParetoNBDModel
+                future_t=0,  # TODO: This is a required parameter if data is provided.
             )
             .mean(("draw", "chain"))
             .values.reshape(mesh_recency.shape)
         )
     else:
         Z = (
-            model.expected_probability_alive(
-                customer_id=np.arange(mesh_recency.size),  # placeholder
-                frequency=mesh_frequency.ravel(),
-                recency=mesh_recency.ravel(),
-                T=max_recency,  # type: ignore
-            )
+            model.expected_probability_alive(data=transaction_data)
             .mean(("draw", "chain"))
             .values.reshape(mesh_recency.shape)
         )

diff --git a/pymc_marketing/mmm/__init__.py b/pymc_marketing/mmm/__init__.py
@@ -12,8 +12,25 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 from pymc_marketing.mmm import base, delayed_saturated_mmm, preprocessing, validating
-from pymc_marketing.mmm.base import MMM, BaseMMM
-from pymc_marketing.mmm.delayed_saturated_mmm import DelayedSaturatedMMM
+from pymc_marketing.mmm.base import (
+    BaseValidateMMM,
+    MMMModelBuilder,
+)
+from pymc_marketing.mmm.components.adstock import (
+    AdstockTransformation,
+    DelayedAdstock,
+    GeometricAdstock,
+    WeibullAdstock,
+)
+from pymc_marketing.mmm.components.saturation import (
+    HillSaturation,
+    LogisticSaturation,
+    MichaelisMentenSaturation,
+    SaturationTransformation,
+    TanhSaturation,
+    TanhSaturationBaselined,
+)
+from pymc_marketing.mmm.delayed_saturated_mmm import MMM, DelayedSaturatedMMM
 from pymc_marketing.mmm.preprocessing import (
     preprocessing_method_X,
     preprocessing_method_y,
@@ -26,10 +43,21 @@
     "preprocessing",
     "validating",
     "MMM",
-    "BaseMMM",
+    "MMMModelBuilder",
+    "BaseValidateMMM",
     "DelayedSaturatedMMM",
     "preprocessing_method_X",
     "preprocessing_method_y",
     "validation_method_X",
     "validation_method_y",
+    "AdstockTransformation",
+    "DelayedAdstock",
+    "GeometricAdstock",
+    "WeibullAdstock",
+    "SaturationTransformation",
+    "MichaelisMentenSaturation",
+    "HillSaturation",
+    "LogisticSaturation",
+    "TanhSaturation",
+    "TanhSaturationBaselined",
 ]