将二维和一维数组 NumPy 相乘

Question

I'm somewhat new to python and numpy, so maybe some of you can help me here.

I have a 1D numpy array called z, and some 2D matrices called X0 , Y0, and SLM.

I want to create a 3D array (a stack of 2D matrices), by doing this operation, trying to avoid a for loop:

for index in range(len(z)):
    3D_array[index] = SLM * np.exp( z[index] * (X0+Y0) )

So far I've tried:

3D_array = SLM[np.newaxis, :] * np.exp( z[:, np.newaxis, np.newaxis] * (X0+Y0))

This performs right, but its slow for my purposes. The arrays are big (z is size 200, and all 2D matrices are 1024x1024)

Do you know any faster implementation for this?

Answer 1

Good question, what you're looking for is np.tensordot , see example below:

import numpy as np

x0 = np.floor(np.random.rand(3, 3)*10)
y0 = np.floor(np.random.rand(3, 3)*10)

z = np.array([1, 2, 3])

SLM = 1

array = SLM*np.exp(np.tensordot(z, x0+y0, axes=0))

Answer 2

# Assuming SLM contains only positive values this is the fastest way 
%timeit np.log(slm[np.newaxis, :, :] + 1e-9) + z[:, np.newaxis, np.newaxis] * (x + y )
1.67 s ± 11 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

 %timeit slm[np.newaxis, :, :] * np.exp(z[:, np.newaxis, np.newaxis] * (x + y ))
4.15 s ± 689 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit slm[np.newaxis, :, :] * np.exp(np.tensordot(z, x + y, axes=0))
3.9 s ± 43.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)