pre-calculating and storing arrays

Say we have a grid, and a function that maps the grid position to a given value (or vector). For example,

$$f(x,y) = ax^3e^{-by},$$

where $a,b$ are constants. And let the grid run from $0$ to $L_x,L_y$.

Something we can do, if $x=idx$ and $y=jdy, is

$$f(x,y) = adx^3i^3e^{-bdyj} = Ai^3B^j,$$

where $A=adx^3$ and $B=e^{-b}dy$. The idea is avoid calculating the same thing more than once.

Another function could be

$$g(x,y) = ax^3 + e^{-by} = Ai^3 + B^j.$$

In this notebook, we propose a methodology that pays off when the calculation of the arrays is expensive.

import numpy as np
import os

Nx, Ny = 200, 100
dx, dy = 0.01, 0.01
a = 12; b = 11;
A = a * dx; B = np.exp(-b)

For function $f$, we pre-calculate the 2D array:

# define the function
def f(ix,iy):
    return A*ix**3*B**iy

# store the array on disk (with a and b values on the file name)
filename = f'array_f_{a}_{b}.npz'

if not os.path.exists(filename): # if file does not exist, create it:
    # calculate the array
    print("calculating array...")
    array_f = np.empty((Nx,Ny))
    for i in range(Nx):
        for j in range(Ny):
            array_f[i][j] = f(i,j)
    print("storing array...")
    np.savez_compressed(filename, array_f=array_f)
else: # if it exists, load it
    print("array exists on disk, loading...")
    data = np.load(filename)
    array_f = data['array_f']

# redefine the function
def f(ix,iy):
    return array_f[ix,iy]

array exists on disk, loading...

In the case of g function, notice how the $x$ and $y$ parts are additive. It is better to calculate them separately and store them in 1D arrays.

# define the function
def gx(ix):
    return A*ix**3
def gy(iy):
    return B**iy

# store the arrays on disk (with a and b values on the file name)
filename2x = f'array_gx_{a}_{b}.npz'
filename2y = f'array_gy_{a}_{b}.npz'

if not os.path.exists(filename2x) and not os.path.exists(filename2y): # if files do not exist, create them:
    # calculate the arrays
    print("calculating arrays...")
    array_gx = np.empty(Nx)
    array_gy = np.empty(Ny)
    for i in range(Nx):
        array_gx[i] = gx(i)
    for i in range(Ny):
        array_gy[i] = gy(i)
        
    print("storing arrays...")
    np.savez_compressed(filename2x, array_gx=array_gx)
    np.savez_compressed(filename2y, array_gy=array_gy)
else: # if they exist, load them
    print("arrays exist on disk, loading...")
    data2x = np.load(filename2x)
    data2y = np.load(filename2y)
    array_gx = data2x['array_gx']
    array_gy = data2y['array_gy']

def g(ix,iy):
    return array_gx[ix] + array_gy[iy]

arrays exist on disk, loading...

print(f(20,30))
print(g(20,30))

4.624791998345352e-141
960.0