convolve a higher dimensional cube with a 2d array

Question

I am convolving a 3D datacube of the shape [nLambda,nX,nY] with a filter of the shape [nLambda,3]. I managed to get this to work for this situation, but I need to repeat this process on a large amount of cubes that are saved as [nt,nLambda,nX,nY]. I want to extend the code that I have to cope with that, but I keep messing up the tiling. Does anyone know how to do this?

My current program does:

datacube.shape         = [21,100,100]
filters.shape          = [21,3]
data_collapsed.shape   = [100,100,3]   

I want it to do

datacube.shape         = [10,21,100,100]
filters.shape          = [21,3]
data_collapsed.shape   = [10,100,100,3]     

The code that I use for 3D cubes

nl,nx,ny = datacube.shape
filter_rgb = np.tile(filters, (ny,nx,1,1))
filter_rgb = np.swapaxes(filter_rgb, 0,2)
data_rgb = np.tile(datacube,(3,1,1,1))
data_rgb = np.swapaxes(data_rgb,0,-1)
data_filtered = data_rgb * filter_rgb
data_collapsed = np.sum(data_filtered, axis=0)

Answer 1

You can use np.tensordot.

For datacube as 4D case -

data_collapsed = np.tensordot(datacube,filters,axes=(1,0))

For 3D case -

data_collapsed = np.tensordot(datacube,filters,axes=(0,0))

Related post to understand tensordot.