Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP: Add DifferentialEquation: API for Bayesian inference of ODEs #3578

Closed
wants to merge 52 commits into from
Closed
Show file tree
Hide file tree
Changes from 19 commits
Commits
Show all changes
52 commits
Select commit Hold shift + click to select a range
e4ecd6f
Commit ODE capabilities. Add unit test for pm.ode
Dpananos Jul 30, 2019
5520ebf
merged upstream
Dpananos Jul 31, 2019
8d9743f
Commit ODE capabilities. Add unit test for pm.ode
Dpananos Jul 30, 2019
5b88886
re run existing ODE notebook to comare against pymc3.ode
Dpananos Aug 1, 2019
ce7dd20
Merge branch 'gsoc_ode' of https://github.com/Dpananos/pymc3 into gso…
Dpananos Aug 1, 2019
58a8352
use service xvfb
aloctavodia Jul 30, 2019
21cf56c
Fixed float32 precision errors
lucianopaz Jul 31, 2019
463bd7a
SMC stabilize covariance matrix (#3573)
aloctavodia Jul 31, 2019
76851e2
WIP: Documentation cleanup (#3575)
rpgoldman Aug 1, 2019
cd140a1
Fix RST bugs Luciano Paz caught.
rpgoldman Aug 1, 2019
8526a5b
Commit ODE capabilities. Add unit test for pm.ode
Dpananos Jul 30, 2019
101fe72
re run existing ODE notebook to comare against pymc3.ode
Dpananos Aug 1, 2019
90f6e50
Add notebook for DifferentialEquation
Dpananos Aug 2, 2019
6facc4c
Merge branch 'gsoc_ode' of https://github.com/Dpananos/pymc3 into gso…
Dpananos Aug 2, 2019
dca3584
Change test conditions
Dpananos Aug 2, 2019
70e2fd4
remove pytest decorator
Dpananos Aug 2, 2019
abc1016
Added citations in notebook.
Dpananos Aug 2, 2019
5f2d63f
An errant / might be cause of formatting issue
Dpananos Aug 2, 2019
11e9ca5
added a word to SIR example.
Dpananos Aug 3, 2019
6ecc0a5
Add tests, edit notebook
Dpananos Aug 3, 2019
54861b2
add blurb about rehsaping and writing func arg
Dpananos Aug 3, 2019
186f9f2
remove some cells from notebook
Dpananos Aug 3, 2019
5c9ca9e
t0 now optional argument in DiffEq constructor
Dpananos Aug 3, 2019
16ce3d4
merge origin/gsoc_ode
Dpananos Aug 3, 2019
bc7d480
Need import of theano.tensor
Dpananos Aug 3, 2019
94c79fa
Merge branch 'master' of https://github.com/pymc-devs/pymc3 into gsoc…
Dpananos Aug 4, 2019
977238c
typos in notebook. ODEGradOp now in __init__
Dpananos Aug 4, 2019
24730cc
fixed ValueError raised when mask is empty (#3580)
nokados Aug 5, 2019
3746666
Refactor smc out of step methods (#3579)
aloctavodia Aug 6, 2019
7998ebb
Clean up docstring for `sample`.
rpgoldman Aug 13, 2019
c1b9272
Merge pull request #3586 from rpgoldman/sample-docstring
fonnesbeck Aug 14, 2019
86efff3
rename parameters, hopefully passes on float32
Dpananos Aug 14, 2019
cb2add3
Commit ODE capabilities. Add unit test for pm.ode
Dpananos Jul 30, 2019
de5ac44
re run existing ODE notebook to comare against pymc3.ode
Dpananos Aug 1, 2019
1509507
SMC stabilize covariance matrix (#3573)
aloctavodia Jul 31, 2019
2694a46
WIP: Documentation cleanup (#3575)
rpgoldman Aug 1, 2019
1cecac4
Fix RST bugs Luciano Paz caught.
rpgoldman Aug 1, 2019
1f888a9
re run existing ODE notebook to comare against pymc3.ode
Dpananos Aug 1, 2019
711287a
Add notebook for DifferentialEquation
Dpananos Aug 2, 2019
997ed9c
Change test conditions
Dpananos Aug 2, 2019
9c166ec
remove pytest decorator
Dpananos Aug 2, 2019
d5914d8
Added citations in notebook.
Dpananos Aug 2, 2019
cd04ca4
An errant / might be cause of formatting issue
Dpananos Aug 2, 2019
84ae5c8
added a word to SIR example.
Dpananos Aug 3, 2019
8ec0b38
Add tests, edit notebook
Dpananos Aug 3, 2019
c980b14
add blurb about rehsaping and writing func arg
Dpananos Aug 3, 2019
9777262
remove some cells from notebook
Dpananos Aug 3, 2019
ce30d31
t0 now optional argument in DiffEq constructor
Dpananos Aug 3, 2019
e51eaf4
Need import of theano.tensor
Dpananos Aug 3, 2019
d5704b7
typos in notebook. ODEGradOp now in __init__
Dpananos Aug 4, 2019
5041c34
rename parameters, hopefully passes on float32
Dpananos Aug 14, 2019
3ea3ffc
rerun previous ODE notebook
Dpananos Aug 14, 2019
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
321 changes: 143 additions & 178 deletions docs/source/notebooks/ODE_parameter_estimation.ipynb

Large diffs are not rendered by default.

608 changes: 608 additions & 0 deletions docs/source/notebooks/bayesian_estimation_of_ode_parameters.ipynb

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions pymc3/ode/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
from . import utils
from .ode import DifferentialEquation
183 changes: 183 additions & 0 deletions pymc3/ode/ode.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,183 @@
import numpy as np
from pymc3.ode.utils import augment_system, ODEGradop
import scipy
import theano
import theano.tensor as tt
THEANO_FLAG = 'compute_test_value=ignore'


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 iun 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__ = ()

def __init__(self, func, t0, times, n_states, n_odeparams):
Dpananos marked this conversation as resolved.
Show resolved Hide resolved

if not callable(func):
raise ValueError("Argument func must be callable.")
if np.any(np.diff(times)<0):
raise ValueError("The values in times must be monotonically increasing or monotonically decreasing; repeated values are allowed.")
if n_states<1:
raise ValueError('Argument n_states must be at least 1.')
if n_odeparams<0:
raise ValueError('Argument n_states must be non-negative.')

#Public
self.func = func
self.t0 = t0
self.times = 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

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=tuple([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
else:
return self._simulate(np.array(parameters))

def state(self, x):

y, sens = self._cached_simulate(np.array(x, dtype=np.float64))
self._cached_y, self._cached_sens, self._cached_parameters = y, sens, x
return y.ravel()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this ravel destroys the shape of the prediction. It's safer to keep it 2D.
(also to keep the tensor that is returned to the user 2D)

Copy link
Contributor Author

@Dpananos Dpananos Aug 4, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a breaking change. Here is the traceback when I try to rerun the notebook.


TypeError                                 Traceback (most recent call last)
<ipython-input-3-16b9d87cbe1d> in <module>
     13     #If we know one of the parameter values, we can simply pass the value.
     14     #No need to specify a prior.
---> 15     ode_solution = ode_model(odeparams = [gamma, 9.8], y0 = [0]).reshape(yobs.shape)
     16 
     17     Y = pm.Normal('Y', mu = ode_solution, sd = sigma, observed = yobs)

~/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/tensor/var.py in reshape(self, shape, ndim)
    319                                  str(type(ndim)))
    320 
--> 321         return theano.tensor.basic.reshape(self, shape, ndim=ndim)
    322 
    323     def dimshuffle(self, *pattern):

~/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/tensor/basic.py in reshape(x, newshape, ndim)
   5064                 "argument to 'reshape' to avoid this problem." % newshape)
   5065     op = Reshape(ndim)
-> 5066     rval = op(x, newshape)
   5067     return rval
   5068 

~/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/gof/op.py in __call__(self, *inputs, **kwargs)
    623             for i, ins in enumerate(node.inputs):
    624                 try:
--> 625                     storage_map[ins] = [self._get_test_value(ins)]
    626                     compute_map[ins] = [True]
    627                 except AttributeError:

~/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/gof/op.py in _get_test_value(cls, v)
    560             # ensure that the test value is correct
    561             try:
--> 562                 ret = v.type.filter(v.tag.test_value)
    563             except Exception as e:
    564                 # Better error message.

~/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/tensor/type.py in filter(self, data, strict, allow_downcast)
    176             raise TypeError("Wrong number of dimensions: expected %s,"
    177                             " got %s with shape %s." % (self.ndim, data.ndim,
--> 178                                                         data.shape))
    179         if not data.flags.aligned:
    180             try:

TypeError: For compute_test_value, one input test value does not have the requested type.
 
Backtrace when that variable is created:

  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 2881, in _run_cell
    return runner(coro)
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/IPython/core/async_helpers.py", line 68, in _pseudo_sync_runner
    coro.send(None)
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3058, in run_cell_async
    interactivity=interactivity, compiler=compiler, result=result)
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3249, in run_ast_nodes
    if (await self.run_code(code, result,  async_=asy)):
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/IPython/core/interactiveshell.py", line 3326, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-3-16b9d87cbe1d>", line 15, in <module>
    ode_solution = ode_model(odeparams = [gamma, 9.8], y0 = [0]).reshape(yobs.shape)
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/theano/gof/op.py", line 615, in __call__
    node = self.make_node(*inputs, **kwargs)
  File "/Users/demetri/anaconda3/envs/gsoc/lib/python3.6/site-packages/pymc3/ode/ode.py", line 155, in make_node
    return theano.Apply(self, [x], [x.type()])

The error when converting the test value to that variable type:
Wrong number of dimensions: expected 1, got 2 with shape (20, 1).

Any idea why this might be?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this after you removed the ravel? There's still the reshape that should be be removable..

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this is after making the suggested change.
No dice. Removing the .reshape(yobs.shape) results in the same error.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could be the reason why I had infer_shape in the other Op...


def numpy_vsp(self, x, g):

numpy_sens = self._cached_simulate(np.array(x, dtype=np.float64))[1].reshape((self.n_states * len(self.times), len(x)))
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} paramteres 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} paramteres 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)
x = tt.concatenate([odeparams, y0])

return theano.Apply(self, [x], [x.type()])

def perform(self, node, inputs_storage, output_storage):

x = inputs_storage[0]
out = output_storage[0]

# get the numerical solution of ODE states
out[0] = self.state(x)

def grad(self, inputs, output_grads):

x = inputs[0]
g = output_grads[0]

# pass the VSP when asked for gradient
grad_op = ODEGradop(self.numpy_vsp)
grad_op_apply = grad_op(x, g)

return [grad_op_apply]
81 changes: 81 additions & 0 deletions pymc3/ode/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
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)

# Parameter(s). Should be vector to allow for generaliztion to multiparameter
# systems of ODEs
t_p = tt.vector('p', dtype=theano.config.floatX)

# Time. Allow for non-automonous systems of ODEs to be analyzed
t_t = tt.scalar('t', dtype=theano.config.floatX)

# 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 = dydp_vec.reshape((n, m))

# Stack the results of the ode_func
# TODO: Does this behave the same of ODE is scalar?
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
Loading