How do I replace values along z-axis in Numpy 3D array based on 2D index and 1D value vector

Question

I'm struggling to understand array indexing it seems.

What is given:

I do have a 3d array like so:

a_3d = np.zeros((3,3,3))

An 2d index array:

a_2d_index = np.array([[0,0,1], [0,0,0], [0,1,1]]).astype('bool')

And values to place into 3d array at x,y location:

a_1d_fill = np.array([10,20,30])

Now, I do want to use a_2d_index to find the locations in a_3d and vertically place the a_1d_fill at this x,y position...

The final outcome should look like this:

a_3d := [[[0,0, 10],
          [0,0,  0],
          [0,10,10]],
         [[0,0, 20],
          [0,0,  0],
          [0,20,20]],
         [[0,0, 30],
          [0,0,  0],
          [0,30,30]]]

This will be used on very large array so memory efficiency and speed are crucial... (little copying, preferably in-place modification)

Answer 1

In [26]: a_3d = np.zeros((3,3,3), dtype=int)

In [27]: a_2d_index = np.array([[0,0,1], [0,0,0], [0,1,1]]).astype('bool')

In [28]: a_1d_fill = np.array([10,20,30])

In [29]: a_3d[:,a_2d_index] = a_1d_fill[:,np.newaxis]

In [30]: a_3d
Out[30]:
array([[[ 0,  0, 10],
        [ 0,  0,  0],
        [ 0, 10, 10]],

       [[ 0,  0, 20],
        [ 0,  0,  0],
        [ 0, 20, 20]],

       [[ 0,  0, 30],
        [ 0,  0,  0],
        [ 0, 30, 30]]])

Answer 2

Maybe a little bit of explanation and intuition might help future readers to better understand slicing. I'm (shamelessly) using Stuart Berg's solution to explain the users to get an intuitive feel for slicing.

Array & mask definition:

In [57]: a_3d        # array with unique numbers for better understanding.
Out[57]: 
array([[[13, 14, 15],
        [23, 24, 25],
        [33, 34, 35]],

       [[ 3,  4,  5],
        [ 6,  7,  8],
        [ 9, 10, 11]],

       [[ 0,  1,  2],
        [30, 31, 32],
        [40, 41, 42]]], dtype=uint8)

Understand the above 3D array as three slices of 3x3 array, with the top most array having index 0 and center one having index 1 and the bottom one having index 2.

Boolean Mask:

In [58]: a_2d_index
Out[58]: 
array([[False, False,  True],
       [False, False, False],
       [False,  True,  True]], dtype=bool)

Now, let's slice the array a_3d using the boolean mask a_2d_index

In [68]: a_3d[:, a_2d_index]   # : means apply the mask to all the slices(here 3)
Out[68]: 
array([[15, 34, 35],
       [ 5, 10, 11],
       [ 2, 41, 42]], dtype=uint8)

Okay, now we get the above result. Why and how? So, imagine like this: we take the boolean-mask array and just superimpose on the array a_3d over each of its slices.

Now, wherever we have boolean value True in the mask, that corresponding element in array a_3d will (glow?) and contribute to the result array. This happens with each slice as we place the mask over each slice (since we've used :, we do this for all the slices in array a_3d).

So Yay, slicing is done! Now, we want to assign new values (where the mask has boolean value True).

Assigning new values:

In [69]: a_1d_fill
Out[69]: array([10, 20, 30])

This is the 1D array that we have. But, the mask is 2D array. So, we change this to 2D using np.newaxis

In [70]: a_2d_fill = a_1d_fill[:, np.newaxis]

In [71]: a_2d_fill
Out[71]: 
array([[10],
       [20],
       [30]])


In [73]: a_3d[:, a_2d_index] = a_2d_fill

Now, this assignment will copy the value 10 to indices in the first slice where the boolean mask has True value, then copy value 20 to indices in second slice where the boolean mask has value True and finally copy value 30 to indices in third slice where the boolean mask has value True.

And the final result looks like:

In [74]: a_3d
Out[74]: 
array([[[13, 14, 10],
        [23, 24, 25],
        [33, 10, 10]],

       [[ 3,  4, 20],
        [ 6,  7,  8],
        [ 9, 20, 20]],

       [[ 0,  1, 30],
        [30, 31, 32],
        [40, 30, 30]]], dtype=uint8)

Phew! That was really long. So, in short, the result array should have values 10, 20, & 30 in first, second, and third slices (where boolean mask has value True) respectively.

---

P.S.: Slicing provides views or references of the original array. So, when we change values using slicing, this will affect the original array. Thus, this is an in-place modification.