PBPK_public/pythonScripts/cModel.py

import numpy
import json
import os
import scipy.interpolate
#for partial function specializations
import functools
import function
import importlib
importlib.reload(function)

class model:
   def __init__(self):
      self.compartments={}
      self.seJ={}
        
   def add_source(self,compartmentName,formula):
      self.compartments[compartmentName]['source']=formula

   def add_compartment(self,compartmentName):
      self.compartments[compartmentName]={}
      self.compartments[compartmentName]['targets']={}
      self.compartments[compartmentName]['sensTargets']={}

   def getTimeUnit(self):

      try: 
         return self.mod['timeUnit']
      except KeyError:
         return 's'

   def bind(self,src,target,qName,pcName,useVolume=0):
      

      #establish a flow from source compartment to the target
      
      #the source equation (where we subtract the current)
      #in fact, this is the diagonal element
      #get volume names
      srcVName=self.getVolumePar(src,useVolume)
      #generate coupling object (w/derivatives)    
      pSrc=self.couplingObject(-1,qName,pcName,srcVName)
      #this includes derivatives and value!
      self.addValueObject(src,src,pSrc)
     
      #special target which is not part of calculation
      if target=='dump':
         return

      #the target equation (where we add the current)
      #get volume names
      targetVName=self.getVolumePar(target,useVolume)
      #generate coupling object
      pTarget=self.couplingObject(1,qName,pcName,targetVName)
      #equation is for target compartment, but binding for source
      self.addValueObject(target,src,pTarget)
   
   def addValueObject(self,targetName,srcName,cObject):
      #always binds equation id and a variable
      targetList=self.compartments[targetName]['targets']
      addValue(targetList,srcName,cObject["value"])
      der=cObject["derivatives"]
      for d in der:
         targetSE=self.getSEJ_comp(d,targetName)
         addValue(targetSE,srcName,der[d])

                    
   def couplingObject(self,sign, qParName, pcParName, vParName):
        
      qPar=self.get(qParName)
      pcPar=self.get(pcParName)
      vPar=self.get(vParName)
      q=qPar['value']
      pc=pcPar['value']
      v=vPar['value']
      DPC=pcPar['derivatives']
      DQ=qPar['derivatives']
      DV=vPar['derivatives']

      if any(['function' in qPar,'function' in pcPar, 'function' in vPar]):
         fq=function.to(q)
         fpc=function.to(pc)
         fv=function.to(v)
         f=lambda t,q=fq,pc=fpc,v=fv,s=sign:s*q(t)/v(t)/pc(t)
         dfdPC=lambda t,f=f,pc=fpc:-f(t)/pc(t)
         dPC=function.generate(dfdPC,DPC)
         dfdQ=lambda t,f=f,q=fq: f(t)/q(t)
         dQ=function.generate(dfdQ,DQ)
         dfdV=lambda t,f=f,v=fv: -f(t)/v(t)
         dV=function.generate(dfdV,DV)
         return function.Object(f,[dPC,dQ,dV])
      else:
         f=sign*q/v/pc
         return derivedObject(sign*q/v/pc,\
            [{'df':-f/pc,'D':DPC},\
            {'df':sign/v/pc,'D':DQ},\
            {'df':-f/v,'D':DV}])
      #derivatives is the combination of the above
            

   def getVolumePar(self,compartment,useVolume=1):
      #returnis volume name, if found and useVolume is directed, 
      #or a standard parameter one
      if not useVolume:
         return "one"
      try:
         return self.mod["volumes"][compartment]
         #parV=self.mod["parameters"][parVName]
      except KeyError:
         pass

      return "one"


   def build(self):
      comps=self.compartments
      self.n=len(comps)
      self.fu=[lambda t:0]*self.n
      self.lut={c:i for (i,c) in zip(range(self.n),comps.keys())}
      self.dM={}
      self.fM=numpy.zeros((self.n,self.n))
      for c in comps:
         comp=comps[c]
         if 'source' in comp:
             self.fu[self.lut[c]]=parseFunction(comp['source'])
                            
         for t in comp['targets']:
            arr=comp['targets'][t]
            
            if function.contains(arr):
               try:
                  self.dM[self.lut[c]][self.lut[t]]=\
                     function.sumArray(arr)
               except (KeyError,TypeError):
                  self.dM[self.lut[c]]={}
                  self.dM[self.lut[c]][self.lut[t]]=\
                     function.sumArray(arr)
            else:
               #just set once
               self.fM[self.lut[c],self.lut[t]]=sum(arr)

      #build SE part
      self.buildSE()

   def buildSE(self):
      #check which parameterst to include
      parList=[]
      pars=self.parSetup['parameters']
      for parName in self.seJ:
         #print(par)
         
         par=pars[parName]
         usePar=calculateDerivative(par)
         #print('[{}]: {}'.format(usePar,par))
         if not usePar:
            continue
         parList.append(parName)
      
      #print(parList)
      self.m=len(parList)
      self.lutSE={c:i for (i,c) in zip(range(self.m),parList)}
      w=self.getWeights(self.lutSE)
      w=numpy.sqrt(w)

      self.qSS={}
      self.SS=numpy.zeros((self.m,self.n,self.n))
      for parName in parList:
         sources=self.seJ[parName]
         for compartment in sources:
            targets=sources[compartment]
            for t in targets:
               k=self.lutSE[parName]
               i=self.lut[compartment]
               j=self.lut[t]
               #print('[{} {} {}] {}'.format(parName,compartment,t,targets[t]))
               arr=targets[t]
               if not function.contains(arr):
                  self.SS[k,i,j]=w[k]*sum(arr)
                  continue
               ft=function.sumArray(arr,w[k])
               try:
                  self.qSS[k][i][j]=ft
               except (KeyError,TypeError):
                  try:
                     self.qSS[k][i]={}
                     self.qSS[k][i][j]=ft
                  except (KeyError,TypeError):
                     self.qSS[k]={}
                     self.qSS[k][i]={}
                     self.qSS[k][i][j]=ft


                 #use fM to build static part of fJ
      N=self.n*(self.m+1)
      self.fJ=numpy.zeros((N,N))
      for i in range(self.m+1):
         self.fJ[i*self.n:(i+1)*self.n,i*self.n:(i+1)*self.n]=self.fM
    
   


   def inspect(self):
      comps=self.compartments
      pars=self.parSetup['parameters']
      #pars=self.mod['parameters']
      
      tu=self.getTimeUnit()
      print('Time unit: {}'.format(tu))
      print('Compartments')
      for c in comps:
         print('{}/{}:'.format(c,self.lut[c]))
         comp=comps[c]
         if 'source' in comp:
             print('\tsource\n\t\t{}'.format(comp['source']))
         print('\ttargets')
         for t in comp['targets']:
             print('\t\t{}[{},{}]: {}'.format(t,self.lut[c],self.lut[t],\
               comp['targets'][t]))
      print('Flows')
      for f in self.flows:
         fName=self.flows[f]
         fParName=self.mod['flows'][fName]
         fPar=pars[fParName]
         print('\t{}[{}]:{} [{}]'.format(f,fName,fParName,self.get(fParName)))

      print('Volumes')
      for v in self.mod['volumes']:
         vParName=self.mod['volumes'][v]
         vPar=pars[vParName]
         print('\t{}:{} [{}]'.format(v,vParName,self.get(vParName)))

      print('Partition coefficients')
      for pc in self.mod['partitionCoefficients']:
         pcParName=self.mod['partitionCoefficients'][pc]
         pcPar=pars[pcParName]
         print('\t{}:{} [{}]'.format(pc,pcParName,self.get(pcParName)))


      print('SE parameters')
      for p in self.seJ:
         print(p)
         sources=self.seJ[p]
         for compartment in sources:
             targets=sources[compartment]
             for t in targets:
                 print('\t SE bind {}/{}:{}'.format(compartment,t,targets[t]))
    
   def parse(self,setupFile,parameterFile):
                    
      with open(setupFile,'r') as f:
         self.mod=json.load(f)

      
      with open(parameterFile,'r') as f:
         self.parSetup=json.load(f)

      for m in self.mod['compartments']:
         self.add_compartment(m)

      self.add_default_parameters()
      #standard parameters such as one,zero etc.
      for s in self.mod['sources']:
         #src=mod['sources'][s]
         self.add_source(s,self.mod['sources'][s])
      self.flows={}
      #pars=self.mod['parameters']
      pars=self.parSetup['parameters']
      for f in self.mod['flows']:
         #skip comments
         if f.find(':')<0:
             continue
         
         comps=f.split(':')
         c0=splitVector(comps[0])
         c1=splitVector(comps[1])
         for x in c0:
             for y in c1:
                 pairName='{}:{}'.format(x,y)
                 self.flows[pairName]=f
                 
      for b in self.mod['bindings']['diffusion']:
         #whether to scale transfer constants to organ volume
         #default is true, but changing here will assume no scaling
         useVolume=1
         comps=b.split('->')
         try:
             pcParName=self.mod['partitionCoefficients'][b]
         except KeyError:
             pcParName="one"
        
         kParName=self.mod['bindings']['diffusion'][b]
         #operate with names to allow for value/function/derived infrastructure
         self.bind(comps[0],comps[1],kParName,pcParName,useVolume)
         
      for q in self.mod['bindings']['flow']:
         comps=q.split('->')
         srcs=splitVector(comps[0])
         tgts=splitVector(comps[1])
         for cs in srcs:
             for ct in tgts:
                 #get partition coefficient
                 try:
                     pcParName=self.mod['partitionCoefficients'][cs]
                 except KeyError:
                     pcParName="one"
                 
                 #get flow (direction could be reversed)
                 try:
                     qName=self.flows['{}:{}'.format(cs,ct)]
                 except KeyError:
                     qName=self.flows['{}:{}'.format(ct,cs)]
                 
                 flowParName=self.mod['flows'][qName]
                 #flowPar=pars[flowParName]
                 
                 self.bind(cs,ct,flowParName,pcParName,1)
                 
      self.build()
   
   def add_default_parameters(self):
      pars=self.parSetup['parameters']
      pars['one']={'value':1}
      pars['zero']={'value':0}
   

   def M(self,t):
      for i in self.dM:
         for j in self.dM[i]:
            self.fM[i,j]=self.dM[i][j](t)
      #create an array and fill it with outputs of function at t
      return self.fM

   def u(self,t):
      ub=[f(t) for f in self.fu]
      return numpy.array(ub)

   def jacobiFull(self,t):

      #update jacobi created during build phase with time dependent values
      for i in self.dM:
         for j in self.dM[i]:
            for k in range(system.m+1):
               self.fJ[k*system.n+i,k*system.n+j]=self.dM[i][j](t)
      return self.fJ



   def fSS(self,t):
      for k in self.qSS:
         for i in self.qSS[k]:
            for j in self.qSS[k][i]:
               #print('[{},{},{}] {}'.format(k,i,j,self.qSS[k][i][j]))
               self.SS[k,i,j]=(self.qSS[k][i][j])(t)
      return self.SS
 
                     
   def fSY(self,y,t):
      #M number of sensitivity parameters
      #N number of equations
      #fSS is MxNxN

      #assume a tabulated solution y(t) at t spaced intervals

      qS=self.fSS(t).dot(y)
      #qS is MxN
      #but NxM is expected, so do a transpose

      #for simultaneous calculation, a Nx(M+1) matrix is expected
      tS=numpy.zeros((self.n,self.m+1))
      #columns from 2..M+1 are the partial derivatives 
      tS[:,1:]=numpy.transpose(qS)
      #first column is the original function
      tS[:,0]=self.u(t)
      return tS
    
   def fS(self,t):
   #M number of sensitivity parameters
   #N number of equations
   #fSS is MxNxN
        
   #assume a tabulated solution y(t) at t spaced intervals
        
      qS=self.fSS(t).dot(self.getY(t))
      return numpy.transpose(qS)
                     
   def getSEJ(self,parName):
      #find the sensitivity (SE) derivative of Jacobi with 
      #respect to parameter  
         try:
            return self.seJ[parName]
         except KeyError:
            self.seJ[parName]={}
            return self.seJ[parName]
    
   def getSEJ_comp(self,parName,compartmentName):
      #find equation dictating concentration in compartmentName 
      #for jacobi-parameter derivative
      seJ=self.getSEJ(parName)

      try:
         return seJ[compartmentName]
      except KeyError:
         seJ[compartmentName]={}
         return seJ[compartmentName]

   def setY(self,t,y):
      self.tck=[None]*self.n
      for i in range(self.n):
         self.tck[i] = scipy.interpolate.splrep(t, y[:,i], s=0)
    
   def getY(self,t):
      fY=numpy.zeros(self.n)
      for i in range(self.n):
         fY[i]=scipy.interpolate.splev(t, self.tck[i], der=0)
      return fY
   
   def getWeight(self,parName):
      pars=self.parSetup['parameters']
      par=pars[parName]
      #self.get parses the units
      v=self.get(parName)["value"]
      if par['dist']=='lognormal':
         #this is sigma^2_lnx
         sln2=numpy.log(par["cv"]*par["cv"]+1)
         #have to multiplied by value to get the derivative 
         #with respect to lnx
         return sln2*v*v
      else:
         #for Gaussian, cv is sigma/value; get sigma by value multiplication
         return par["cv"]*par["cv"]*v*v
      
   def getMax(lutSE):
      fm=-1
      for x in lutSE:
         if int(lutSE[x])>fm:
            fm=lutSE[x]
      return fm

   def getWeights(self,lutSE):
      #pars=self.parSetup['parameters']
      wts=numpy.zeros((model.getMax(lutSE)+1))
      for parName in lutSE:
         j=lutSE[parName]
         wts[j]=self.getWeight(parName)
      return wts

   def getDerivatives(self,se,i):
      #return latest point derivatives
      fse=se[-1][i]
      #fse is an m-vector
      return fse*fse

  

   def calculateUncertainty(self,se):
      
      s2out=numpy.zeros(se.shape[1:])
      se2=numpy.multiply(se,se)
      #w=self.getWeights(self.lutSE)
      w=numpy.ones((self.m))
      return numpy.sqrt(numpy.dot(se2,w))


   def get(self,parName):
      pars=self.parSetup['parameters']
      par=pars[parName]
      par['name']=parName
      if "value" in par:
         return self.getValue(par)
      if "function" in par:
         return self.getFunction(par)
      if "derived" in par:
         return self.getDerived(par)
      print('Paramter {} not found!'.format(parName))

   def getValue(self,par):

      v=par["value"]
      parName=par['name']
      #convert to seconds
      try:
         parUnits=par['unit'].split('/')
      except (KeyError,IndexError):
         #no unit given
         return valueObject(v,parName)
     
      timeUnit=self.getTimeUnit()

      try:
         if parUnits[1]==timeUnit:
            return valueObject(v,parName)
      except IndexError:
         #no / in unit name
         return valueObject(v,parName)

      if parUnits[1]=='min' and timeUnit=='s':
         return valueObject(v/60,parName)
      
      if parUnits[1]=='s' and timeUnit=='min':
         return valueObject(60*v,parName)

      if parUnits[1]=='day' and timeUnit=='min':
         return valueObject(v/24/60,parName)


      if parUnits[1]=='hour' and timeUnit=='min':
         return valueObject(v/60,parName)
      #no idea what to do
      return valueObject(v,parName)

   def getFunction(self,par):
      fName=par['function']
      #print('[{}]: getFunction({})'.format(par['name'],par['function']))
      df=self.parSetup['functions'][fName]
      skip=['type']
      par1={x:self.get(df[x]) for x in df if x not in skip}
      if df['type']=='linearGrowth':
         #print(par1)
         return function.linearGrowth(par1)
      if df['type']=='linearGrowthFixedSlope':
         return function.linearGrowthFixedSlope(par1)
         
      print('Function {} not found!'.format(fName))

   def getDerived(self,par):
      dName=par['derived']
      d=self.parSetup['derivedParameters'][dName]
      #print('Derived [{}]: type {}'.format(dName,d['type']))
      if d['type']=='product':
         #print('{}*{}'.format(d['a'],d['b']))
         pA=self.get(d['a'])
         a=pA['value']
         DA=pA['derivatives']
         pB=self.get(d['b'])
         b=pB['value']
         DB=pB['derivatives']
         #even more generic -> df/dp=[df/dA*dA/dp+df/dB*dB/dp]
         if any(['function' in pA,'function' in pB]):
            fa=function.to(a)
            fb=function.to(b)
            f=lambda t,a=fa,b=fb:a(t)*b(t)
            dfdA=lambda t,b=fb: b(t)
            dfdB=lambda t,a=fa: a(t)
            dA=function.generate(dfdA,DA)
            dB=function.generate(dfdB,DB)
            return function.Object(f,[dA,dB])
         else:
            return derivedObject(a*b,[{'df':b,'D':DA},{'df':a,'D':DB}])

      if d['type']=='power':
         #print('{}^{}'.format(d['a'],d['n']))
         pA=self.get(d['a'])
         a=pA['value']
         DA=pA['derivatives']
         pN=self.get(d['n'])
         n=pN['value']
         DN=pN['derivatives']
         if any(['function' in pA,'function' in pN]):
            fa=function.to(a)
            fn=function.to(n)
            f=lambda t,a=fa,n=fn:numpy.power(a(t),n(t))
            dfdA=lambda t,n=fn,f=f,a=fa:n(t)*f(t)/a(t)
            dfdN=lambda t,a=fa,f=f:numpy.log(a(t))*f(t)
            dA=function.generate(dfdA,DA)
            dN=function.generate(dfdN,DN)
            return function.Object(f,[dA,dN])
         else:
            f=numpy.power(a,n)
            return derivedObject(f,[{'df':n*f/a,'D':DA},{'df':f*numpy.log(a),'D':DN}])
      
      if d['type']=='ratio':
         #print('{}/{}'.format(d['a'],d['b']))
         pA=self.get(d['a'])
         a=pA['value']
         DA=pA['derivatives']
         pB=self.get(d['b'])
         b=pB['value']
         DB=pB['derivatives']
         #even more generic -> df/dp=[df/dA*dA/dp+df/dB*dB/dp]
         if any(['function' in pA,'function' in pB]):
            fa=function.to(a)
            fb=function.to(b)
            f=lambda t,a=fa,b=fb,:a(t)/b(t)
            dfdA=lambda t,f=f,a=fa: f(t)/a(t)
            dfdB=lambda t,f=f,b=fb: -f(t)/b(t)
            dA=function.generate(dfdA,DA)
            dB=function.generate(dfdB,DB)
            return function.Object(f,[dA,dB])
         else:
            return derivedObject(a/b,[{'df':1/b,'D':DA},{'df':-a/b/b,'D':DB}])
      if d['type']=='sum':
         #print('{}+{}'.format(d['a'],d['b']))
         pA=self.get(d['a'])
         a=pA['value']
         DA=pA['derivatives']
         pB=self.get(d['b'])
         b=pB['value']
         DB=pB['derivatives']
         #even more generic -> df/dp=[df/dA*dA/dp+df/dB*dB/dp]
         if any(['function' in pA,'function' in pB]):
            fa=function.to(a)
            fb=function.to(b)
            f=lambda t,a=fa,b=fb,:a(t)+b(t)
            dfdA=lambda t: 1
            dfdB=lambda t: 1
            dA=function.generate(dfdA,DA)
            dB=function.generate(dfdB,DB)
            return function.Object(f,[dA,dB])
         else:
            return derivedObject(a+b,[{'df':1,'D':DA},{'df':1,'D':DB}])


def calculateDerivative(par):
   #add derivatives if dist(short for distribution) is specified
   return "dist" in par    

def sumValues(dArray,x):
   s=0
   for a in dArray:
      try:
         s+=a[x]
      except KeyError:
         continue
   return s


def valueObject(v,parName):
   #convert everything to functions
   d0={parName:1}
   return {'value':v,'derivatives':{parName:1}}

def derivedObject(f,ders):
   o={'value':f}
   DD=[]
   for d in ders:
      df=d['df']
      D=d['D']
      DD.append({x:df*D[x] for x in D})
   allKeys=[]
   for x in DD:
      allKeys.extend(x.keys()) 
   allKeys=list(set(allKeys))
   o['derivatives']={x:sumValues(DD,x) for x in allKeys}
   return o
 

def splitVector(v):
   if v.find('(')<0:
      return [v]
   return v[1:-1].split(',')

def parseFunction(formula):
   if formula['name']=='exponential':
      c0=formula['constant']
      k=formula['k']
      return lambda t,c=c0,k=k:c*numpy.exp(k*t)
   if formula['name']=='constant':
      c0=formula['value']
      return lambda t,c0=c0:c0
   if formula['name']=='Heavyside':
      t0=formula['limit']
      v=formula['value']
      return lambda t,v=v,t0=t0:v if t<t0 else 0
   return lambda t:1

def addValue(qdict,compName,v):
   #add function to compName of dictionary qdict, 
   #check if compName exists and handle the potential error
   #lambda functions can't be summed directly, so qdict is a list
   #that will be merged at matrix generation time
   try:
      qdict[compName].append(v)
   except KeyError:
      qdict[compName]=[v]

   #also add derivatives
   #
   #   for d in dTarget:
   #      ctarget=self.getSEJ_comp(d,target)
   #      addValue(ctarget,target,dTarget[d])




def get(timeUnit,par):
   v=par["value"]
#convert to seconds
   try:
      parUnits=par['unit'].split('/')
   except (KeyError,IndexError):
      #no unit given
      return v
   
   try:
      if parUnits[1]==timeUnit:
         return v
   except IndexError:
      #no / in unit name
      return v
   if parUnits[1]=='min' and timeUnit=='s':
      return v/60
   
   if parUnits[1]=='s' and timeUnit=='min':
      return 60*v

   #no idea what to do
   return v