diff --git a/conda-envs/environment-dev-py37.yml b/conda-envs/environment-dev-py37.yml index 7263fb4245..90e22cf6b7 100644 --- a/conda-envs/environment-dev-py37.yml +++ b/conda-envs/environment-dev-py37.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/environment-dev-py38.yml b/conda-envs/environment-dev-py38.yml index 21858097ce..69fd28473b 100644 --- a/conda-envs/environment-dev-py38.yml +++ b/conda-envs/environment-dev-py38.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/environment-dev-py39.yml b/conda-envs/environment-dev-py39.yml index 0b007a5aa2..6a0bfc917d 100644 --- a/conda-envs/environment-dev-py39.yml +++ b/conda-envs/environment-dev-py39.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/environment-test-py37.yml b/conda-envs/environment-test-py37.yml index c5a4a14b47..964939f493 100644 --- a/conda-envs/environment-test-py37.yml +++ b/conda-envs/environment-test-py37.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/environment-test-py38.yml b/conda-envs/environment-test-py38.yml index 5a2b1ee762..cec0b1ffa2 100644 --- a/conda-envs/environment-test-py38.yml +++ b/conda-envs/environment-test-py38.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/environment-test-py39.yml b/conda-envs/environment-test-py39.yml index 942c05c03d..50d87383c6 100644 --- a/conda-envs/environment-test-py39.yml +++ b/conda-envs/environment-test-py39.yml @@ -4,7 +4,7 @@ channels: - conda-forge - defaults dependencies: -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools diff --git a/conda-envs/windows-environment-dev-py38.yml b/conda-envs/windows-environment-dev-py38.yml index 513fe2e2af..ffac106acf 100644 --- a/conda-envs/windows-environment-dev-py38.yml +++ b/conda-envs/windows-environment-dev-py38.yml @@ -4,7 +4,7 @@ channels: - defaults dependencies: # base dependencies (see install guide for Windows) -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.4 - cachetools>=4.2.1 diff --git a/conda-envs/windows-environment-test-py38.yml b/conda-envs/windows-environment-test-py38.yml index 47374aa425..6ca7b966dd 100644 --- a/conda-envs/windows-environment-test-py38.yml +++ b/conda-envs/windows-environment-test-py38.yml @@ -4,7 +4,7 @@ channels: - defaults dependencies: # base dependencies (see install guide for Windows) -- aeppl>=0.0.17 +- aeppl=0.0.17 - aesara>=2.2.6 - arviz>=0.11.2 - cachetools diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 9c3bd37a91..43a66fe1e5 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -406,8 +406,9 @@ class Dirichlet(Continuous): Parameters ---------- - a: array - Concentration parameters (a > 0). + a: float array + Concentration parameters (a > 0). The number of categories is given by the + length of the last axis. """ rv_op = dirichlet @@ -507,13 +508,12 @@ class Multinomial(Discrete): Parameters ---------- - n: int or array - Number of trials (n > 0). If n is an array its shape must be (N,) with - N = p.shape[0] - p: one- or two-dimensional array - Probability of each one of the different outcomes. Elements must - be non-negative and sum to 1 along the last axis. They will be - automatically rescaled otherwise. + n: int + Total counts in each replicate (n > 0). + p: float array + Probability of each one of the different outcomes (0 <= p <= 1). The number of + categories is given by the length of the last axis. Elements are expected to sum + to 1 along the last axis, and they will be automatically rescaled otherwise. """ rv_op = multinomial @@ -626,17 +626,12 @@ class DirichletMultinomial(Discrete): Parameters ---------- - n : int or array - Total counts in each replicate. If n is an array its shape must be (N,) - with N = a.shape[0] - - a : one- or two-dimensional array - Dirichlet parameter. Elements must be strictly positive. - The number of categories is given by the length of the last axis. + n : int + Total counts in each replicate (n > 0). - shape : integer tuple - Describes shape of distribution. For example if n=array([5, 10]), and - a=array([1, 1, 1]), shape should be (2, 3). + a : float array + Dirichlet concentration parameters (a > 0). The number of categories is given by + the length of the last axis. """ rv_op = dirichlet_multinomial @@ -661,15 +656,10 @@ def logp(value, n, a): ------- TensorVariable """ - if value.ndim >= 1: - n = at.shape_padright(n) - if a.ndim > 1: - a = at.shape_padleft(a) - - sum_a = a.sum(axis=-1, keepdims=True) + sum_a = a.sum(axis=-1) const = (gammaln(n + 1) + gammaln(sum_a)) - gammaln(n + sum_a) series = gammaln(value + a) - (gammaln(value + 1) + gammaln(a)) - result = const + series.sum(axis=-1, keepdims=True) + result = const + series.sum(axis=-1) # Bounds checking to confirm parameters and data meet all constraints # and that each observation value_i sums to n_i. @@ -678,13 +668,10 @@ def logp(value, n, a): value >= 0, a > 0, n >= 0, - at.eq(value.sum(axis=-1, keepdims=True), n), + at.eq(value.sum(axis=-1), n), broadcast_conditions=False, ) - def _distr_parameters_for_repr(self): - return ["n", "a"] - class _OrderedMultinomial(Multinomial): r""" diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index b19f24543b..a4c1ee472d 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -48,9 +48,9 @@ def polyagamma_cdf(*args, **kwargs): from aesara.tensor.random.op import RandomVariable from aesara.tensor.var import TensorVariable from numpy import array, inf, log -from numpy.testing import assert_allclose, assert_almost_equal, assert_equal +from numpy.testing import assert_almost_equal, assert_equal from scipy import integrate -from scipy.special import erf, logit +from scipy.special import erf, gammaln, logit import pymc as pm @@ -327,31 +327,21 @@ def f3(a, b, c): raise ValueError("Dont know how to integrate shape: " + str(shape)) -def multinomial_logpdf(value, n, p): +def _dirichlet_multinomial_logpmf(value, n, a): if value.sum() == n and (0 <= value).all() and (value <= n).all(): - logpdf = scipy.special.gammaln(n + 1) - logpdf -= scipy.special.gammaln(value + 1).sum() - logpdf += logpow(p, value).sum() - return logpdf - else: - return -inf - - -def dirichlet_multinomial_logpmf(value, n, a): - value, n, a = (np.asarray(x) for x in [value, n, a]) - assert value.ndim == 1 - assert n.ndim == 0 - assert a.shape == value.shape - gammaln = scipy.special.gammaln - if value.sum() == n and (0 <= value).all() and (value <= n).all(): - sum_a = a.sum(axis=-1) + sum_a = a.sum() const = gammaln(n + 1) + gammaln(sum_a) - gammaln(n + sum_a) series = gammaln(value + a) - gammaln(value + 1) - gammaln(a) - return const + series.sum(axis=-1) + return const + series.sum() else: return -inf +dirichlet_multinomial_logpmf = np.vectorize( + _dirichlet_multinomial_logpmf, signature="(n),(),(n)->()" +) + + def beta_mu_sigma(value, mu, sigma): kappa = mu * (1 - mu) / sigma ** 2 - 1 if kappa > 0: @@ -473,8 +463,12 @@ def discrete_weibull_logpmf(value, q, beta): ) -def dirichlet_logpdf(value, a): - return floatX((-betafn(a) + logpow(value, a - 1).sum(-1)).sum()) +def _dirichlet_logpdf(value, a): + # scipy.stats.dirichlet.logpdf suffers from numerical precision issues + return -betafn(a) + logpow(value, a - 1).sum() + + +dirichlet_logpdf = np.vectorize(_dirichlet_logpdf, signature="(n),(n)->()") def categorical_logpdf(value, p): @@ -2111,182 +2105,74 @@ def test_lkj(self, x, eta, n, lp): @pytest.mark.parametrize("n", [1, 2, 3]) def test_dirichlet(self, n): - self.check_logp(Dirichlet, Simplex(n), {"a": Vector(Rplus, n)}, dirichlet_logpdf) - - @pytest.mark.parametrize("dist_shape", [(1, 2), (2, 4, 3)]) - def test_dirichlet_with_batch_shapes(self, dist_shape): - a = np.ones(dist_shape) - with pm.Model() as model: - d = pm.Dirichlet("d", a=a) - - # Generate sample points to test - d_value = d.tag.value_var - d_point = d.eval().astype("float64") - d_point /= d_point.sum(axis=-1)[..., None] - - if hasattr(d_value.tag, "transform"): - d_point_trans = d_value.tag.transform.forward( - at.as_tensor(d_point), *d.owner.inputs - ).eval() - else: - d_point_trans = d_point - - pymc_res = logpt(d, d_point_trans, jacobian=False, sum=False).eval() - scipy_res = np.empty_like(pymc_res) - for idx in np.ndindex(a.shape[:-1]): - scipy_res[idx] = scipy.stats.dirichlet(a[idx]).logpdf(d_point[idx]) - - assert_almost_equal(pymc_res, scipy_res) - - def test_dirichlet_shape(self): - a = at.as_tensor_variable(np.r_[1, 2]) - dir_rv = Dirichlet.dist(a) - assert dir_rv.shape.eval() == (2,) - - with pytest.warns(DeprecationWarning), aesara.change_flags(compute_test_value="ignore"): - dir_rv = Dirichlet.dist(at.vector()) - - def test_dirichlet_2D(self): self.check_logp( Dirichlet, - MultiSimplex(2, 2), - {"a": Vector(Vector(Rplus, 2), 2)}, + Simplex(n), + {"a": Vector(Rplus, n)}, dirichlet_logpdf, ) - @pytest.mark.parametrize("n", [2, 3]) - def test_multinomial(self, n): - self.check_logp( - Multinomial, Vector(Nat, n), {"p": Simplex(n), "n": Nat}, multinomial_logpdf - ) - @pytest.mark.parametrize( - "p, size, n", + "a", [ - [[0.25, 0.25, 0.25, 0.25], (4,), 2], - [[0.25, 0.25, 0.25, 0.25], (1, 4), 3], - # 3: expect to fail - # [[.25, .25, .25, .25], (10, 4)], - [[0.25, 0.25, 0.25, 0.25], (10, 1, 4), 5], - # 5: expect to fail - # [[[.25, .25, .25, .25]], (2, 4), [7, 11]], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (2, 4), 13], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (1, 2, 4), [23, 29]], - [ - [[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], - (10, 2, 4), - [31, 37], - ], - [[[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]], (2, 4), [17, 19]], + ([2, 3, 5]), + ([[2, 3, 5], [9, 19, 3]]), + (np.abs(np.random.randn(2, 2, 4)) + 1), ], ) - def test_multinomial_random(self, p, size, n): - p = np.asarray(p) - with Model() as model: - m = Multinomial("m", n=n, p=p, size=size) - - assert m.eval().shape == size + p.shape - - def test_multinomial_vec(self): - vals = np.array([[2, 4, 4], [3, 3, 4]]) - p = np.array([0.2, 0.3, 0.5]) - n = 10 - - with Model() as model_single: - Multinomial("m", n=n, p=p) - - with Model() as model_many: - Multinomial("m", n=n, p=p, size=2) + @pytest.mark.parametrize("size", [2, (1, 2), (2, 4, 3)]) + def test_dirichlet_vectorized(self, a, size): + a = floatX(np.array(a)) - assert_almost_equal( - scipy.stats.multinomial.logpmf(vals, n, p), - np.asarray([model_single.fastlogp({"m": val}) for val in vals]), - decimal=4, - ) - - assert_almost_equal( - scipy.stats.multinomial.logpmf(vals, n, p), - logp(model_many.m, vals).eval().squeeze(), - decimal=4, - ) + dir = pm.Dirichlet.dist(a=a, size=size) + vals = dir.eval() assert_almost_equal( - sum(model_single.fastlogp({"m": val}) for val in vals), - model_many.fastlogp({"m": vals}), + dirichlet_logpdf(vals, a), + pm.logp(dir, vals).eval(), decimal=4, + err_msg=f"vals={vals}", ) - def test_multinomial_vec_1d_n(self): - vals = np.array([[2, 4, 4], [4, 3, 4]]) - p = np.array([0.2, 0.3, 0.5]) - ns = np.array([10, 11]) - - with Model() as model: - Multinomial("m", n=ns, p=p) - - assert_almost_equal( - sum(multinomial_logpdf(val, n, p) for val, n in zip(vals, ns)), - model.fastlogp({"m": vals}), - decimal=4, - ) - - def test_multinomial_vec_1d_n_2d_p(self): - vals = np.array([[2, 4, 4], [4, 3, 4]]) - ps = np.array([[0.2, 0.3, 0.5], [0.9, 0.09, 0.01]]) - ns = np.array([10, 11]) - - with Model() as model: - Multinomial("m", n=ns, p=ps) - - assert_almost_equal( - sum(multinomial_logpdf(val, n, p) for val, n, p in zip(vals, ns, ps)), - model.fastlogp({"m": vals}), - decimal=4, + @pytest.mark.parametrize("n", [2, 3]) + def test_multinomial(self, n): + self.check_logp( + Multinomial, + Vector(Nat, n), + {"p": Simplex(n), "n": Nat}, + lambda value, n, p: scipy.stats.multinomial.logpmf(value, n, p), ) - def test_multinomial_vec_2d_p(self): - vals = np.array([[2, 4, 4], [3, 3, 4]]) - ps = np.array([[0.2, 0.3, 0.5], [0.3, 0.3, 0.4]]) - n = 10 + @pytest.mark.parametrize("n", [(10), ([10, 11]), ([[5, 6], [10, 11]])]) + @pytest.mark.parametrize( + "p", + [ + ([0.2, 0.3, 0.5]), + ([[0.2, 0.3, 0.5], [0.9, 0.09, 0.01]]), + (np.abs(np.random.randn(2, 2, 4))), + ], + ) + @pytest.mark.parametrize("size", [1, 2, (2, 3)]) + def test_multinomial_vectorized(self, n, p, size): + n = intX(np.array(n)) + p = floatX(np.array(p)) + p /= p.sum(axis=-1, keepdims=True) - with Model() as model: - Multinomial("m", n=n, p=ps) + mn = pm.Multinomial.dist(n=n, p=p, size=size) + vals = mn.eval() assert_almost_equal( - sum(multinomial_logpdf(val, n, p) for val, p in zip(vals, ps)), - model.fastlogp({"m": vals}), + scipy.stats.multinomial.logpmf(vals, n, p), + pm.logp(mn, vals).eval(), decimal=4, + err_msg=f"vals={vals}", ) - def test_batch_multinomial(self): - n = 10 - vals = intX(np.zeros((4, 5, 3))) - p = floatX(np.zeros_like(vals)) - inds = np.random.randint(vals.shape[-1], size=vals.shape[:-1])[..., None] - np.put_along_axis(vals, inds, n, axis=-1) - np.put_along_axis(p, inds, 1, axis=-1) - - dist = Multinomial.dist(n=n, p=p) - logp_mn = at.exp(pm.logp(dist, vals)).eval() - assert_almost_equal( - logp_mn, - np.ones(vals.shape[:-1]), - decimal=select_by_precision(float64=6, float32=3), - ) - - dist = Multinomial.dist(n=n, p=p, size=2) - sample = dist.eval() - assert_allclose(sample, np.stack([vals, vals], axis=0)) - def test_multinomial_zero_probs(self): # test multinomial accepts 0 probabilities / observations: - value = aesara.shared(np.array([0, 0, 100], dtype=int)) - logp = pm.Multinomial.logp(value=value, n=100, p=at.constant([0.0, 0.0, 1.0])) - logp_fn = aesara.function(inputs=[], outputs=logp) - assert logp_fn() >= 0 - - value.set_value(np.array([50, 50, 0], dtype=int)) - assert np.isneginf(logp_fn()) + mn = pm.Multinomial.dist(n=100, p=[0.0, 0.0, 1.0]) + assert pm.logp(mn, np.array([0, 0, 100])).eval() >= 0 + assert pm.logp(mn, np.array([50, 50, 0])).eval() == -np.inf @pytest.mark.parametrize("n", [2, 3]) def test_dirichlet_multinomial(self, n): @@ -2314,105 +2200,30 @@ def test_dirichlet_multinomial_matches_beta_binomial(self): decimal=select_by_precision(float64=6, float32=3), ) - def test_dirichlet_multinomial_vec(self): - vals = np.array([[2, 4, 4], [3, 3, 4]]) - a = np.array([0.2, 0.3, 0.5]) - n = 10 - - with Model() as model_single: - DirichletMultinomial("m", n=n, a=a) - - with Model() as model_many: - DirichletMultinomial("m", n=n, a=a, size=2) - - assert_almost_equal( - np.asarray([dirichlet_multinomial_logpmf(val, n, a) for val in vals]), - np.asarray([model_single.fastlogp({"m": val}) for val in vals]), - decimal=4, - ) - - assert_almost_equal( - np.asarray([dirichlet_multinomial_logpmf(val, n, a) for val in vals]), - logp(model_many.m, vals).eval().squeeze(), - decimal=4, - ) - - assert_almost_equal( - sum(model_single.fastlogp({"m": val}) for val in vals), - model_many.fastlogp({"m": vals}), - decimal=4, - ) - - def test_dirichlet_multinomial_vec_1d_n(self): - vals = np.array([[2, 4, 4], [4, 3, 4]]) - a = np.array([0.2, 0.3, 0.5]) - ns = np.array([10, 11]) - - with Model() as model: - DirichletMultinomial("m", n=ns, a=a) - - assert_almost_equal( - sum(dirichlet_multinomial_logpmf(val, n, a) for val, n in zip(vals, ns)), - model.fastlogp({"m": vals}), - decimal=4, - ) - - def test_dirichlet_multinomial_vec_1d_n_2d_a(self): - vals = np.array([[2, 4, 4], [4, 3, 4]]) - as_ = np.array([[0.2, 0.3, 0.5], [0.9, 0.09, 0.01]]) - ns = np.array([10, 11]) - - with Model() as model: - DirichletMultinomial("m", n=ns, a=as_) - - assert_almost_equal( - sum(dirichlet_multinomial_logpmf(val, n, a) for val, n, a in zip(vals, ns, as_)), - model.fastlogp({"m": vals}), - decimal=4, - ) - - def test_dirichlet_multinomial_vec_2d_a(self): - vals = np.array([[2, 4, 4], [3, 3, 4]]) - as_ = np.array([[0.2, 0.3, 0.5], [0.3, 0.3, 0.4]]) - n = 10 + @pytest.mark.parametrize("n", [(10), ([10, 11]), ([[5, 6], [10, 11]])]) + @pytest.mark.parametrize( + "a", + [ + ([0.2, 0.3, 0.5]), + ([[0.2, 0.3, 0.5], [0.9, 0.09, 0.01]]), + (np.abs(np.random.randn(2, 2, 4))), + ], + ) + @pytest.mark.parametrize("size", [1, 2, (2, 3)]) + def test_dirichlet_multinomial_vectorized(self, n, a, size): + n = intX(np.array(n)) + a = floatX(np.array(a)) - with Model() as model: - DirichletMultinomial("m", n=n, a=as_) + dm = pm.DirichletMultinomial.dist(n=n, a=a, size=size) + vals = dm.eval() assert_almost_equal( - sum(dirichlet_multinomial_logpmf(val, n, a) for val, a in zip(vals, as_)), - model.fastlogp({"m": vals}), + dirichlet_multinomial_logpmf(vals, n, a), + pm.logp(dm, vals).eval(), decimal=4, + err_msg=f"vals={vals}", ) - def test_batch_dirichlet_multinomial(self): - # Test that DM can handle a 3d array for `a` - - # Create an almost deterministic DM by setting a to 0.001, everywhere - # except for one category / dimension which is given the value of 1000 - n = 5 - vals = np.zeros((4, 5, 3), dtype="int32") - a = np.zeros_like(vals, dtype=aesara.config.floatX) + 0.001 - inds = np.random.randint(vals.shape[-1], size=vals.shape[:-1])[..., None] - np.put_along_axis(vals, inds, n, axis=-1) - np.put_along_axis(a, inds, 1000, axis=-1) - - dist = DirichletMultinomial.dist(n=n, a=a) - - # Logp should be approx -9.98004998e-06 - dist_logp = logp(dist, vals).eval() - expected_logp = np.full_like(dist_logp, fill_value=-9.98004998e-06) - assert_almost_equal( - dist_logp, - expected_logp, - decimal=select_by_precision(float64=6, float32=3), - ) - - # Samples should be equal given the almost deterministic DM - dist = DirichletMultinomial.dist(n=n, a=a, size=2) - sample = dist.eval() - assert_allclose(sample, np.stack([vals, vals], axis=0)) - @aesara.config.change_flags(compute_test_value="raise") def test_categorical_bounds(self): with Model(): diff --git a/pymc/tests/test_distributions_random.py b/pymc/tests/test_distributions_random.py index 425f2ccc5e..0216de75fe 100644 --- a/pymc/tests/test_distributions_random.py +++ b/pymc/tests/test_distributions_random.py @@ -255,18 +255,6 @@ class TestGaussianRandomWalk(BaseTestCases.BaseTestCase): default_shape = (1,) -@pytest.mark.xfail(reason="This distribution has not been refactored for v4") -class TestZeroInflatedNegativeBinomial(BaseTestCases.BaseTestCase): - distribution = pm.ZeroInflatedNegativeBinomial - params = {"mu": 1.0, "alpha": 1.0, "psi": 0.3} - - -@pytest.mark.xfail(reason="This distribution has not been refactored for v4") -class TestZeroInflatedBinomial(BaseTestCases.BaseTestCase): - distribution = pm.ZeroInflatedBinomial - params = {"n": 10, "p": 0.6, "psi": 0.3} - - class BaseTestDistribution(SeededTest): """ This class provides a base for tests that new RandomVariables are correctly diff --git a/requirements-dev.txt b/requirements-dev.txt index 81509ea957..c5fbff991b 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,7 +1,7 @@ # This file is auto-generated by scripts/generate_pip_deps_from_conda.py, do not modify. # See that file for comments about the need/usage of each dependency. -aeppl>=0.0.17 +aeppl==0.0.17 aesara>=2.2.6 arviz>=0.11.4 cachetools>=4.2.1 diff --git a/requirements.txt b/requirements.txt index d80b9c8a46..75da478008 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -aeppl>=0.0.17 +aeppl==0.0.17 aesara>=2.2.6 arviz>=0.11.4 cachetools>=4.2.1