diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst index be9716599..d48c29435 100644 --- a/docs/source/distributions.rst +++ b/docs/source/distributions.rst @@ -245,6 +245,14 @@ TruncatedCauchy :show-inheritance: :member-order: bysource +TruncatedDistribution +--------------------- +.. autoclass:: numpyro.distributions.continuous.TruncatedDistribution + :members: + :undoc-members: + :show-inheritance: + :member-order: bysource + TruncatedNormal --------------- .. autoclass:: numpyro.distributions.continuous.TruncatedNormal diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py index cc5fbec9a..9f2cb662c 100644 --- a/numpyro/distributions/__init__.py +++ b/numpyro/distributions/__init__.py @@ -26,6 +26,7 @@ Pareto, StudentT, TruncatedCauchy, + TruncatedDistribution, TruncatedNormal, TruncatedPolyaGamma, Uniform @@ -121,6 +122,7 @@ 'StudentT', 'TransformedDistribution', 'TruncatedCauchy', + 'TruncatedDistribution', 'TruncatedNormal', 'TruncatedPolyaGamma', 'Uniform', diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py index f4b87cef6..acfd7112e 100644 --- a/numpyro/distributions/continuous.py +++ b/numpyro/distributions/continuous.py @@ -26,12 +26,12 @@ # POSSIBILITY OF SUCH DAMAGE. -from jax import lax, ops +from jax import lax, ops, tree_map import jax.nn as nn import jax.numpy as jnp import jax.random as random from jax.scipy.linalg import cho_solve, solve_triangular -from jax.scipy.special import gammaln, log_ndtr, logsumexp, multigammaln, ndtr, ndtri +from jax.scipy.special import betainc, expit, gammaln, logit, log_ndtr, logsumexp, multigammaln, ndtr, ndtri from numpyro.distributions import constraints from numpyro.distributions.distribution import Distribution, TransformedDistribution @@ -81,6 +81,9 @@ def variance(self): total = self.concentration1 + self.concentration0 return self.concentration1 * self.concentration0 / (total ** 2 * (total + 1)) + def cdf(self, value): + return betainc(self.concentration1, self.concentration0, value) + class Cauchy(Distribution): arg_constraints = {'loc': constraints.real, 'scale': constraints.positive} @@ -109,6 +112,13 @@ def mean(self): def variance(self): return jnp.full(self.batch_shape, jnp.nan) + def cdf(self, value): + scaled = (value - self.loc) / self.scale + return jnp.arctan(scaled) / jnp.pi + 0.5 + + def icdf(self, q): + return self.loc + self.scale * jnp.tan(jnp.pi * (q - 0.5)) + class Dirichlet(Distribution): arg_constraints = {'concentration': constraints.independent(constraints.positive, 1)} @@ -412,6 +422,14 @@ def mean(self): def variance(self): return jnp.broadcast_to(2 * self.scale ** 2, self.batch_shape) + def cdf(self, value): + scaled = (value - self.loc) / self.scale + return 0.5 - 0.5 * jnp.sign(scaled) * jnp.expm1(-jnp.abs(scaled)) + + def icdf(self, q): + a = q - 0.5 + return self.loc - self.scale * jnp.sign(a) * jnp.log1p(-2 * jnp.abs(a)) + class LKJ(TransformedDistribution): r""" @@ -659,6 +677,44 @@ def tree_flatten(self): return super(TransformedDistribution, self).tree_flatten() +class Logistic(Distribution): + arg_constraints = {'loc': constraints.real, 'scale': constraints.positive} + support = constraints.real + reparametrized_params = ['loc', 'scale'] + + def __init__(self, loc=0., scale=1., validate_args=None): + self.loc, self.scale = promote_shapes(loc, scale) + batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale)) + super(Logistic, self).__init__(batch_shape, validate_args=validate_args) + + def sample(self, key, sample_shape=()): + assert is_prng_key(key) + z = random.logistic(key, shape=sample_shape + self.batch_shape + self.event_shape) + return self.loc + z * self.scale + + @validate_sample + def log_prob(self, value): + log_exponent = (self.loc - value) / self.scale + log_denominator = jnp.log(self.scale) + 2 * nn.softplus(log_exponent) + return log_exponent - log_denominator + + @property + def mean(self): + return jnp.broadcast_to(self.loc, self.batch_shape) + + @property + def variance(self): + var = (self.scale ** 2) * (jnp.pi ** 2) / 3 + return jnp.broadcast_to(var, self.batch_shape) + + def cdf(self, value): + scaled = (value - self.loc) / self.scale + return expit(scaled) + + def icdf(self, q): + return self.loc + self.scale * logit(q) + + def _batch_mahalanobis(bL, bx): if bL.shape[:-1] == bx.shape: # no need to use the below optimization procedure @@ -966,6 +1022,10 @@ def log_prob(self, value): value_scaled = (value - self.loc) / self.scale return -0.5 * value_scaled ** 2 - normalize_term + def cdf(self, value): + scaled = (value - self.loc) / self.scale + return ndtr(scaled) + def icdf(self, q): return self.loc + self.scale * ndtri(q) @@ -1049,6 +1109,241 @@ def variance(self): var = jnp.where(self.df <= 1, jnp.nan, var) return jnp.broadcast_to(var, self.batch_shape) + def cdf(self, value): + # Ref: https://en.wikipedia.org/wiki/Student's_t-distribution#Related_distributions + # X^2 ~ F(1, df) -> df / (df + X^2) ~ Beta(df/2, 0.5) + scaled = (value - self.loc) / self.scale + scaled_squared = scaled * scaled + beta_value = self.df / (self.df + scaled_squared) + # when scaled < 0, returns 0.5 * Beta(df/2, 0.5).cdf(beta_value) + # when scaled > 0, returns 1 - 0.5 * Beta(df/2, 0.5).cdf(beta_value) + return 0.5 * (1 + jnp.sign(scaled) * (1 - betainc(0.5 * self.df, 0.5, beta_value))) + + def icdf(self, q): + # scipy.special.betaincinv is not avaiable yet in JAX + # upstream issue: https://github.com/google/jax/issues/2399 + raise NotImplementedError + + +class LeftTruncatedDistribution(Distribution): + arg_constraints = {"low": constraints.real} + reparametrized_params = ["low"] + supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT) + + def __init__(self, base_dist, low=0., validate_args=None): + assert isinstance(base_dist, self.supported_types) + assert base_dist.support is constraints.real, \ + "The base distribution should be univariate and have real support." + batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low)) + self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist) + self.low, = promote_shapes(low, shape=batch_shape) + self._support = constraints.greater_than(low) + super().__init__(batch_shape, validate_args=validate_args) + + @constraints.dependent_property(is_discrete=False, event_dim=0) + def support(self): + return self._support + + @lazy_property + def _tail_prob_at_low(self): + # if low < loc, returns cdf(low); otherwise returns 1 - cdf(low) + loc = self.base_dist.loc + sign = jnp.where(loc >= self.low, 1., -1.) + return self.base_dist.cdf(loc - sign * (loc - self.low)) + + @lazy_property + def _tail_prob_at_high(self): + # if low < loc, returns cdf(high) = 1; otherwise returns 1 - cdf(high) = 0 + return jnp.where(self.low < self.base_dist.loc, 1., 0.) + + def sample(self, key, sample_shape=()): + assert is_prng_key(key) + u = random.uniform(key, sample_shape + self.batch_shape) + loc = self.base_dist.loc + sign = jnp.where(loc >= self.low, 1., -1.) + return (1 - sign) * loc + sign * self.base_dist.icdf( + (1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high) + + @validate_sample + def log_prob(self, value): + sign = jnp.where(self.base_dist.loc >= self.low, 1., -1.) + return self.base_dist.log_prob(value) - \ + jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low)) + + def tree_flatten(self): + base_flatten, base_aux = self.base_dist.tree_flatten() + if isinstance(self._support.lower_bound, (int, float)): + return base_flatten, (type(self.base_dist), base_aux, self._support.lower_bound) + else: + return (base_flatten, self.low), (type(self.base_dist), base_aux) + + @classmethod + def tree_unflatten(cls, aux_data, params): + if len(aux_data) == 2: + base_flatten, low = params + base_cls, base_aux = aux_data + else: + base_flatten = params + base_cls, base_aux, low = aux_data + base_dist = base_cls.tree_unflatten(base_aux, base_flatten) + return cls(base_dist, low=low) + + +class RightTruncatedDistribution(Distribution): + arg_constraints = {"high": constraints.real} + reparametrized_params = ["high"] + supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT) + + def __init__(self, base_dist, high=0., validate_args=None): + assert isinstance(base_dist, self.supported_types) + assert base_dist.support is constraints.real, \ + "The base distribution should be univariate and have real support." + batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(high)) + self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist) + self.high, = promote_shapes(high, shape=batch_shape) + self._support = constraints.less_than(high) + super().__init__(batch_shape, validate_args=validate_args) + + @constraints.dependent_property(is_discrete=False, event_dim=0) + def support(self): + return self._support + + @lazy_property + def _cdf_at_high(self): + return self.base_dist.cdf(self.high) + + def sample(self, key, sample_shape=()): + assert is_prng_key(key) + u = random.uniform(key, sample_shape + self.batch_shape) + return self.base_dist.icdf(u * self._cdf_at_high) + + @validate_sample + def log_prob(self, value): + return self.base_dist.log_prob(value) - jnp.log(self._cdf_at_high) + + def tree_flatten(self): + base_flatten, base_aux = self.base_dist.tree_flatten() + if isinstance(self._support.upper_bound, (int, float)): + return base_flatten, (type(self.base_dist), base_aux, self._support.upper_bound) + else: + return (base_flatten, self.high), (type(self.base_dist), base_aux) + + @classmethod + def tree_unflatten(cls, aux_data, params): + if len(aux_data) == 2: + base_flatten, high = params + base_cls, base_aux = aux_data + else: + base_flatten = params + base_cls, base_aux, high = aux_data + base_dist = base_cls.tree_unflatten(base_aux, base_flatten) + return cls(base_dist, high=high) + + +class TwoSidedTruncatedDistribution(Distribution): + arg_constraints = {"low": constraints.dependent, "high": constraints.dependent} + reparametrized_params = ["low", "high"] + supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT) + + def __init__(self, base_dist, low=0., high=1., validate_args=None): + assert isinstance(base_dist, self.supported_types) + assert base_dist.support is constraints.real, \ + "The base distribution should be univariate and have real support." + batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low), jnp.shape(high)) + self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist) + self.low, = promote_shapes(low, shape=batch_shape) + self.high, = promote_shapes(high, shape=batch_shape) + self._support = constraints.interval(low, high) + super().__init__(batch_shape, validate_args=validate_args) + + @constraints.dependent_property(is_discrete=False, event_dim=0) + def support(self): + return self._support + + @lazy_property + def _tail_prob_at_low(self): + # if low < loc, returns cdf(low); otherwise returns 1 - cdf(low) + loc = self.base_dist.loc + sign = jnp.where(loc >= self.low, 1., -1.) + return self.base_dist.cdf(loc - sign * (loc - self.low)) + + @lazy_property + def _tail_prob_at_high(self): + # if low < loc, returns cdf(high); otherwise returns 1 - cdf(high) + loc = self.base_dist.loc + sign = jnp.where(loc >= self.low, 1., -1.) + return self.base_dist.cdf(loc - sign * (loc - self.high)) + + def sample(self, key, sample_shape=()): + assert is_prng_key(key) + u = random.uniform(key, sample_shape + self.batch_shape) + + # NB: we use a more numerically stable formula for a symmetric base distribution + # A = icdf(cdf(low) + (cdf(high) - cdf(low)) * u) = icdf[(1 - u) * cdf(low) + u * cdf(high)] + # will suffer by precision issues when low is large; + # If low < loc: + # A = icdf[(1 - u) * cdf(low) + u * cdf(high)] + # Else + # A = 2 * loc - icdf[(1 - u) * cdf(2*loc-low)) + u * cdf(2*loc - high)] + loc = self.base_dist.loc + sign = jnp.where(loc >= self.low, 1., -1.) + return (1 - sign) * loc + sign * self.base_dist.icdf( + (1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high) + + @validate_sample + def log_prob(self, value): + # NB: we use a more numerically stable formula for a symmetric base distribution + # if low < loc + # cdf(high) - cdf(low) = as-is + # if low > loc + # cdf(high) - cdf(low) = cdf(2 * loc - low) - cdf(2 * loc - high) + sign = jnp.where(self.base_dist.loc >= self.low, 1., -1.) + return self.base_dist.log_prob(value) - \ + jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low)) + + def tree_flatten(self): + base_flatten, base_aux = self.base_dist.tree_flatten() + if isinstance(self._support.lower_bound, (int, float)) and \ + isinstance(self._support.upper_bound, (int, float)): + return base_flatten, (type(self.base_dist), base_aux, + self._support.lower_bound, self._support.upper_bound) + else: + return (base_flatten, self.low, self.high), (type(self.base_dist), base_aux) + + @classmethod + def tree_unflatten(cls, aux_data, params): + if len(aux_data) == 2: + base_flatten, low, high = params + base_cls, base_aux = aux_data + else: + base_flatten = params + base_cls, base_aux, low, high = aux_data + base_dist = base_cls.tree_unflatten(base_aux, base_flatten) + return cls(base_dist, low=low, high=high) + + +def TruncatedDistribution(base_dist, low=None, high=None, validate_args=None): + """ + A function to generate a truncated distribution. + + :param base_dist: The base distribution to be truncated. This should be a univariate + distribution. Currently, only the following distributions are supported: + Cauchy, Laplace, Logistic, Normal, and StudentT. + :param low: the value which is used to truncate the base distribution from below. + Setting this parameter to None to not truncate from below. + :param high: the value which is used to truncate the base distribution from above. + Setting this parameter to None to not truncate from above. + """ + if high is None: + if low is None: + return base_dist + else: + return LeftTruncatedDistribution(base_dist, low=low, validate_args=validate_args) + elif low is None: + return RightTruncatedDistribution(base_dist, high=high, validate_args=validate_args) + else: + return TwoSidedTruncatedDistribution(base_dist, low=low, high=high, validate_args=validate_args) + class _BaseTruncatedCauchy(Distribution): # NB: this is a truncated cauchy with low=0, scale=1 @@ -1252,37 +1547,6 @@ def infer_shapes(low=(), high=()): return batch_shape, event_shape -class Logistic(Distribution): - arg_constraints = {'loc': constraints.real, 'scale': constraints.positive} - support = constraints.real - reparametrized_params = ['loc', 'scale'] - - def __init__(self, loc=0., scale=1., validate_args=None): - self.loc, self.scale = promote_shapes(loc, scale) - batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale)) - super(Logistic, self).__init__(batch_shape, validate_args=validate_args) - - def sample(self, key, sample_shape=()): - assert is_prng_key(key) - z = random.logistic(key, shape=sample_shape + self.batch_shape + self.event_shape) - return self.loc + z * self.scale - - @validate_sample - def log_prob(self, value): - log_exponent = (self.loc - value) / self.scale - log_denominator = jnp.log(self.scale) + 2 * nn.softplus(log_exponent) - return log_exponent - log_denominator - - @property - def mean(self): - return jnp.broadcast_to(self.loc, self.batch_shape) - - @property - def variance(self): - var = (self.scale ** 2) * (jnp.pi ** 2) / 3 - return jnp.broadcast_to(var, self.batch_shape) - - class TruncatedPolyaGamma(Distribution): truncation_point = 2.5 num_log_prob_terms = 7 diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py index 599424af0..898f093c0 100644 --- a/numpyro/distributions/distribution.py +++ b/numpyro/distributions/distribution.py @@ -391,6 +391,24 @@ def infer_shapes(cls, *args, **kwargs): event_shape = () return batch_shape, event_shape + def cdf(self, value): + """ + The cummulative distribution function of this distribution. + + :param value: samples from this distribution. + :return: output of the cummulative distribution function evaluated at `value`. + """ + raise NotImplementedError + + def icdf(self, q): + """ + The inverse cumulative distribution function of this distribution. + + :param q: quantile values, should belong to [0, 1]. + :return: the samples whose cdf values equals to `q`. + """ + raise NotImplementedError + class ExpandedDistribution(Distribution): arg_constraints = {} diff --git a/test/test_distributions.py b/test/test_distributions.py index d8b746c38..76602bd5f 100644 --- a/test/test_distributions.py +++ b/test/test_distributions.py @@ -60,6 +60,27 @@ def _lowrank_mvn_to_scipy(loc, cov_fac, cov_diag): return osp.multivariate_normal(mean=mean, cov=cov) +def _truncnorm_to_scipy(loc, scale, low, high): + if low is None: + a = -np.inf + else: + a = (low - loc) / scale + if high is None: + b = np.inf + else: + b = (high - loc) / scale + return osp.truncnorm(a, b, loc=loc, scale=scale) + + +def _TruncatedNormal(loc, scale, low, high): + return dist.TruncatedDistribution(dist.Normal(loc, scale), low, high) + + +_TruncatedNormal.arg_constraints = {} +_TruncatedNormal.reparametrized_params = [] +_TruncatedNormal.infer_shapes = lambda *args: (lax.broadcast_shapes(*args), ()) + + class _ImproperWrapper(dist.ImproperUniform): def sample(self, key, sample_shape=()): transform = biject_to(self.support) @@ -103,7 +124,8 @@ def sample(self, key, sample_shape=()): dist.Uniform: lambda a, b: osp.uniform(a, b - a), dist.Logistic: lambda loc, scale: osp.logistic(loc=loc, scale=scale), dist.VonMises: lambda loc, conc: osp.vonmises(loc=np.array(loc, dtype=np.float64), - kappa=np.array(conc, dtype=np.float64)) + kappa=np.array(conc, dtype=np.float64)), + _TruncatedNormal: _truncnorm_to_scipy, } CONTINUOUS = [ @@ -178,6 +200,12 @@ def sample(self, key, sample_shape=()): T(dist.TruncatedNormal, -1., 0., 1.), T(dist.TruncatedNormal, 1., -1., jnp.array([1., 2.])), T(dist.TruncatedNormal, jnp.array([-2., 2.]), jnp.array([0., 1.]), jnp.array([[1.], [2.]])), + T(_TruncatedNormal, -1., 2., 1., 5.), + T(_TruncatedNormal, jnp.array([-1., 4.]), 2., None, 5.), + T(_TruncatedNormal, -1., jnp.array([2., 3.]), 1., None), + T(_TruncatedNormal, -1., 2., jnp.array([-6., 4.]), jnp.array([-4., 6.])), + T(_TruncatedNormal, jnp.array([0., 1.]), jnp.array([[1.], [2.]]), None, jnp.array([-2., 2.])), + T(dist.continuous.TwoSidedTruncatedDistribution, dist.Laplace(0., 1.), -2., 3.), T(dist.Uniform, 0., 2.), T(dist.Uniform, 1., jnp.array([2., 3.])), T(dist.Uniform, jnp.array([0., 0.]), jnp.array([[2.], [3.]])), @@ -346,6 +374,7 @@ def test_dist_shape(jax_dist, sp_dist, params, prepend_shape): @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS + DISCRETE + DIRECTIONAL) def test_infer_shapes(jax_dist, sp_dist, params, prepend_shape): shapes = tuple(getattr(p, "shape", ()) for p in params) + shapes = tuple(x() if callable(x) else x for x in shapes) try: expected_batch_shape, expected_event_shape = jax_dist.infer_shapes(*shapes) except NotImplementedError: @@ -519,6 +548,40 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit): assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5) +@pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS) +def test_cdf_and_icdf(jax_dist, sp_dist, params): + d = jax_dist(*params) + if d.event_dim > 0: + pytest.skip('skip testing cdf/icdf methods of multivariate distributions') + samples = d.sample(key=random.PRNGKey(0), sample_shape=(100,)) + quantiles = random.uniform(random.PRNGKey(1), (100,) + d.shape()) + try: + if d.shape() == (): + rtol = 1e-3 if jax_dist is dist.StudentT else 1e-5 + assert_allclose(jax.vmap(jax.grad(d.cdf))(samples), + jnp.exp(d.log_prob(samples)), atol=1e-5, rtol=rtol) + assert_allclose(jax.vmap(jax.grad(d.icdf))(quantiles), + jnp.exp(-d.log_prob(d.icdf(quantiles))), atol=1e-5, rtol=rtol) + assert_allclose(d.cdf(d.icdf(quantiles)), quantiles, atol=1e-5, rtol=1e-5) + assert_allclose(d.icdf(d.cdf(samples)), samples, atol=1e-5, rtol=1e-5) + except NotImplementedError: + pass + + # test against scipy + if not sp_dist: + pytest.skip('no corresponding scipy distn.') + sp_dist = sp_dist(*params) + try: + actual_cdf = d.cdf(samples) + expected_cdf = sp_dist.cdf(samples) + assert_allclose(actual_cdf, expected_cdf, atol=1e-5, rtol=1e-5) + actual_icdf = d.icdf(quantiles) + expected_icdf = sp_dist.ppf(quantiles) + assert_allclose(actual_icdf, expected_icdf, atol=1e-5, rtol=1e-4) + except NotImplementedError: + pass + + @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS) def test_gof(jax_dist, sp_dist, params): if "Improper" in jax_dist.__name__: @@ -713,6 +776,8 @@ def fn(*args): eps = 1e-3 for i in range(len(params)): + if isinstance(params[i], dist.Distribution): # skip taking grad w.r.t. base_dist + continue if params[i] is None or jnp.result_type(params[i]) in (jnp.int32, jnp.int64): continue actual_grad = jax.grad(fn, i)(*params) @@ -734,6 +799,8 @@ def fn(*args): def test_mean_var(jax_dist, sp_dist, params): if jax_dist is _ImproperWrapper: pytest.skip("Improper distribution does not has mean/var implemented") + if jax_dist in (_TruncatedNormal, dist.continuous.TwoSidedTruncatedDistribution): + pytest.skip("Truncated distributions do not has mean/var implemented") n = 20000 if jax_dist in [dist.LKJ, dist.LKJCholesky] else 200000 d_jax = jax_dist(*params) @@ -808,6 +875,8 @@ def test_mean_var(jax_dist, sp_dist, params): (2, 3), ]) def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape): + if jax_dist is _TruncatedNormal: + pytest.skip("_TruncatedNormal is a function, not a class") dist_args = [p for p in inspect.getfullargspec(jax_dist.__init__)[0][1:]] valid_params, oob_params = list(params), list(params) @@ -816,6 +885,8 @@ def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape): for i in range(len(params)): if jax_dist in (_ImproperWrapper, dist.LKJ, dist.LKJCholesky) and dist_args[i] != "concentration": continue + if jax_dist is dist.continuous.TwoSidedTruncatedDistribution and dist_args[i] == "base_dist": + continue if jax_dist is dist.GaussianRandomWalk and dist_args[i] == "num_steps": continue if params[i] is None: