Merge pull request #301 from pints-team/issue-298-sensitivities

add sensitivities for fitzhugh nagumo model

pints/toy/_fitzhugh_nagumo_model.py

@@ -21,11 +21,37 @@ class FitzhughNagumoModel(pints.ForwardModel):
     Has two states, and three phenomenological parameters: ``a`` , ``b``,
     ``c``. All states are visible
 
+    .. math::
+        \\frac{d \mathbf{y}}{dt} = \\mathbf{f}(\\mathbf{y},\\mathbf{p},t)
+
+    where
+
+    .. math::
+        \\mathbf{y} &= (V,R)\\\\
+        \\mathbf{p} &= (a,b,c)
+
+    The RHS, jacobian and change in RHS with the parameters are given by
+
+    .. math::
+        \\mathbf{f}(\\mathbf{y},\\mathbf{p},t) &= \\left[\\begin{matrix}
+                    c \\left(R - V^{3}/3+V\\right)\\\\
+                    - \\frac{1}{c} \\left(R b + V - a\\right)\\end{matrix}
+                    \\right]\\\\
+        \\frac{\partial \mathbf{f}}{\partial \mathbf{y}} &=
+        \\left[\\begin{matrix} c \\left(1- V^{2}\\right) & c \\\\
+                    - \\frac{1}{c} & - \\frac{b}{c}\\end{matrix}\\right] \\\\
+        \\frac{\partial \mathbf{f}}{\partial \mathbf{p}} &=
+                        \\left[\\begin{matrix}0 & 0 & R - V^{3}/3 + V\\\\
+                        \\frac{1}{c} & - \\frac{R}{c} &
+                        \\frac{1}{c^{2}} \\left(R b + V - a\\right)
+                        \\end{matrix}\\right]
+
     Arguments:
 
     ``y0``
-        The system's initial
+        The system's initial state
     """
+
     def __init__(self, y0=None):
         super(FitzhughNagumoModel, self).__init__()
 
@@ -46,7 +72,34 @@ def n_outputs(self):
         return 2
 
     def simulate(self, parameters, times):
-        """ See :meth:`pints.ForwardModel.simulate()`. """
+        return self._simulate(parameters, times, False)
+
+    def simulate_with_sensitivities(self, parameters, times):
+        return self._simulate(parameters, times, True)
+
+    def _simulate(self, parameters, times, sensitivities):
+        """
+        Private helper function that either simulates the model with
+        sensitivities (`sensitivities == true`) or without
+        (`sensitivities == false`)
+
+        Arguments:
+
+        ``parameters``
+            The three phenomenological parameters: ``a`` , ``b``, ``c``.
+
+        ``times``
+            The times at which to calculate the model output / sensitivities
+
+        ``sensitivities``
+            If set to `true` the function returns the model outputs and
+            sensitivities `(values,sensitivities)`. If set to `false` the
+            function only returns the model outputs `values`. See
+            :meth:`pints.ForwardModel.simulate()` and
+            :meth:`pints.ForwardModel.simulate_with_sensitivities()` for
+            details.
+        """
+
         a, b, c = [float(x) for x in parameters]
 
         times = np.asarray(times)
@@ -59,4 +112,42 @@ def r(y, t, p):
             dR_dt = (V - a + b * R) / -c
             return dV_dt, dR_dt
 
-        return odeint(r, self._y0, times, (parameters,))
+        if sensitivities:
+            def jac(y):
+                V, R = y
+                ret = np.empty((2, 2))
+                ret[0, 0] = c * (1 - V**2)
+                ret[0, 1] = c
+                ret[1, 0] = -1 / c
+                ret[1, 1] = -b / c
+                return ret
+
+            def dfdp(y):
+                V, R = y
+                ret = np.empty((2, 3))
+                ret[0, 0] = 0
+                ret[0, 1] = 0
+                ret[0, 2] = R - V**3 / 3 + V
+                ret[1, 0] = 1 / c
+                ret[1, 1] = -R / c
+                ret[1, 2] = (R * b + V - a) / c**2
+                return ret
+
+            def rhs(y_and_dydp, t, p):
+                y = y_and_dydp[0:2]
+                dydp = y_and_dydp[2:].reshape((2, 3))
+
+                dydt = r(y, t, p)
+                d_dydp_dt = np.matmul(jac(y), dydp) + dfdp(y)
+
+                return np.concatenate((dydt, d_dydp_dt.reshape(-1)))
+
+            y0 = np.zeros(8)
+            y0[0:2] = self._y0
+            result = odeint(rhs, y0, times, (parameters,))
+            values = result[:, 0:2]
+            dvalues_dp = result[:, 2:].reshape((len(times), 2, 3))
+            return values, dvalues_dp
+        else:
+            values = odeint(r, self._y0, times, (parameters,))
+            return values

test/test_fitzhugh_nagumo_model.py

@@ -16,6 +16,7 @@ class TestFitzhughNagumoModel(unittest.TestCase):
     """
     Tests if the Fitzhugh-Nagumo toy model runs.
     """
+
     def test_run(self):
 
         model = pints.toy.FitzhughNagumoModel()
@@ -25,12 +26,16 @@ def test_run(self):
         x = [1, 1, 1]
         times = [1, 2, 3, 4]
         values = model.simulate(x, times)
-        self.assertEqual(values.shape, (4, 2))
+        self.assertEqual(values.shape, (len(times), 2))
+
+        values, dvalues_dp = model.simulate_with_sensitivities(x, times)
+        self.assertEqual(values.shape, (len(times), 2))
+        self.assertEqual(dvalues_dp.shape, (len(times), 2, 3))
 
         # Test alternative starting position
         model = pints.toy.FitzhughNagumoModel([0.1, 0.1])
         values = model.simulate(x, times)
-        self.assertEqual(values.shape, (4, 2))
+        self.assertEqual(values.shape, (len(times), 2))
 
         # Test errors
         times = [-1, 2, 3, 4]