From ae84163a5480f06ffebb8bac0f47c259c199ab6e Mon Sep 17 00:00:00 2001 From: Colt Allen <10178857+ColtAllen@users.noreply.github.com> Date: Mon, 27 May 2024 12:28:35 -0600 Subject: [PATCH] Add `BetaGeoBetaBinom` Distribution Block (#431) * init commit * removed scan * Fix logp * Remove print statement * Add test for logp notimplemented errors * docstrings * dev notebook added * updated to vectorize_graph * import order * update oldest pymc_version * Update ci.yml pymc version * Update pyproject.toml pymc version * WIP sample prior testing * sample prior compared against lifetimes * increase rtol * remove commented code, add logp reference * fix latex docstring * notebook testing and misc edits * revert latex in docstring * add ruff ignore comment --------- Co-authored-by: Ricardo Vieira --- .../clv/dev/beta_geo_beta_binom.ipynb | 1577 +++++++++++++++++ pymc_marketing/clv/distributions.py | 194 +- tests/clv/test_distributions.py | 194 +- 3 files changed, 1956 insertions(+), 9 deletions(-) create mode 100644 docs/source/notebooks/clv/dev/beta_geo_beta_binom.ipynb diff --git a/docs/source/notebooks/clv/dev/beta_geo_beta_binom.ipynb b/docs/source/notebooks/clv/dev/beta_geo_beta_binom.ipynb new file mode 100644 index 00000000..d434df4a --- /dev/null +++ b/docs/source/notebooks/clv/dev/beta_geo_beta_binom.ipynb @@ -0,0 +1,1577 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 9, + "id": "5e06e043-4631-47ae-a658-a9a928ff15e5", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "from lifetimes import BetaGeoBetaBinomFitter\n", + "from lifetimes.datasets import load_donations, load_cdnow_summary_data_with_monetary_value\n", + "\n", + "import pymc as pm\n", + "from pymc_marketing.clv.distributions import BetaGeoBetaBinom" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "7b7431fc-1bad-41b6-8ec4-ec843a81709d", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Sampling: [alpha, beta, beta_geo_beta_binom, delta, gamma]\n" + ] + }, + { + "data": { + "text/html": [ + "\n", + "
\n", + "
\n", + "
arviz.InferenceData
\n", + "
\n", + " \n", + "
\n", + " " + ], + "text/plain": [ + "Inference data with groups:\n", + "\t> prior\n", + "\t> constant_data" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "T_true = 6\n", + "alpha_true = 1.204\n", + "beta_true = 0.750\n", + "gamma_true = 0.657\n", + "delta_true = 2.783\n", + "\n", + "with pm.Model():\n", + " alpha = pm.Normal(name=\"alpha\", mu=alpha_true, sigma=1e-3)\n", + " beta = pm.Normal(name=\"beta\", mu=beta_true, sigma=1e-3)\n", + " gamma = pm.Normal(name=\"gamma\", mu=gamma_true, sigma=1e-3)\n", + " delta = pm.Normal(name=\"delta\", mu=delta_true, sigma=1e-3)\n", + "\n", + " T = pm.MutableData(name=\"T\", value=np.array(T_true))\n", + "\n", + " BetaGeoBetaBinom(\n", + " name=\"beta_geo_beta_binom\",\n", + " alpha=alpha,\n", + " beta=beta,\n", + " gamma=gamma,\n", + " delta=delta,\n", + " T=T,\n", + " )\n", + " prior = pm.sample_prior_predictive(samples=10)\n", + "\n", + "# recency = prior[\"prior\"][\"beta_geo_beta_binom\"][0][:,0]\n", + "# #frequency = prior[\"prior\"][\"beta_geo_beta_binom\"][1]\n", + "# recency.mean()\n", + "prior" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "8ff42e71-fc43-4d45-8446-7205e3d37bce", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([ -3.94031398, -10.25427751, -6.82582822])" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "value = np.array([[1.5, 1], [5.3, 4], [6, 2]])\n", + "alpha = 0.55\n", + "beta = 10.58\n", + "gamma = 0.61\n", + "delta = 11.67\n", + "T = 12\n", + "\n", + "BetaGeoBetaBinomFitter._loglikelihood((alpha, beta, gamma, delta), value[..., 1], value[..., 0], T)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "371bb7a1-9f5c-4bdf-81b1-a9504261badb", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "discrete_noncontract_df = load_donations()\n", + "\n", + "periods = 6\n", + "bgbb = BetaGeoBetaBinomFitter().fit(discrete_noncontract_df['frequency'].values,\n", + " discrete_noncontract_df['recency'].values,\n", + " discrete_noncontract_df['periods'].values,\n", + " discrete_noncontract_df['weights'].values)\n", + "bgbb" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "f8052f01-5aca-48fd-917e-1f2bbeb6326f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "['conditional_expected_number_of_purchases_up_to_time', 'conditional_probability_alive', 'expected_number_of_transactions_in_first_n_periods', 'fit', 'load_model', 'save_model', 'summary']\n" + ] + } + ], + "source": [ + "method_list = [method for method in dir(BetaGeoBetaBinomFitter) if not method.startswith('_')]\n", + "print(method_list)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "f4f56b83-f830-4610-8f4f-e67ff242967e", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0 0.072863\n", + "1 0.085696\n", + "2 0.314238\n", + "3 0.593853\n", + "4 0.839396\n", + "5 1.021689\n", + "6 1.147885\n", + "7 0.119121\n", + "8 0.536111\n", + "9 1.057604\n", + "10 1.443042\n", + "11 1.668817\n", + "12 0.223595\n", + "13 1.034572\n", + "14 1.804703\n", + "15 2.189749\n", + "16 0.583192\n", + "17 2.030024\n", + "18 2.710681\n", + "19 1.812942\n", + "20 3.231612\n", + "21 3.752544\n", + "dtype: float64" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# equation 13 in paper\n", + "bgbb.conditional_expected_number_of_purchases_up_to_time(5,\n", + " discrete_noncontract_df['frequency'],\n", + " discrete_noncontract_df['recency'],\n", + " discrete_noncontract_df['periods'])" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "7ea5dc42-160f-4f96-9e16-a97f01dd4bdc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0 0.070072\n", + "1 0.045012\n", + "2 0.165056\n", + "3 0.311927\n", + "4 0.440900\n", + "5 0.536651\n", + "6 0.602936\n", + "7 0.043038\n", + "8 0.193695\n", + "9 0.382108\n", + "10 0.521365\n", + "11 0.602936\n", + "12 0.061566\n", + "13 0.284864\n", + "14 0.496916\n", + "15 0.602936\n", + "16 0.129719\n", + "17 0.451538\n", + "18 0.602936\n", + "19 0.338249\n", + "20 0.602936\n", + "21 0.602936\n", + "dtype: float64" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# equation 11 in paper\n", + "bgbb.conditional_probability_alive(10,\n", + " discrete_noncontract_df['frequency'],\n", + " discrete_noncontract_df['recency'],\n", + " discrete_noncontract_df['periods'])" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "96bcab46-7279-400a-82c1-e8b509ece774", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
model
frequency
03454.881979
11888.654343
21348.882330
31113.378533
41017.922413
51027.166049
61253.114353
\n", + "
" + ], + "text/plain": [ + " model\n", + "frequency \n", + "0 3454.881979\n", + "1 1888.654343\n", + "2 1348.882330\n", + "3 1113.378533\n", + "4 1017.922413\n", + "5 1027.166049\n", + "6 1253.114353" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# TODO: write and test (8) as a replacement. Compare against just aggregating means across the exploded DF \n", + "# TODO: Can the arviz functions in the BetaGeoBetaBinom distribution block preclude the need for this?\n", + "# TODO: Replace this with (9) or (10) in a future PR, since that expression can predict interval ranges\n", + "\n", + "# equation 7 in paper, but that's for probabilities. should it be 8 for predicting mean n?\n", + "# yeah, this function should be renamed for clarity. \n", + "# it distributes customers in the dataset across n transaction opportunies\n", + "# it works better as an evaluation function, since it assumes a fixed customer population size\n", + "# if n > n_periods, it will keep right on predicting. This may be a bug\n", + "bgbb.expected_number_of_transactions_in_first_n_periods(n=6)" + ] + }, + { + "cell_type": "markdown", + "id": "9d55e986-d1f2-4c0d-8c25-3e289e90d5fe", + "metadata": {}, + "source": [ + "### Expected transactions in N periods\n", + "This expression will blow up to inf with large values of n (n=167 in this example). Recalculating on the log scale will allow for larger values, but this isn't possible if gamma < 1 because term1 will be negative." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "2e82f5b4-1b4a-4477-843b-58cbd411d348", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "average of 1.938137499995133 purchases expected in 5 opportunities\n" + ] + } + ], + "source": [ + "from scipy import special\n", + "from numpy import log,exp\n", + "\n", + "n = 5\n", + "alpha,beta,delta,gamma = bgbb._unload_params('alpha','beta','delta','gamma')\n", + "\n", + "# add a larger gamma value for testing\n", + "#gamma = .9\n", + "\n", + "log_scale = False\n", + "\n", + "if not log_scale:\n", + " term1 = alpha/(alpha+beta)*delta/(gamma-1)\n", + " term2 = 1-(special.gamma(gamma+delta))/special.gamma(gamma+delta+n)*(special.gamma(1+delta+n))/special.gamma(1+delta)\n", + " expected_purchases_n_periods = term1 * term2\n", + "else:\n", + " term1 = log(alpha/(alpha+beta)) + log(delta/(gamma-1))\n", + " term2 = special.gammaln(gamma+delta) - special.gammaln(gamma+delta+n) + special.gammaln(1+delta+n) - special.gammaln(1+delta)\n", + " expected_purchases_n_periods = exp(term1) - exp(term2)\n", + "\n", + "print(f'average of {expected_purchases_n_periods} purchases expected in {n} opportunities')" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "5186cf4d-710d-4e85-bef9-b66ccced5586", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[1.203522393608101, 0.7497163581757842, 2.783441982887136, 0.6567181695498788]" + ] + }, + "execution_count": 17, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "bgbb._unload_params('alpha','beta','delta','gamma')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "80d11cc8-98fb-426e-89b2-693f0a8d22fa", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pymc_marketing/clv/distributions.py b/pymc_marketing/clv/distributions.py index c47bb505..196f9e2d 100644 --- a/pymc_marketing/clv/distributions.py +++ b/pymc_marketing/clv/distributions.py @@ -15,14 +15,13 @@ import pymc as pm import pytensor.tensor as pt from pymc.distributions.continuous import PositiveContinuous -from pymc.distributions.dist_math import check_parameters +from pymc.distributions.dist_math import betaln, check_parameters +from pymc.distributions.distribution import Discrete +from pytensor import scan +from pytensor.graph import vectorize_graph from pytensor.tensor.random.op import RandomVariable -__all__ = [ - "ContContract", - "ContNonContract", - "ParetoNBD", -] +__all__ = ["ContContract", "ContNonContract", "ParetoNBD", "BetaGeoBetaBinom"] class ContNonContractRV(RandomVariable): @@ -391,7 +390,7 @@ class ParetoNBD(PositiveContinuous): \end{align} ======== =============================================== - Support :math:`t_j > 0` for :math:`j = 1, \dots, x` + Support :math:`t_j >= 0` for :math:`j = 1, \dots, x` Mean :math:`\mathbb{E}[X(t) | r, \alpha, s, \beta] = \frac{r\beta}{\alpha(s-1)}[1-(\frac{\beta}{\beta + t})^{s-1}]` ======== =============================================== @@ -476,3 +475,184 @@ def logp(value, r, alpha, s, beta, T): beta > 0, msg="r > 0, alpha > 0, s > 0, beta > 0", ) + + +class BetaGeoBetaBinomRV(RandomVariable): + name = "beta_geo_beta_binom" + ndim_supp = 1 + ndims_params = [0, 0, 0, 0, 0] + 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 == (): + size = np.broadcast_shapes( + alpha.shape, beta.shape, gamma.shape, delta.shape, T.shape + ) + + alpha = np.broadcast_to(alpha, size) + beta = np.broadcast_to(beta, size) + gamma = np.broadcast_to(gamma, size) + delta = np.broadcast_to(delta, size) + T = np.broadcast_to(T, size) + + output = np.zeros(shape=(*size, 2)) + + purchase_prob = rng.beta(a=alpha, b=beta, size=size) + churn_prob = rng.beta(a=delta, b=gamma, size=size) + + def sim_data(purchase_prob, churn_prob, T): + t_x = 0 + x = 0 + active = True + recency = 0 + + while t_x <= T and active: + t_x += 1 + active = rng.binomial(1, churn_prob) + purchase = rng.binomial(1, purchase_prob) + if active and purchase: + recency = t_x + x += 1 + return np.array( + [ + recency if x > 0 else T, + x, + ], + ) + + for index in np.ndindex(*size): + output[index] = sim_data(purchase_prob[index], churn_prob[index], T[index]) + + return output + + def _supp_shape_from_params(*args, **kwargs): + return (2,) + + +beta_geo_beta_binom = BetaGeoBetaBinomRV() + + +class BetaGeoBetaBinom(Discrete): + r""" + Population-level distribution class for a discrete, non-contractual, Beta-Geometric/Beta-Binomial process, + based on equation(5) from Fader, et al. in [1]_. + + .. math:: + + \mathbb{L}(\alpha, \beta, \gamma, \delta | x, t_x, n) &= + \frac{B(\alpha+x,\beta+n-x)}{B(\alpha,\beta)} + \frac{B(\gamma,\delta+n)}{B(\gamma,\delta)} \\ + &+ \sum_{i=0}^{n-t_x-1}\frac{B(\alpha+x,\beta+t_x-x+i)}{B(\alpha,\beta)} \\ + &\cdot \frac{B(\gamma+1,\delta+t_x+i)}{B(\gamma,\delta)} + + ======== =============================================== + Support :math:`t_j >= 0` for :math:`j = 1, \dots,x` + Mean :math:`\mathbb{E}[X(n) | \alpha, \beta, \gamma, \delta] = (\frac{\alpha}{\alpha+\beta})(\frac{\delta}{\gamma-1}) \cdot{1-\frac{\Gamma(\gamma+\delta)}{\Gamma(\gamma+\delta+n)}\frac{\Gamma(1+\delta+n)}{\Gamma(1+ \delta)}}` + ======== =============================================== + + References + ---------- + .. [1] Fader, Peter S., Bruce G.S. Hardie, and Jen Shang (2010), + "Customer-Base Analysis in a Discrete-Time Noncontractual Setting," + Marketing Science, 29 (6), 1086-1108. https://www.brucehardie.com/papers/020/fader_et_al_mksc_10.pdf + + """ # noqa: E501 + + rv_op = beta_geo_beta_binom + + @classmethod + def dist(cls, alpha, beta, gamma, delta, T, **kwargs): + return super().dist([alpha, beta, gamma, delta, T], **kwargs) + + def logp(value, alpha, beta, gamma, delta, T): + t_x = pt.atleast_1d(value[..., 0]) + x = pt.atleast_1d(value[..., 1]) + scalar_case = t_x.type.broadcastable == (True,) + + for param in (t_x, x, alpha, beta, gamma, delta, T): + if param.type.ndim > 1: + raise NotImplementedError( + f"BetaGeoBetaBinom logp only implemented for vector parameters, got ndim={param.type.ndim}" + ) + if scalar_case: + if param.type.broadcastable == (False,): + raise NotImplementedError( + f"Parameter {param} cannot be larger than scalar value" + ) + + # Broadcast all the parameters so they are sequences. + # Potentially inefficient, but otherwise ugly logic needed to unpack arguments in the scan function, + # since sequences always precede non-sequences. + _, alpha, beta, gamma, delta, T = pt.broadcast_arrays( + t_x, alpha, beta, gamma, delta, T + ) + + def logp_customer_died(t_x_i, x_i, alpha_i, beta_i, gamma_i, delta_i, T_i): + i = pt.scalar("i", dtype=int) + died = pt.lt(t_x_i + i, T_i) + + unnorm_logprob_customer_died_at_tx_plus_i = betaln( + alpha_i + x_i, beta_i + t_x_i - x_i + i + ) + betaln(gamma_i + died, delta_i + t_x_i + i) + + # Maximum prevents invalid T - t_x values from crashing logp + i_vec = pt.arange(pt.maximum(T_i - t_x_i, 0) + 1) + unnorm_logprob_customer_died_at_tx_plus_i_vec = vectorize_graph( + unnorm_logprob_customer_died_at_tx_plus_i, replace={i: i_vec} + ) + + return pt.logsumexp(unnorm_logprob_customer_died_at_tx_plus_i_vec) + + unnorm_logp, _ = scan( + fn=logp_customer_died, + outputs_info=[None], + sequences=[t_x, x, alpha, beta, gamma, delta, T], + ) + + logp = unnorm_logp - betaln(alpha, beta) - betaln(gamma, delta) + + logp = pt.switch( + pt.or_( + pt.or_( + pt.lt(t_x, 0), + pt.lt(x, 0), + ), + pt.gt(t_x, T), + ), + -np.inf, + logp, + ) + + if value.ndim == 1: + logp = pt.specify_shape(logp, 1).squeeze(0) + + return check_parameters( + logp, + alpha > 0, + beta > 0, + gamma > 0, + delta > 0, + msg="alpha > 0, beta > 0, gamma > 0, delta > 0", + ) diff --git a/tests/clv/test_distributions.py b/tests/clv/test_distributions.py index 4562db5f..c868d3c1 100644 --- a/tests/clv/test_distributions.py +++ b/tests/clv/test_distributions.py @@ -13,12 +13,21 @@ # limitations under the License. import numpy as np import pymc as pm +import pytensor +import pytensor.tensor as pt import pytest +from lifetimes import BetaGeoBetaBinomFitter as BGBBF from lifetimes import ParetoNBDFitter as PF +from lifetimes.generate_data import beta_geometric_beta_binom_model from numpy.testing import assert_almost_equal from pymc import Model -from pymc_marketing.clv.distributions import ContContract, ContNonContract, ParetoNBD +from pymc_marketing.clv.distributions import ( + BetaGeoBetaBinom, + ContContract, + ContNonContract, + ParetoNBD, +) class TestContNonContract: @@ -265,4 +274,185 @@ def test_pareto_nbd_sample_prior( ) prior = pm.sample_prior_predictive(samples=100) - assert prior["prior"]["pareto_nbd"][0].shape == (100,) + expected_size # noqa: RUF005 + assert prior["prior"]["pareto_nbd"][0].shape == (100, *expected_size) + + +class TestBetaGeoBetaBinom: + @pytest.mark.parametrize("batch_shape", [(), (5,)]) + def test_logp_matches_lifetimes(self, batch_shape): + rng = np.random.default_rng(269) + + alpha = pm.draw( + pm.Gamma.dist(mu=1.2, sigma=3, shape=batch_shape), random_seed=rng + ) + beta = pm.draw( + pm.Gamma.dist(mu=0.75, sigma=3, shape=batch_shape), random_seed=rng + ) + gamma = pm.draw( + pm.Gamma.dist(mu=0.657, sigma=3, shape=(1,) * len(batch_shape)), + random_seed=rng, + ) + delta = pm.draw(pm.Gamma.dist(mu=2.783, sigma=3), random_seed=rng) + T = pm.draw(pm.DiscreteUniform.dist(1, 10, shape=batch_shape), random_seed=rng) + + t_x = pm.draw(pm.DiscreteUniform.dist(0, T, shape=batch_shape), random_seed=rng) + x = pm.draw(pm.DiscreteUniform.dist(0, t_x, shape=batch_shape), random_seed=rng) + value = np.concatenate([t_x[..., None], x[..., None]], axis=-1) + + dist = BetaGeoBetaBinom.dist(alpha, beta, gamma, delta, T) + np.testing.assert_allclose( + pm.logp(dist, value).eval(), + BGBBF._loglikelihood((alpha, beta, gamma, delta), x, t_x, T), + ) + + def test_logp_matches_excel(self): + # Expected logp values can be found in excel file in http://brucehardie.com/notes/010/ + # Spreadsheet: Parameter estimate + + alpha = 1.204 + beta = 0.750 + gamma = 0.657 + delta = 2.783 + T = 6 + + x = np.array([6, 5, 4, 3, 2, 1, 5, 4, 3, 2, 1, 4, 3, 2, 1, 3, 2, 1, 2, 1, 1, 0]) + t_x = np.array( + [6, 6, 6, 6, 6, 6, 5, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3, 2, 2, 1, 0] + ) + expected_logp = np.array( + [ + -2.18167018824, + -4.29485034929, + -5.38473334360, + -5.80915881601, + -5.65172964525, + -4.88370164695, + -3.71682127437, + -5.09558227343, + -5.61576884108, + -5.50636893346, + -4.76723821904, + -3.84829625138, + -5.05936147828, + -5.19562191019, + -4.57070931973, + -3.52745257839, + -4.51620272962, + -4.22465969453, + -3.01199924784, + -3.58817880928, + -2.28882847451, + -1.16751622367, + ] + ) + + value = np.concatenate([t_x[:, None], x[:, None]], axis=-1) + dist = BetaGeoBetaBinom.dist( + alpha=alpha, + beta=beta, + gamma=gamma, + delta=delta, + T=T, + ) + np.testing.assert_allclose( + pm.logp(dist, value).eval(), + expected_logp, + rtol=1e-3, + ) + + def test_invalid_value_logp(self): + beta_geo_beta_binom = BetaGeoBetaBinom.dist( + alpha=1.20, beta=0.75, gamma=0.66, delta=2.78, T=6 + ) + value = pt.vector("value", shape=(2,)) + logp = pm.logp(beta_geo_beta_binom, value) + + logp_fn = pytensor.function([value], logp) + assert logp_fn(np.array([3, -1])) == -np.inf + assert logp_fn(np.array([-1, 1.5])) == -np.inf + assert logp_fn(np.array([11, 1.5])) == -np.inf + + def test_notimplemented_logp(self): + dist = BetaGeoBetaBinom.dist(alpha=1, beta=1, gamma=2, delta=2, T=10) + invalid_value = np.broadcast_to([1, 3], (4, 3, 2)) + with pytest.raises(NotImplementedError): + pm.logp(dist, invalid_value) + + invalid_dist = BetaGeoBetaBinom.dist( + alpha=np.ones( + 5, + ), + beta=1, + gamma=2, + delta=2, + T=10, + ) + value = np.array([1, 3]) + with pytest.raises(NotImplementedError): + pm.logp(invalid_dist, value) + + @pytest.mark.parametrize( + "alpha_size, beta_size, gamma_size, delta_size, beta_geo_beta_binom_size, expected_size", + [ + (None, None, None, None, None, (2,)), + ((5,), None, None, None, None, (5, 2)), + (None, (5,), None, None, (5,), (5, 2)), + (None, None, (5, 1), (1, 3), (5, 3), (5, 3, 2)), + (None, None, None, None, (5, 3), (5, 3, 2)), + ], + ) + def test_beta_geo_beta_binom_sample_prior( + self, + alpha_size, + beta_size, + gamma_size, + delta_size, + beta_geo_beta_binom_size, + expected_size, + ): + # Declare simulation params + T_true = 60 + alpha_true = 1.204 + beta_true = 0.750 + gamma_true = 0.657 + delta_true = 2.783 + + # Generate simulated data from lifetimes + # this does not have a random seed, so rtol must be higher + lt_bgbb = beta_geometric_beta_binom_model( + N=T_true, + alpha=alpha_true, + beta=beta_true, + gamma=gamma_true, + delta=delta_true, + size=1000, + ) + lt_frequency = lt_bgbb["frequency"].values + lt_recency = lt_bgbb["recency"].values + + with Model(): + alpha = pm.Normal(name="alpha", mu=alpha_true, sigma=1e-4, size=alpha_size) + beta = pm.Normal(name="beta", mu=beta_true, sigma=1e-4, size=beta_size) + gamma = pm.Normal(name="gamma", mu=gamma_true, sigma=1e-4, size=gamma_size) + delta = pm.Normal(name="delta", mu=delta_true, sigma=1e-4, size=delta_size) + + T = pm.MutableData(name="T", value=np.array(T_true)) + + BetaGeoBetaBinom( + name="beta_geo_beta_binom", + alpha=alpha, + beta=beta, + gamma=gamma, + delta=delta, + T=T, + size=beta_geo_beta_binom_size, + ) + prior = pm.sample_prior_predictive(samples=1000) + prior = prior["prior"]["beta_geo_beta_binom"][0] + recency = prior[:, 0] + frequency = prior[:, 1] + + assert prior.shape == (1000, *expected_size) + + np.testing.assert_allclose(lt_frequency.mean(), recency.mean(), rtol=0.84) + np.testing.assert_allclose(lt_recency.mean(), frequency.mean(), rtol=0.84)