PBPK_public/pythonScripts/ivp.py

import numpy
import scipy.integrate
#scipy solve_ivp interface using system of n equations
#infrastructure for sensitivity analysis of y with respect
#to model parameters p also integrated
#
#General form of equation:
#  dy/dt=M*y+u
#where y is an n-vector and M is nxn Jacobi matrix 
#
#implementation of system has to provide attributes:
# - n, number of variables
# - m, number of parameters
# - M(t), nxn, a time dependent Jacobi matrix,
# - u(t), 1xn, a time dependent source term
# - fS(t) mxn, a time dependent array of derivatives 
#           fS(t)[p,:] = d(M*y)/dp (t)
#   where y is internal interpolated y-only ivp solution
# - fSY(t,y), mxn, a time dependent array of derivatives:
#           fS(t)[p,:] = d(M*y)/dp (t)
#   where current y is supplied at each step

def dfdy(t,y,system):
    dfdy=system.M(t).dot(y)+system.u(t)
    return dfdy

def jacobi(t,y,system):
    return system.M(t)

#SE post calculation
def dfdyS(t,S,system):
    #unwrap S to NxM where M is number of parameters
    mS=numpy.reshape(S,(system.n,system.m))
    mOut=system.M(t).dot(mS)+system.fS(t)
    return numpy.ravel(mOut)

def jacobiSE(t,S,system):
    N=system.n*(system.m)
    fJ=numpy.zeros((N,N))
    #print('fJ shape {}'.format(fJ.shape))
    for i in range(system.m):
        fJ[i*system.n:(i+1)*system.n,i*system.n:(i+1)*system.n]=system.M(t)
    return fJ

#SE simultaeneous calculation
def dfdySFull(t,S,system):
    #unwrap S to NxM where M is number of parameters
    mS=numpy.reshape(S,(system.n,system.m+1))
    #system.fS(y,t) is NxM matrix where M are parameters
    y=mS[:,0]
    mOut=system.M(t).dot(mS)+system.fSY(y,t)
    try:
      system.iPrint+=1
    except AttributeError:
      system.iPrint=0
    if system.iPrint>1000:
       print('At t={:.2f}'.format(t),flush=True)
       system.iPrint=0
    return numpy.ravel(mOut)

def jacobiSEFull(t,S,system):
    N=system.n*(system.m+1)
    fJ=numpy.zeros((N,N))
    #print('fJ shape {}'.format(fJ.shape))
    for i in range(system.m+1):
        fJ[i*system.n:(i+1)*system.n,i*system.n:(i+1)*system.n]=system.M(t)
    return fJ

def solveSimultaneous(model,tmax,atol,rtol,method='LSODA',t0=0,y0=numpy.array([]), Sin=numpy.array([])):
    S1=numpy.zeros((model.n,model.m+1))
    if Sin.size!=0:
       S1[:,1:]=Sin
    if y0.size==0:
       y0=numpy.zeros(model.n)
    #set initial condition
    S1[:,0]=y0
    S1=S1.ravel()
    sol=scipy.integrate.solve_ivp(dfdySFull,[t0, tmax],S1, args=(model,), jac=jacobiSEFull,
                                  method=method, atol=atol, rtol=rtol)
    t=sol.t
    sFull=numpy.reshape(numpy.transpose(sol.y),(len(t),model.n,model.m+1))
    s1=sFull[:,:,1:]
    ysol=sFull[:,:,0]
    se=model.calculateUncertainty(ysol,s1)
    print('Done simultaneous LSODA SE')
    return t,ysol,se,s1

def solveSequential(model,tmax,atol,rtol,method='LSODA',t0=0,y0=numpy.array([]), S1=numpy.array([])):
   if y0.size==0:
      y0=numpy.zeros(model.n)
   model.iPrint=0
   solIVP=scipy.integrate.solve_ivp(dfdy,[t0, tmax],y0, args=(model,), jac=jacobi,
                                  method=method, atol=atol, rtol=rtol)
   #y is n x nt (odeint nt x n)
   sol=numpy.transpose(solIVP.y)
   t=solIVP.t
   print('shape (y) {}'.format(sol.shape))
   model.setY(t,sol)
   if S1.size==0:
      S1=numpy.zeros((model.n,model.m))
   S1=S1.ravel()
    
   solIVPSE=scipy.integrate.solve_ivp(dfdyS,[0, tmax],S1, args=(model,), jac=jacobiSE,
                               method=method, atol=atol, rtol=rtol)
   sraw=numpy.reshape(numpy.transpose(solIVPSE.y),(len(solIVPSE.t),model.n,model.m))
   #interpolate on t
   s1=numpy.zeros((len(t),model.n,model.m))
   for i in range(model.n):
      for j in range(model.m):
         tck = scipy.interpolate.splrep(solIVPSE.t, sraw[:,i,j], s=0)
         s1[:,i,j]=scipy.interpolate.splev(t, tck, der=0)

   se=model.calculateUncertainty(sol,s1)
   return t,sol,se,s1

def solveSimultaneousOdeint(model,tmax,nt=201,t0=0,y0=numpy.array([]), Sin=numpy.array([])):
   t = numpy.linspace(t0,tmax, nt)
   if y0.size==0:
      y0=numpy.zeros(model.n)
   
   S1=numpy.zeros((model.n,model.m+1))
   if Sin.size!=0:
      S1[:,1:]=Sin
   #set initial condition
   S1[:,0]=y0
   S1=S1.ravel()
   solSE1=scipy.integrate.odeint(dfdySFull, S1, t, args=(model,),Dfun=jacobiSEFull,tfirst=True)
   sFull=numpy.reshape(solSE1,(len(t),model.n,model.m+1))
   s1=sFull[:,:,1:]
   sol=sFull[:,:,0]
   se=sys.calculateUncertainty(sol,s1)
   print('Done simultaneous SE')
   return t,sol,se,s1

def solveSequentialOdeint(model,tmax,nt=201,t0=0,y0=numpy.array([]),S1=numpy.array([])):
   t = numpy.linspace(t0,tmax, nt)
   if y0.size==0:
      y0=numpy.zeros(sys.n)
   sol = scipy.integrate.odeint(dfdy, y0=y0, t=t, args=(model,),Dfun=jacobi,tfirst=True)
   print('shape (y) {}'.format(sol.shape))

   model.setY(t,sol)
   if S1.size==0:
      S1=numpy.zeros((model.n,model.m))
   S1=S1.ravel()
   solSE=scipy.integrate.odeint(dfdyS, S1, t, args=(model,),Dfun=jacobiSE,tfirst=True)
   s1=numpy.reshape(solSE,(len(t),model.n,model.m))
   se=model.calculateUncertainty(sol,s1)
   print('Done sequential SE')
   return t,sol,se,s1