diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py index f1e03cde60..57504b342d 100644 --- a/pymc3/distributions/continuous.py +++ b/pymc3/distributions/continuous.py @@ -20,8 +20,11 @@ from . import transforms from pymc3.util import get_variable_name from .special import log_i0 -from ..math import invlogit, logit -from .dist_math import bound, logpow, gammaln, betaln, std_cdf, alltrue_elemwise, SplineWrapper, i0e +from ..math import invlogit, logit, logdiffexp +from .dist_math import ( + bound, logpow, gammaln, betaln, std_cdf, alltrue_elemwise, + SplineWrapper, i0e, normal_lcdf, normal_lccdf +) from .distribution import Continuous, draw_values, generate_samples __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta', @@ -51,10 +54,21 @@ def __init__(self, transform=transforms.logodds, *args, **kwargs): class BoundedContinuous(Continuous): """Base class for bounded continuous distributions""" - def __init__(self, transform='interval', *args, **kwargs): + def __init__(self, transform='auto', lower=None, upper=None, + *args, **kwargs): + + lower = tt.as_tensor_variable(lower) if lower is not None else None + upper = tt.as_tensor_variable(upper) if upper is not None else None - if transform == 'interval': - transform = transforms.interval(self.lower, self.upper) + if transform == 'auto': + if lower is None and upper is None: + transform = None + elif lower is not None and upper is None: + transform = transforms.lowerbound(lower) + elif lower is None and upper is not None: + transform = transforms.upperbound(upper) + else: + transform = transforms.interval(lower, upper) super(BoundedContinuous, self).__init__( transform=transform, *args, **kwargs) @@ -173,7 +187,8 @@ def __init__(self, lower=0, upper=1, *args, **kwargs): self.mean = (upper + lower) / 2. self.median = self.mean - super(Uniform, self).__init__(*args, **kwargs) + super(Uniform, self).__init__( + lower=lower, upper=upper, *args, **kwargs) def random(self, point=None, size=None): """ @@ -446,7 +461,7 @@ def _repr_latex_(self, name=None, dist=None): get_variable_name(sd)) -class TruncatedNormal(Continuous): +class TruncatedNormal(BoundedContinuous): R""" Univariate truncated normal log-likelihood. @@ -521,34 +536,29 @@ class TruncatedNormal(Continuous): """ def __init__(self, mu=0, sd=None, tau=None, lower=None, upper=None, - transform='infer', *args, **kwargs): + transform='auto', *args, **kwargs): tau, sd = get_tau_sd(tau=tau, sd=sd) self.sd = tt.as_tensor_variable(sd) self.tau = tt.as_tensor_variable(tau) self.lower = tt.as_tensor_variable(lower) if lower is not None else lower self.upper = tt.as_tensor_variable(upper) if upper is not None else upper - self.mu = tt.as_tensor_variable(mu) + self.mu = tt.as_tensor_variable(mu) - # Calculate mean - pdf_a, pdf_b, cdf_a, cdf_b = self._get_boundary_parameters() - z = cdf_b - cdf_a - self.mean = self.mu + (pdf_a+pdf_b) / z * self.sd + if self.lower is None and self.upper is None: + self._defaultval = mu + elif self.lower is None and self.upper is not None: + self._defaultval = self.upper - 1. + elif self.lower is not None and self.upper is None: + self._defaultval = self.lower + 1. + else: + self._defaultval = (self.lower + self.upper) / 2 assert_negative_support(sd, 'sd', 'TruncatedNormal') assert_negative_support(tau, 'tau', 'TruncatedNormal') - if transform == 'infer': - if lower is None and upper is None: - transform = None - elif lower is not None and upper is not None: - transform = transforms.interval(lower, upper) - elif upper is not None: - transform = transforms.upperbound(upper) - else: - transform = transforms.lowerbound(lower) - super(TruncatedNormal, self).__init__( - transform=transform, *args, **kwargs) + defaults=('_defaultval',), transform=transform, + lower=lower, upper=upper, *args, **kwargs) def random(self, point=None, size=None): """ @@ -592,89 +602,57 @@ def logp(self, value): ------- TensorVariable """ - sd = self.sd - tau = self.tau mu = self.mu - a = self.lower - b = self.upper - - # In case either a or b are not specified, normalization terms simplify to 1.0 and 0.0 - # https://en.wikipedia.org/wiki/Truncated_normal_distribution - norm_left, norm_right = 1.0, 0.0 - - # Define normalization - if b is not None: - norm_left = self._cdf((b - mu) / sd) - - if a is not None: - norm_right = self._cdf((a - mu) / sd) - - f = self._pdf((value - mu) / sd) / sd / (norm_left - norm_right) + sd = self.sd - return bound(tt.log(f), value >= a, value <= b, sd > 0) + norm = self._normalization() + logp = Normal.dist(mu=mu, sd=sd).logp(value) - norm - def _cdf(self, value): - """ - Calculate cdf of standard normal distribution + bounds = [sd > 0] + if self.lower is not None: + bounds.append(value >= self.lower) + if self.upper is not None: + bounds.append(value <= self.upper) + return bound(logp, *bounds) - Parameters - ---------- - value : numeric - Value(s) for which log-probability is calculated. If the log probabilities for multiple - values are desired the values must be provided in a numpy array or theano tensor + def _normalization(self): + mu, sd = self.mu, self.sd - Returns - ------- - TensorVariable - """ - return 0.5 * (1.0 + tt.erf(value / tt.sqrt(2))) + if self.lower is None and self.upper is None: + return 0. - def _pdf(self, value): - """ - Calculate pdf of standard normal distribution + if self.lower is not None and self.upper is not None: + lcdf_a = normal_lcdf(mu, sd, self.lower) + lcdf_b = normal_lcdf(mu, sd, self.upper) + lsf_a = normal_lccdf(mu, sd, self.lower) + lsf_b = normal_lccdf(mu, sd, self.upper) - Parameters - ---------- - value : numeric - Value(s) for which log-probability is calculated. If the log probabilities for multiple - values are desired the values must be provided in a numpy array or theano tensor + return tt.switch( + self.lower > 0, + logdiffexp(lsf_a, lsf_b), + logdiffexp(lcdf_b, lcdf_a), + ) - Returns - ------- - TensorVariable - """ - return 1.0 / tt.sqrt(2 * np.pi) * tt.exp(-0.5 * (value ** 2)) + if self.lower is not None: + return normal_lccdf(mu, sd, self.lower) + else: + return normal_lcdf(mu, sd, self.upper) def _repr_latex_(self, name=None, dist=None): if dist is None: dist = self - sd = dist.sd - mu = dist.mu - a = dist.a - b = dist.b name = r'\text{%s}' % name - return r'${} \sim \text{{TruncatedNormal}}(\mathit{{mu}}={},~\mathit{{sd}}={},a={},b={})$'.format(name, - get_variable_name(mu), - get_variable_name(sd), - get_variable_name(a), - get_variable_name(b)) - - def _get_boundary_parameters(self): - """ - Calcualte values of cdf and pdf at boundary points a and b - - Returns - ------- - pdf(a), pdf(b), cdf(a), cdf(b) if a,b defined, otherwise 0.0, 0.0, 0.0, 1.0 - """ - # pdf = 0 at +-inf - pdf_a = self._pdf(self.lower) if not self.lower is None else 0.0 - pdf_b = self._pdf(self.upper) if not self.upper is None else 0.0 - - # b-> inf, cdf(b) = 1.0, a->-inf, cdf(a) = 0 - cdf_a = self._cdf(self.lower) if not self.lower is None else 0.0 - cdf_b = self._cdf(self.upper) if not self.upper is None else 1.0 - return pdf_a, pdf_b, cdf_a, cdf_b + return ( + r'${} \sim \text{{TruncatedNormal}}(' + '\mathit{{mu}}={},~\mathit{{sd}}={},a={},b={})$' + .format( + name, + get_variable_name(self.mu), + get_variable_name(self.sd), + get_variable_name(self.lower), + get_variable_name(self.upper), + ) + ) class HalfNormal(PositiveContinuous): @@ -3054,7 +3032,8 @@ def __init__(self, lower=0, upper=1, c=0.5, self.lower = lower = tt.as_tensor_variable(lower) self.upper = upper = tt.as_tensor_variable(upper) - super(Triangular, self).__init__(*args, **kwargs) + super(Triangular, self).__init__(lower=lower, upper=upper, + *args, **kwargs) def random(self, point=None, size=None): """ @@ -3553,7 +3532,8 @@ def __init__(self, x_points, pdf_points, *args, **kwargs): self.lower = lower = tt.as_tensor_variable(x_points[0]) self.upper = upper = tt.as_tensor_variable(x_points[-1]) - super(Interpolated, self).__init__(*args, **kwargs) + super(Interpolated, self).__init__(lower=lower, upper=upper, + *args, **kwargs) interp = InterpolatedUnivariateSpline( x_points, pdf_points, k=1, ext='zeros') diff --git a/pymc3/distributions/dist_math.py b/pymc3/distributions/dist_math.py index dc0684b371..7671f5b7fa 100644 --- a/pymc3/distributions/dist_math.py +++ b/pymc3/distributions/dist_math.py @@ -87,6 +87,25 @@ def std_cdf(x): return .5 + .5 * tt.erf(x / tt.sqrt(2.)) +def normal_lcdf(mu, sigma, x): + """Compute the log of the cumulative density function of the normal.""" + z = (x - mu) / sigma + return tt.switch( + tt.lt(z, -1.0), + tt.log(tt.erfcx(-z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2., + tt.log1p(-tt.erfc(z / tt.sqrt(2.)) / 2.) + ) + + +def normal_lccdf(mu, sigma, x): + z = (x - mu) / sigma + return tt.switch( + tt.gt(z, 1.0), + tt.log(tt.erfcx(z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2., + tt.log1p(-tt.erfc(-z / tt.sqrt(2.)) / 2.) + ) + + def sd2rho(sd): """ `sd -> rho` theano converter diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py index ca6db26bd5..c5b8040ecb 100644 --- a/pymc3/distributions/distribution.py +++ b/pymc3/distributions/distribution.py @@ -6,7 +6,10 @@ from theano import function import theano from ..memoize import memoize -from ..model import Model, get_named_nodes_and_relations, FreeRV, ObservedRV, MultiObservedRV +from ..model import ( + Model, get_named_nodes_and_relations, FreeRV, + ObservedRV, MultiObservedRV +) from ..vartypes import string_types __all__ = ['DensityDist', 'Distribution', 'Continuous', 'Discrete', diff --git a/pymc3/examples/censored_data.py b/pymc3/examples/censored_data.py index 51c9a0da2e..534b67d6b3 100644 --- a/pymc3/examples/censored_data.py +++ b/pymc3/examples/censored_data.py @@ -1,7 +1,3 @@ -import numpy as np -import matplotlib.pyplot as plt -import pymc3 as pm -import theano.tensor as tt ''' Data can be left, right, or interval censored. @@ -31,26 +27,11 @@ To establish a baseline they compare to an uncensored model of uncensored data. ''' +import numpy as np +import matplotlib.pyplot as plt +import pymc3 as pm - -# Helper functions -def normal_lcdf(mu, sigma, x): - z = (x - mu) / sigma - return tt.switch( - tt.lt(z, -1.0), - tt.log(tt.erfcx(-z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2., - tt.log1p(-tt.erfc(z / tt.sqrt(2.)) / 2.) - ) - - -def normal_lccdf(mu, sigma, x): - z = (x - mu) / sigma - return tt.switch( - tt.gt(z, 1.0), - tt.log(tt.erfcx(z / tt.sqrt(2.)) / 2.) - tt.sqr(z) / 2., - tt.log1p(-tt.erfc(-z / tt.sqrt(2.)) / 2.) - ) - +from pymc3.distributions.dist_math import normal_lcdf, normal_lccdf # Produce normally distributed samples np.random.seed(123) @@ -103,6 +84,7 @@ def censored_right_likelihood(mu, sigma, n_right_censored, upper_bound): def censored_left_likelihood(mu, sigma, n_left_censored, lower_bound): return n_left_censored * normal_lcdf(mu, sigma, lower_bound) + with pm.Model() as unimputed_censored_model: mu = pm.Normal('mu', mu=0., sd=(high - low) / 2.) sigma = pm.HalfNormal('sigma', sd=(high - low) / 2.) @@ -115,7 +97,7 @@ def censored_left_likelihood(mu, sigma, n_left_censored, lower_bound): right_censored = pm.Potential( 'right_censored', censored_right_likelihood(mu, sigma, n_right_censored, high) - ) + ) left_censored = pm.Potential( 'left_censored', censored_left_likelihood(mu, sigma, n_left_censored, low) @@ -144,5 +126,6 @@ def run(n=1500): pm.plot_posterior(trace[-1000:], varnames=['mu', 'sigma']) plt.show() + if __name__ == '__main__': run() diff --git a/pymc3/math.py b/pymc3/math.py index efac2e3a51..0ea46f0036 100644 --- a/pymc3/math.py +++ b/pymc3/math.py @@ -114,8 +114,13 @@ def logsumexp(x, axis=None): def logaddexp(a, b): diff = b - a return tt.switch(diff > 0, - b + tt.log1p(tt.exp(-diff)), - a + tt.log1p(tt.exp(diff))) + b + tt.log1p(tt.exp(-diff)), + a + tt.log1p(tt.exp(diff))) + + +def logdiffexp(a, b): + """log(exp(a) - exp(b))""" + return a + log1mexp(a - b) def invlogit(x, eps=sys.float_info.epsilon): diff --git a/pymc3/tests/test_distributions.py b/pymc3/tests/test_distributions.py index 264753304d..f778f10e93 100644 --- a/pymc3/tests/test_distributions.py +++ b/pymc3/tests/test_distributions.py @@ -7,17 +7,19 @@ from ..vartypes import continuous_types from ..model import Model, Point, Potential, Deterministic from ..blocking import DictToVarBijection, DictToArrayBijection, ArrayOrdering -from ..distributions import (DensityDist, Categorical, Multinomial, VonMises, Dirichlet, - MvStudentT, MvNormal, MatrixNormal, ZeroInflatedPoisson, - ZeroInflatedNegativeBinomial, Constant, Poisson, Bernoulli, Beta, - BetaBinomial, HalfStudentT, StudentT, Weibull, Pareto, - InverseGamma, Gamma, Cauchy, HalfCauchy, Lognormal, Laplace, - NegativeBinomial, Geometric, Exponential, ExGaussian, Normal, TruncatedNormal, - Flat, LKJCorr, Wald, ChiSquared, HalfNormal, DiscreteUniform, - Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull, - Gumbel, Logistic, OrderedLogistic, LogitNormal, Interpolated, - ZeroInflatedBinomial, HalfFlat, AR1, KroneckerNormal, Rice, - Kumaraswamy) +from ..distributions import ( + DensityDist, Categorical, Multinomial, VonMises, Dirichlet, + MvStudentT, MvNormal, MatrixNormal, ZeroInflatedPoisson, + ZeroInflatedNegativeBinomial, Constant, Poisson, Bernoulli, Beta, + BetaBinomial, HalfStudentT, StudentT, Weibull, Pareto, + InverseGamma, Gamma, Cauchy, HalfCauchy, Lognormal, Laplace, + NegativeBinomial, Geometric, Exponential, ExGaussian, Normal, TruncatedNormal, + Flat, LKJCorr, Wald, ChiSquared, HalfNormal, DiscreteUniform, + Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull, + Gumbel, Logistic, OrderedLogistic, LogitNormal, Interpolated, + ZeroInflatedBinomial, HalfFlat, AR1, KroneckerNormal, Rice, + Kumaraswamy +) from ..distributions import continuous from pymc3.theanof import floatX @@ -535,15 +537,21 @@ def test_normal(self): ) def test_truncated_normal(self): - # Rplusbig domain is specified for eveything, to avoid silly cases such as - # {'mu': array(-2.1), 'a': array(-100.), 'b': array(0.01), 'value': array(0.), 'sd': array(0.01)} - # TruncatedNormal: pdf = 0.0, logpdf = -inf - # Scipy's answer: pdf = 0.0, logpdf = -22048.413!!! - self.pymc3_matches_scipy(TruncatedNormal, R, {'mu': R, 'sd': Rplusbig, 'lower':-Rplusbig, 'upper':Rplusbig}, - lambda value, mu, sd, lower, upper: sp.truncnorm.logpdf( - value, (lower-mu)/sd, (upper-mu)/sd, loc=mu, scale=sd), - decimal=select_by_precision(float64=6, float32=1) - ) + def scipy_logp(value, mu, sd, lower, upper): + return sp.truncnorm.logpdf( + value, (lower-mu)/sd, (upper-mu)/sd, loc=mu, scale=sd) + + args = {'mu': array(-2.1), 'lower': array(-100.), 'upper': array(0.01), + 'sd': array(0.01)} + val = TruncatedNormal.dist(**args).logp(0.) + assert_allclose(val.eval(), scipy_logp(value=0, **args)) + + self.pymc3_matches_scipy( + TruncatedNormal, R, + {'mu': R, 'sd': Rplusbig, 'lower': -Rplusbig, 'upper': Rplusbig}, + scipy_logp, + decimal=select_by_precision(float64=6, float32=1) + ) def test_half_normal(self): self.pymc3_matches_scipy(HalfNormal, Rplus, {'sd': Rplus},