Add LogitNormal distribution and tests (#2877)

* Add LogitNormal disribution and tests

Co-authored-by: denadai2 <[email protected]>
Co-authored-by: aloctavodia <[email protected]>

* reduce precision for float32

Author: aloctavodia

RELEASE-NOTES.md

@@ -14,6 +14,7 @@
 - Add new 'pairplot' function, for plotting scatter or hexbin matrices of sampled parameters.
   Optionally it can plot divergences.
 - Plots of discrete distributions in the docstrings
+- Add logitnormal distribution
 
 ### Fixes
 

docs/source/api/distributions/continuous.rst

@@ -30,6 +30,7 @@ Continuous
    Triangular
    Gumbel
    Logistic
+   LogitNormal
    Interpolated
 
 .. automodule:: pymc3.distributions.continuous

pymc3/distributions/__init__.py

@@ -26,6 +26,7 @@
 from .continuous import Triangular
 from .continuous import Gumbel
 from .continuous import Logistic
+from .continuous import LogitNormal
 from .continuous import Interpolated
 
 from .discrete import Binomial
@@ -145,6 +146,7 @@
            'DiscreteWeibull',
            'Gumbel',
            'Logistic',
+           'LogitNormal',
            'Interpolated',
            'Bound',
            ]

pymc3/distributions/continuous.py

@@ -10,21 +10,23 @@
 import numpy as np
 import theano.tensor as tt
 from scipy import stats
+from scipy.special import expit
 from scipy.interpolate import InterpolatedUnivariateSpline
 import warnings
 
 from pymc3.theanof import floatX
 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
 from .distribution import Continuous, draw_values, generate_samples
 
 __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'Beta', 'Exponential',
            'Laplace', 'StudentT', 'Cauchy', 'HalfCauchy', 'Gamma', 'Weibull',
            'HalfStudentT', 'Lognormal', 'ChiSquared', 'HalfNormal', 'Wald',
            'Pareto', 'InverseGamma', 'ExGaussian', 'VonMises', 'SkewNormal',
-           'Triangular', 'Gumbel', 'Logistic', 'Interpolated']
+           'Triangular', 'Gumbel', 'Logistic', 'LogitNormal', 'Interpolated']
 
 
 class PositiveContinuous(Continuous):
@@ -2262,6 +2264,86 @@ def _repr_latex_(self, name=None, dist=None):
                                                                 get_variable_name(s))
 
 
+class LogitNormal(UnitContinuous):
+    R"""
+    Logit-Normal log-likelihood.
+
+    The pdf of this distribution is
+
+    .. math::
+       f(x \mid \mu, \tau) =
+           \frac{1}{x(1-x)} \sqrt{\frac{\tau}{2\pi}}
+           \exp\left\{ -\frac{\tau}{2} (logit(x)-\mu)^2 \right\}
+
+
+    .. plot::
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import scipy.stats as st
+        from scipy.special import logit
+        plt.style.use('seaborn-darkgrid')
+        x = np.linspace(0.0001, 0.9999, 500)
+        mus = [0., 0., 0., 1.]
+        sds = [0.3, 1., 2., 1.]
+        for mu, sd in  zip(mus, sds):
+            pdf = st.norm.pdf(logit(x), loc=mu, scale=sd) * 1/(x * (1-x))
+            plt.plot(x, pdf, label=r'$\mu$ = {}, $\sigma$ = {}'.format(mu, sd))
+            plt.legend(loc=1)
+        plt.show()
+
+    ========  ==========================================
+    Support   :math:`x \in (0, 1)`
+    Mean      no analytical solution
+    Variance  no analytical solution
+    ========  ==========================================
+
+    Parameters
+    ----------
+    mu : float
+        Location parameter.
+    sd : float
+        Scale parameter (sd > 0).
+    tau : float
+        Scale parameter (tau > 0).
+    """
+
+    def __init__(self, mu=0, sd=None, tau=None, **kwargs):
+        self.mu = mu = tt.as_tensor_variable(mu)
+        tau, sd = get_tau_sd(tau=tau, sd=sd)
+        self.sd = tt.as_tensor_variable(sd)
+        self.tau = tau = tt.as_tensor_variable(tau)
+
+        self.median = invlogit(mu)
+        assert_negative_support(sd, 'sd', 'LogitNormal')
+        assert_negative_support(tau, 'tau', 'LogitNormal')
+
+        super(LogitNormal, self).__init__(**kwargs)
+
+    def random(self, point=None, size=None, repeat=None):
+        mu, _, sd = draw_values([self.mu, self.tau, self.sd], point=point)
+        return expit(generate_samples(stats.norm.rvs, loc=mu, scale=sd, dist_shape=self.shape,
+                                      size=size))
+
+    def logp(self, value):
+        sd = self.sd
+        mu = self.mu
+        tau = self.tau
+        return bound(-0.5 * tau * (logit(value) - mu) ** 2
+                     + 0.5 * tt.log(tau / (2. * np.pi))
+                     - tt.log(value * (1 - value)), value > 0, value < 1, tau > 0)
+
+    def _repr_latex_(self, name=None, dist=None):
+        if dist is None:
+            dist = self
+        sd = dist.sd
+        mu = dist.mu
+        name = r'\text{%s}' % name
+        return r'${} \sim \text{{LogitNormal}}(\mathit{{mu}}={},~\mathit{{sd}}={})$'.format(name,
+                                                                get_variable_name(mu),
+                                                                get_variable_name(sd))
+
+
 class Interpolated(Continuous):
     R"""
     Univariate probability distribution defined as a linear interpolation

pymc3/tests/test_distributions.py

@@ -14,8 +14,8 @@
                              NegativeBinomial, Geometric, Exponential, ExGaussian, Normal,
                              Flat, LKJCorr, Wald, ChiSquared, HalfNormal, DiscreteUniform,
                              Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull,
-                             Gumbel, Logistic, Interpolated, ZeroInflatedBinomial, HalfFlat, AR1,
-                             KroneckerNormal)
+                             Gumbel, Logistic, LogitNormal, Interpolated, ZeroInflatedBinomial,
+                             HalfFlat, AR1, KroneckerNormal)
 from ..distributions import continuous
 from pymc3.theanof import floatX
 from numpy import array, inf, log, exp
@@ -27,6 +27,7 @@
 from scipy import integrate
 import scipy.stats.distributions as sp
 import scipy.stats
+from scipy.special import logit
 import theano
 import theano.tensor as tt
 from ..math import kronecker
@@ -1035,6 +1036,12 @@ def test_logistic(self):
                                  lambda value, mu, s: sp.logistic.logpdf(value, mu, s),
                                  decimal=select_by_precision(float64=6, float32=1))
 
+    def test_logitnormal(self):
+        self.pymc3_matches_scipy(LogitNormal, Unit, {'mu': R, 'sd': Rplus},
+                                 lambda value, mu, sd: (sp.norm.logpdf(logit(value), mu, sd)
+                                                        - (np.log(value) + np.log1p(-value))),
+                                 decimal=select_by_precision(float64=6, float32=1))
+
     def test_multidimensional_beta_construction(self):
         with Model():
             Beta('beta', alpha=1., beta=1., shape=(10, 20))

pymc3/tests/test_distributions_random.py

@@ -4,6 +4,7 @@
 import numpy as np
 import numpy.testing as npt
 import scipy.stats as st
+from scipy.special import expit
 from scipy import linalg
 import numpy.random as nr
 import theano
@@ -299,6 +300,16 @@ class TestGumbel(BaseTestCases.BaseTestCase):
     params = {'mu': 0., 'beta': 1.}
 
 
+class TestLogistic(BaseTestCases.BaseTestCase):
+    distribution = pm.Logistic
+    params = {'mu': 0., 's': 1.}
+
+
+class TestLogitNormal(BaseTestCases.BaseTestCase):
+    distribution = pm.LogitNormal
+    params = {'mu': 0., 'sd': 1.}
+
+
 class TestBinomial(BaseTestCases.BaseTestCase):
     distribution = pm.Binomial
     params = {'n': 5, 'p': 0.5}
@@ -668,6 +679,16 @@ def ref_rand(size, mu, beta):
             return st.gumbel_r.rvs(loc=mu, scale=beta, size=size)
         pymc3_random(pm.Gumbel, {'mu': R, 'beta': Rplus}, ref_rand=ref_rand)
 
+    def test_logistic(self):
+        def ref_rand(size, mu, s):
+            return st.logistic.rvs(loc=mu, scale=s, size=size)
+        pymc3_random(pm.Logistic, {'mu': R, 's': Rplus}, ref_rand=ref_rand)
+
+    def test_logitnormal(self):
+        def ref_rand(size, mu, sd):
+            return expit(st.norm.rvs(loc=mu, scale=sd, size=size))
+        pymc3_random(pm.LogitNormal, {'mu': R, 'sd': Rplus}, ref_rand=ref_rand)
+
     @pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
     def test_interpolated(self):
         for mu in R.vals: