From 0e31efa300654ad3fcea164437f00659c6693ebf Mon Sep 17 00:00:00 2001
From: Will Dean <57733339+wd60622@users.noreply.github.com>
Date: Wed, 22 May 2024 09:25:42 +0200
Subject: [PATCH] Correct BetaGeo docstring example (#693)

* fix example in docstring

* change docstring
---
 pymc_marketing/clv/models/beta_geo.py | 53 ++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 10 deletions(-)

diff --git a/pymc_marketing/clv/models/beta_geo.py b/pymc_marketing/clv/models/beta_geo.py
index 12f2b63ce..26c72e940 100644
--- a/pymc_marketing/clv/models/beta_geo.py
+++ b/pymc_marketing/clv/models/beta_geo.py
@@ -43,7 +43,6 @@ class BetaGeoModel(CLVModel):
     Methods below are adapted from the BetaGeoFitter class from the lifetimes package
     (see https://github.com/CamDavidsonPilon/lifetimes/).
 
-
     Parameters
     ----------
     data: pd.DataFrame
@@ -97,19 +96,28 @@ class BetaGeoModel(CLVModel):
         model.fit()
         print(model.fit_summary())
 
-        # Estimating the expected number of purchases for a randomly chosen
-        # individual in a future time period of length t
+    Predicting the customer-specific number of purchases for a future time
+    interval of length t given their previous frequency, recency, and T
+
+    .. code-block:: python
+
+        future_t = [2, 5, 7, 10]
+
         expected_num_purchases = model.expected_num_purchases(
-            t=[2, 5, 7, 10],
+            t=future_t,
+            customer_id=data["customer_id"],
+            frequency=data["frequenc"y],
+            recency=data["recency"],
+            T=data["T"],
         )
 
-        # Predicting the customer-specific number of purchases for a future
-        # time interval of length t given their previous frequency and recency
+    Estimating the expected number of purchases for a randomly chosen
+    individual in a future time period of length t
+
+    .. code-block:: python
+
         expected_num_purchases_new_customer = model.expected_num_purchases_new_customer(
-            t=[5, 5, 5, 5],
-            frequency=[5, 2, 1, 8],
-            recency=[7, 4, 2.5, 11],
-            T=[10, 8, 10, 22],
+            t=future_t,
         )
 
     References
@@ -255,6 +263,22 @@ def expected_num_purchases(
                 \frac{\alpha + T}{\alpha + T + t}
             \right)^{r + x}
            }
+
+        Parameters
+        ----------
+        customer_id : np.ndarray | pd.Series
+            Customer identifier. Not required to be from data.
+        t : np.ndarray | pd.Series | TensorVariable
+            Time interval for which to predict the number of future purchases.
+        frequency : np.ndarray | pd.Series | TensorVariable
+            Number of repeat purchases.
+        recency : np.ndarray | pd.Series | TensorVariable
+            Time between the first and the last purchase.
+
+        Returns
+        -------
+        xr.DataArray
+            Expected number of purchases in the next time interval of length :math:`t`.
         """
         t = np.asarray(t)
         if t.size != 1:
@@ -340,6 +364,15 @@ def expected_num_purchases_new_customer(
                 \text{hyp2f1}\left(r, b; a + b - 1; \frac{t}{\alpha + t}\right)
             \right]
 
+        Parameters
+        ----------
+        t : np.ndarray | pd.Series
+            Time interval for which to predict the number of future purchases.
+
+        Returns
+        -------
+        xr.DataArray
+            Expected number of purchases in the next time interval of length :math:`t`.
         """
         t = np.asarray(t)
         if t.size != 1: