Optimize system of equation function organizer with memory bottleneck in python

Question

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...

Answer 1

The original post asserts that

    sys = np.empty_like(v)

takes significant time.

I find no evidence for that, in the OP nor in these tests: (See new timing results below.)

$ python -m timeit \
   -s 'import numpy as np; shp = (1000, 1000, 1000); v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
20000 loops, best of 5: 19.2 usec per loop
$
$ python -m timeit \
   -s 'import numpy as np; shp = (10,) * 9; v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
10000 loops, best of 5: 20.4 usec per loop
$
$ python -m timeit \
   -s 'import numpy as np; shp = (2,) * 29; v = np.ones(shp)' \
      'assert shp == np.empty_like(v).shape'
20000 loops, best of 5: 11 usec per loop

Bumping that last one up to 2 ** 30 takes us to "22 usec per loop", which is due to beginning to page to backing store on this 8 GiB laptop rather than a reflection on the BLAS / numpy routines.

The .empty_like(...) operation appears to take negligible time, even when operating on a billion entries. That is, empirically it takes O(1) time, independent of size of v.

---

I have explored a subset of the design space, and documented what I found. It is a large space. I encourage the sharing of additional observations.

---

EDIT

Based on cProfile describing the nested loops as using the bulk of the cycles, I immediately broke this out:

@njit(fastmath=True)
def energy_balance(t_n, sys, l):
    ...

I wanted to use timeit.timeit(), but alas that's not compatible with JIT'ing. So I have this at the top:

@njit(fastmath=True)
def el(v):
    return np.empty_like(v)


# @njit(fastmath=True)
def edp7(v):
    t_n = np.reshape(v, (n_points_1, n_points_2))  # Variable declaration

    # Model
    sys = el(v)
    assert len(sys) == len(v) == 11_000_000
    print("")
    number = 10_000
    elapsed = timeit(lambda: np.empty_like(v), number=number)
    print(f"elapsed: {elapsed:.3f} s, {elapsed/number:.6f} s/loop")

If you adjust the lambda, you'll see the elapsed time explodes 100x when we call the trivial el wrapper instead of empty_like directly. But calling the wrapper is "safe" (fast) as long as it's not JITed.

Also, here's a second clue. I wound up with this at the bottom:

if __name__ == "__main__":
    edp7(initial_guess.copy())

    t0 = time.time()
    compile_1 = edp7(initial_guess)
    print(f"total:  {time.time() - t0}:.3f s")

The idea was to give numba a chance to JIT prior to the timing run. But I seem to have done something terrible, as these timing figures show:

elapsed: 0.018 s, 0.000002 s/loop

elapsed: 42.895 s, 0.004290 s/loop
total:  43.702 s

The timeit figure just went to pieces the second time through.

OTOH if both calls make an initial_guess.copy() then things go smoothly:

elapsed: 0.020 s, 0.000002 s/loop

elapsed: 0.019 s, 0.000002 s/loop
total:  0.823 s

There's plenty of free RAM on the laptop during the run. I cannot explain this behavior. But it is possible to skirt around the edge of it and avoid it.

---

It doesn't really hurt to re-assign v like this

    v = np.float64(v)
    sys = np.empty_like(v)

and it helps numba to be quite definite about the inferred type. (Put another way, then the developer will be certain inference did the right thing.)

If you can get away with it, assign np.float32(v) to save on memory, which implies saving on memory bandwidth. Fewer bytes moved means less elapsed time.