Numpy 2D spatial mask to be filled with specific values from a 2D array to form a 3D structure

Question

I'm quite new to programming in general, but I could not figure this problem out until now.

I've got a two-dimensional numpy array mask, lets say mask.shape is (3800,3500)which is filled with 0s and 1s representing a spatial resolution of a 2D image, where a 1 represents a visible pixel and 0 represents background.
I've got a second two-dimensional array data of data.shape is (909,x) where x is exactly the amount of 1s in the first array. I now want to replace each 1 in the first array with a vector of length 909 from the second array. Resulting in a final 3D array of shape(3800,3500,909) which is basically a 2D x by y image where select pixels have a spectrum of 909 values in z direction.

I tried

mask_vector = mask.flatten
ones = np.ones((909,1))
mask_909 = mask_vector.dot(ones) #results in a 13300000 by 909 2d array
count = 0
for i in mask_vector:
    if i == 1:
        mask_909[i,:] = data[:,count]
        count += 1

result = mask_909.reshape((3800,3500,909))

This results in a viable 3D array giving a 2D picture when doing plt.imshow(result.mean(axis=2)) But the values are still only 1s and 0s not the wanted spectral data in z direction. I also tried using np.where but broadcasting fails as the two 2D arrays have clearly different shapes.
Has anybody got a solution? I am sure that there must be an easy way...

Answer 1

Basically, you simply need to use np.where to locate the 1s in your mask array. Then initialize your result array to zero and replace the third dimension with your data using the outputs of np.where:

import numpy as np

m, n, k = 380, 350, 91
mask = np.round(np.random.rand(m, n))
x = np.sum(mask == 1)
data = np.random.rand(k, x)

result = np.zeros((m, n, k))
row, col = np.where(mask == 1)
result[row,col] = data.transpose()