From f03976c94cfbbcc588386b35faf06688eb5f1293 Mon Sep 17 00:00:00 2001 From: kc611 Date: Wed, 29 Sep 2021 22:48:06 +0530 Subject: [PATCH] Refactor LKJ Cholesky Cov and LKJCorr for V4 (WIP) * `sd_dist` is now a standard parameter (with a prior specified outside of the distribution) --- pymc/distributions/multivariate.py | 373 ++++++++++-------------- pymc/tests/sampler_fixtures.py | 2 +- pymc/tests/test_distributions.py | 28 +- pymc/tests/test_distributions_random.py | 104 ++++++- pymc/tests/test_idata_conversion.py | 3 +- pymc/tests/test_mixture.py | 2 +- 6 files changed, 273 insertions(+), 239 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 28440a3e85..2041496f6c 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -43,7 +43,12 @@ from pymc.aesaraf import floatX, intX from pymc.distributions import transforms -from pymc.distributions.continuous import ChiSquared, Normal, assert_negative_support +from pymc.distributions.continuous import ( + BoundedContinuous, + ChiSquared, + Normal, + assert_negative_support, +) from pymc.distributions.dist_math import ( betaln, check_parameters, @@ -58,6 +63,7 @@ to_tuple, ) from pymc.math import kron_diag, kron_dot +from pymc.util import UNSET __all__ = [ "MvNormal", @@ -1085,67 +1091,117 @@ def _lkj_normalizing_constant(eta, n): return result -class _LKJCholeskyCov(Continuous): - r"""Underlying class for covariance matrix with LKJ distributed correlations. - See docs for LKJCholeskyCov function for more details on how to use it in models. - """ +# TODO: Can the `sd_dist` be a separate distribution that we take as an input? +# TODO: Explain that this takes draws from flatten upper triangular matrix, change name as well! +class _LKJCholeskyCovRV(RandomVariable): - def __init__(self, eta, n, sd_dist, *args, **kwargs): - self.n = at.as_tensor_variable(n) - self.eta = at.as_tensor_variable(eta) + name = "_lkjcholeskycov" + ndim_supp = 1 + ndims_params = (0, 0, 1) + dtype = "floatX" + _print_name = ("_lkjcholeskycov", "\\operatorname{_lkjcholeskycov}") - if "transform" in kwargs and kwargs["transform"] is not None: - raise ValueError("Invalid parameter: transform.") - if "shape" in kwargs: - raise ValueError("Invalid parameter: shape.") + def _shape_from_params(self, dist_params, **kwargs): + n = dist_params[1] + dist_shape = ((n * (n + 1)) // 2,) + return dist_shape - shape = n * (n + 1) // 2 + def rng_fn(self, rng, eta, n, sd_dist, size): - if sd_dist.shape.ndim not in [0, 1]: - raise ValueError("Invalid shape for sd_dist.") + dist_shape = ((n * (n + 1)) // 2,) - def transform_params(rv_var): - _, _, _, n, eta = rv_var.owner.inputs - return np.arange(1, n + 1).cumsum() - 1 + # We will always provide random with an integer size and then reshape + # the output to get the correct size + if size is not None: + orig_size = size + size = np.prod(size) + else: + size = 1 + orig_size = None + + P = np.full((size, n, n), np.eye(n)) + # original implementation in R see: + # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r + beta = eta - 1.0 + n / 2.0 + r12 = 2.0 * stats.beta.rvs(a=beta, b=beta, size=(size,)) - 1.0 + P[..., 0, 1] = r12 + P[..., 1, 1] = np.sqrt(1.0 - r12 ** 2) + for mp1 in range(2, n): + beta -= 0.5 + y = stats.beta.rvs(a=mp1 / 2.0, b=beta, size=(size,)) + z = stats.norm.rvs(loc=0, scale=1, size=(size,) + (mp1,)) + z = z / np.sqrt(np.einsum("ij,ij->j", z, z)) + P[..., 0:mp1, mp1] = np.sqrt(y[..., np.newaxis]) * z + P[..., mp1, mp1] = np.sqrt(1.0 - y) + C = np.einsum("...ji,...jk->...ik", P, P) + + D = sd_dist + C *= D[..., :, np.newaxis] * D[..., np.newaxis, :] + tril_idx = np.tril_indices(n, k=0) + + samples = np.linalg.cholesky(C)[..., tril_idx[0], tril_idx[1]] + + if orig_size is None: + samples = samples[0] + else: + samples = np.reshape(samples, orig_size + dist_shape) - transform = transforms.CholeskyCovPacked(transform_params) + return samples - kwargs["shape"] = shape - kwargs["transform"] = transform - super().__init__(*args, **kwargs) - self.sd_dist = sd_dist - self.diag_idxs = transform.diag_idxs +_ljk_cholesky_cov = _LKJCholeskyCovRV() - self.mode = floatX(np.zeros(shape)) - self.mode[self.diag_idxs] = 1 - def logp(self, x): +class _LKJCholeskyCov(Continuous): + r"""Underlying class for covariance matrix with LKJ distributed correlations. + See docs for LKJCholeskyCov function for more details on how to use it in models. + """ + rv_op = _ljk_cholesky_cov + + def __new__(cls, *args, **kwargs): + transform = kwargs.get("transform", UNSET) + if transform is UNSET: + kwargs["transform"] = cls.default_transform() + return super().__new__(cls, *args, **kwargs) + + @classmethod + def default_transform(cls): + def transform_params(rv_var): + _, _, _, _, n, _ = rv_var.owner.inputs + return at.arange(1, n + 1).cumsum() - 1 + + return transforms.CholeskyCovPacked(transform_params) + + @classmethod + def dist(cls, eta, n, sd_dist, **kwargs): + eta = at.as_tensor_variable(floatX(eta)) + n = at.at_least_1d(at.as_tensor_variable(intX(n))) + sd_dist = at.as_tensor_variable(floatX(sd_dist)) + return super().dist([eta, n, sd_dist], **kwargs) + + def logp(value, eta, n, sd_dist): """ Calculate log-probability of Covariance matrix with LKJ distributed correlations at specified value. Parameters ---------- - x: numeric + value: numeric Value for which log-probability is calculated. Returns ------- TensorVariable """ - n = self.n - eta = self.eta - diag_idxs = self.diag_idxs - cumsum = at.cumsum(x ** 2) + diag_idxs = at.cumsum(at.arange(1, n + 1)) - 1 + cumsum = at.cumsum(value ** 2) variance = at.zeros(n) - variance = at.inc_subtensor(variance[0], x[0] ** 2) + variance = at.inc_subtensor(variance[0], value[0] ** 2) variance = at.inc_subtensor(variance[1:], cumsum[diag_idxs[1:]] - cumsum[diag_idxs[:-1]]) sd_vals = at.sqrt(variance) - logp_sd = self.sd_dist.logp(sd_vals).sum() - corr_diag = x[diag_idxs] / sd_vals + corr_diag = value[diag_idxs] / sd_vals logp_lkj = (2 * eta - 3 + n - at.arange(n)) * at.log(corr_diag) logp_lkj = at.sum(logp_lkj) @@ -1158,111 +1214,10 @@ def logp(self, x): norm = _lkj_normalizing_constant(eta, n) - return norm + logp_lkj + logp_sd + det_invjac + return norm + logp_lkj + det_invjac - def _random(self, n, eta, size=1): - eta_sample_shape = (size,) + eta.shape - P = np.eye(n) * np.ones(eta_sample_shape + (n, n)) - # original implementation in R see: - # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r - beta = eta - 1.0 + n / 2.0 - r12 = 2.0 * stats.beta.rvs(a=beta, b=beta, size=eta_sample_shape) - 1.0 - P[..., 0, 1] = r12 - P[..., 1, 1] = np.sqrt(1.0 - r12 ** 2) - for mp1 in range(2, n): - beta -= 0.5 - y = stats.beta.rvs(a=mp1 / 2.0, b=beta, size=eta_sample_shape) - z = stats.norm.rvs(loc=0, scale=1, size=eta_sample_shape + (mp1,)) - z = z / np.sqrt(np.einsum("ij,ij->j", z, z)) - P[..., 0:mp1, mp1] = np.sqrt(y[..., np.newaxis]) * z - P[..., mp1, mp1] = np.sqrt(1.0 - y) - C = np.einsum("...ji,...jk->...ik", P, P) - D = np.atleast_1d(self.sd_dist.random(size=P.shape[:-2])) - if D.shape in [tuple(), (1,)]: - D = self.sd_dist.random(size=P.shape[:-1]) - elif D.ndim < C.ndim - 1: - D = [D] + [self.sd_dist.random(size=P.shape[:-2]) for _ in range(n - 1)] - D = np.moveaxis(np.array(D), 0, C.ndim - 2) - elif D.ndim == C.ndim - 1: - if D.shape[-1] == 1: - D = [D] + [self.sd_dist.random(size=P.shape[:-2]) for _ in range(n - 1)] - D = np.concatenate(D, axis=-1) - elif D.shape[-1] != n: - raise ValueError( - "The size of the samples drawn from the " - "supplied sd_dist.random have the wrong " - "size. Expected {} but got {} instead.".format(n, D.shape[-1]) - ) - else: - raise ValueError( - "Supplied sd_dist.random generates samples with " - "too many dimensions. It must yield samples " - "with 0 or 1 dimensions. Got {} instead".format(D.ndim - C.ndim - 2) - ) - C *= D[..., :, np.newaxis] * D[..., np.newaxis, :] - tril_idx = np.tril_indices(n, k=0) - return np.linalg.cholesky(C)[..., tril_idx[0], tril_idx[1]] - - def random(self, point=None, size=None): - """ - Draw random values from Covariance matrix with LKJ - distributed correlations. - Parameters - ---------- - point: dict, optional - Dict of variable values on which random values are to be - conditioned (uses default point if not specified). - size: int, optional - Desired size of random sample (returns one sample if not - specified). - - Returns - ------- - array - """ - # # Get parameters and broadcast them - # n, eta = draw_values([self.n, self.eta], point=point, size=size) - # broadcast_shape = np.broadcast(n, eta).shape - # # We can only handle cov matrices with a constant n per random call - # n = np.unique(n) - # if len(n) > 1: - # raise RuntimeError("Varying n is not supported for LKJCholeskyCov") - # n = int(n[0]) - # dist_shape = ((n * (n + 1)) // 2,) - # # We make sure that eta and the drawn n get their shapes broadcasted - # eta = np.broadcast_to(eta, broadcast_shape) - # # We change the size of the draw depending on the broadcast shape - # sample_shape = broadcast_shape + dist_shape - # if size is not None: - # if not isinstance(size, tuple): - # try: - # size = tuple(size) - # except TypeError: - # size = (size,) - # if size == sample_shape: - # size = None - # elif size == broadcast_shape: - # size = None - # elif size[-len(sample_shape) :] == sample_shape: - # size = size[: len(size) - len(sample_shape)] - # elif size[-len(broadcast_shape) :] == broadcast_shape: - # size = size[: len(size) - len(broadcast_shape)] - # # We will always provide _random with an integer size and then reshape - # # the output to get the correct size - # if size is not None: - # _size = np.prod(size) - # else: - # _size = 1 - # samples = self._random(n, eta, size=_size) - # if size is None: - # samples = samples[0] - # else: - # samples = np.reshape(samples, size + sample_shape) - # return samples - - -def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=True, *args, **kwargs): +def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=True, store_in_trace=True, *args, **kwargs): r"""Wrapper function for covariance matrix with LKJ distributed correlations. This defines a distribution over Cholesky decomposed covariance @@ -1282,7 +1237,7 @@ def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=Tru Dimension of the covariance matrix (n > 1). sd_dist: pm.Distribution A distribution for the standard deviations. - compute_corr: bool, default=False + compute_corr: bool, default=True If `True`, returns three values: the Cholesky decomposition, the correlations and the standard deviations of the covariance matrix. Otherwise, only returns the packed Cholesky decomposition. Defaults to `False` to ensure backwards @@ -1399,7 +1354,6 @@ def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=Tru packed_chol = _LKJCholeskyCov(name, eta=eta, n=n, sd_dist=sd_dist) if not compute_corr: return packed_chol - else: chol = pm.expand_packed_triangular(n, packed_chol, lower=True) # compute covariance matrix @@ -1415,7 +1369,46 @@ def LKJCholeskyCov(name, eta, n, sd_dist, compute_corr=False, store_in_trace=Tru return chol, corr, stds -class LKJCorr(Continuous): +class LKJCorrRV(RandomVariable): + name = "lkjcorr" + ndim_supp = 1 + ndims_params = [0, 0] + dtype = "floatX" + _print_name = ("LKJCorrRV", "\\operatorname{LKJCorrRV}") + + def _infer_shape(self, size, dist_params, param_shapes=None): + n = dist_params[0] + shape = tuple(size) + (n,) + return shape + + @classmethod + def rng_fn(cls, rng, n, eta, size=None): + # TODO: rng is not being used by the stat.sebta.rvs! + size = 1 if size is None else size + size = size if isinstance(size, tuple) else (size,) + # original implementation in R see: + # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r + beta = eta - 1.0 + n / 2.0 + r12 = 2.0 * stats.beta.rvs(a=beta, b=beta, size=size) - 1.0 + P = np.eye(n)[:, :, np.newaxis] * np.ones(size) + P[0, 1] = r12 + P[1, 1] = np.sqrt(1.0 - r12 ** 2) + for mp1 in range(2, n): + beta -= 0.5 + y = stats.beta.rvs(a=mp1 / 2.0, b=beta, size=size) + z = stats.norm.rvs(loc=0, scale=1, size=(mp1,) + size) + z = z / np.sqrt(np.einsum("ij,ij->j", z, z)) + P[0:mp1, mp1] = np.sqrt(y) * z + P[mp1, mp1] = np.sqrt(1.0 - y) + C = np.einsum("ji...,jk...->...ik", P, P) + triu_idx = np.triu_indices(n, k=1) + return C[..., triu_idx[0], triu_idx[1]] + + +lkjcorr = LKJCorrRV() + + +class LKJCorr(BoundedContinuous): r""" The LKJ (Lewandowski, Kurowicka and Joe) log-likelihood. @@ -1457,86 +1450,17 @@ class LKJCorr(Continuous): 100(9), pp.1989-2001. """ - def __init__(self, eta=None, n=None, p=None, transform="interval", *args, **kwargs): - if (p is not None) and (n is not None) and (eta is None): - warnings.warn( - "Parameters to LKJCorr have changed: shape parameter n -> eta " - "dimension parameter p -> n. Please update your code. " - "Automatically re-assigning parameters for backwards compatibility.", - FutureWarning, - ) - self.n = p - self.eta = n - eta = self.eta - n = self.n - elif (n is not None) and (eta is not None) and (p is None): - self.n = n - self.eta = eta - else: - raise ValueError( - "Invalid parameter: please use eta as the shape parameter and " - "n as the dimension parameter." - ) - - shape = n * (n - 1) // 2 - self.mean = floatX(np.zeros(shape)) - - if transform == "interval": - transform = transforms.interval(-1, 1) - - super().__init__(shape=shape, transform=transform, *args, **kwargs) - warnings.warn( - "Parameters in LKJCorr have been rename: shape parameter n -> eta " - "dimension parameter p -> n. Please double check your initialization.", - FutureWarning, - ) - self.tri_index = np.zeros([n, n], dtype="int32") - self.tri_index[np.triu_indices(n, k=1)] = np.arange(shape) - self.tri_index[np.triu_indices(n, k=1)[::-1]] = np.arange(shape) + bound_args_indices = [-1, 1] + rv_op = lkjcorr - def _random(self, n, eta, size=None): - size = size if isinstance(size, tuple) else (size,) - # original implementation in R see: - # https://github.com/rmcelreath/rethinking/blob/master/R/distributions.r - beta = eta - 1.0 + n / 2.0 - r12 = 2.0 * stats.beta.rvs(a=beta, b=beta, size=size) - 1.0 - P = np.eye(n)[:, :, np.newaxis] * np.ones(size) - P[0, 1] = r12 - P[1, 1] = np.sqrt(1.0 - r12 ** 2) - for mp1 in range(2, n): - beta -= 0.5 - y = stats.beta.rvs(a=mp1 / 2.0, b=beta, size=size) - z = stats.norm.rvs(loc=0, scale=1, size=(mp1,) + size) - z = z / np.sqrt(np.einsum("ij,ij->j", z, z)) - P[0:mp1, mp1] = np.sqrt(y) * z - P[mp1, mp1] = np.sqrt(1.0 - y) - C = np.einsum("ji...,jk...->...ik", P, P) - triu_idx = np.triu_indices(n, k=1) - return C[..., triu_idx[0], triu_idx[1]] - - def random(self, point=None, size=None): - """ - Draw random values from LKJ distribution. - - Parameters - ---------- - point: dict, optional - Dict of variable values on which random values are to be - conditioned (uses default point if not specified). - size: int, optional - Desired size of random sample (returns one sample if not - specified). + @classmethod + def dist(cls, n, eta, *args, **kwargs): + n = at.as_tensor_variable(intX(n)) + eta = at.as_tensor_variable(floatX(eta)) - Returns - ------- - array - """ - # n, eta = draw_values([self.n, self.eta], point=point, size=size) - # size = 1 if size is None else size - # samples = generate_samples(self._random, n, eta, broadcast_shape=(size,)) - # return samples + return super().dist([n, eta], *args, **kwargs) - def logp(self, x): + def logp(x, n, eta): """ Calculate log-probability of LKJ distribution at specified value. @@ -1550,10 +1474,15 @@ def logp(self, x): ------- TensorVariable """ - n = self.n - eta = self.eta - - X = x[self.tri_index] + # TODO: What is this n.data? + _n = n.data + _eta = eta.data + shape = _n * (_n - 1) // 2 + tri_index = np.zeros([_n, _n], dtype="int32") + tri_index[np.triu_indices(_n, k=1)] = np.arange(shape) + tri_index[np.triu_indices(_n, k=1)[::-1]] = np.arange(shape) + + X = at.take(x, tri_index) X = at.fill_diagonal(X, 1) result = _lkj_normalizing_constant(eta, n) diff --git a/pymc/tests/sampler_fixtures.py b/pymc/tests/sampler_fixtures.py index ce5f4f0490..5923a759eb 100644 --- a/pymc/tests/sampler_fixtures.py +++ b/pymc/tests/sampler_fixtures.py @@ -121,7 +121,7 @@ class LKJCholeskyCovFixture(KnownCDF): def make_model(cls): with pm.Model() as model: sd_mu = np.array([1, 2, 3, 4, 5]) - sd_dist = pm.LogNormal.dist(mu=sd_mu, sigma=sd_mu / 10.0, size=5) + sd_dist = pm.LogNormal("sd_dist", mu=sd_mu, sigma=sd_mu / 10.0, size=5) chol_packed = pm.LKJCholeskyCov("chol_packed", eta=3, n=5, sd_dist=sd_dist) chol = pm.expand_packed_triangular(5, chol_packed, lower=True) cov = at.dot(chol, chol.T) diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index 9b377825ed..a8bf3965f4 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -141,7 +141,6 @@ def get_lkj_cases(): return [ (tri, 1, 3, 1.5963125911388549), (tri, 3, 3, -7.7963493376312742), - (tri, 0, 3, -np.inf), (np.array([1.1, 0.0, -0.7]), 1, 3, -np.inf), (np.array([0.7, 0.0, -1.1]), 1, 3, -np.inf), ] @@ -2094,7 +2093,6 @@ def test_wishart(self, n): ) @pytest.mark.parametrize("x,eta,n,lp", LKJ_CASES) - @pytest.mark.xfail(reason="Distribution not refactored yet") def test_lkj(self, x, eta, n, lp): with Model() as model: LKJCorr("lkj", eta=eta, n=n, transform=None) @@ -3311,3 +3309,29 @@ def test_censored_invalid_dist(self): match="The dist dist was already registered in the current model", ): x = pm.Censored("x", registered_dist, lower=None, upper=None) + + +# TODO: Finish this test +def test_lkjcholeskycov_sampling_shape(): + samples = 5 + D = 15 + dist_shape = ((D * (D + 1)) // 2,) + + with pm.Model() as model: + sd_dist = pm.HalfCauchy("sd_dist", beta=2.5, shape=1) + packedL = pm.LKJCholeskyCov("packedL", eta=2, n=D, sd_dist=sd_dist, compute_corr=False) + + with model: + trace = pm.sample( + tune=5, + draws=samples, + chains=1, + return_inferencedata=False, + compute_convergence_checks=False, + ) + + # TODO: assert something about the logp (separate test?) + # pt = model.recompute_initial_point(0) + # logp = model.compile_logp(pt) + + assert trace["packedL"].shape == (samples,) + dist_shape diff --git a/pymc/tests/test_distributions_random.py b/pymc/tests/test_distributions_random.py index 8afa50fd8e..e8bbbc6d2b 100644 --- a/pymc/tests/test_distributions_random.py +++ b/pymc/tests/test_distributions_random.py @@ -47,7 +47,11 @@ def random_polyagamma(*args, **kwargs): from pymc.distributions.discrete import _OrderedLogistic, _OrderedProbit from pymc.distributions.dist_math import clipped_beta_rvs from pymc.distributions.logprob import logp -from pymc.distributions.multivariate import _OrderedMultinomial, quaddist_matrix +from pymc.distributions.multivariate import ( + _LKJCholeskyCov, + _OrderedMultinomial, + quaddist_matrix, +) from pymc.distributions.shape_utils import to_tuple from pymc.tests.helpers import SeededTest, select_by_precision from pymc.tests.test_distributions import ( @@ -365,8 +369,9 @@ def check_rv_size(self): sizes_expected = self.sizes_expected or [(), (), (1,), (1,), (5,), (4, 5), (2, 4, 2)] for size, expected in zip(sizes_to_check, sizes_expected): pymc_rv = self.pymc_dist.dist(**self.pymc_dist_params, size=size) - actual = tuple(pymc_rv.shape.eval()) - assert actual == expected, f"size={size}, expected={expected}, actual={actual}" + expected_symbolic = tuple(pymc_rv.shape.eval()) + actual = pymc_rv.eval().shape + assert actual == expected_symbolic == expected # test multi-parameters sampling for univariate distributions (with univariate inputs) if ( @@ -386,8 +391,9 @@ def check_rv_size(self): ] for size, expected in zip(sizes_to_check, sizes_expected): pymc_rv = self.pymc_dist.dist(**params, size=size) - actual = tuple(pymc_rv.shape.eval()) - assert actual == expected + expected_symbolic = tuple(pymc_rv.shape.eval()) + actual = pymc_rv.eval().shape + assert actual == expected_symbolic == expected def validate_tests_list(self): assert len(self.tests_to_run) == len( @@ -1546,6 +1552,84 @@ class TestZeroInflatedNegativeBinomial(BaseTestDistribution): tests_to_run = ["check_pymc_params_match_rv_op"] +class TestLKJCov(BaseTestDistribution): + pymc_dist = _LKJCholeskyCov + pymc_dist_params = {"eta": 1.0, "n": 2, "sd_dist": np.array([0.5, 2.0])} + expected_rv_op_params = {"eta": 1.0, "n": 2, "sd_dist": np.array([0.5, 2.0])} + + sizes_to_check = [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)] + sizes_expected = [ + (3,), + (3,), + (1, 3), + (1, 3), + (5, 3), + (4, 5, 3), + (2, 4, 2, 3), + ] + + tests_to_run = [ + "check_pymc_params_match_rv_op", + "check_rv_size", + "check_expected_draws", + ] + + def check_expected_draws(self): + rng = aesara.shared(self.get_random_state(reset=True)) + x = _LKJCholeskyCov.dist(n=2, eta=10_000, sd_dist=np.array([1.0, 2.0]), rng=rng) + assert np.all(np.abs(x.eval() - np.array([1.0, 0, 2.0])) < 0.1) + + # TODO: Enable or remove this test + def check_expected_draws_slow(self): + def ref_rand(size, n, eta): + shape = n * (n - 1) // 2 + beta = eta - 1 + n / 2 + return (st.beta.rvs(size=(size, shape), a=beta, b=beta) - 0.5) * 2 + + pymc_random( + pm.LKJCorr, + {"n": Domain([2, 10, 50], edges=[0, -1]), "eta": Domain([1.0, 10.0, 100.0])}, + size=1000, + ref_rand=ref_rand, + ) + + +class TestLKJCorr(BaseTestDistribution): + pymc_dist = pm.LKJCorr + pymc_dist_params = {"n": 3, "eta": 1.0} + expected_rv_op_params = {"n": 3, "eta": 1.0} + + sizes_to_check = [None, (), 1, (1,), 5, (4, 5), (2, 4, 2)] + sizes_expected = [ + (3,), + (3,), + (1, 3), + (1, 3), + (5, 3), + (4, 5, 3), + (2, 4, 2, 3), + ] + + tests_to_run = [ + "check_pymc_params_match_rv_op", + "test_lkj", + "check_rv_size", + ] + + def test_lkj(self): + def ref_rand(size, n, eta): + shape = n * (n - 1) // 2 + beta = eta - 1 + n / 2 + return (st.beta.rvs(size=(size, shape), a=beta, b=beta) - 0.5) * 2 + + pymc3_random( + pm.LKJCorr, + {"n": Domain([2, 10, 50], edges=[0, -1]), "eta": Domain([1.0, 10.0, 100.0])}, + size=1000, + ref_rand=ref_rand, + ) + + class TestOrderedLogistic(BaseTestDistribution): pymc_dist = _OrderedLogistic pymc_dist_params = {"eta": 0, "cutpoints": np.array([-2, 0, 2])} @@ -2241,7 +2325,6 @@ def generate_shapes(include_params=False): return data -@pytest.mark.skip(reason="This test is covered by Aesara") class TestMvNormal(SeededTest): @pytest.mark.parametrize( ["sample_shape", "dist_shape", "mu_shape", "param"], @@ -2261,7 +2344,7 @@ def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param): def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape): with pm.Model() as model: mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape) - sd_dist = pm.Exponential.dist(1.0, shape=3) + sd_dist = pm.Exponential("sd_dist", 1.0, shape=3) chol, corr, stds = pm.LKJCholeskyCov( "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True ) @@ -2278,7 +2361,7 @@ def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape): def test_with_cov_rv(self, sample_shape, dist_shape, mu_shape): with pm.Model() as model: mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape) - sd_dist = pm.Exponential.dist(1.0, shape=3) + sd_dist = pm.Exponential("sd_dist", 1.0, shape=3) chol, corr, stds = pm.LKJCholeskyCov( "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True ) @@ -2325,7 +2408,6 @@ def test_issue_3706(self): assert prior_pred["X"].shape == (1, N, 2) -@pytest.mark.xfail(reason="This distribution has not been refactored for v4") def test_matrix_normal_random_with_random_variables(): """ This test checks for shape correctness when using MatrixNormal distribution @@ -2337,7 +2419,7 @@ def test_matrix_normal_random_with_random_variables(): mu_0 = np.zeros((D, K)) lambd = 1.0 with pm.Model() as model: - sd_dist = pm.HalfCauchy.dist(beta=2.5) + sd_dist = pm.HalfCauchy("sd_dist", beta=2.5) packedL = pm.LKJCholeskyCov("packedL", eta=2, n=D, sd_dist=sd_dist) L = pm.expand_packed_triangular(D, packedL, lower=True) Sigma = pm.Deterministic("Sigma", L.dot(L.T)) # D x D covariance @@ -2372,7 +2454,7 @@ def test_with_np_arrays(self, sample_shape, dist_shape, mu_shape, param): def test_with_chol_rv(self, sample_shape, dist_shape, mu_shape): with pm.Model() as model: mu = pm.Normal("mu", 0.0, 1.0, shape=mu_shape) - sd_dist = pm.Exponential.dist(1.0, shape=3) + sd_dist = pm.Exponential("sd_dist", 1.0, shape=3) chol, corr, stds = pm.LKJCholeskyCov( "chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True ) diff --git a/pymc/tests/test_idata_conversion.py b/pymc/tests/test_idata_conversion.py index b846831f52..754110ff6e 100644 --- a/pymc/tests/test_idata_conversion.py +++ b/pymc/tests/test_idata_conversion.py @@ -332,14 +332,13 @@ def test_missing_data_model(self): assert inference_data.log_likelihood["y_observed"].shape == (2, 100, 3) @pytest.mark.xfal(reason="Multivariate partial observed RVs not implemented for V4") - @pytest.mark.xfail(reason="LKJCholeskyCov not refactored for v4") def test_mv_missing_data_model(self): data = ma.masked_values([[1, 2], [2, 2], [-1, 4], [2, -1], [-1, -1]], value=-1) model = pm.Model() with model: mu = pm.Normal("mu", 0, 1, size=2) - sd_dist = pm.HalfNormal.dist(1.0) + sd_dist = pm.HalfNormal("sd_dist", 1.0) chol, *_ = pm.LKJCholeskyCov("chol_cov", n=2, eta=1, sd_dist=sd_dist, compute_corr=True) y = pm.MvNormal("y", mu=mu, chol=chol, observed=data) inference_data = pm.sample(100, chains=2, return_inferencedata=True) diff --git a/pymc/tests/test_mixture.py b/pymc/tests/test_mixture.py index 3262f83b8f..d76a6b6800 100644 --- a/pymc/tests/test_mixture.py +++ b/pymc/tests/test_mixture.py @@ -368,7 +368,7 @@ def build_toy_dataset(N, K): mu.append(pm.Normal("mu%i" % i, 0, 10, shape=D)) packed_chol.append( pm.LKJCholeskyCov( - "chol_cov_%i" % i, eta=2, n=D, sd_dist=pm.HalfNormal.dist(2.5) + "chol_cov_%i" % i, eta=2, n=D, sd_dist=pm.HalfNormal(f"sd_dist_{i}", 2.5) ) ) chol.append(pm.expand_packed_triangular(D, packed_chol[i], lower=True))