The Code
So i have the following function that just gathers a system of equations so it can be input into the standard solvers of scipy
def EDP7 (v):
T_N = np.reshape(v, (N_points_1, N_points_2)) #Variable declaration
#Model
sys = np.empty_like(v)
l=0
##Boundary Condition x = 0
for i in range(N_points_1):
BC1 = T_N[i,0] - T_x_0
sys[l] = BC1
l+=1
##Boundary Condition x = L
for i in range(N_points_1):
BC2 = T_N[i,-1] - T_x_L
sys[l] = BC2
l+=1
##Initial Condition
for j in range(1,N_points_2-1):
IC1 = T_N[0,j] - T_t_0
sys[l] = IC1
l+=1
## Energy Balance
for i in range(N_points_1-1):
for j in range (1,N_points_2 -1):
EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt
sys[l] = EB
l+=1
return sys
Parameters here:
t_max = 25
x_max = 1
N_points_1 = 1000000
N_points_2 = 11
dt = t_max/(N_points_1-1)
dx = x_max/(N_points_2-1)
N_variables = N_points_1 * N_points_2
alpha = 0.01
T_x_0 = 60
T_x_L = 25
T_t_0 = 25
initial_guess = T_t_0*np.ones((N_points_1,N_points_2))
initial_guess[:,0] = T_x_0
initial_guess[:,-1] = T_x_L
initial_guess = initial_guess.flatten()
and i found out for problems with large number of variables(what i am interested in), like N=1e6-1e7, there is a significant time spent simply creating the system of equation array
sys = np.empty_like(v)
About half the time of the function spent on it(tests below).
What i Tried
Well i figured out i could simply declare the vector outside and pass it into the function as a parameter
syst = np.empty_like(initial_guess)
@njit(fastmath=True)
def EDP8 (v,system):
sys = system
T_N = np.reshape(v, (N_points_1, N_points_2))
...
32.3 ms ± 245 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) -Declaration inside
21.9 ms ± 91.6 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) - Declaration outside
Improvement of about 33%
However, when passed into the scipy solvers, the function refuses to leave the initial guess position
Solution_T7 = optimize.least_squares(EDP8,initial_guess, args = [syst], verbose = 2)
Iteration Total nfev Cost Cost reduction Step norm Optimality
0 1 3.0625e+04 0.00e+00
`gtol` termination condition is satisfied.
Function evaluations 1, initial cost 3.0625e+04, final cost 3.0625e+04, first-order optimality 0.00e+00.
print(Solution_7.x)
[60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25.
25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25.
25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25.
25. 60. 25. 25. 25. 25. 25. 25. 25. 25. 25. 25. 60. 25. 25. 25. 25. 25.
...]
This is the initial_guess array that i declared. I suppose that something in the machinations of scipy(i tested it with many other solvers with the same results) prevents the iterations with the syst variable declared outside, maybe because sys becomes a global variable? Also, this is not caused by the njit decorator, the same thing happens without it.
In any case, is there any way to remove the array declaration bottleneck of the original function?
Edits
This is the original full code
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from scipy.optimize import fsolve
from scipy import optimize
import pandas as pd
from numba import njit
import time
t_max = 25
x_max = 1
N_points_1 = 1000000
N_points_2 = 11
dt = t_max/(N_points_1-1)
dx = x_max/(N_points_2-1)
N_variables = N_points_1 * N_points_2
alpha = 0.01
T_x_0 = 60
T_x_L = 25
T_t_0 = 25
initial_guess = T_t_0*np.ones((N_points_1,N_points_2))
initial_guess[:,0] = T_x_0
initial_guess[:,-1] = T_x_L
initial_guess = initial_guess.flatten()
@njit(fastmath=True)
def EDP7 (v):
T_N = np.reshape(v, (N_points_1, N_points_2)) #Variable declaration
#Model
sys = np.empty_like(v)
l=0
##Boundary Condition x = 0
for i in range(N_points_1):
BC1 = T_N[i,0] - T_x_0
sys[l] = BC1
l+=1
##Boundary Condition x = L
for i in range(N_points_1):
BC2 = T_N[i,-1] - T_x_L
sys[l] = BC2
l+=1
##Initial Condition
for j in range(1,N_points_2-1):
IC1 = T_N[0,j] - T_t_0
sys[l] = IC1
l+=1
## Energy Balance
for i in range(N_points_1-1):
for j in range (1,N_points_2 -1):
EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt #partial(T_N,t_1,1)[i,j] - alpha*partial_2(T_N,x,2)[i,j]
sys[l] = EB
l+=1
return sys
compile_1 = EDP7(initial_guess)
And the modified one
syst = np.empty_like(initial_guess)
@njit(fastmath=True)
def EDP8 (v,system):
sys = system
T_N = np.reshape(v, (N_points_1, N_points_2))
#Model
l=0
##Boundary Condition x = 0
for i in range(N_points_1):
BC1 = T_N[i,0] - T_x_0
sys[l] = BC1
l+=1
##Boundary Condition x = L
for i in range(N_points_1):
BC2 = T_N[i,-1] - T_x_L
sys[l] = BC2
l+=1
##Initial Condition
for j in range(1,N_points_2-1):
IC1 = T_N[0,j] - T_t_0
sys[l] = IC1
l+=1
## Energy Balance
for i in range(N_points_1-1):
for j in range (1,N_points_2 -1):
EB = alpha*(T_N[i,j+1]-2*T_N[i,j]+T_N[i,j-1])/dx**2 - (T_N[i+1,j]-T_N[i,j])/dt #partial(T_N,t_1,1)[i,j] - alpha*partial_2(T_N,x,2)[i,j]
sys[l] = EB
l+=1
return sys
compile_2 = EDP8(initial_guess,syst)
And the tests that i just now, did.
%timeit EDP7(initial_guess)
%timeit EDP8(initial_guess,syst)
47.2 ms ± 1.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
23.3 ms ± 633 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
This is already a bigger difference than the OP. Since the only difference between these codes is the outside declaration, i suppose it was because of it...