pymc-extras 0.2.0__py3-none-any.whl

tests/distributions/test_discrete.py

@@ -0,0 +1,210 @@
+#   Copyright 2023 The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+import numpy as np
+import pymc as pm
+import pytensor
+import pytensor.tensor as pt
+import pytest
+import scipy.stats
+
+from pymc.logprob.utils import ParameterValueError
+from pymc.testing import (
+    BaseTestDistributionRandom,
+    Domain,
+    I,
+    Rplus,
+    assert_support_point_is_expected,
+    check_logp,
+    discrete_random_tester,
+)
+from pytensor import config
+
+from pymc_extras.distributions import (
+    BetaNegativeBinomial,
+    GeneralizedPoisson,
+    Skellam,
+)
+
+
+class TestGeneralizedPoisson:
+    class TestRandomVariable(BaseTestDistributionRandom):
+        pymc_dist = GeneralizedPoisson
+        pymc_dist_params = {"mu": 4.0, "lam": 1.0}
+        expected_rv_op_params = {"mu": 4.0, "lam": 1.0}
+        tests_to_run = [
+            "check_pymc_params_match_rv_op",
+            "check_rv_size",
+        ]
+
+        def test_random_matches_poisson(self):
+            discrete_random_tester(
+                dist=self.pymc_dist,
+                paramdomains={"mu": Rplus, "lam": Domain([0], edges=(None, None))},
+                ref_rand=lambda mu, lam, size: scipy.stats.poisson.rvs(mu, size=size),
+            )
+
+        @pytest.mark.parametrize("mu", (2.5, 20, 50))
+        def test_random_lam_expected_moments(self, mu):
+            lam = np.array([-0.9, -0.7, -0.2, 0, 0.2, 0.7, 0.9])
+            dist = self.pymc_dist.dist(mu=mu, lam=lam, size=(10_000, len(lam)))
+            draws = dist.eval()
+
+            expected_mean = mu / (1 - lam)
+            np.testing.assert_allclose(draws.mean(0), expected_mean, rtol=1e-1)
+
+            expected_std = np.sqrt(mu / (1 - lam) ** 3)
+            np.testing.assert_allclose(draws.std(0), expected_std, rtol=1e-1)
+
+    def test_logp_matches_poisson(self):
+        # We are only checking this distribution for lambda=0 where it's equivalent to Poisson.
+        mu = pt.scalar("mu")
+        lam = pt.scalar("lam")
+        value = pt.vector("value", dtype="int64")
+
+        logp = pm.logp(GeneralizedPoisson.dist(mu, lam), value)
+        logp_fn = pytensor.function([value, mu, lam], logp)
+
+        test_value = np.array([0, 1, 2, 30])
+        for test_mu in (0.01, 0.1, 0.9, 1, 1.5, 20, 100):
+            np.testing.assert_allclose(
+                logp_fn(test_value, test_mu, lam=0),
+                scipy.stats.poisson.logpmf(test_value, test_mu),
+                rtol=1e-7 if config.floatX == "float64" else 1e-5,
+            )
+
+        # Check out-of-bounds values
+        value = pt.scalar("value")
+        logp = pm.logp(GeneralizedPoisson.dist(mu, lam), value)
+        logp_fn = pytensor.function([value, mu, lam], logp)
+
+        logp_fn(-1, mu=5, lam=0) == -np.inf
+        logp_fn(9, mu=5, lam=-1) == -np.inf
+
+        # Check mu/lam restrictions
+        with pytest.raises(ParameterValueError):
+            logp_fn(1, mu=1, lam=2)
+
+        with pytest.raises(ParameterValueError):
+            logp_fn(1, mu=0, lam=0)
+
+        with pytest.raises(ParameterValueError):
+            logp_fn(1, mu=1, lam=-1)
+
+    def test_logp_lam_expected_moments(self):
+        mu = 30
+        lam = np.array([-0.9, -0.7, -0.2, 0, 0.2, 0.7, 0.9])
+        with pm.Model():
+            x = GeneralizedPoisson("x", mu=mu, lam=lam)
+            trace = pm.sample(chains=1, draws=10_000, random_seed=96).posterior
+
+        expected_mean = mu / (1 - lam)
+        np.testing.assert_allclose(trace["x"].mean(("chain", "draw")), expected_mean, rtol=1e-1)
+
+        expected_std = np.sqrt(mu / (1 - lam) ** 3)
+        np.testing.assert_allclose(trace["x"].std(("chain", "draw")), expected_std, rtol=1e-1)
+
+    @pytest.mark.parametrize(
+        "mu, lam, size, expected",
+        [
+            (50, [-0.6, 0, 0.6], None, np.floor(50 / (1 - np.array([-0.6, 0, 0.6])))),
+            ([5, 50], -0.1, (4, 2), np.full((4, 2), np.floor(np.array([5, 50]) / 1.1))),
+        ],
+    )
+    def test_moment(self, mu, lam, size, expected):
+        with pm.Model() as model:
+            GeneralizedPoisson("x", mu=mu, lam=lam, size=size)
+        assert_support_point_is_expected(model, expected)
+
+
+class TestBetaNegativeBinomial:
+    """
+    Wrapper class so that tests of experimental additions can be dropped into
+    PyMC directly on adoption.
+    """
+
+    def test_logp(self):
+        """
+
+        Beta Negative Binomial logp function test values taken from R package as
+        there is currently no implementation in scipy.
+        https://github.com/scipy/scipy/issues/17330
+
+        The test values can be generated in R with the following code:
+
+        .. code-block:: r
+
+            library(extraDistr)
+
+            create.test.rows <- function(alpha, beta, r, x) {
+                logp <- dbnbinom(x, alpha, beta, r, log=TRUE)
+                paste0("(", paste(alpha, beta, r, x, logp, sep=", "), ")")
+            }
+
+            x <- c(0, 1, 250, 5000)
+            print(create.test.rows(1, 1, 1, x), quote=FALSE)
+            print(create.test.rows(1, 1, 10, x), quote=FALSE)
+            print(create.test.rows(1, 10, 1, x), quote=FALSE)
+            print(create.test.rows(10, 1, 1, x), quote=FALSE)
+            print(create.test.rows(10, 10, 10, x), quote=FALSE)
+
+        """
+        alpha, beta, r, value = pt.scalars("alpha", "beta", "r", "value")
+        logp = pm.logp(BetaNegativeBinomial.dist(alpha, beta, r), value)
+        logp_fn = pytensor.function([value, alpha, beta, r], logp)
+
+        tests = [
+            # 1, 1, 1
+            (1, 1, 1, 0, -0.693147180559945),
+            (1, 1, 1, 1, -1.79175946922805),
+            (1, 1, 1, 250, -11.0548820266432),
+            (1, 1, 1, 5000, -17.0349862828565),
+            # 1, 1, 10
+            (1, 1, 10, 0, -2.39789527279837),
+            (1, 1, 10, 1, -2.58021682959232),
+            (1, 1, 10, 250, -8.82261694534392),
+            (1, 1, 10, 5000, -14.7359968760473),
+            # 1, 10, 1
+            (1, 10, 1, 0, -2.39789527279837),
+            (1, 10, 1, 1, -2.58021682959232),
+            (1, 10, 1, 250, -8.82261694534418),
+            (1, 10, 1, 5000, -14.7359968760446),
+            # 10, 1, 1
+            (10, 1, 1, 0, -0.0953101798043248),
+            (10, 1, 1, 1, -2.58021682959232),
+            (10, 1, 1, 250, -43.5891148758123),
+            (10, 1, 1, 5000, -76.2953173311091),
+            # 10, 10, 10
+            (10, 10, 10, 0, -5.37909807285049),
+            (10, 10, 10, 1, -4.17512526852455),
+            (10, 10, 10, 250, -21.781591505836),
+            (10, 10, 10, 5000, -53.4836799634603),
+        ]
+        for test_alpha, test_beta, test_r, test_value, expected_logp in tests:
+            np.testing.assert_allclose(
+                logp_fn(test_value, test_alpha, test_beta, test_r), expected_logp
+            )
+
+
+class TestSkellam:
+    def test_logp(self):
+        # Scipy Skellam underflows to -inf earlier than PyMC
+        Rplus_small = Domain([0, 0.01, 0.1, 0.9, 0.99, 1, 1.5, 2, 10, np.inf])
+        # Suppress warnings coming from Scipy logpmf implementation
+        with np.errstate(divide="ignore", invalid="ignore"):
+            check_logp(
+                Skellam,
+                I,
+                {"mu1": Rplus_small, "mu2": Rplus_small},
+                lambda value, mu1, mu2: scipy.stats.skellam.logpmf(value, mu1, mu2),
+            )

tests/distributions/test_discrete_markov_chain.py

@@ -0,0 +1,258 @@
+import numpy as np
+import pymc as pm
+
+# general imports
+import pytensor.tensor as pt
+import pytest
+
+from pymc.distributions import Categorical
+from pymc.distributions.shape_utils import change_dist_size
+from pymc.logprob.utils import ParameterValueError
+from pymc.sampling.mcmc import assign_step_methods
+
+from pymc_extras.distributions.timeseries import (
+    DiscreteMarkovChain,
+    DiscreteMarkovChainGibbsMetropolis,
+)
+
+
+def transition_probability_tests(steps, n_states, n_lags, n_draws, atol):
+    P = np.full((n_states,) * (n_lags + 1), 1 / n_states)
+    x0 = pm.Categorical.dist(p=np.ones(n_states) / n_states)
+
+    chain = DiscreteMarkovChain.dist(
+        P=pt.as_tensor_variable(P), init_dist=x0, steps=steps, n_lags=n_lags
+    )
+
+    draws = pm.draw(chain, n_draws, random_seed=172)
+
+    # Test x0 is uniform over n_states
+    for i in range(n_lags):
+        assert np.allclose(
+            np.histogram(draws[:, ..., i], bins=n_states)[0] / n_draws, 1 / n_states, atol=atol
+        )
+
+    n_grams = [[tuple(row[i : i + n_lags + 1]) for i in range(len(row) - n_lags)] for row in draws]
+    freq_table = np.zeros((n_states,) * (n_lags + 1))
+
+    for row in n_grams:
+        for ngram in row:
+            freq_table[ngram] += 1
+    freq_table /= freq_table.sum(axis=-1)[:, None]
+
+    # Test continuation probabilities match P
+    assert np.allclose(P, freq_table, atol=atol)
+
+
+class TestDiscreteMarkovRV:
+    def test_fail_if_P_not_square(self):
+        P = pt.eye(3, 2)
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
+        with pytest.raises(ParameterValueError):
+            pm.logp(chain, np.zeros((3,))).eval()
+
+    def test_fail_if_P_not_valid(self):
+        P = pt.zeros((3, 3))
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
+        with pytest.raises(ParameterValueError):
+            pm.logp(chain, np.zeros((3,))).eval()
+
+    def test_high_dimensional_P(self):
+        P = pm.Dirichlet.dist(a=pt.ones(3), size=(3, 3, 3))
+        n_lags = 3
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+        chain = DiscreteMarkovChain.dist(P=P, steps=10, init_dist=x0, n_lags=n_lags)
+        draws = pm.draw(chain, 10)
+        logp = pm.logp(chain, draws)
+
+    def test_default_init_dist_warns_user(self):
+        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
+
+        with pytest.warns(UserWarning):
+            DiscreteMarkovChain.dist(P=P, steps=3)
+
+    def test_logp_shape(self):
+        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+        # Test with steps
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
+        draws = pm.draw(chain, 5)
+        logp = pm.logp(chain, draws).eval()
+
+        assert logp.shape == (5,)
+
+        # Test with shape
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3,))
+        draws = pm.draw(chain, 5)
+        logp = pm.logp(chain, draws).eval()
+
+        assert logp.shape == (5,)
+
+    def test_logp_with_default_init_dist(self):
+        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+        value = np.array([0, 1, 2])
+        logp_expected = np.log((1 / 3) * 0.5 * 0.3)
+
+        # Test dist directly
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
+        logp_eval = pm.logp(chain, value).eval()
+        np.testing.assert_allclose(logp_eval, logp_expected, rtol=1e-6)
+
+        # Test via Model
+        with pm.Model() as m:
+            DiscreteMarkovChain("chain", P=P, init_dist=x0, steps=3)
+        model_logp_eval = m.compile_logp()({"chain": value})
+        np.testing.assert_allclose(model_logp_eval, logp_expected, rtol=1e-6)
+
+    def test_logp_with_user_defined_init_dist(self):
+        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
+        x0 = pm.Categorical.dist(p=[0.2, 0.6, 0.2])
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=3)
+
+        logp = pm.logp(chain, [0, 1, 2]).eval()
+        assert logp == np.log(0.2 * 0.5 * 0.3)
+
+    def test_moment_function(self):
+        P_np = np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]])
+
+        x0_np = np.array([0, 1, 0])
+
+        P = pt.as_tensor_variable(P_np)
+        x0 = pm.Categorical.dist(p=x0_np.tolist())
+        n_steps = 3
+
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, steps=n_steps)
+
+        chain_np = np.empty(shape=n_steps + 1, dtype="int8")
+        chain_np[0] = np.argmax(x0_np)
+        for i in range(n_steps):
+            state = chain_np[i]
+            chain_np[i + 1] = np.argmax(P_np[state])
+
+        dmc_chain = pm.distributions.distribution.support_point(chain).eval()
+
+        assert np.allclose(dmc_chain, chain_np)
+
+    def test_define_steps_via_shape_arg(self):
+        P = pt.full((3, 3), 1 / 3)
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3,))
+        assert chain.eval().shape == (3,)
+
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(3, 2))
+        assert chain.eval().shape == (3, 2)
+
+    def test_define_steps_via_dim_arg(self):
+        coords = {"steps": [1, 2, 3]}
+
+        with pm.Model(coords=coords):
+            P = pt.full((3, 3), 1 / 3)
+            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+            chain = DiscreteMarkovChain("chain", P=P, init_dist=x0, dims=["steps"])
+
+        assert chain.eval().shape == (3,)
+
+    def test_dims_when_steps_are_defined(self):
+        coords = {"steps": [1, 2, 3, 4]}
+
+        with pm.Model(coords=coords):
+            P = pt.full((3, 3), 1 / 3)
+            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+            chain = DiscreteMarkovChain("chain", P=P, steps=3, init_dist=x0, dims=["steps"])
+
+        assert chain.eval().shape == (4,)
+
+    def test_multiple_dims_with_steps(self):
+        coords = {"steps": [1, 2, 3], "mc_chains": [1, 2, 3]}
+
+        with pm.Model(coords=coords):
+            P = pt.full((3, 3), 1 / 3)
+            x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+            chain = DiscreteMarkovChain(
+                "chain", P=P, steps=2, init_dist=x0, dims=["steps", "mc_chains"]
+            )
+
+        assert chain.eval().shape == (3, 3)
+
+    def test_mutiple_dims_with_steps_and_init_dist(self):
+        coords = {"steps": [1, 2, 3], "mc_chains": [1, 2, 3]}
+
+        with pm.Model(coords=coords):
+            P = pt.full((3, 3), 1 / 3)
+            x0 = pm.Categorical.dist(p=[0.1, 0.1, 0.8], size=(3,))
+            chain = DiscreteMarkovChain(
+                "chain", P=P, init_dist=x0, steps=2, dims=["steps", "mc_chains"]
+            )
+
+        assert chain.eval().shape == (3, 3)
+
+    def test_multiple_lags_with_data(self):
+        with pm.Model():
+            P = pt.full((3, 3, 3), 1 / 3)
+            x0 = pm.Categorical.dist(p=[0.1, 0.1, 0.8], size=2)
+            data = pm.draw(x0, 100)
+
+            chain = DiscreteMarkovChain("chain", P=P, init_dist=x0, n_lags=2, observed=data)
+
+        assert chain.eval().shape == (100, 2)
+
+    def test_random_draws(self):
+        transition_probability_tests(steps=3, n_states=2, n_lags=1, n_draws=2500, atol=0.05)
+        transition_probability_tests(steps=3, n_states=2, n_lags=3, n_draws=7500, atol=0.05)
+
+    def test_change_size_univariate(self):
+        P = pt.as_tensor_variable(np.array([[0.1, 0.5, 0.4], [0.3, 0.4, 0.3], [0.9, 0.05, 0.05]]))
+        x0 = pm.Categorical.dist(p=np.ones(3) / 3)
+
+        chain = DiscreteMarkovChain.dist(P=P, init_dist=x0, shape=(100, 5))
+
+        new_rw = change_dist_size(chain, new_size=(7,))
+        assert tuple(new_rw.shape.eval()) == (7, 5)
+
+        new_rw = change_dist_size(chain, new_size=(4, 3), expand=True)
+        assert tuple(new_rw.shape.eval()) == (4, 3, 100, 5)
+
+    def test_mcmc_sampling(self):
+        with pm.Model(coords={"step": range(100)}) as model:
+            init_dist = Categorical.dist(p=[0.5, 0.5])
+            markov_chain = DiscreteMarkovChain(
+                "markov_chain",
+                P=[[0.1, 0.9], [0.1, 0.9]],
+                init_dist=init_dist,
+                shape=(100,),
+                dims="step",
+            )
+
+            _, assigned_step_methods = assign_step_methods(model)
+            assert assigned_step_methods[DiscreteMarkovChainGibbsMetropolis] == [
+                model.rvs_to_values[markov_chain]
+            ]
+
+            # Sampler needs no tuning
+            idata = pm.sample(
+                tune=0, chains=4, draws=250, progressbar=False, compute_convergence_checks=False
+            )
+
+        np.testing.assert_allclose(
+            idata.posterior["markov_chain"].isel(step=0).mean(("chain", "draw")),
+            0.5,
+            atol=0.05,
+        )
+
+        np.testing.assert_allclose(
+            idata.posterior["markov_chain"].isel(step=slice(1, None)).mean(("chain", "draw")),
+            0.9,
+            atol=0.05,
+        )
+
+        assert pm.stats.ess(idata, method="tail").min() > 950