Raise NotImplementedError for SplineWrapper gradient operation (#2211)

* Make spacing more compliant to PEP8

* Raise NotImplementedError from SplineWrapper grad operation

Fixes #2209

* Instantiate SplineWrapper recursively to simplify the architecture

* Create spline derivatives lazily

* Move grad_op to a separate property

* Add tests for SplineWrapper

* Fix style issues

Author: a-rodin

pymc3/distributions/continuous.py

@@ -17,7 +17,10 @@
 from . import transforms
 from pymc3.util import get_variable_name
 
-from .dist_math import bound, logpow, gammaln, betaln, std_cdf, i0, i1, alltrue_elemwise, DifferentiableSplineWrapper
+from .dist_math import (
+    bound, logpow, gammaln, betaln, std_cdf, i0,
+    i1, alltrue_elemwise, SplineWrapper
+)
 from .distribution import Continuous, draw_values, generate_samples, Bound
 
 __all__ = ['Uniform', 'Flat', 'Normal', 'Beta', 'Exponential', 'Laplace',
@@ -1625,7 +1628,7 @@ def __init__(self, x_points, pdf_points, transform='interval',
         Z = interp.integral(x_points[0], x_points[-1])
 
         self.Z = tt.as_tensor_variable(Z)
-        self.interp_op = DifferentiableSplineWrapper(interp)
+        self.interp_op = SplineWrapper(interp)
         self.x_points = x_points
         self.pdf_points = pdf_points / Z
         self.cdf_points = interp.antiderivative()(x_points) / Z

pymc3/distributions/dist_math.py

@@ -369,7 +369,8 @@ def conjugate_solve_triangular(outer, inner):
             grad = tt.triu(s + s.T) - tt.diag(tt.diagonal(s))
         return [tt.switch(ok, grad, floatX(np.nan))]
 
-class SplineWrapper (theano.Op):
+
+class SplineWrapper(theano.Op):
     """
     Creates a theano operation from scipy.interpolate.UnivariateSpline
     """
@@ -381,22 +382,24 @@ class SplineWrapper (theano.Op):
     def __init__(self, spline):
         self.spline = spline
 
+    @property
+    def grad_op(self):
+        if not hasattr(self, '_grad_op'):
+            try:
+                self._grad_op = SplineWrapper(self.spline.derivative())
+            except ValueError:
+                self._grad_op = None
+
+        if self._grad_op is None:
+            raise NotImplementedError('Spline of order 0 is not differentiable')
+        return self._grad_op
+
     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)]
+
+        return [x_grad * self.grad_op(x)]

pymc3/tests/test_dist_math.py

@@ -4,13 +4,13 @@
 import theano
 import theano.tests.unittest_tools as utt
 import pymc3 as pm
-from scipy import stats
+from scipy import stats, interpolate
 import pytest
 
 from ..theanof import floatX
 from ..distributions import Discrete
 from ..distributions.dist_math import (
-    bound, factln, alltrue_scalar, MvNormalLogp)
+    bound, factln, alltrue_scalar, MvNormalLogp, SplineWrapper)
 
 
 def test_bound():
@@ -176,3 +176,20 @@ def test_hessian(self):
         logp = MvNormalLogp()(cov, delta)
         g_cov, g_delta = tt.grad(logp, [cov, delta])
         tt.grad(g_delta.sum() + g_cov.sum(), [delta, cov])
+
+
+class TestSplineWrapper(object):
+    def test_grad(self):
+        x = np.linspace(0, 1, 100)
+        y = x * x
+        spline = SplineWrapper(interpolate.InterpolatedUnivariateSpline(x, y, k=1))
+        utt.verify_grad(spline, [0.5])
+
+    def test_hessian(self):
+        x = np.linspace(0, 1, 100)
+        y = x * x
+        spline = SplineWrapper(interpolate.InterpolatedUnivariateSpline(x, y, k=1))
+        x_var = tt.dscalar('x')
+        g_x, = tt.grad(spline(x_var), [x_var])
+        with pytest.raises(NotImplementedError):
+            tt.grad(g_x, [x_var])