2

I have a Python3 code. I want to parallelize a very long for loop, in which a function (solve_ivp) integrates a system of differential equations. The loop is over the initial conditions, contained in the rows of a matrix called "esemble".

I have already tried with Numba but the problem is that the function from scipy called "solve_ivp" is not supported by Numba.

Do you know how I can parallelize this code in order to make it faster?

N.B: I need to keep track of all the trajectory (as noted by t_eval=tt in the arguments of solve_ivp).

mm=100000 #number of initial conditions
m=500000  #interval of integration
step=0.0001  #integration step
N=36
F=8
#----------------------------------------
#differential equations to be integrated
def lorenz96(t,x):
"""Lorenz 96 model."""
# Compute state derivatives
d = np.zeros(N)

# First the 3 edge cases: i=1,2,N
d[0] = (x[1] - x[N-2]) * x[N-1] - x[0]
d[1] = (x[2] - x[N-1]) * x[0] - x[1]
d[N-1] = (x[0] - x[N-3]) * x[N-2] - x[N-1]
# Then the general case
for i in range(2, N-1):
    d[i] = (x[i+1] - x[i-2]) * x[i-1] - x[i]
# Add the forcing term
d = d + F               

# Return the state derivatives
return d
#----------------------------------------
#x0 part (initial transient)

x0 = np.zeros(N)
for i in range(0, N):
    x0[i]=np.random.uniform()*(2*0.4) 
esemble = np.zeros((mm,N))

#time lag iniziale
lag = int((50.0)/step)
tlag = np.arange(0.0, float(lag*step), step)    #time steps

b=0
x = solve_ivp(lorenz96,[0.0,float(lag*step)], x0, t_eval=tlag)

for j in range(0,N):
    x0[j]=x.y[j,lag-1] #end lag as CI
    esemble[0,j] = x.y[j,lag-1]

#---------------------------------------
#Creating the esemble of initial conditions from a long run of the differential system
m_es = int((10.0)/step) # distanza membri esemble
t_es = np.arange(0.0, float(m_es*step), step)

for ll in range(1,mm): 

    x = solve_ivp(lorenz96,[0.0,float(m_es*step)], x0, t_eval=t_es)
    for j in range(0,N):        
        x0[j]=x.y[j,m_es-1] #end as CI for next cycle
        esemble[ll,j] = x.y[j,m_es-1]

### Heading---This the loop to be parallelized ##

Gamma_psi=np.zeros((2,m))#array which will contain the result of the integration along the time interval (0, m*step)

for i in range(mm): #pick the ensemble member I

    xx0 = np.zeros(N) 

    for j in range(0,N):
        xx0[j]=esemble[i,j] #pick the start for each ensemble member I

    #FORWARD m STEPS
    tt = np.arange(0.0, float(m*step), step) #interval of integration   
    xx00 = solve_ivp(lorenz96,[0.0,float(m*step)], xx0, t_eval=tt)

    #saving result for just two dynamical variables
    loc=1
    for p in range(0,m):    #cycle over time steps 

        Gamma_psi[0,p] = Gamma_psi[0,p]+ (xx00.y[loc,p])
        Gamma_psi[1,p] = Gamma_psi[1,p]+ (xx00.y[loc-1,p] )

Gamma_psi = Gamma_psi/(float(mm)) #NB
#---------------------------------------------
Umberto Tomasini
  • 31
  • 3
  • 1
    You are completely right, I add now the Python tag! – Umberto Tomasini May 05 '20 at 09:09
  • 1
    Please provide a [minimal reproducible example](https://stackoverflow.com/help/minimal-reproducible-example). At the moment, it isn't obvious what the shape of `esemble` (ensemble?) is or how you use/store the result, `xx00`. – FiddleStix May 05 '20 at 09:50
  • I tried to add some other parts of the code in order to make it reproducible, I hope it is better now. Thank you very much – Umberto Tomasini May 05 '20 at 11:56
  • 1
    See https://stackoverflow.com/a/34301676/3088138. But check the usefulness of your parameters. A step size of 1e-4 for a standard test problem is likely too small, think about reducing it to 1e-3 or even 1e-2. Note that this step size of the sample point sequence does not influence the internal step size control, that is only controlled by the atol and rtol parameters. – Lutz Lehmann May 05 '20 at 11:56
  • Thank you very much, now I look into it. I resorted to such a small time step because I will Fourier transform these results, and I will need a high density of frequencies. To achieve that, I thought to reduce the time step. Do you have another method in mind to make the frequencies denser? – Umberto Tomasini May 05 '20 at 12:11
  • 1
    @UmbertoTomasini could you also add example values for esemble, N, x0, etc.? They don't have to be realistic, just enough to get the code to run. – FiddleStix May 05 '20 at 14:26
  • @FiddleStix, done it, thank you! Anyway I have implemented what was reported in the link by Lutz Lehmann (thank you very much), and it seems to me that it works! – Umberto Tomasini May 05 '20 at 15:25
  • 1
    You could also use Numba on your `lorenz96(t,x)` function. A low-level-callable isn't supported by solve_ivp but nevertheless this should provide a significant speedup. – max9111 May 05 '20 at 17:52
  • @max9111 thank you very much! It speeds up a bit indeed! – Umberto Tomasini May 06 '20 at 12:24

0 Answers0