Numpy insert matrix values with matrix index

I have the following code which creates a 4D grid matrix and I am looking to insert the rolled 2D vals matrix into this grid.

import numpy as np

k = 100
x = 20
y = 10
z = 3

grid = np.zeros((y, k, x, z))
insert_map = np.random.randint(low=0, high=y, size=(5, k, x))
vals = np.random.random((5, k))

for i in range(x):
    grid[insert_map[:, :, i], i, 0] = np.roll(vals, i)

If vals would be a 1D array and I would use a 1D insert_map array as a reference it would work, however using it in multiple dimensions seems to be an issue and it raises error:

ValueError: shape mismatch: value array of shape (5,100)  could not be broadcast to indexing result of shape (5,100,3)

I’m confused as to why it’s saying that error as grid[insert_map[:, :, i], i, 0] should in my mind give a (5, 100) insert location for the y and k portion of the grid array and then fixes the x and z portion with i and 0?

Is there any way to insert the 2D (5, 100) rolled vals matrix into the 4D (10, 100, 20, 3) grid matrix by 2D indexing?

Answer

grid is (y, k, x, z)

insert_map is (5, k, x). insert_map[:, :, i] is then (5,k).

grid[insert_map[:, :, i], i, 0] will then be (5,k,z). insert_map[] only indexes the first y dimension.

vals is (5,k); roll doesn’t change that.

np.roll(vals, i)[...,None] can broadcast to fill the z dimension, if that’s what you want.

Your insert_map can’t select values along the k dimension. It’s values, created by the randint are valid the y dimension.

If the i and 0 are supposed to apply to the last two dimensions, you still need an index for the k dimension. Possibilities are:

grid[insert_map[:, j, i], j, i, 0]
grid[insert_map[:, :, i], 0, i, 0]
grid[insert_map[:, :, i], :, i, 0]
grid[insert_map[:, :, i], np.arange(k), i, 0]