Merge branch 'main' into bass-model

.github/labeler.yml

@@ -26,7 +26,6 @@ CLV:
     - any-glob-to-any-file:
       - pymc_marketing/prior.py
 
-
 streamlit:
   - changed-files:
     - any-glob-to-any-file:
@@ -41,3 +40,8 @@ mlflow:
   - changed-files:
     - any-glob-to-any-file:
       - pymc_marketing/mlflow.py
+
+ModelBuilder:
+  - changed-files:
+    - any-glob-to-any-file:
+      - pymc_marketing/model_builder.py

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
           - --exclude=docs/
           - --exclude=scripts/
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.8.4
+    rev: v0.8.6
     hooks:
       - id: ruff
         types_or: [python, pyi, jupyter]
@@ -21,7 +21,7 @@ repos:
         types_or: [python, pyi, jupyter]
         exclude: ^docs/source/notebooks/clv/dev/
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.14.0
+    rev: v1.14.1
     hooks:
       - id: mypy
         args: [--ignore-missing-imports]

docs/source/notebooks/index.md

@@ -46,4 +46,5 @@ customer_choice/mv_its_unsaturated
 
 general/model_configuration
 general/other_nuts_samplers
+general/prior_predictive
 :::

environment.yml

@@ -15,7 +15,7 @@ dependencies:
 - pydantic
 - preliz
 # NOTE: Keep minimum pymc version in sync with ci.yml `OLDEST_PYMC_VERSION`
-- pymc>=5.12.0,<5.16.0
+- pymc>=5.20.0
 - scikit-learn>=1.1.1
 - seaborn>=0.12.2
 - xarray

pymc_marketing/clv/distributions.py

@@ -13,8 +13,9 @@
 #   limitations under the License.
 """Distributions for the CLV module."""
 
+from functools import reduce
+
 import numpy as np
-import pymc as pm
 import pytensor.tensor as pt
 from pymc.distributions.continuous import PositiveContinuous
 from pymc.distributions.dist_math import betaln, check_parameters
@@ -28,26 +29,16 @@
 
 class ContNonContractRV(RandomVariable):
     name = "continuous_non_contractual"
-    ndim_supp = 1
-    ndims_params = [0, 0, 0, 0]
+    signature = "(),(),()->(2)"
     dtype = "floatX"
     _print_name = ("ContNonContract", "\\operatorname{ContNonContract}")
 
-    def make_node(self, rng, size, dtype, lam, p, T):
-        T = pt.as_tensor_variable(T)
-
-        return super().make_node(rng, size, dtype, lam, p, T)
+    def __call__(self, lam, p, T, size=None, **kwargs):
+        return super().__call__(lam, p, T, size=size, **kwargs)
 
     @classmethod
     def rng_fn(cls, rng, lam, p, T, size):
-        size = pm.distributions.shape_utils.to_tuple(size)
-
-        # TODO: broadcast sizes
-        lam = np.asarray(lam)
-        p = np.asarray(p)
-        T = np.asarray(T)
-
-        if size == ():
+        if size is None:
             size = np.broadcast_shapes(lam.shape, p.shape, T.shape)
 
         lam = np.broadcast_to(lam, size)
@@ -74,9 +65,6 @@ def rng_fn(cls, rng, lam, p, T, size):
 
         return np.stack([t_x, x], axis=-1)
 
-    def _supp_shape_from_params(*args, **kwargs):
-        return (2,)
-
 
 continuous_non_contractual = ContNonContractRV()
 
@@ -129,13 +117,14 @@ def logp(value, lam, p, T):
         )
 
         logp = pt.switch(
-            pt.any(
-                (
+            reduce(
+                pt.bitwise_or,
+                [
                     pt.and_(pt.ge(t_x, 0), zero_observations),
                     pt.lt(t_x, 0),
                     pt.lt(x, 0),
                     pt.gt(t_x, T),
-                ),
+                ],
             ),
             -np.inf,
             logp,
@@ -152,29 +141,16 @@ def logp(value, lam, p, T):
 
 class ContContractRV(RandomVariable):
     name = "continuous_contractual"
-    ndim_supp = 1
-    ndims_params = [0, 0, 0, 0]
+    signature = "(),(),()->(3)"
     dtype = "floatX"
     _print_name = ("ContinuousContractual", "\\operatorname{ContinuousContractual}")
 
-    def make_node(self, rng, size, dtype, lam, p, T):
-        T = pt.as_tensor_variable(T)
-
-        return super().make_node(rng, size, dtype, lam, p, T)
-
     def __call__(self, lam, p, T, size=None, **kwargs):
         return super().__call__(lam, p, T, size=size, **kwargs)
 
     @classmethod
     def rng_fn(cls, rng, lam, p, T, size):
-        size = pm.distributions.shape_utils.to_tuple(size)
-
-        # To do: broadcast sizes
-        lam = np.asarray(lam)
-        p = np.asarray(p)
-        T = np.asarray(T)
-
-        if size == ():
+        if size is None:
             size = np.broadcast_shapes(lam.shape, p.shape, T.shape)
 
         lam = np.broadcast_to(lam, size)
@@ -254,24 +230,15 @@ def logp(value, lam, p, T):
         )
 
         logp = pt.switch(
-            pt.any(pt.or_(pt.lt(t_x, 0), zero_observations)),
-            -np.inf,
-            logp,
-        )
-        logp = pt.switch(
-            pt.all(
-                pt.or_(pt.eq(churn, 0), pt.eq(churn, 1)),
-            ),
-            logp,
-            -np.inf,
-        )
-        logp = pt.switch(
-            pt.any(
-                (
+            reduce(
+                pt.bitwise_or,
+                [
+                    zero_observations,
                     pt.lt(t_x, 0),
                     pt.lt(x, 0),
                     pt.gt(t_x, T),
-                ),
+                    pt.bitwise_not(pt.bitwise_or(pt.eq(churn, 0), pt.eq(churn, 1))),
+                ],
             ),
             -np.inf,
             logp,
@@ -289,34 +256,16 @@ def logp(value, lam, p, T):
 
 class ParetoNBDRV(RandomVariable):
     name = "pareto_nbd"
-    ndim_supp = 1
-    ndims_params = [0, 0, 0, 0, 0]
+    signature = "(),(),(),(),()->(2)"
     dtype = "floatX"
     _print_name = ("ParetoNBD", "\\operatorname{ParetoNBD}")
 
-    def make_node(self, rng, size, dtype, r, alpha, s, beta, T):
-        r = pt.as_tensor_variable(r)
-        alpha = pt.as_tensor_variable(alpha)
-        s = pt.as_tensor_variable(s)
-        beta = pt.as_tensor_variable(beta)
-        T = pt.as_tensor_variable(T)
-
-        return super().make_node(rng, size, dtype, r, alpha, s, beta, T)
-
     def __call__(self, r, alpha, s, beta, T, size=None, **kwargs):
         return super().__call__(r, alpha, s, beta, T, size=size, **kwargs)
 
     @classmethod
     def rng_fn(cls, rng, r, alpha, s, beta, T, size):
-        size = pm.distributions.shape_utils.to_tuple(size)
-
-        r = np.asarray(r)
-        alpha = np.asarray(alpha)
-        s = np.asarray(s)
-        beta = np.asarray(beta)
-        T = np.asarray(T)
-
-        if size == ():
+        if size is None:
             size = np.broadcast_shapes(
                 r.shape, alpha.shape, s.shape, beta.shape, T.shape
             )
@@ -357,9 +306,6 @@ def sim_data(lam, mu, T):
 
         return output
 
-    def _supp_shape_from_params(*args, **kwargs):
-        return (2,)
-
 
 pareto_nbd = ParetoNBDRV()
 
@@ -489,34 +435,16 @@ def logp(value, r, alpha, s, beta, T):
 
 class BetaGeoBetaBinomRV(RandomVariable):
     name = "beta_geo_beta_binom"
-    ndim_supp = 1
-    ndims_params = [0, 0, 0, 0, 0]
+    signature = "(),(),(),(),()->(2)"
     dtype = "floatX"
     _print_name = ("BetaGeoBetaBinom", "\\operatorname{BetaGeoBetaBinom}")
 
-    def make_node(self, rng, size, dtype, alpha, beta, gamma, delta, T):
-        alpha = pt.as_tensor_variable(alpha)
-        beta = pt.as_tensor_variable(beta)
-        gamma = pt.as_tensor_variable(gamma)
-        delta = pt.as_tensor_variable(delta)
-        T = pt.as_tensor_variable(T)
-
-        return super().make_node(rng, size, dtype, alpha, beta, gamma, delta, T)
-
     def __call__(self, alpha, beta, gamma, delta, T, size=None, **kwargs):
         return super().__call__(alpha, beta, gamma, delta, T, size=size, **kwargs)
 
     @classmethod
     def rng_fn(cls, rng, alpha, beta, gamma, delta, T, size) -> np.ndarray:
-        size = pm.distributions.shape_utils.to_tuple(size)
-
-        alpha = np.asarray(alpha)
-        beta = np.asarray(beta)
-        gamma = np.asarray(gamma)
-        delta = np.asarray(delta)
-        T = np.asarray(T)
-
-        if size == ():
+        if size is None:
             size = np.broadcast_shapes(
                 alpha.shape, beta.shape, gamma.shape, delta.shape, T.shape
             )
@@ -557,9 +485,6 @@ def sim_data(purchase_prob, churn_prob, T):
 
         return output
 
-    def _supp_shape_from_params(*args, **kwargs):
-        return (2,)
-
 
 beta_geo_beta_binom = BetaGeoBetaBinomRV()
 

pymc_marketing/clv/models/basic.py

@@ -26,7 +26,6 @@
 from pymc.backends import NDArray
 from pymc.backends.base import MultiTrace
 from pymc.model.core import Model
-from xarray import Dataset
 
 from pymc_marketing.model_builder import ModelBuilder
 from pymc_marketing.model_config import ModelConfig, parse_model_config
@@ -256,23 +255,6 @@ def default_sampler_config(self) -> dict:
     def _serializable_model_config(self) -> dict:
         return self.model_config
 
-    @property
-    def fit_result(self) -> Dataset:
-        """Get the fit result."""
-        if self.idata is None or "posterior" not in self.idata:
-            raise RuntimeError("The model hasn't been fit yet, call .fit() first")
-        return self.idata["posterior"]
-
-    @fit_result.setter
-    def fit_result(self, res: az.InferenceData) -> None:
-        if self.idata is None:
-            self.idata = res
-        elif "posterior" in self.idata:
-            warnings.warn("Overriding pre-existing fit_result", stacklevel=1)
-            self.idata.posterior = res
-        else:
-            self.idata.posterior = res
-
     def fit_summary(self, **kwargs):
         """Compute the summary of the fit result."""
         res = self.fit_result