Demo of numpy einsum matrix multiplication / summation functionality

einsum_demo.py

import numpy as np

# dimensions (age,metallicity,wavelength)
a = 2
Z = 3
wl = 4

weights = np.random.rand(a,Z,1)
grid = np.random.randint(1,10,(a,Z,wl))

print(weights.shape)
print(grid.shape)

def grid_sum(grid,weights):
    spectra = grid * weights
    return np.sum(spectra,axis=(0,1))

summed_spectra = grid_sum(grid,weights)
print(summed_spectra)

summed_spectra = np.einsum('ijk,ijl->k',grid,weights)
print(summed_spectra)


%timeit grid_sum(grid,weights)
%timeit np.einsum('ijk,ijl->k',grid,weights)


## scale up the dimensions to test scaling
a = 10
Z = 20
wl = 6000

weights = np.random.rand(a,Z,1)
grid = np.random.randint(1,10,(a,Z,wl))

%timeit grid_sum(grid,weights)
%timeit np.einsum('ijk,ijl->k',grid,weights)

# can also do this for *multiple* object simultaneously
a = 2
Z = 3
wl = 4
N = 100
weights = np.random.rand(a,Z,N)
grid = np.random.randint(1,10,(a,Z,wl))

summed_spectra = np.einsum('ijk,ijl->kl',grid,weights)
print(summed_spectra)