diff --git a/pymc/distributions/mixture.py b/pymc/distributions/mixture.py index 8afc5b2f1e..3cfda1a837 100644 --- a/pymc/distributions/mixture.py +++ b/pymc/distributions/mixture.py @@ -56,10 +56,8 @@ ] -class MarginalMixtureRV(SymbolicRandomVariable): - """A placeholder used to specify a distribution for a mixture sub-graph.""" - - _print_name = ("MarginalMixture", "\\operatorname{MarginalMixture}") +class _BaseMixtureRV(SymbolicRandomVariable): + """Base class SymbolicRandomVariable for Mixture and Hurdle RVs.""" @classmethod def rv_op(cls, weights, *components, size=None): @@ -139,7 +137,7 @@ def rv_op(cls, weights, *components, size=None): comps_s = ",".join(f"({s})" for _ in components) extended_signature = f"[rng],(w),{comps_s}->[rng],({s})" - return MarginalMixtureRV( + return cls( inputs=[mix_indexes_rng, weights, *components], outputs=[mix_indexes_rng_next, mix_out], extended_signature=extended_signature, @@ -161,117 +159,8 @@ def update(self, node: Apply): return {node.inputs[0]: node.outputs[0]} -class Mixture(Distribution): - R""" - Mixture distribution. - - Often used to model subpopulation heterogeneity - - .. math:: f(x \mid w, \theta) = \sum_{i = 1}^n w_i f_i(x \mid \theta_i) - - ======== ============================================ - Support :math:`\cup_{i = 1}^n \textrm{support}(f_i)` - Mean :math:`\sum_{i = 1}^n w_i \mu_i` - ======== ============================================ - - Parameters - ---------- - w : tensor_like of float - w >= 0 and w <= 1 - the mixture weights - comp_dists : iterable of unnamed distributions or single batched distribution - Distributions should be created via the `.dist()` API. If a single distribution - is passed, the last size dimension (not shape) determines the number of mixture - components (e.g. `pm.Poisson.dist(..., size=components)`) - :math:`f_1, \ldots, f_n` - - .. warning:: comp_dists will be cloned, rendering them independent of the ones passed as input. - - - Examples - -------- - .. code-block:: python - - # Mixture of 2 Poisson variables - with pm.Model() as model: - w = pm.Dirichlet("w", a=np.array([1, 1])) # 2 mixture weights - - lam1 = pm.Exponential("lam1", lam=1) - lam2 = pm.Exponential("lam2", lam=1) - - # As we just need the logp, rather than add a RV to the model, we need to call `.dist()` - # These two forms are equivalent, but the second benefits from vectorization - components = [ - pm.Poisson.dist(mu=lam1), - pm.Poisson.dist(mu=lam2), - ] - # `shape=(2,)` indicates 2 mixture components - components = pm.Poisson.dist(mu=pm.math.stack([lam1, lam2]), shape=(2,)) - - like = pm.Mixture("like", w=w, comp_dists=components, observed=data) - - - .. code-block:: python - - # Mixture of Normal and StudentT variables - with pm.Model() as model: - w = pm.Dirichlet("w", a=np.array([1, 1])) # 2 mixture weights - - mu = pm.Normal("mu", 0, 1) - - components = [ - pm.Normal.dist(mu=mu, sigma=1), - pm.StudentT.dist(nu=4, mu=mu, sigma=1), - ] - - like = pm.Mixture("like", w=w, comp_dists=components, observed=data) - - - .. code-block:: python - - # Mixture of (5 x 3) Normal variables - with pm.Model() as model: - # w is a stack of 5 independent size 3 weight vectors - # If shape was `(3,)`, the weights would be shared across the 5 replication dimensions - w = pm.Dirichlet("w", a=np.ones(3), shape=(5, 3)) - - # Each of the 3 mixture components has an independent mean - mu = pm.Normal("mu", mu=np.arange(3), sigma=1, shape=3) - - # These two forms are equivalent, but the second benefits from vectorization - components = [ - pm.Normal.dist(mu=mu[0], sigma=1, shape=(5,)), - pm.Normal.dist(mu=mu[1], sigma=1, shape=(5,)), - pm.Normal.dist(mu=mu[2], sigma=1, shape=(5,)), - ] - components = pm.Normal.dist(mu=mu, sigma=1, shape=(5, 3)) - - # The mixture is an array of 5 elements - # Each element can be thought of as an independent scalar mixture of 3 - # components with different means - like = pm.Mixture("like", w=w, comp_dists=components, observed=data) - - - .. code-block:: python - - # Mixture of 2 Dirichlet variables - with pm.Model() as model: - w = pm.Dirichlet("w", a=np.ones(2)) # 2 mixture weights - - # These two forms are equivalent, but the second benefits from vectorization - components = [ - pm.Dirichlet.dist(a=[1, 10, 100], shape=(3,)), - pm.Dirichlet.dist(a=[100, 10, 1], shape=(3,)), - ] - components = pm.Dirichlet.dist(a=[[1, 10, 100], [100, 10, 1]], shape=(2, 3)) - - # The mixture is an array of 3 elements - # Each element comes from only one of the two core Dirichlet components - like = pm.Mixture("like", w=w, comp_dists=components, observed=data) - """ - - rv_type = MarginalMixtureRV - rv_op = MarginalMixtureRV.rv_op +class _BaseMixtureDistribution(Distribution): + """Base class distribution for Mixture and Hurdle distributions.""" @classmethod def dist(cls, w, comp_dists, **kwargs): @@ -298,8 +187,6 @@ def dist(cls, w, comp_dists, **kwargs): # Check that components are not associated with a registered variable in the model components_ndim_supp = set() for dist in comp_dists: - # TODO: Allow these to not be a RandomVariable as long as we can call `ndim_supp` on them - # and resize them if not isinstance(dist, TensorVariable) or not isinstance( dist.owner.op, RandomVariable | SymbolicRandomVariable ): @@ -318,8 +205,8 @@ def dist(cls, w, comp_dists, **kwargs): return super().dist([w, *comp_dists], **kwargs) -@_change_dist_size.register(MarginalMixtureRV) -def change_marginal_mixture_size(op, dist, new_size, expand=False): +@_change_dist_size.register(_BaseMixtureRV) +def change_mixture_size(op, dist, new_size, expand=False): rng, weights, *components = dist.owner.inputs if expand: @@ -333,39 +220,32 @@ def change_marginal_mixture_size(op, dist, new_size, expand=False): old_size = components[0].shape[:size_dims] new_size = tuple(new_size) + tuple(old_size) - return Mixture.rv_op(weights, *components, size=new_size) + return op.rv_op(weights, *components, size=new_size) -@_logprob.register(MarginalMixtureRV) -def marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs): - (value,) = values +@_support_point.register(_BaseMixtureRV) +def mixture_support_point(op, rv, rng, weights, *components): + ndim_supp = components[0].owner.op.ndim_supp + weights = pt.shape_padright(weights, ndim_supp) + mix_axis = -ndim_supp - 1 - # single component if len(components) == 1: - # Need to broadcast value across mixture axis - mix_axis = -components[0].owner.op.ndim_supp - 1 - components_logp = logp(components[0], pt.expand_dims(value, mix_axis)) + support_point_components = support_point(components[0]) + else: - components_logp = pt.stack( - [logp(component, value) for component in components], - axis=-1, + support_point_components = pt.stack( + [support_point(component) for component in components], + axis=mix_axis, ) - mix_logp = pt.logsumexp(pt.log(weights) + components_logp, axis=-1) - - mix_logp = check_parameters( - mix_logp, - 0 <= weights, - weights <= 1, - pt.isclose(pt.sum(weights, axis=-1), 1), - msg="0 <= weights <= 1, sum(weights) == 1", - ) - - return mix_logp + mix_support_point = pt.sum(weights * support_point_components, axis=mix_axis) + if components[0].dtype in discrete_types: + mix_support_point = pt.round(mix_support_point) + return mix_support_point -@_logcdf.register(MarginalMixtureRV) -def marginal_mixture_logcdf(op, value, rng, weights, *components, **kwargs): +@_logcdf.register(_BaseMixtureRV) +def mixture_logcdf(op, value, rng, weights, *components, **kwargs): # single component if len(components) == 1: # Need to broadcast value across mixture axis @@ -390,27 +270,6 @@ def marginal_mixture_logcdf(op, value, rng, weights, *components, **kwargs): return mix_logcdf -@_support_point.register(MarginalMixtureRV) -def marginal_mixture_support_point(op, rv, rng, weights, *components): - ndim_supp = components[0].owner.op.ndim_supp - weights = pt.shape_padright(weights, ndim_supp) - mix_axis = -ndim_supp - 1 - - if len(components) == 1: - support_point_components = support_point(components[0]) - - else: - support_point_components = pt.stack( - [support_point(component) for component in components], - axis=mix_axis, - ) - - mix_support_point = pt.sum(weights * support_point_components, axis=mix_axis) - if components[0].dtype in discrete_types: - mix_support_point = pt.round(mix_support_point) - return mix_support_point - - # List of transforms that can be used by Mixture, either because they do not require # special handling or because we have custom logic to enable them. If new default # transforms are implemented, this list and function should be updated @@ -431,8 +290,8 @@ class MixtureTransformWarning(UserWarning): pass -@_default_transform.register(MarginalMixtureRV) -def marginal_mixture_default_transform(op, rv): +@_default_transform.register(_BaseMixtureRV) +def mixture_default_transform(op, rv): def transform_warning(): warnings.warn( f"No safe default transform found for Mixture distribution {rv}. This can " @@ -491,6 +350,151 @@ def mixture_args_fn(rng, weights, *components): return default_transform +class MixtureRV(_BaseMixtureRV): + _print_name = ("Mixture", "\\operatorname{Mixture}") + + +class Mixture(_BaseMixtureDistribution): + R""" + Mixture distribution. + + Often used to model subpopulation heterogeneity + + .. math:: f(x \mid w, \theta) = \sum_{i = 1}^n w_i f_i(x \mid \theta_i) + + ======== ============================================ + Support :math:`\cup_{i = 1}^n \textrm{support}(f_i)` + Mean :math:`\sum_{i = 1}^n w_i \mu_i` + ======== ============================================ + + Parameters + ---------- + w : tensor_like of float + w >= 0 and w <= 1 + the mixture weights + comp_dists : iterable of unnamed distributions or single batched distribution + Distributions should be created via the `.dist()` API. If a single distribution + is passed, the last size dimension (not shape) determines the number of mixture + components (e.g. `pm.Poisson.dist(..., size=components)`) + :math:`f_1, \ldots, f_n` + + .. warning:: comp_dists will be cloned, rendering them independent of the ones passed as input. + + + Examples + -------- + .. code-block:: python + + # Mixture of 2 Poisson variables + with pm.Model() as model: + w = pm.Dirichlet("w", a=np.array([1, 1])) # 2 mixture weights + + lam1 = pm.Exponential("lam1", lam=1) + lam2 = pm.Exponential("lam2", lam=1) + + # As we just need the logp, rather than add a RV to the model, we need to call `.dist()` + # These two forms are equivalent, but the second benefits from vectorization + components = [ + pm.Poisson.dist(mu=lam1), + pm.Poisson.dist(mu=lam2), + ] + # `shape=(2,)` indicates 2 mixture components + components = pm.Poisson.dist(mu=pm.math.stack([lam1, lam2]), shape=(2,)) + + like = pm.Mixture("like", w=w, comp_dists=components, observed=data) + + + .. code-block:: python + + # Mixture of Normal and StudentT variables + with pm.Model() as model: + w = pm.Dirichlet("w", a=np.array([1, 1])) # 2 mixture weights + + mu = pm.Normal("mu", 0, 1) + + components = [ + pm.Normal.dist(mu=mu, sigma=1), + pm.StudentT.dist(nu=4, mu=mu, sigma=1), + ] + + like = pm.Mixture("like", w=w, comp_dists=components, observed=data) + + + .. code-block:: python + + # Mixture of (5 x 3) Normal variables + with pm.Model() as model: + # w is a stack of 5 independent size 3 weight vectors + # If shape was `(3,)`, the weights would be shared across the 5 replication dimensions + w = pm.Dirichlet("w", a=np.ones(3), shape=(5, 3)) + + # Each of the 3 mixture components has an independent mean + mu = pm.Normal("mu", mu=np.arange(3), sigma=1, shape=3) + + # These two forms are equivalent, but the second benefits from vectorization + components = [ + pm.Normal.dist(mu=mu[0], sigma=1, shape=(5,)), + pm.Normal.dist(mu=mu[1], sigma=1, shape=(5,)), + pm.Normal.dist(mu=mu[2], sigma=1, shape=(5,)), + ] + components = pm.Normal.dist(mu=mu, sigma=1, shape=(5, 3)) + + # The mixture is an array of 5 elements + # Each element can be thought of as an independent scalar mixture of 3 + # components with different means + like = pm.Mixture("like", w=w, comp_dists=components, observed=data) + + + .. code-block:: python + + # Mixture of 2 Dirichlet variables + with pm.Model() as model: + w = pm.Dirichlet("w", a=np.ones(2)) # 2 mixture weights + + # These two forms are equivalent, but the second benefits from vectorization + components = [ + pm.Dirichlet.dist(a=[1, 10, 100], shape=(3,)), + pm.Dirichlet.dist(a=[100, 10, 1], shape=(3,)), + ] + components = pm.Dirichlet.dist(a=[[1, 10, 100], [100, 10, 1]], shape=(2, 3)) + + # The mixture is an array of 3 elements + # Each element comes from only one of the two core Dirichlet components + like = pm.Mixture("like", w=w, comp_dists=components, observed=data) + """ + + rv_type = MixtureRV + rv_op = MixtureRV.rv_op + + +@_logprob.register(MixtureRV) +def mixture_logprob(op, values, rng, weights, *components, **kwargs): + (value,) = values + + # single component + if len(components) == 1: + # Need to broadcast value across mixture axis + mix_axis = -components[0].owner.op.ndim_supp - 1 + components_logp = logp(components[0], pt.expand_dims(value, mix_axis)) + else: + components_logp = pt.stack( + [logp(component, value) for component in components], + axis=-1, + ) + + mix_logp = pt.logsumexp(pt.log(weights) + components_logp, axis=-1) + + mix_logp = check_parameters( + mix_logp, + 0 <= weights, + weights <= 1, + pt.isclose(pt.sum(weights, axis=-1), 1), + msg="0 <= weights <= 1, sum(weights) == 1", + ) + + return mix_logp + + class NormalMixture: R""" Normal mixture distribution. @@ -799,34 +803,65 @@ def dist(cls, psi, mu=None, alpha=None, p=None, n=None, **kwargs): ) -def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, max_n_steps=10_000, **kwargs): - """Create a hurdle mixtures (helper function). +class _HurdleRV(_BaseMixtureRV): + _print_name = ("Hurdle", "\\operatorname{Hurdle}") - If name is `None`, this function returns an unregistered variable - In hurdle models, the zeros come from a completely different process than the rest of the data. - In other words, the zeros are not inflated, they come from a different process. - """ - if dtype == "float": - zero = 0.0 - lower = np.finfo(pytensor.config.floatX).eps - elif dtype == "int": - zero = 0 - lower = 1 - else: - raise ValueError("dtype must be 'float' or 'int'") +class _Hurdle(_BaseMixtureDistribution): + rv_type = _HurdleRV + rv_op = _HurdleRV.rv_op - nonzero_p = pt.as_tensor_variable(nonzero_p) - weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1) - comp_dists = [ - DiracDelta.dist(zero), - Truncated.dist(nonzero_dist, lower=lower, max_n_steps=max_n_steps), - ] + @classmethod + def _create(cls, *, name, nonzero_p, nonzero_dist, max_n_steps=10_000, **kwargs): + """Create a hurdle mixture (helper function). + + If name is `None`, this function returns an unregistered variable + + In hurdle models, the zeros come from a completely different process than the rest of the data. + In other words, the zeros are not inflated, they come from a different process. + + Note: this is invalid for discrete nonzero distributions with mass below 0, as we simply truncate[lower=1]. + """ + dtype = nonzero_dist.dtype + + if dtype.startswith("int"): + # Need to truncate the distribution to exclude zero. + # Continuous distributions have "zero" mass at zero (and anywhere else), so can be used as is. + nonzero_dist = Truncated.dist(nonzero_dist, lower=1, max_n_steps=max_n_steps) + elif not dtype.startswith("float"): + raise ValueError(f"nonzero_dist dtype must be 'float' or 'int', got {dtype}") + + nonzero_p = pt.as_tensor_variable(nonzero_p) + weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1) + comp_dists = [ + DiracDelta.dist(np.asarray(0, dtype=dtype)), + nonzero_dist, + ] + + if name is not None: + return cls(name, weights, comp_dists, **kwargs) + else: + return cls.dist(weights, comp_dists, **kwargs) - if name is not None: - return Mixture(name, weights, comp_dists, **kwargs) - else: - return Mixture.dist(weights, comp_dists, **kwargs) + +@_logprob.register(_HurdleRV) +def marginal_hurdle_logprob(op, values, rng, weights, zero_dist, dist, **kwargs): + (value,) = values + + psi = weights[..., 1] + + hurdle_logp = pt.where( + pt.eq(value, 0), + pt.log(1 - psi), + pt.log(psi) + logp(dist, value), + ) + + return check_parameters( + hurdle_logp, + 0 <= psi, + psi <= 1, + msg="0 <= psi <= 1", + ) class HurdlePoisson: @@ -864,14 +899,20 @@ class HurdlePoisson: """ def __new__(cls, name, psi, mu, **kwargs): - return _hurdle_mixture( - name=name, nonzero_p=psi, nonzero_dist=Poisson.dist(mu=mu), dtype="int", **kwargs + return _Hurdle._create( + name=name, + nonzero_p=psi, + nonzero_dist=Poisson.dist(mu=mu), + **kwargs, ) @classmethod def dist(cls, psi, mu, **kwargs): - return _hurdle_mixture( - name=None, nonzero_p=psi, nonzero_dist=Poisson.dist(mu=mu), dtype="int", **kwargs + return _Hurdle._create( + name=None, + nonzero_p=psi, + nonzero_dist=Poisson.dist(mu=mu), + **kwargs, ) @@ -914,21 +955,19 @@ class HurdleNegativeBinomial: """ def __new__(cls, name, psi, mu=None, alpha=None, p=None, n=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=name, nonzero_p=psi, nonzero_dist=NegativeBinomial.dist(mu=mu, alpha=alpha, p=p, n=n), - dtype="int", **kwargs, ) @classmethod def dist(cls, psi, mu=None, alpha=None, p=None, n=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=None, nonzero_p=psi, nonzero_dist=NegativeBinomial.dist(mu=mu, alpha=alpha, p=p, n=n), - dtype="int", **kwargs, ) @@ -963,24 +1002,28 @@ class HurdleGamma: Alternative shape parameter (mu > 0). sigma : tensor_like of float, optional Alternative scale parameter (sigma > 0). + + .. warning:: + HurdleGamma distributions cannot be sampled correctly with MCMC methods, + as this would require a specialized step sampler. They are intended to be used as + observed variables, and/or sampled exclusively with forward methods like + `sample_prior_predictive` and `sample_posterior_predictive`. """ def __new__(cls, name, psi, alpha=None, beta=None, mu=None, sigma=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=name, nonzero_p=psi, nonzero_dist=Gamma.dist(alpha=alpha, beta=beta, mu=mu, sigma=sigma), - dtype="float", **kwargs, ) @classmethod def dist(cls, psi, alpha=None, beta=None, mu=None, sigma=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=None, nonzero_p=psi, nonzero_dist=Gamma.dist(alpha=alpha, beta=beta, mu=mu, sigma=sigma), - dtype="float", **kwargs, ) @@ -1015,23 +1058,27 @@ class HurdleLogNormal: tau : tensor_like of float, optional Scale parameter (tau > 0). (only required if sigma is not specified). Defaults to 1. + + .. warning:: + HurdleLogNormal distributions cannot be sampled correctly with MCMC methods, + as this would require a specialized step sampler. They are intended to be used as + observed variables, and/or sampled exclusively with forward methods like + `sample_prior_predictive` and `sample_posterior_predictive`. """ def __new__(cls, name, psi, mu=0, sigma=None, tau=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=name, nonzero_p=psi, nonzero_dist=LogNormal.dist(mu=mu, sigma=sigma, tau=tau), - dtype="float", **kwargs, ) @classmethod def dist(cls, psi, mu=0, sigma=None, tau=None, **kwargs): - return _hurdle_mixture( + return _Hurdle._create( name=None, nonzero_p=psi, nonzero_dist=LogNormal.dist(mu=mu, sigma=sigma, tau=tau), - dtype="float", **kwargs, ) diff --git a/pymc/distributions/moments/means.py b/pymc/distributions/moments/means.py index f183ace5db..0e3129935e 100644 --- a/pymc/distributions/moments/means.py +++ b/pymc/distributions/moments/means.py @@ -70,7 +70,7 @@ ) from pymc.distributions.discrete import DiscreteUniformRV from pymc.distributions.distribution import DiracDeltaRV -from pymc.distributions.mixture import MarginalMixtureRV +from pymc.distributions.mixture import MixtureRV from pymc.distributions.multivariate import ( CARRV, DirichletMultinomialRV, @@ -300,7 +300,7 @@ def lognormal_mean(op, rv, rng, size, mu, sigma): return maybe_resize(pt.exp(mu + 0.5 * sigma**2), size) -@_mean.register(MarginalMixtureRV) +@_mean.register(MixtureRV) def marginal_mixture_mean(op, rv, rng, weights, *components): ndim_supp = components[0].owner.op.ndim_supp weights = pt.shape_padright(weights, ndim_supp) diff --git a/tests/distributions/test_distribution.py b/tests/distributions/test_distribution.py index d7e2bbd0a1..4059465443 100644 --- a/tests/distributions/test_distribution.py +++ b/tests/distributions/test_distribution.py @@ -92,8 +92,14 @@ def test_all_distributions_have_support_points(): dists = (getattr(dist_module, dist) for dist in dist_module.__all__) dists = (dist for dist in dists if isinstance(dist, DistributionMeta)) + generic_func = _support_point.dispatch(object) missing_support_points = { - dist for dist in dists if getattr(dist, "rv_type", None) not in _support_point.registry + dist + for dist in dists + if ( + getattr(dist, "rv_type", None) is not None + and _support_point.dispatch(dist.rv_type) is generic_func + ) } # Ignore super classes diff --git a/tests/distributions/test_mixture.py b/tests/distributions/test_mixture.py index 7fd00bcb5a..28a09744b6 100644 --- a/tests/distributions/test_mixture.py +++ b/tests/distributions/test_mixture.py @@ -49,13 +49,12 @@ Poisson, StickBreakingWeights, Triangular, - Truncated, Uniform, ZeroInflatedBinomial, ZeroInflatedNegativeBinomial, ZeroInflatedPoisson, ) -from pymc.distributions.mixture import MixtureTransformWarning +from pymc.distributions.mixture import MixtureTransformWarning, _Hurdle from pymc.distributions.shape_utils import change_dist_size, to_tuple from pymc.distributions.transforms import _default_transform from pymc.logprob.basic import logp @@ -1563,26 +1562,24 @@ def test_zero_inflated_dists_dtype_and_broadcast(self, dist, non_psi_args): assert x.eval().shape == (3,) -class TestHurdleMixtures: +class TestHurdleDistributions: @staticmethod - def check_hurdle_mixture_graph(dist): + def check_hurdle_graph(dist): # Assert it's a mixture - assert isinstance(dist.owner.op, Mixture) + assert isinstance(dist.owner.op, _Hurdle) # Extract the distribution for zeroes and nonzeroes zero_dist, nonzero_dist = dist.owner.inputs[-2:] # Assert ops are of the right type assert isinstance(zero_dist.owner.op, DiracDelta) - assert isinstance(nonzero_dist.owner.op, Truncated) - return zero_dist, nonzero_dist def test_hurdle_poisson_graph(self): # There's nothing special in these values psi, mu = 0.3, 4 dist = HurdlePoisson.dist(psi=psi, mu=mu) - _, nonzero_dist = self.check_hurdle_mixture_graph(dist) + _, nonzero_dist = self.check_hurdle_graph(dist) # Assert the truncated distribution is of the right type assert isinstance(nonzero_dist.owner.op.base_rv_op, Poisson) @@ -1593,7 +1590,7 @@ def test_hurdle_poisson_graph(self): def test_hurdle_negativebinomial_graph(self): psi, p, n = 0.2, 0.6, 10 dist = HurdleNegativeBinomial.dist(psi=psi, p=p, n=n) - _, nonzero_dist = self.check_hurdle_mixture_graph(dist) + _, nonzero_dist = self.check_hurdle_graph(dist) assert isinstance(nonzero_dist.owner.op.base_rv_op, NegativeBinomial) assert nonzero_dist.owner.inputs[-4].data == n @@ -1602,22 +1599,24 @@ def test_hurdle_negativebinomial_graph(self): def test_hurdle_gamma_graph(self): psi, alpha, beta = 0.25, 3, 4 dist = HurdleGamma.dist(psi=psi, alpha=alpha, beta=beta) - _, nonzero_dist = self.check_hurdle_mixture_graph(dist) + _, nonzero_dist = self.check_hurdle_graph(dist) # Under the hood it uses the shape-scale parametrization of the Gamma distribution. # So the second value is the reciprocal of the rate (i.e. 1 / beta) - assert isinstance(nonzero_dist.owner.op.base_rv_op, Gamma) - assert nonzero_dist.owner.inputs[-4].data == alpha - assert nonzero_dist.owner.inputs[-3].eval() == 1 / beta + assert isinstance(nonzero_dist.owner.op, Gamma) + alpha_param, reciprocal_beta_param = nonzero_dist.owner.op.dist_params(nonzero_dist.owner) + assert alpha_param.data == alpha + assert reciprocal_beta_param.eval() == 1 / beta def test_hurdle_lognormal_graph(self): psi, mu, sigma = 0.1, 2, 2.5 dist = HurdleLogNormal.dist(psi=psi, mu=mu, sigma=sigma) - _, nonzero_dist = self.check_hurdle_mixture_graph(dist) + _, nonzero_dist = self.check_hurdle_graph(dist) - assert isinstance(nonzero_dist.owner.op.base_rv_op, LogNormal) - assert nonzero_dist.owner.inputs[-4].data == mu - assert nonzero_dist.owner.inputs[-3].data == sigma + assert isinstance(nonzero_dist.owner.op, LogNormal) + mu_param, sigma_param = nonzero_dist.owner.op.dist_params(nonzero_dist.owner) + assert mu_param.data == mu + assert sigma_param.data == sigma @pytest.mark.parametrize( "dist, psi, non_psi_args", @@ -1699,11 +1698,7 @@ def logp_fn(value, psi, alpha, beta): if value == 0: return np.log(1 - psi) else: - return ( - np.log(psi) - + st.gamma.logpdf(value, alpha, scale=1.0 / beta) - - np.log(1 - st.gamma.cdf(np.finfo(float).eps, alpha, scale=1.0 / beta)) - ) + return np.log(psi) + st.gamma.logpdf(value, alpha, scale=1.0 / beta) check_logp(HurdleGamma, Rplus, {"psi": Unit, "alpha": Rplusbig, "beta": Rplusbig}, logp_fn) @@ -1712,10 +1707,6 @@ def logp_fn(value, psi, mu, sigma): if value == 0: return np.log(1 - psi) else: - return ( - np.log(psi) - + st.lognorm.logpdf(value, sigma, 0, np.exp(mu)) - - np.log(1 - st.lognorm.cdf(np.finfo(float).eps, sigma, 0, np.exp(mu))) - ) + return np.log(psi) + st.lognorm.logpdf(value, sigma, 0, np.exp(mu)) check_logp(HurdleLogNormal, Rplus, {"psi": Unit, "mu": R, "sigma": Rplusbig}, logp_fn) diff --git a/tests/test_printing.py b/tests/test_printing.py index 917a2d1ee2..879b25b96b 100644 --- a/tests/test_printing.py +++ b/tests/test_printing.py @@ -141,11 +141,11 @@ def setup_class(self): r"beta ~ Normal(0, 10)", r"Z ~ MultivariateNormal(f(), f())", r"nb_with_p_n ~ NegativeBinomial(10, nbp)", - r"zip ~ MarginalMixture(f(), DiracDelta(0), Poisson(5))", + r"zip ~ Mixture(f(), DiracDelta(0), Poisson(5))", r"w ~ Dirichlet()", ( - r"nested_mix ~ MarginalMixture(w, " - r"MarginalMixture(f(), DiracDelta(0), Poisson(5)), " + r"nested_mix ~ Mixture(w, " + r"Mixture(f(), DiracDelta(0), Poisson(5)), " r"Censored(Bernoulli(0.5), -1, 1))" ), r"Y_obs ~ Normal(mu, sigma)", @@ -159,9 +159,9 @@ def setup_class(self): r"beta ~ Normal", r"Z ~ MultivariateNormal", r"nb_with_p_n ~ NegativeBinomial", - r"zip ~ MarginalMixture", + r"zip ~ Mixture", r"w ~ Dirichlet", - r"nested_mix ~ MarginalMixture", + r"nested_mix ~ Mixture", r"Y_obs ~ Normal", r"pot ~ Potential", r"pred ~ Deterministic", @@ -173,11 +173,11 @@ def setup_class(self): r"$\text{beta} \sim \operatorname{Normal}(0,~10)$", r"$\text{Z} \sim \operatorname{MultivariateNormal}(f(),~f())$", r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}(10,~\text{nbp})$", - r"$\text{zip} \sim \operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))$", + r"$\text{zip} \sim \operatorname{Mixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))$", r"$\text{w} \sim \operatorname{Dirichlet}(\text{})$", ( - r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}(\text{w}," - r"~\operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))," + r"$\text{nested\_mix} \sim \operatorname{Mixture}(\text{w}," + r"~\operatorname{Mixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))," r"~\operatorname{Censored}(\operatorname{Bernoulli}(0.5),~-1,~1))$" ), r"$\text{Y\_obs} \sim \operatorname{Normal}(\text{mu},~\text{sigma})$", @@ -191,9 +191,9 @@ def setup_class(self): r"$\text{beta} \sim \operatorname{Normal}$", r"$\text{Z} \sim \operatorname{MultivariateNormal}$", r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}$", - r"$\text{zip} \sim \operatorname{MarginalMixture}$", + r"$\text{zip} \sim \operatorname{Mixture}$", r"$\text{w} \sim \operatorname{Dirichlet}$", - r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}$", + r"$\text{nested\_mix} \sim \operatorname{Mixture}$", r"$\text{Y\_obs} \sim \operatorname{Normal}$", r"$\text{pot} \sim \operatorname{Potential}$", r"$\text{pred} \sim \operatorname{Deterministic}", @@ -276,7 +276,7 @@ def test_model_latex_repr_mixture_model(): "\\begin{array}{rcl}", "\\text{w} &\\sim & " "\\operatorname{Dirichlet}(\\text{})\\\\\\text{mix} &\\sim & " - "\\operatorname{MarginalMixture}(\\text{w},~\\operatorname{Normal}(0,~5),~\\operatorname{StudentT}(7,~0,~1))", + "\\operatorname{Mixture}(\\text{w},~\\operatorname{Normal}(0,~5),~\\operatorname{StudentT}(7,~0,~1))", "\\end{array}", "$$", ]