Bump PyMC dependency

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()
 

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
     "pandas",
     "pydantic>=2.1.0",
     # NOTE: Used as minimum pymc version with ci.yml `OLDEST_PYMC_VERSION`
-    "pymc>=5.13.0,<5.16.0",
+    "pymc>=5.19.1",
     "scikit-learn>=1.1.1",
     "seaborn>=0.12.2",
     "xarray>=2024.1.0",