Improve truncnormal test and fix example

Author: aseyboldt

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()

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},