Solving ODEs in PYMC3

Question

Here I aim to estimate the parameters (gama and omega) of a damped harmonic oscillator given by

dX^2/dt^2+gamma*dX/dt+(2*pi*omega)^2*X=0. (We can add white gaussian noise to the system.)

---

 import pymc
 import numpy as np
 import scipy.io as sio
 import matplotlib.pyplot as plt; 
 from scipy.integrate import odeint

 #import data
 xdata = sio.loadmat('T.mat')['T'][0]  #time
 ydata1 = sio.loadmat('V1.mat')['V1'][0]  #  V2=dV1/dt, (X=V1),
 ydata2 = sio.loadmat('V2.mat')['V2'][0]  # dV2/dt=-(2pi*omega)^2*V1-gama*V2

 #time span for solving the equations
 npts= 500   
 dt=0.01
 Tspan=5.0
 time = np.linspace(0,Tspan,npts+1) 

 #initial condition
 V0 = [1.0, 1.0]

# Priors for unknown model parameters
sigma = pymc.Uniform('sigma',  0.0, 100.0)
gama= pymc.Uniform('gama', 0.0, 20.0) 
omega=pymc.Uniform('omega',0.0, 20.0)

#Solve the equations
@pymc.deterministic
def DHOS(gama=gama, omega=omega):
V1= np.zeros(npts+1)
V2= np.zeros(npts+1)
V1[0] = V0[0]
V2[0] = V0[1]
for i in range(1,npts+1):
    V1[i]=  V1[i-1] + dt*V2[i-1]; 
    V2[i] = V2[i-1] + dt*(-((2*np.pi*omega)**2)*V1[i-1]-gama*V2[i-1]); 
return [V1, V2]

#or we can use odeint
#@pymc.deterministic
#def DHS( gama=gama, omega=omega):
#    def DOS_func(y, time):
#        V1, V2 = y[0], y[1]
#        dV1dt = V2
#        dV2dt = -((2*np.pi*omega)**2)* V1 -gama*V2 
#        dydt = [dV1dt, dV2dt]
#        return dydt
#    soln = odeint(DOS_func,V0, time)
#    V1, V2 = soln[:,0], soln[:,1]
#    return V1, V2


#  value of outcome (observations)
V1 = pymc.Lambda('V1', lambda DHOS=DHOS: DHOS[0])
V2 = pymc.Lambda('V2', lambda DHOS=DHOS: DHOS[1])


# liklihood of observations
Yobs1 = pymc.Normal('Yobs1', mu=V1, tau=1.0/sigma**2, value=ydata1, observed=True)
Yobs2 = pymc.Normal('Yobs2', mu=V2, tau=1.0/sigma**2, value=ydata2, observed=True)

By saving the above code as DampedOscil_model.py, then we are able to run PYMC as follows

import pymc
import DampedOscil_model


MDL = pymc.MCMC(DampedOscil_model, db='pickle')
MDL.sample(iter=1e4, burn=1e2, thin=2)

gama_trace=MDL.trace('gama')[- 1000:]
omega_trace=MDL.trace('omega')[-1000:]

gama=MDL.gama.value
omega=MDL.omega.value

---

And it works well (See below).

The true signal constructed by gama_true=2.0 and omega_est=1.5 versus the estimated signal. The estimated parameter values are gama_est=2.04 and omega_est=1.49

Now I would convert this code to PYMC3 to use NUTS and ADVI.

import matplotlib.pyplot as plt
import scipy.io as sio
import pandas as pd
import numpy as np

import pymc3 as pm
import theano.tensor as tt
import theano

from pymc3 import Model, Normal, HalfNormal, Uniform
from pymc3 import NUTS, find_MAP, sample, Slice, traceplot, summary
from pymc3 import Deterministic  
from scipy.optimize import fmin_powell


#import data
xdata = sio.loadmat('T.mat')['T'][0]  #time
ydata1 = sio.loadmat('V1.mat')['V1'][0]  #  V2=dV1/dt, (X=V1),
ydata2 = sio.loadmat('V2.mat')['V2'][0]  # dV2/dt=-(2pi*omega)^2*V1-gama*V2

#time span for solving the equations
npts= 500   
dt=0.01
Tspan=5.0
time = np.linspace(0,Tspan,npts+1) 

niter=10000
burn=niter//2;

with pm.Model() as model:

 #Priors for unknown model parameters
 sigma = pm.HalfNormal('sigma', sd=1)
 gama= pm.Uniform('gama', 0.0, 20.0)
 omega=pm.Uniform('omega',0.0, 20.0)

 #initial condition
 V0 = [1.0, 1.0]

#Solve the equations     
# do I need to use theano.tensor here?!
 @theano.compile.ops.as_op(itypes=[tt.dscalar, tt.dscalar],otypes=[tt.dvector])  
def DHOS(gama=gama, omega=omega):
V1= np.zeros(npts+1)
V2= np.zeros(npts+1)
V1[0] = V0[0]
V2[0] = V0[1]
for i in range(1,npts+1):
    V1[i]=  V1[i-1] + dt*V2[i-1]; 
    V2[i] = V2[i-1] + dt*(-((2*np.pi*1)**2)*V1[i-1]-gama*V2[i-1]); 
return V1,V2


 V1 = pm.Deterministic('V1', DHOS[0])
 V2 = pm.Deterministic('V2', DHOS[1])


 start = pm.find_MAP(fmin=fmin_powell, disp=True)
 step=pm.NUTS()
 trace=pm.sample(niter, step, start=start, progressbar=False)

 traceplot(trace);

 Summary=pm.df_summary(trace[-1000:])

  gama_trace = trace.get_values('gama', burn)
  omega_trace = trace.get_values('omega', burn)

For this code I get the following error: V1 = pm.Deterministic('V1', DHOS[0])

      TypeError: 'FromFunctionOp' object does not support indexing

Briefly, I wonder to know how can I to convert the following part of PYMC code to PYMC3.

@pymc.deterministic
def DOS(gama=gama, omega=omega):
V1= np.zeros(npts+1)
V2= np.zeros(npts+1)
V1[0] = V0[0]
V2[0] = V0[1]
for i in range(1,npts+1):
    V1[i]=  V1[i-1] + dt*V2[i-1]; 
    V2[i] = V2[i-1] + dt*(-((2*np.pi*omega)**2)*V1[i-1]-gama*V2[i-1]); 
return [V1, V2]


V1 = pymc.Lambda('V1', lambda DOS=DOS: DOS[0]) 
V2 = pymc.Lambda('V2', lambda DOS=DOS: DOS[1])

The problem is, first, the argumentation of Deterministic function is different in PYMC3 from PYMC, secondly, there in no Lambda function in PYMC3.

I appreciate your help in solving ODEs in PYMC3 to solve parameter estimation task in biological systems (estimating the equation parameters from data).

Thanks a lot in advance for your help.

Kind Regards,

Meysam

Answer 1

I would suggest, and have successfully implemented, using a 'black box' method for interfacing with PYMC3. In this case what that means is calculating the log-liklihood yourself and then using PYMC3 to sample it. This requires writing your functions in a way that Theano and PYMC3 can interface with them.

This is outlined in a notebook on the PYMC3 page, which uses cython as an example.

Here is a bit shorter sample of what needs to be done.

First you can load your data and set-up any parameters you need such as your time steps etc.

import pymc3 as pm
import numpy as np
import theano
import theano.tensor as tt


 #import data
 xdata = sio.loadmat('T.mat')['T'][0]  #time
 ydata1 = sio.loadmat('V1.mat')['V1'][0]  #  V2=dV1/dt, (X=V1),
 ydata2 = sio.loadmat('V2.mat')['V2'][0]  # dV2/dt=-(2pi*omega)^2*V1-gama*V2

 #time span for solving the equations
 npts= 500   
 dt=0.01
 Tspan=5.0
 time = np.linspace(0,Tspan,npts+1) 

 #initial condition
 V0 = [1.0, 1.0]

Then you define a data generating function just as before but you don't need to use any decorators from PYMC for this. The output of this function should be whatever you need to compare to your data to calculate the likelihood.

def DHOS(theta):
    gama,omega=theta
    V1= np.zeros(npts+1)
    V2= np.zeros(npts+1)
    V1[0] = V0[0]
    V2[0] = V0[1]
    for i in range(1,npts+1):
        V1[i]=  V1[i-1] + dt*V2[i-1]; 
        V2[i] = V2[i-1] + dt*(-((2*np.pi*omega)**2)*V1[i-1]-gama*V2[i-1]); 
    return [V1, V2]

Next you write a function that calls the previous function and calculates the likelihood using whatever distribution you want, in this a normal distribution.

def my_loglike(theta,data,sigma):
    """
    A Gaussian log-likelihood function for a model with parameters given in theta
    """

    model = DHOS(theta) #V1 and V2 from the DHOS function

    #Here data = [ydata1,ydata2] to compare with model
    #sigma is either the same shape as model or a scalar
    #which corresponds to the uncertainty on the data. 

    return -(0.5)*sum((data - model)**2/sigma**2)

From here you have now have to define a Theano class so that it can interface with PYMC3.

# define a theano Op for our likelihood function
class LogLike(tt.Op):

    """
    Specify what type of object will be passed and returned to the Op when it is
    called. In our case we will be passing it a vector of values (the parameters
    that define our model) and returning a single "scalar" value (the
    log-likelihood)
    """
    itypes = [tt.dvector] # expects a vector of parameter values when called
    otypes = [tt.dscalar] # outputs a single scalar value (the log likelihood)

    def __init__(self, loglike, data, sigma):
        """
        Initialise the Op with various things that our log-likelihood function
        requires. Below are the things that are needed in this particular
        example.

        Parameters
        ----------
        loglike:
            The log-likelihood (or whatever) function we've defined
        data:
            The "observed" data that our log-likelihood function takes in
        x:
            The dependent variable (aka 'x') that our model requires
        sigma:
            The noise standard deviation that our function requires.
        """

        # add inputs as class attributes
        self.likelihood = loglike
        self.data = data
        self.sigma = sigma

    def perform(self, node, inputs, outputs):
        # the method that is used when calling the Op
        theta, = inputs  # this will contain my variables

        # call the log-likelihood function
        logl = self.likelihood(theta, self.data, self.sigma)

        outputs[0][0] = array(logl) # output the log-likelihood

Finally you can use PYMC3 to build your model and sample accordingly.

ndraws = 10000  # number of draws from the distribution
nburn = 1000   # number of "burn-in points" (which we'll discard)

# create our Op
logl = LogLike(my_loglike, rdat_sim, 10)

# use PyMC3 to sampler from log-likelihood
with pm.Model():
    gama= pm.Uniform('gama', 0.0, 20.0)
    omega=pm.Uniform('omega',0.0, 20.0)


    # convert m and c to a tensor vector
    theta = tt.as_tensor_variable([gama,omega])

    # use a DensityDist (use a lamdba function to "call" the Op)
    pm.DensityDist('likelihood', lambda v: logl(v), observed={'v': theta})

    trace = pm.sample(ndraws, tune=nburn, discard_tuned_samples=True)

And you can use the internal plotting to see the results of the sampling

_ = pm.traceplot(trace)

This was just adapted from the example notebook in the link, and as mentioned there if you want to use NUTS you need gradient information, which you do not have given you custom function. In the link it talks about how to sample the gradient and construct it so you can pass it into the sampler, but I have not shown that here.

Additionally if you want to use solve_ivp (or odeint or another solver), all you have to do is change the DHOS function as you normally would to invoke the solver. The rest of the code should be portable to whatever problem you, or anyone else, need.