I use jitclass to run faster my model. I find the compilation time quite long. On my computer it takes about 7-8s to compile the one I give below as a minimal example. I have other longer models which take above 20s to compile. I'm seeking some advice to accelerate the compilation time? Maybe I have written the jitclass the wrong way. I mean compare to some guvectorize functions I have, which are quite long, their compilation time is almost instantaneous.
import numpy as np
import random
from numba import jitclass
from numba import boolean, int32, float64,uint8
import time
spec = [('A' ,float64[:]),
('B' ,float64[:]),
('a' ,float64[:]),
('b',float64[:]),
('C1' ,float64[:]),
('C2' ,float64[:]),
('C3' ,float64[:]),
('C4' ,float64[:]),
('v0' ,float64[:]),
('e0',float64[:]),
('r',float64[:]),
('m',float64[:]),
('s',float64[:]),
('coef',float64[:]),
('dt',float64),
('NbODEs',int32),
('NbNMMs',int32),
('dydx' ,float64[:,:]),
('dydx1' ,float64[:,:]),
('dydx2' ,float64[:,:]),
('dydx3' ,float64[:,:]),
('y' ,float64[:,:]),
('yt' ,float64[:,:]),
('bruit',float64[:]),
('LFP',float64[:]),
('pulseOutputP',float64[:]),
('pulseOutputINGabaA',float64[:]),
('PPSEpyrP',float64[:]),
('PPSIgabaA',float64[:]),
('ExtI_P_to_P',float64[:]),
('ExtI_P_to_I',float64[:]),
('ExtI_I_to_P',float64[:]),
('ExtI_I_to_I',float64[:]),
('k_P',float64[:]),
('k_G',float64[:]),
('Stim',float64[:]),
('Pre_Post',int32),
('CM_P_to_P' ,float64[:,:]),
('CM_P_to_I' ,float64[:,:]),
('CM_I_to_P' ,float64[:,:]),
('CM_I_to_I' ,float64[:,:]),
('DelayMat' ,float64[:,:]),
('Delay_in_index_Mat' ,int32[:,:]),
('EmptyMat' ,boolean[:]),
('Nb_NMM_m1' ,int32),
('history_pulseP' ,float64[:,:]),
('history_pulseI' ,float64[:,:]),
('max_delay_index' ,int32),
('pulseP_delay_Mat' ,float64[:,:]),
('pulseI_delay_Mat' ,float64[:,:]),]
@jitclass(spec)
class pop_JansenRit:
"""Cellular and tissular model *Red3*"""
def __init__(self,):
self.dt = 1./2048.
self.NbODEs = 14
self.NbNMMs = 1
self.Nb_NMM_m1=1
self.init_vector()
self.init_vector_param()
self.Pre_Post = True
def init_vector(self):
self.dydx = np.zeros((self.NbODEs,self.NbNMMs))
self.dydx1 = np.zeros((self.NbODEs,self.NbNMMs))
self.dydx2 = np.zeros((self.NbODEs,self.NbNMMs))
self.dydx3 = np.zeros((self.NbODEs,self.NbNMMs))
self.y =np.zeros((self.NbODEs,self.NbNMMs))
self.yt =np.zeros((self.NbODEs,self.NbNMMs))
#external input
#to pyramide
self.ExtI_P_to_P = np.zeros((self.NbNMMs))
self.ExtI_P_to_I = np.zeros((self.NbNMMs))
#to GABA A slow
self.ExtI_I_to_P = np.zeros((self.NbNMMs))
self.ExtI_I_to_I = np.zeros((self.NbNMMs))
self.bruit = np.zeros((self.NbNMMs))
self.LFP=np.zeros((self.NbNMMs))
self.pulseOutputP=np.zeros((self.NbNMMs))
self.pulseOutputINGabaA=np.zeros((self.NbNMMs))
self.PPSEpyrP=np.zeros((self.NbNMMs))
self.PPSIgabaA=np.zeros((self.NbNMMs))
self.Stim = np.zeros((self.NbNMMs))
self.CM_P_to_P = np.zeros((self.NbNMMs, self.NbNMMs))
self.CM_P_to_I = np.zeros((self.NbNMMs, self.NbNMMs))
self.CM_I_to_P = np.zeros((self.NbNMMs, self.NbNMMs))
self.CM_I_to_I = np.zeros((self.NbNMMs, self.NbNMMs))
self.DelayMat = np.zeros((self.NbNMMs, self.NbNMMs))
self.Delay_in_index_Mat = np.zeros((self.NbNMMs, self.NbNMMs), dtype=np.int32)
self.EmptyMat = np.full((4), True)
self.history_pulseP = np.zeros((self.NbNMMs, 1))
self.history_pulseI = np.zeros((self.NbNMMs, 1))
self.max_delay_index = 0
self.pulseP_delay_Mat = np.zeros((self.NbNMMs, self.NbNMMs))
self.pulseI_delay_Mat = np.zeros((self.NbNMMs, self.NbNMMs))
def init_vector_param(self):
self.A = np.ones((self.NbNMMs)) * 3.25
self.B = np.ones((self.NbNMMs)) * 22.
self.a = np.ones((self.NbNMMs)) * 100.
self.b = np.ones((self.NbNMMs)) * 50.
self.C1 = np.ones((self.NbNMMs)) * 135.*1.
self.C2 = np.ones((self.NbNMMs)) * 135.*0.8
self.C3 = np.ones((self.NbNMMs)) * 135.*0.25
self.C4 = np.ones((self.NbNMMs)) * 135.*0.25
self.v0 = np.ones((self.NbNMMs)) * 6.
self.e0 = np.ones((self.NbNMMs)) * 5.
self.r = np.ones((self.NbNMMs)) * 0.56
self.m = np.ones((self.NbNMMs)) * 3.
self.s = np.ones((self.NbNMMs)) * 1.
self.coef = np.ones((self.NbNMMs)) * 30.
self.k_P = np.ones((self.NbNMMs)) * 1.
self.k_G = np.ones((self.NbNMMs)) * 1.
def random_seeded(self,seed):
random.seed(int(seed))
def sigm(self,v):
return self.e0/(1+np.exp( self.r*(self.v0-v)))
def noise(self):
for i in range(self.NbNMMs):
self.bruit[i] = self.coef[i] * np.random.normal(self.m[i],self.s[i])
def PSP(self,y0,y1,y2,V,v):
return (V*v*y0 - 2*v*y2 - v*v*y1)
def rk4(self):
self.noise() #p(t) est le bruit d'entree dans le modele
self.yt = self.y+0. #y origine
self.dydx1=self.derivT()
self.y = self.yt + self.dydx1 * self.dt / 2
self.dydx2=self.derivT()
self.y = self.yt + self.dydx2 * self.dt / 2
self.dydx3=self.derivT()
self.y = self.yt + self.dydx3 * self.dt
self.derivT()
self.y =self.yt + self.dt/6. *(self.dydx1+2*self.dydx2+2*self.dydx3+self.dydx)
def mp(self):
self.noise() #p(t) est le bruit d'entree dans le modele
self.yt = self.y+0. #y origine
self.dydx1=self.derivT()
self.y = self.yt + self.dydx1 * self.dt / 2
self.derivT()
self.y = self.yt + self.dydx * self.dt
def Eul(self):
self.noise()
self.dydx1=self.derivT()
self.y += (self.dydx1 * self.dt)
def Eul_Time(self,N):
self.init_vector()
lfp = np.zeros((N,self.NbNMMs))
for k in range(N):
self.rk4()
lfp[k,:]= self.LFP
return lfp
def derivT(self , ):
# Equation Pyramides P#
self.LFP = self.y[1,:]-self.y[2,:]
self.pulseOutputP= self.sigm(self.k_P * self.Stim * self.Pre_Post +
self.LFP +
self.y[6,:] - self.y[8,:] )
self.dydx[0,:] = self.y[3,:]
self.dydx[3,:] = self.PSP(self.k_P * self.Stim * (not self.Pre_Post) + self.pulseOutputP,self.y[0,:],self.y[3,:],self.A,self.a)
# Equation Pyramides P'#
self.dydx[1,:] = self.y[4,:]
self.dydx[4,:] =self.PSP(self.k_P * self.Stim * (not self.Pre_Post) + self.bruit +self.C2*self.sigm(self.k_P * self.Stim * self.Pre_Post + self.C1* self.y[0,:]),self.y[1,:],self.y[4,:],self.A,self.a)
# Equations IN GABAa slow#
self.pulseOutputINGabaA = self.sigm(self.k_G * self.Stim * self.Pre_Post +
self.C3 * self.y[0,:]+
self.y[10,:] - self.y[12,:] )
self.dydx[2,:] = self.y[5,:]
self.dydx[5,:] = self.PSP(self.k_G * self.Stim * (not self.Pre_Post) + self.C4 *self.pulseOutputINGabaA,self.y[2,:],self.y[5,:],self.B,self.b)
self.dydx[6,:] = self.y[7,:]
self.dydx[7,:] = self.PSP(self.ExtI_P_to_P,self.y[6,:],self.y[7,:],self.A,self.a)
self.dydx[8,:] = self.y[9,:]
self.dydx[9,:] = self.PSP(self.ExtI_I_to_P,self.y[8,:],self.y[9,:],self.B,self.b)
self.dydx[10,:] = self.y[11,:]
self.dydx[11,:] = self.PSP(self.ExtI_P_to_I,self.y[10,:],self.y[11,:],self.A,self.a)
self.dydx[12,:] = self.y[13,:]
self.dydx[13,:] = self.PSP(self.ExtI_I_to_I,self.y[12,:],self.y[13,:],self.B,self.b)
self.PPSEpyrP =self.y[0,:]
self.PPSIgabaA=self.y[2,:]
return self.dydx+0.
def NonNullMat(self):
if np.max(self.CM_P_to_P)==0.:
self.EmptyMat[0]=False
else:
self.EmptyMat[0]=True
if np.max(self.CM_P_to_I)==0.:
self.EmptyMat[1]=False
else:
self.EmptyMat[1]=True
if np.max(self.CM_I_to_P)==0.:
self.EmptyMat[2]=False
else:
self.EmptyMat[2]=True
if np.max(self.CM_I_to_I)==0.:
self.EmptyMat[3]=False
else:
self.EmptyMat[3]=True
def apply_connectivity_Mat(self):
if self.EmptyMat[0]:
self.ExtI_P_to_P = np.dot(self.CM_P_to_P ,self.pulseOutputP)
if self.EmptyMat[1]:
self.ExtI_P_to_I = np.dot(self.CM_P_to_I ,self.pulseOutputP)
if self.EmptyMat[2]:
self.ExtI_I_to_P = np.dot(self.CM_I_to_P ,self.pulseOutputINGabaA)
if self.EmptyMat[3]:
self.ExtI_I_to_I = np.dot(self.CM_I_to_I ,self.pulseOutputINGabaA)
def apply_connectivity_Mat_delay(self):
if self.EmptyMat[0]:
self.ExtI_P_to_P =np.sum(self.CM_P_to_P *self.pulseP_delay_Mat,axis=1)
if self.EmptyMat[1]:
self.ExtI_P_to_I =np.sum(self.CM_P_to_I *self.pulseP_delay_Mat,axis=1)
if self.EmptyMat[2]:
self.ExtI_I_to_P =np.sum(self.CM_I_to_P *self.pulseI_delay_Mat,axis=1)
if self.EmptyMat[3]:
self.ExtI_I_to_I =np.sum(self.CM_I_to_I *self.pulseI_delay_Mat,axis=1)
def convert_delay_in_index(self):
self.Delay_in_index_Mat = (self.DelayMat / self.dt).astype(np.int32)
self.max_delay_index = np.max(self.Delay_in_index_Mat)+1
self.history_pulseP = np.zeros((self.NbNMMs, self.max_delay_index))
self.history_pulseI = np.zeros((self.NbNMMs, self.max_delay_index))
def compute_pulse_delayed(self):
if np.any(self.EmptyMat[:2]):
self.history_pulseP[:, 1:] = self.history_pulseP[:, 0:-1]
self.history_pulseP[:, 0]=self.pulseOutputP
for idx2 in range(self.NbNMMs):
for idx in range(self.NbNMMs):
self.pulseP_delay_Mat[idx2, idx] = self.history_pulseP[idx,self.Delay_in_index_Mat[idx2,idx]]
if np.any(self.EmptyMat[2:4]):
self.history_pulseI[:, 1:] = self.history_pulseI[:, 0:-1]
self.history_pulseI[:, 0] = self.pulseOutputINGabaA
for idx2 in range(self.NbNMMs):
for idx in range(self.NbNMMs):
self.pulseI_delay_Mat[idx2, idx] = self.history_pulseI[idx, self.Delay_in_index_Mat[idx2, idx]]
t0 = time.time()
pop = pop_JansenRit()
print(time.time() - t0)
from Why my jitclass class is quite long to compile?
No comments:
Post a Comment