Add Differential Equation API

pymc3/__init__.py

@@ -8,6 +8,7 @@
 from .math import logaddexp, logsumexp, logit, invlogit, expand_packed_triangular, probit, invprobit
 from .model import *
 from .model_graph import model_to_graphviz
+from . import ode
 from .stats import *
 from .sampling import *
 from .step_methods import *

pymc3/ode/__init__.py

@@ -0,0 +1,2 @@
+from . import utils
+from .ode import DifferentialEquation

pymc3/ode/ode.py

@@ -0,0 +1,184 @@
+import numpy as np
+import scipy
+import theano
+import theano.tensor as tt
+from ..ode.utils import augment_system, ODEGradop
+
+
+class DifferentialEquation(theano.Op):
+    """
+    Specify an ordinary differential equation
+
+    .. math::
+        \dfrac{dy}{dt} = f(y,t,p) \quad y(t_0) = y_0
+
+    Parameters
+    ----------
+
+    func : callable
+        Function specifying the differential equation
+    t0 : float
+        Time corresponding to the initial condition
+    times : array
+        Array of times at which to evaluate the solution of the differential equation.
+    n_states : int
+        Dimension of the differential equation.  For scalar differential equations, n_states=1.
+        For vector valued differential equations, n_states = number of differential equations in the system.
+    n_odeparams : int
+        Number of parameters in the differential equation.
+
+    .. code-block:: python
+
+        def odefunc(y, t, p):
+            #Logistic differential equation
+            return p[0]*y[0]*(1-y[0])
+
+        times=np.arange(0.5, 5, 0.5)
+
+        ode_model = DifferentialEquation(func=odefunc, t0=0, times=times, n_states=1, n_odeparams=1)
+    """
+
+    __props__ = ("func", "t0", "times", "n_states", "n_odeparams")
+
+    def __init__(self, func, times, n_states, n_odeparams, t0=0):
+        if not callable(func):
+            raise ValueError("Argument func must be callable.")
+        if n_states < 1:
+            raise ValueError("Argument n_states must be at least 1.")
+        if n_odeparams <= 0:
+            raise ValueError("Argument n_odeparams must be positive.")
+
+        # Public
+        self.func = func
+        self.t0 = t0
+        self.times = tuple(times)
+        self.n_states = n_states
+        self.n_odeparams = n_odeparams
+
+        # Private
+        self._n = n_states
+        self._m = n_odeparams + n_states
+
+        self._augmented_times = np.insert(times, t0, 0)
+        self._augmented_func = augment_system(func, self._n, self._m)
+        self._sens_ic = self._make_sens_ic()
+
+        self._cached_y = None
+        self._cached_sens = None
+        self._cached_parameters = None
+
+        self._grad_op = ODEGradop(self._numpy_vsp)
+
+    def _make_sens_ic(self):
+        # The sensitivity matrix will always have consistent form.
+        # If the first n_odeparams entries of the parameters vector in the simulate call
+        # correspond to ode paramaters, then the first n_odeparams columns in
+        # the sensitivity matrix will be 0
+        sens_matrix = np.zeros((self._n, self._m))
+
+        # If the last n_states entrues of the paramters vector in the simulate call
+        # correspond to initial conditions of the system,
+        # then the last n_states columns of the sensitivity matrix should form
+        # an identity matrix
+        sens_matrix[:, -self.n_states :] = np.eye(self.n_states)
+
+        # 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
+        Args:
+            Y (vector): current state and current gradient state
+            t (scalar): current time
+            p (vector): parameters
+        Returns:
+            derivatives (vector): derivatives of state and gradient
+        """
+
+        dydt, ddt_dydp = self._augmented_func(Y[: self._n], t, p, Y[self._n :])
+        derivatives = np.concatenate([dydt, ddt_dydp])
+        return derivatives
+
+    def _simulate(self, parameters):
+        # Initial condition comprised of state initial conditions and raveled
+        # sensitivity matrix
+        y0 = np.concatenate([parameters[self.n_odeparams :], self._sens_ic])
+
+        # perform the integration
+        sol = scipy.integrate.odeint(
+            func=self._system, y0=y0, t=self._augmented_times, args=(parameters,)
+        )
+        # The solution
+        y = sol[1:, : self.n_states]
+
+        # The sensitivities, reshaped to be a sequence of matrices
+        sens = sol[1:, self.n_states :].reshape(len(self.times), self._n, self._m)
+
+        return y, sens
+
+    def _cached_simulate(self, parameters):
+        if np.array_equal(np.array(parameters), self._cached_parameters):
+
+            return self._cached_y, self._cached_sens
+
+        return self._simulate(np.array(parameters))
+
+    def _state(self, parameters):
+        y, sens = self._cached_simulate(np.array(parameters))
+        self._cached_y, self._cached_sens, self._cached_parameters = y, sens, parameters
+        return y.ravel()
+
+    def _numpy_vsp(self, parameters, g):
+        _, sens = self._cached_simulate(np.array(parameters))
+
+        # Each element of sens is an nxm sensitivity matrix
+        # There is one sensitivity matrix per time step, making sens a (len(times), n_states, len(parameter))
+        # dimensional array.  Reshaping the sens array in this way is like stacking each of the elements of sens on top
+        # of one another.
+        numpy_sens = sens.reshape((self.n_states * len(self.times), len(parameters)))
+        # The dot product here is equivalent to np.einsum('ijk,jk', sens, g)
+        # if sens was not reshaped and if g had the same shape as yobs
+        return numpy_sens.T.dot(g)
+
+    def make_node(self, odeparams, y0):
+        if len(odeparams) != self.n_odeparams:
+            raise ValueError(
+                "odeparams has too many or too few parameters.  Expected {a} parameter(s) but got {b}".format(
+                    a=self.n_odeparams, b=len(odeparams)
+                )
+            )
+        if len(y0) != self.n_states:
+            raise ValueError(
+                "y0 has too many or too few parameters.  Expected {a} parameter(s) but got {b}".format(
+                    a=self.n_states, b=len(y0)
+                )
+            )
+
+        if np.ndim(odeparams) > 1:
+            odeparams = np.ravel(odeparams)
+        if np.ndim(y0) > 1:
+            y0 = np.ravel(y0)
+
+        odeparams = tt.as_tensor_variable(odeparams)
+        y0 = tt.as_tensor_variable(y0)
+        parameters = tt.concatenate([odeparams, y0])
+        return theano.Apply(self, [parameters], [parameters.type()])
+
+    def perform(self, node, inputs_storage, output_storage):
+        parameters = inputs_storage[0]
+        out = output_storage[0]
+        # get the numerical solution of ODE states
+        out[0] = self._state(parameters)
+
+    def grad(self, inputs, output_grads):
+        x = inputs[0]
+        g = output_grads[0]
+        # pass the VSP when asked for gradient
+        grad_op_apply = self._grad_op(x, g)
+
+        return [grad_op_apply]

pymc3/ode/utils.py

@@ -0,0 +1,79 @@
+import numpy as np
+import theano
+import theano.tensor as tt
+
+
+def augment_system(ode_func, n, m):
+    """Function to create augmented system.
+
+    Take a function which specifies a set of differential equations and return
+    a compiled function which allows for computation of gradients of the
+    differential equation's solition with repsect to the parameters.
+
+    Args:
+        ode_func (function): Differential equation.  Returns array-like
+        n: Number of rows of the sensitivity matrix
+        m: Number of columns of the sensitivity matrix
+
+    Returns:
+        system (function): Augemted system of differential equations.
+
+    """
+
+    # Present state of the system
+    t_y = tt.vector("y", dtype=theano.config.floatX)
+    t_y.tag.test_value = np.zeros((n,))
+    # Parameter(s).  Should be vector to allow for generaliztion to multiparameter
+    # systems of ODEs.  Is m dimensional because it includes all ode parameters as well as initical conditions
+    t_p = tt.vector("p", dtype=theano.config.floatX)
+    t_p.tag.test_value = np.zeros((m,))
+    # Time.  Allow for non-automonous systems of ODEs to be analyzed
+    t_t = tt.scalar("t", dtype=theano.config.floatX)
+    t_t.tag.test_value = 2.459
+
+    # Present state of the gradients:
+    # 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=theano.config.floatX)
+    dydp_vec.tag.test_value = np.zeros(n * m)
+
+    dydp = dydp_vec.reshape((n, m))
+
+    # Stack the results of the ode_func
+    f_tensor = tt.stack(ode_func(t_y, t_t, t_p))
+
+    # Now compute gradients
+    J = tt.jacobian(f_tensor, t_y)
+
+    Jdfdy = tt.dot(J, dydp)
+
+    grad_f = tt.jacobian(f_tensor, t_p)
+
+    # This is the time derivative of dydp
+    ddt_dydp = (Jdfdy + grad_f).flatten()
+
+    system = theano.function(
+        inputs=[t_y, t_t, t_p, dydp_vec],
+        outputs=[f_tensor, ddt_dydp],
+        on_unused_input="ignore",
+    )
+
+    return system
+
+
+class ODEGradop(theano.Op):
+    def __init__(self, numpy_vsp):
+        self._numpy_vsp = numpy_vsp
+
+    def make_node(self, x, g):
+
+        x = theano.tensor.as_tensor_variable(x)
+        g = theano.tensor.as_tensor_variable(g)
+        node = theano.Apply(self, [x, g], [g.type()])
+        return node
+
+    def perform(self, node, inputs_storage, output_storage):
+        x = inputs_storage[0]
+        g = inputs_storage[1]
+        out = output_storage[0]
+        out[0] = self._numpy_vsp(x, g)  # get the numerical VSP