Adding a distribution class for the continuous contractual setting (#36)

* Adding continuous contractual distribution + notebook

* Fix logp method w/ Ricardo

* Removing test file for debugging

* Adding tests for ContContract

* Fix prior sampling test

pymmmc/__init__.py

@@ -1,7 +1,8 @@
 __version__ = "0.0.1"
 
-from pymmmc.distributions import ContNonContract
+from pymmmc.distributions import ContContract, ContNonContract
 
 __all__ = [
     "ContNonContract",
+    "ContContract",
 ]

pymmmc/distributions.py

@@ -143,3 +143,152 @@ def logp(value, lam, p, T, T0):
             at.all(T0 < T),
             msg="lam > 0, 0 <= p <= 1, T0 < T",
         )
+
+
+class ContContractRV(RandomVariable):
+    name = "continuous_contractual"
+    ndim_supp = 1
+    ndims_params = [0, 0, 0, 0]
+    dtype = "floatX"
+    _print_name = ("ContinuousContractual", "\\operatorname{ContinuousContractual}")
+
+    def make_node(self, rng, size, dtype, lam, p, T, T0):
+
+        T = at.as_tensor_variable(T)
+        T0 = at.as_tensor_variable(T0)
+
+        return super().make_node(rng, size, dtype, lam, p, T, T0)
+
+    def __call__(self, lam, p, T, T0=0, size=None, **kwargs):
+        return super().__call__(lam, p, T, T0, size=size, **kwargs)
+
+    @classmethod
+    def rng_fn(cls, rng, lam, p, T, T0, size) -> np.array:
+
+        size = pm.distributions.shape_utils.to_tuple(size)
+
+        # To do: broadcast sizes
+        lam = np.asarray(lam)
+        p = np.asarray(p)
+        T = np.asarray(T)
+        T0 = np.asarray(T0)
+
+        if size is None:
+            size = np.broadcast_shapes(lam.shape, p.shape, T.shape, T0.shape)
+
+        lam = np.broadcast_to(lam, size)
+        p = np.broadcast_to(p, size)
+        T = np.broadcast_to(T, size)
+        T0 = np.broadcast_to(T0, size)
+
+        output = np.zeros(shape=size + (3,))
+
+        def sim_data(lam, p, T, T0):
+            t = 0
+            n = 0
+
+            while True:
+                wait = rng.exponential(scale=1 / lam)
+                dropout = rng.binomial(n=1, p=p)
+
+                if t + wait > T:
+                    break
+                else:
+                    t += wait
+                    n += 1
+
+                    if dropout == 1:
+                        break
+
+            return np.array(
+                [
+                    t,
+                    n,
+                    dropout,
+                ],
+            )
+
+        for index in np.ndindex(*size):
+            output[index] = sim_data(lam[index], p[index], T[index], T0[index])
+
+        return output
+
+    def _supp_shape_from_params(*args, **kwargs):
+        return (3,)
+
+
+continuous_contractual = ContContractRV()
+
+
+class ContContract(PositiveContinuous):
+    r"""
+    Distribution class of a continuous contractual data-generating process,
+    that is where purchases can occur at any time point (continuous) and
+    churning/dropping out is explicit (contractual).
+    .. math:
+
+        f(\lambda, p | d, x, t_1, \dots, t_x, T)
+        = f(\lambda, p | t_x, T) = (1 - p)^{x-1} \lambda^x \exp(-\lambda t_x)
+        p^d \left\{(1-p)\exp(-\lambda*(T - t_x))\right\}^{1 - d}
+
+    ========  ===============================================
+    Support   :math:`t_j > 0` for :math:`j = 1, \dots, x`
+    Mean      :math:`\mathbb{E}[X(t) | \lambda, p, d] = \frac{1}{p}`
+                    `- \frac{1}{p}\exp\left(-\lambda p \min(t, T)\right)`
+    """
+    rv_op = continuous_contractual
+
+    @classmethod
+    def dist(cls, lam, p, T, T0, **kwargs):
+        return super().dist([lam, p, T, T0], **kwargs)
+
+    def logp(value, lam, p, T, T0):
+        t_x = value[..., 0]
+        x = value[..., 1]
+        churn = value[..., 2]
+
+        zero_observations = at.eq(x, 0)
+
+        logp = (x - 1) * at.log(1 - p) + x * at.log(lam) - lam * t_x
+        logp += churn * at.log(p) + (1 - churn) * (
+            at.log(1 - p) - lam * ((T - T0) - t_x)
+        )
+
+        logp = at.switch(
+            zero_observations,
+            -lam * (T - T0),
+            logp,
+        )
+
+        logp = at.switch(
+            at.any(at.or_(at.lt(t_x, 0), at.lt(x, 0))),
+            -np.inf,
+            logp,
+        )
+        logp = at.switch(
+            at.all(
+                at.or_(at.eq(churn, 0), at.eq(churn, 1)),
+            ),
+            logp,
+            -np.inf,
+        )
+        logp = at.switch(
+            at.any(
+                (
+                    at.lt(t_x, T0),
+                    at.lt(x, 0),
+                    at.gt(t_x, T),
+                ),
+            ),
+            -np.inf,
+            logp,
+        )
+
+        return check_parameters(
+            logp,
+            lam > 0,
+            0 <= p,
+            p <= 1,
+            at.all(T0 < T),
+            msg="lam > 0, 0 <= p <= 1, T0 < T",
+        )

pymmmc/tests/test_distributions.py

@@ -5,7 +5,7 @@
 from pymc import Model
 from pymc.tests.helpers import select_by_precision
 
-from pymmmc.distributions import ContNonContract
+from pymmmc.distributions import ContContract, ContNonContract
 
 
 class TestContNonContract:
@@ -78,3 +78,77 @@ def test_continuous_non_contractual_sample_prior(
             prior = pm.sample_prior_predictive(samples=100)
 
         assert prior["prior"]["cnc"][0].shape == (100,) + expected_size
+
+
+class TestContContract:
+    @pytest.mark.parametrize(
+        "value, lam, p, T, T0, logp",
+        [
+            (np.array([6.3, 5, 1]), 0.3, 0.15, 12, 2, -10.45705972),
+            (
+                np.array([6.3, 5, 1]),
+                np.array([0.3, 0.2]),
+                0.15,
+                12,
+                2,
+                np.array([-10.45705972, -11.85438527]),
+            ),
+            (
+                np.array([[6.3, 5, 1], [5.3, 4, 0]]),
+                np.array([0.3, 0.2]),
+                0.15,
+                12,
+                2,
+                np.array([-10.45705972, -9.08782737]),
+            ),
+            (
+                np.array([6.3, 5, 0]),
+                0.3,
+                np.full((5, 3), 0.15),
+                12,
+                2,
+                np.full(shape=(5, 3), fill_value=-9.83245867),
+            ),
+        ],
+    )
+    def test_continuous_contractual(self, value, lam, p, T, T0, logp):
+        with Model():
+            cc = ContContract("cc", lam=lam, p=p, T=T, T0=T0)
+        pt = {"cc": value}
+
+        assert_almost_equal(
+            pm.logp(cc, value).eval(),
+            logp,
+            decimal=select_by_precision(float64=6, float32=2),
+            err_msg=str(pt),
+        )
+
+    def test_continuous_contractual_invalid(self):
+        cc = ContContract.dist(lam=0.8, p=0.15, T=10, T0=2)
+        assert pm.logp(cc, np.array([-1, 3, 1])).eval() == -np.inf
+        assert pm.logp(cc, np.array([1.5, -1, 1])).eval() == -np.inf
+        assert pm.logp(cc, np.array([1.5, 0, 1])).eval() == -np.inf
+        assert pm.logp(cc, np.array([11, 3, 1])).eval() == -np.inf
+        assert pm.logp(cc, np.array([1.5, 3, 0.5])).eval() == -np.inf
+        assert pm.logp(cc, np.array([1.5, 3, -1])).eval() == -np.inf
+
+    # TODO: test broadcasting of parameters, including T and T0
+    @pytest.mark.parametrize(
+        "lam_size, p_size, cc_size, expected_size",
+        [
+            (None, None, None, (3,)),
+            ((7,), None, (7,), (7, 3)),
+            ((7, 1), (1, 5), (7, 5), (7, 5, 3)),
+            (None, None, (7, 5), (7, 5, 3)),
+        ],
+    )
+    def test_continuous_contractual_sample_prior(
+        self, lam_size, p_size, cc_size, expected_size
+    ):
+        with Model():
+            lam = pm.Gamma(name="lam", alpha=1, beta=1, size=lam_size)
+            p = pm.Beta(name="p", alpha=1.0, beta=1.0, size=p_size)
+            ContContract(name="cc", lam=lam, p=p, T=10, T0=2, size=cc_size)
+            prior = pm.sample_prior_predictive(samples=100)
+
+        assert prior["prior"]["cc"][0].shape == (100,) + expected_size