ODE code style and details (#3687)

* don't use built-in as variable name

* add test for DtypeError

* move make_sens_ic to utils
+ use ones instead of zeros for test values
+ use make_sens_ic for setting the dydp test_value
+ n_states and n_theta args for augment_system function

* also pass n_theta instead of n_p

* enable all test on float32

* state the correct return shape in the docstring

Author: michaelosthege

pymc3/ode/ode.py

@@ -3,7 +3,7 @@
 import scipy
 import theano
 import theano.tensor as tt
-from ..ode.utils import augment_system
+from ..ode import utils
 from ..exceptions import ShapeError, DtypeError
 
 _log = logging.getLogger('pymc3')
@@ -71,38 +71,13 @@ def __init__(self, func, times, *, n_states, n_theta, t0=0):
 
         # Private
         self._augmented_times = np.insert(times, 0, t0).astype(floatX)
-        self._augmented_func = augment_system(func, self.n_states, self.n_p)
-        self._sens_ic = self._make_sens_ic()
+        self._augmented_func = utils.augment_system(func, self.n_states, self.n_theta)
+        self._sens_ic = utils.make_sens_ic(self.n_states, self.n_theta, floatX)
 
         # Cache symbolic sensitivities by the hash of inputs
         self._apply_nodes = {}
         self._output_sensitivities = {}
 
-    def _make_sens_ic(self):
-        """
-        The sensitivity matrix will always have consistent form. (n_states, n_states + n_theta)
-
-        If the first n_states entries of the parameters vector in the simulate call
-        correspond to initial conditions of the system,
-        then the first n_states columns of the sensitivity matrix should form
-        an identity matrix.
-
-        If the last n_theta entries of the parameters vector in the simulate call
-        correspond to ode paramaters, then the last n_theta columns in
-        the sensitivity matrix will be 0.        
-        """
-
-        # Initialize the sensitivity matrix to be 0 everywhere
-        sens_matrix = np.zeros((self.n_states, self.n_states + self.n_theta), dtype=floatX)
-
-        # Slip in the identity matrix in the appropirate place
-        sens_matrix[:,:self.n_states] = np.eye(self.n_states, dtype=floatX)
-
-        # We need the sensitivity matrix to be a vector (see augmented_function)
-        # Ravel and return
-        dydp = sens_matrix.ravel()
-        return dydp
-
     def _system(self, Y, t, p):
         """This is the function that will be passed to odeint. Solves both ODE and sensitivities.
 
@@ -151,10 +126,10 @@ def __call__(self, y0, theta, return_sens=False, **kwargs):
         y0 = tt.cast(tt.unbroadcast(tt.as_tensor_variable(y0), 0), floatX)
         theta = tt.cast(tt.unbroadcast(tt.as_tensor_variable(theta), 0), floatX)
         inputs = [y0, theta]
-        for i, (input, itype) in enumerate(zip(inputs, self._itypes)):
-            if not input.type == itype:
-                raise ValueError('Input {} of type {} does not have the expected type of {}'.format(i, input.type, itype))
-
+        for i, (input_val, itype) in enumerate(zip(inputs, self._itypes)):
+            if not input_val.type == itype:
+                raise ValueError('Input {} of type {} does not have the expected type of {}'.format(i, input_val.type, itype))
+        
         # use default implementation to prepare symbolic outputs (via make_node)
         states, sens = super(theano.Op, self).__call__(y0, theta, **kwargs)
 

pymc3/ode/utils.py

@@ -3,7 +3,47 @@
 import theano.tensor as tt
 
 
-def augment_system(ode_func, n, m):
+def make_sens_ic(n_states, n_theta, floatX):
+        """
+        The sensitivity matrix will always have consistent form. (n_states, n_states + n_theta)
+
+        If the first n_states entries of the parameters vector in the simulate call
+        correspond to initial conditions of the system,
+        then the first n_states columns of the sensitivity matrix should form
+        an identity matrix.
+
+        If the last n_theta entries of the parameters vector in the simulate call
+        correspond to ode paramaters, then the last n_theta columns in
+        the sensitivity matrix will be 0.
+
+        Parameters
+        ----------
+        n_states : int
+            Number of state variables in the ODE
+        n_theta : int
+            Number of ODE parameters
+        floatX : str
+            dtype to be used for the array
+
+        Returns
+        -------
+        dydp : array
+            1D-array of shape (n_states * (n_states + n_theta),), representing the initial condition of the sensitivities
+        """
+
+        # Initialize the sensitivity matrix to be 0 everywhere
+        sens_matrix = np.zeros((n_states, n_states + n_theta), dtype=floatX)
+
+        # Slip in the identity matrix in the appropirate place
+        sens_matrix[:,:n_states] = np.eye(n_states, dtype=floatX)
+
+        # We need the sensitivity matrix to be a vector (see augmented_function)
+        # Ravel and return
+        dydp = sens_matrix.ravel()
+        return dydp
+
+
+def augment_system(ode_func, n_states, n_theta):
     """
     Function to create augmented system.
 
@@ -17,10 +57,10 @@ def augment_system(ode_func, n, m):
     ----------
     ode_func : function
         Differential equation.  Returns array-like.
-    n : int
+    n_states : int
         Number of rows of the sensitivity matrix. (n_states)
-    m : int
-        Number of columns of the sensitivity matrix. (n_states + n_theta)
+    n_theta : int
+        Number of ODE parameters
 
     Returns
     -------
@@ -30,11 +70,11 @@ def augment_system(ode_func, n, m):
 
     # Present state of the system
     t_y = tt.vector("y", dtype='float64')
-    t_y.tag.test_value = np.zeros((n,), dtype='float64')
+    t_y.tag.test_value = np.ones((n_states,), dtype='float64')
     # Parameter(s).  Should be vector to allow for generaliztion to multiparameter
     # systems of ODEs.  Is m dimensional because it includes all initial conditions as well as ode parameters
     t_p = tt.vector("p", dtype='float64')
-    t_p.tag.test_value = np.zeros((m,), dtype='float64')
+    t_p.tag.test_value = np.ones((n_states + n_theta,), dtype='float64')
     # Time.  Allow for non-automonous systems of ODEs to be analyzed
     t_t = tt.scalar("t", dtype='float64')
     t_t.tag.test_value = 2.459
@@ -43,12 +83,12 @@ def augment_system(ode_func, n, m):
     # Will always be 0 unless the parameter is the inital condition
     # Entry i,j is partial of y[i] wrt to p[j]
     dydp_vec = tt.vector("dydp", dtype='float64')
-    dydp_vec.tag.test_value = np.zeros(n * m, dtype='float64')
+    dydp_vec.tag.test_value = make_sens_ic(n_states, n_theta, 'float64')
 
-    dydp = dydp_vec.reshape((n, m))
+    dydp = dydp_vec.reshape((n_states, n_states + n_theta))
 
     # Get symbolic representation of the ODEs by passing tensors for y, t and theta
-    yhat = ode_func(t_y, t_t, t_p[n:])
+    yhat = ode_func(t_y, t_t, t_p[n_states:])
     # Stack the results of the ode_func into a single tensor variable
     if not isinstance(yhat, (list, tuple)):
         yhat = (yhat,)

pymc3/tests/test_ode.py

@@ -8,9 +8,6 @@
 import pytest
 
 
-@pytest.mark.xfail(
-    condition=(theano.config.floatX == "float32"), reason="Fails on float32"
-)
 def test_gradients():
     """Tests the computation of the sensitivities from the theano computation graph"""
 
@@ -53,9 +50,6 @@ def augmented_system(Y, t, p):
     np.testing.assert_allclose(sensitivity, simulated_sensitivity, rtol=1e-5)
 
 
-@pytest.mark.xfail(
-    condition=(theano.config.floatX == "float32"), reason="Fails on float32"
-)
 def test_simulate():
     """Tests the integration in DifferentialEquation"""
 
@@ -81,9 +75,6 @@ def ode_func(y, t, p):
     np.testing.assert_allclose(y, simulated_y, rtol=1e-5)
 
 
-@pytest.mark.xfail(
-    condition=(theano.config.floatX == "float32"), reason="Fails on float32"
-)
 class TestSensitivityInitialCondition(object):
 
     t = np.arange(0, 12, 0.25).reshape(-1, 1)
@@ -104,7 +95,7 @@ def ode_func_1(y, t, p):
         # Sensitivity initial condition for this model should be 1 by 2
         model1_sens_ic = np.array([1, 0])
 
-        np.testing.assert_array_equal(model1_sens_ic, model1._make_sens_ic())
+        np.testing.assert_array_equal(model1_sens_ic, model1._sens_ic)
 
     def test_sens_ic_scalar_2_param(self):
         # Scalar ODE 2 Param
@@ -118,7 +109,7 @@ def ode_func_2(y, t, p):
 
         model2_sens_ic = np.array([1, 0, 0])
 
-        np.testing.assert_array_equal(model2_sens_ic, model2._make_sens_ic())
+        np.testing.assert_array_equal(model2_sens_ic, model2._sens_ic)
 
     def test_sens_ic_vector_1_param(self):
         # Vector ODE 1 Param
@@ -138,7 +129,7 @@ def ode_func_3(y, t, p):
             0, 1, 0
         ])
 
-        np.testing.assert_array_equal(model3_sens_ic, model3._make_sens_ic())
+        np.testing.assert_array_equal(model3_sens_ic, model3._sens_ic)
 
     def test_sens_ic_vector_2_param(self):
         # Vector ODE 2 Param
@@ -158,7 +149,7 @@ def ode_func_4(y, t, p):
             0, 1, 0, 0
         ])
 
-        np.testing.assert_array_equal(model4_sens_ic, model4._make_sens_ic())
+        np.testing.assert_array_equal(model4_sens_ic, model4._sens_ic)
 
     def test_sens_ic_vector_3_params(self):
         # Big System with Many Parameters
@@ -183,12 +174,9 @@ def ode_func_5(y, t, p):
             [0, 0, 1, 0, 0, 0]
         ])
 
-        np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._make_sens_ic())
+        np.testing.assert_array_equal(np.ravel(model5_sens_ic), model5._sens_ic)
 
 
-@pytest.mark.xfail(
-    condition=(theano.config.floatX == "float32"), reason="Fails on float32"
-)
 def test_logp_scalar_ode():
     """Test the computation of the log probability for these models"""
 
@@ -279,9 +267,6 @@ def test_number_of_params(self):
             )
 
 
-@pytest.mark.xfail(
-    condition=(theano.config.floatX == "float32"), reason="Fails on float32"
-)
 class TestDiffEqModel(object):
     def test_op_equality(self):
         """Tests that the equality of mathematically identical Ops evaluates True"""

pymc3/tests/test_util.py

@@ -96,3 +96,15 @@ def test_shape_error(self):
         with pytest.raises(pm.exceptions.ShapeError):
             raise err
         pass
+
+    def test_dtype_error(self):
+        err = pm.exceptions.DtypeError('Without dtypes.')
+        with pytest.raises(pm.exceptions.DtypeError):
+            raise err
+
+        err = pm.exceptions.DtypeError('With shapes.', np.float64, np.float32)
+        assert 'float64' in err.args[0]
+        assert 'float32' in err.args[0]
+        with pytest.raises(pm.exceptions.DtypeError):
+            raise err
+        pass