Merge pull request #2163 from a-rodin/interpolated

Add Interpolated distribution class

Author: Junpeng Lao

docs/source/notebooks/updating_priors.ipynb

@@ -25,6 +25,7 @@
 from .continuous import SkewNormal
 from .continuous import Triangular
 from .continuous import Gumbel
+from .continuous import Interpolated
 
 from .discrete import Binomial
 from .discrete import BetaBinomial
@@ -132,5 +133,6 @@
            'NormalMixture',
            'Triangular',
            'DiscreteWeibull',
-           'Gumbel'
+           'Gumbel',
+           'Interpolated'
            ]

pymc3/distributions/__init__.py

@@ -10,19 +10,20 @@
 import numpy as np
 import theano.tensor as tt
 from scipy import stats
+from scipy.interpolate import InterpolatedUnivariateSpline
 import warnings
 
 from pymc3.theanof import floatX
 from . import transforms
 
-from .dist_math import bound, logpow, gammaln, betaln, std_cdf, i0, i1, alltrue_elemwise
+from .dist_math import bound, logpow, gammaln, betaln, std_cdf, i0, i1, alltrue_elemwise, DifferentiableSplineWrapper
 from .distribution import Continuous, draw_values, generate_samples, Bound
 
 __all__ = ['Uniform', 'Flat', 'Normal', 'Beta', 'Exponential', 'Laplace',
            'StudentT', 'Cauchy', 'HalfCauchy', 'Gamma', 'Weibull',
            'HalfStudentT', 'StudentTpos', 'Lognormal', 'ChiSquared',
            'HalfNormal', 'Wald', 'Pareto', 'InverseGamma', 'ExGaussian',
-           'VonMises', 'SkewNormal']
+           'VonMises', 'SkewNormal', 'Interpolated']
 
 
 class PositiveContinuous(Continuous):
@@ -1389,3 +1390,71 @@ def random(self, point=None, size=None, repeat=None):
     def logp(self, value):
         scaled = (value - self.mu) / self.beta
         return bound(-scaled - tt.exp(-scaled) - tt.log(self.beta), self.beta > 0)
+
+class Interpolated(Continuous):
+    R"""
+    Probability distribution defined as a linear interpolation of
+    of a set of points and values of probability density function
+    evaluated on them.
+
+    The points are not variables, but plain array-like objects, so
+    they are constant and cannot be sampled.
+
+    ========  =========================================
+    Support   :math:`x \in [x_points[0], x_points[-1]]`
+    ========  =========================================
+
+    Parameters
+    ----------
+    x_points : array-like
+        A monotonically growing list of values
+    pdf_points : array-like
+        Probability density function evaluated at points from `x`
+    """
+
+    def __init__(self, x_points, pdf_points, transform='interval',
+                 *args, **kwargs):
+        if transform == 'interval':
+            transform = transforms.interval(x_points[0], x_points[-1])
+        super(Interpolated, self).__init__(transform=transform,
+                                           *args, **kwargs)
+
+        interp = InterpolatedUnivariateSpline(x_points, pdf_points, k=1, ext='zeros')
+        Z = interp.integral(x_points[0], x_points[-1])
+
+        self.Z = tt.as_tensor_variable(Z)
+        self.interp_op = DifferentiableSplineWrapper(interp)
+        self.x_points = x_points
+        self.pdf_points = pdf_points / Z
+        self.cdf_points = interp.antiderivative()(x_points) / Z
+
+        self.median = self._argcdf(0.5)
+
+    def _argcdf(self, p):
+        pdf = self.pdf_points
+        cdf = self.cdf_points
+        x = self.x_points
+
+        index = np.searchsorted(cdf, p) - 1
+        slope = (pdf[index + 1] - pdf[index]) / (x[index + 1] - x[index])
+
+        return x[index] + np.where(
+            np.abs(slope) <= 1e-8,
+            np.where(
+                np.abs(pdf[index]) <= 1e-8,
+                np.zeros(index.shape),
+                (p - cdf[index]) / pdf[index]
+            ),
+            (-pdf[index] + np.sqrt(pdf[index] ** 2 + 2 * slope * (p - cdf[index]))) / slope
+        )
+
+    def _random(self, size=None):
+        return self._argcdf(np.random.uniform(size=size))
+
+    def random(self, point=None, size=None, repeat=None):
+        return generate_samples(self._random,
+                                dist_shape=self.shape,
+                                size=size)
+
+    def logp(self, value):
+        return tt.log(self.interp_op(value) / self.Z)

pymc3/distributions/continuous.py

@@ -364,3 +364,35 @@ def conjugate_solve_triangular(outer, inner):
         else:
             grad = tt.triu(s + s.T) - tt.diag(tt.diagonal(s))
         return [tt.switch(ok, grad, floatX(np.nan))]
+
+class SplineWrapper (theano.Op):
+    """
+    Creates a theano operation from scipy.interpolate.UnivariateSpline
+    """
+
+    __props__ = ('spline',)
+    itypes = [tt.dscalar]
+    otypes = [tt.dscalar]
+
+    def __init__(self, spline):
+        self.spline = spline
+
+    def perform(self, node, inputs, output_storage):
+        x, = inputs
+        output_storage[0][0] = np.asarray(self.spline(x))
+
+class DifferentiableSplineWrapper (SplineWrapper):
+    """
+    Creates a theano operation with defined gradient from
+    scipy.interpolate.UnivariateSpline
+    """
+
+    def __init__(self, spline):
+        super(DifferentiableSplineWrapper, self).__init__(spline)
+        self.spline_grad = SplineWrapper(spline.derivative())
+        self.__props__ += ('spline_grad',)
+
+    def grad(self, inputs, grads):
+        x, = inputs
+        x_grad, = grads
+        return [x_grad * self.spline_grad(x)]

pymc3/distributions/dist_math.py

@@ -13,7 +13,8 @@
                              InverseGamma, Gamma, Cauchy, HalfCauchy, Lognormal, Laplace,
                              NegativeBinomial, Geometric, Exponential, ExGaussian, Normal,
                              Flat, LKJCorr, Wald, ChiSquared, HalfNormal, DiscreteUniform,
-                             Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull, Gumbel)
+                             Bound, Uniform, Triangular, Binomial, SkewNormal, DiscreteWeibull, Gumbel,
+                             Interpolated)
 from ..distributions import continuous
 from pymc3.theanof import floatX
 from numpy import array, inf, log, exp
@@ -791,3 +792,30 @@ def test_gumbel(self):
     def test_multidimensional_beta_construction(self):
         with Model():
             Beta('beta', alpha=1., beta=1., shape=(10, 20))
+
+    def test_interpolated(self):
+        for mu in R.vals:
+            for sd in Rplus.vals:
+                #pylint: disable=cell-var-from-loop
+                xmin = mu - 5 * sd
+                xmax = mu + 5 * sd
+
+                class TestedInterpolated (Interpolated):
+
+                    def __init__(self, **kwargs):
+                        x_points = np.linspace(xmin, xmax, 100000)
+                        pdf_points = sp.norm.pdf(x_points, loc=mu, scale=sd)
+                        super(TestedInterpolated, self).__init__(
+                            x_points=x_points,
+                            pdf_points=pdf_points,
+                            **kwargs
+                        )
+
+                def ref_pdf(value):
+                    return np.where(
+                        np.logical_and(value >= xmin, value <= xmax),
+                        sp.norm.logpdf(value, mu, sd),
+                        -np.inf * np.ones(value.shape)
+                    )
+
+                self.pymc3_matches_scipy(TestedInterpolated, R, {}, ref_pdf)

pymc3/tests/test_distributions.py

@@ -577,6 +577,26 @@ 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_interpolated(self):
+        for mu in R.vals:
+            for sd in Rplus.vals:
+                #pylint: disable=cell-var-from-loop
+                def ref_rand(size):
+                    return st.norm.rvs(loc=mu, scale=sd, size=size)
+
+                class TestedInterpolated (pm.Interpolated):
+
+                    def __init__(self, **kwargs):
+                        x_points = np.linspace(mu - 5 * sd, mu + 5 * sd, 100)
+                        pdf_points = st.norm.pdf(x_points, loc=mu, scale=sd)
+                        super(TestedInterpolated, self).__init__(
+                            x_points=x_points,
+                            pdf_points=pdf_points,
+                            **kwargs
+                        )
+
+                pymc3_random(TestedInterpolated, {}, ref_rand=ref_rand)
+
     @pytest.mark.skip('Wishart random sampling not implemented.\n'
                       'See https://github.com/pymc-devs/pymc3/issues/538')
     def test_wishart(self):