How can I apply the same set of values at various positions in a larger array, specified by masks?

Question

I have a 2D array Y of dimensions N x N, and an array of K binary masks X, each of dimensions M x M (so, X has shape K x M x M). Each binary mask in X has exactly one N x N patch of ones, and the rest of the elements are zero. I would like to create a K x M x M array Z, with elements of Y in the locations specified by X, and zero everywhere else.

For example, if M = 3, N = 2, K = 3, and

X = [[[1., 1., 0.],
     [1., 1., 0.],
     [0., 0., 0.]],

    [[0., 1., 1.],
     [0., 1., 1.],
     [0., 0., 0.]],

    [[0., 0., 0.],
     [0., 1., 1.],
     [0., 1., 1.]]]

and

Y = [[0.1, 0.2],
    [0.3, 0.4]]

then Z should be

Z = [[[0.1, 0.2, 0.],
     [0.3, 0.4, 0.],
     [0., 0., 0.]],

    [[0., 0.1, 0.2],
     [0., 0.3, 0.4],
     [0., 0., 0.]],

    [[0., 0., 0.],
     [0., 0.1, 0.2],
     [0., 0.3, 0.4]]]

I would like to do this using functions either from Numpy or Pytorch. I thought it might be possible to do this using numpy.where(), perhaps as

Z = numpy.zeros((3,3,3))
Z[numpy.where(X == 1)] = Y

but this gives a shape mismatch error.

What would be the most concise and efficient way to do this (apart from explicitly using loops to set the values)?

Answer 1

Try this:

X[X == 1] = np.tile(Y.flatten(), X.shape[0])