PBPK_public/pythonScripts/runSolver.py

import ivp
import os
import cModel
import sys
import json
import numpy
import time
import scipy.interpolate

defaultValues={\
      'method':'LSODA',\
      'atol':1e-4,\
      'rtol':1e-4,\
      't0':0,\
      'tmax':1,\
      'mode':'IVPSimultaneous',\
      'nt':201,
      'tUnit':'day',\
      'network':'labkey-klimt.json',\
      'startFromRef':'None',\
      'project':'Analysis/Run',\
      'schema':'lists',\
      'query':'runs',\
      'view':'withStrings'}

def get(setup,par):
   try:
      return setup[par]
   except KeyError:
      pass
   return defaultValues[par]


def parseSetup(setupFile):
   with open(setupFile,'r') as f:
      setup=json.load(f)
   out={}
   for p in defaultValues:
      out[p]=get(setup,p)
   return out

def getScale(setup):
   tUnit=setup['tUnit']
   if tUnit=='min':
      return 1
   if tUnit=='hour':
      return 60
   if tUnit=='day':
      return 24*60
   if tUnit=='month':
      return 24*60*30
   if tUnit=='year':
      return 24*60*30*365
   return 1



def main(parFiles,jobDir,startDir='NONE'):
   model=cModel.model()
   setupFile=parFiles[0]
   modelFile=parFiles[1]
   parameterFile=parFiles[2]
   model.parse(modelFile,parameterFile)
   setup=parseSetup(setupFile)
   scale=getScale(setup)
   tmax=setup['tmax']*scale
   print('Using {}'.format(jobDir))
   if startDir!='NONE':
      print('Using solution from {}'.format(startDir))
      t0,y0,s0,lut,lutSE=getStartPointFromDir(startDir)
      print('t0={}'.format(t0))
   else:
      t0,y0,s0,lut,lutSE=getStartPoint(setup)

   #we should reorder S1 to match model.lutSE
   #S1 is nxm, so we should reorder columns
   sIn=numpy.zeros(s0.shape)
   if s0.size>0:
      for x in model.lutSE:
         #this could be modified for different ordering of lut as well
         sIn[:,model.lutSE[x]]=s0[:,lutSE[x]]

   start_time=time.time()
   if setup['mode']=='SequentialOdeint':
      t,sol,se,s1=ivp.solveSequentialOdeint(model,tmax,setup['nt'],t0,y0,sIn)
   if setup['mode']=='SimultaneousOdeint':
      t,sol,se,s1=ivp.solveSimultaneousOdeint(model,tmax,setup['nt'],t0,y0,sIn)

   if setup['mode']=='IVP':
      t,sol,se,s1=ivp.solveSequential(model,tmax,\
         nt=setup['nt'],\
         atol=setup['atol'],\
         rtol=setup['rtol'],method=setup['method'],\
         t0=t0,y0=y0,sIn=sIn)
        
   if setup['mode']=='IVPSimultaneous':
      t,sol,se,s1=ivp.solveSimultaneous(model,tmax,\
         nt=setup['nt'],\
         atol=setup['atol'],\
         rtol=setup['rtol'],method=setup['method'],\
         t0=t0,y0=y0,sIn=sIn)

   end_time=time.time()
   print('Time: {:.3f} s'.format(end_time-start_time))
   #store
   #interpolate on s1
   #s1 has shape ft x n x m
   #where ft is LSODE driven number of points
   qt,sOut=interpolate(setup,model,t,s1,t0,tmax)

   numpy.savetxt(os.path.join(jobDir,'t.txt'),t)
   numpy.savetxt(os.path.join(jobDir,'sol.txt'),sol)
   numpy.savetxt(os.path.join(jobDir,'se.txt'),se)
   numpy.savetxt(os.path.join(jobDir,'qt.txt'),qt)
   #this is 3D, so new routines
   write3D(os.path.join(jobDir,'sOut.txt'),sOut,model.lut,model.lutSE)

def interpolate(setup,model,t,s1,t0,tmax):
   #interpolate on s1
   #s1 has shape ft x n x m
   #where ft is LSODE driven number of points
   sOut=numpy.zeros((setup['nt'],model.n,model.m))
   qt=numpy.linspace(t0,tmax,setup['nt'])
   for i in range(model.n):
      for j in range(model.m):
         y=s1[:,i,j]
         f = scipy.interpolate.interp1d(t, y, kind='cubic')
         sOut[:,i,j]=f(qt)
   return qt,sOut

def invertLUT(lut):
   fcode=[None]*len(lut)
   for x in lut:
      fcode[lut[x]]=x
   return fcode

def write3D(fName,a,lut,lutSE):
   with open(fName,'w') as f:
      f.write('#Shape {}\n'.format(a.shape))
      f.write('#Lut {}\n'.format(invertLUT(lut)))
      f.write('#Lut (SE) {}\n'.format(invertLUT(lutSE)))

      i=0
      for data_slice in a:
         f.write('#\t Slice {}\n'.format(i))
         numpy.savetxt(f,data_slice)
         i+=1
      
def parseLUT(line,isSE=False):
   SE=''
   if isSE:
      SE='(SE)'
   header='#Lut {}'.format(SE)
   names=line.replace(header,'').\
      replace('[','').\
      replace(']','').\
      replace(' ','').\
      replace('\n','').\
      replace("'",'').\
      split(',')
   lut={names[i]:i for i in range(len(names))}
   return lut

def parseShape(line):
   shape=line.replace('#Shape ','').\
      replace('\n','').\
      replace('(','').\
      replace(')','').\
      split(',')
   return [int(x) for x in shape]

      
def read3D(fName):
   new_data=numpy.loadtxt(fName)
   #read and parse first three lines to get shape,lut and lutSE
   nLines=3
   with open(fName,'r') as f:
      lines=[f.readline() for i in range(nLines)]

   shape=parseShape(lines[0])
   lut=parseLUT(lines[1])
   lutSE=parseLUT(lines[2],isSE=True)

   return new_data.reshape(shape),lut,lutSE


def loadSolutionFromRef(setup):
   
   if setup['startFromRef']=='None':
      return 0,numpy.array([]),numpy.array([]),numpy.array([]),numpy.array([])

   nixSuite=os.path.join(os.path.expanduser('~'),'software','src','nixSuite')
   sys.path.append(os.path.join(nixSuite,'wrapper'))
   import nixWrapper
   nixWrapper.loadLibrary('labkeyInterface')
      
   import labkeyInterface
   import labkeyFileBrowser

   net=labkeyInterface.labkeyInterface()
   fconfig=os.path.join(os.path.expanduser('~'),'.labkey',setup['network'])
   net.init(fconfig)
   #debug
   net.getCSRF()

   #db=labkeyDatabaseBrowser.labkeyDB(net)
   fb=labkeyFileBrowser.labkeyFileBrowser(net)

   #refFilter={'variable':'runRef','value':setup['startFromRef'],'oper':'eq'}
   #ds=db.selectRows(setup['project'],setup['schema'],setup['query'],[refFilter],setup['view']'WithStrings')
   #ordering requires the latest to be first
   #ref=ds['rows'][0]['runRef']
   #parFileString=ds['rows'][0]['parameterFileString']
   #parFiles=parFileString.split(';')
   ref=setup['startFromRef']
   remoteDir=fb.formatPathURL(setup['project'],'/'.join(['jobs',ref]))
   localDir=os.path.join(os.path.expanduser('~'),'temp')
   inFiles=['t.txt','sol.txt','se.txt','qt.txt','sOut.txt']
   #inFiles.extend(parFiles)
   for f in inFiles:
       localPath=os.path.join(localDir,f)
       remotePath='/'.join([remoteDir,f])
       fb.readFileToFile(remotePath,localPath)

   return loadSolutionFromDir(localDir)

def loadSolutionFromDir(jobDir):
   t=numpy.loadtxt(os.path.join(jobDir,'t.txt'))
   sol=numpy.loadtxt(os.path.join(jobDir,'sol.txt'))
   se=numpy.loadtxt(os.path.join(jobDir,'se.txt'))
   sOut,lut,lutSE=read3D(os.path.join(jobDir,'sOut.txt'))
   qt=numpy.loadtxt(os.path.join(jobDir,'qt.txt'))
   return t,sol,se,sOut,qt,lut,lutSE

def getStartPointFromDir(jobDir):
   t,sol,se,sOut,qt,lut,lutSE=loadSolutionFromDir(jobDir)
   return t[-1],sol[-1],sOut[-1],lut,lutSE

def getStartPoint(setup):
   if setup['startFromRef']=='None':
      return 0,numpy.array([]),numpy.array([]),[],[]
   t,sol,se,sOut,qt,lut,lutSE=loadSolutionFromRef(setup)
   return t[-1],sol[-1],sOut[-1],lut,lutSE

if __name__=="__main__":
    parFiles=sys.argv[1].split(';')
    jobDir=sys.argv[2]
    main(parFiles,jobDir)