Python: Stretch 2D array into 3D based on corresponding array with 3rd-dimension index

Question

Lets say I have some 2D array a = np.ones((3,3))

I want to stretch this array into 3 dimensions. I have array b, the same size as a, that provides the index in the 3rd dimension that each corresponding element in a needs to go too.

I also have 3D array c that is filled with NaNs. This is the array that the information from a should be put into. The remaining blank spaces that do not get "filled: can remain NaNs.

>>> a = np.ones((3,3))
>>> b = np.random.randint(0,3,(3,3))
>>> c = np.empty((3,3,3))*np.nan
>>> 
>>> a
array([[ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.]])
>>> b
array([[2, 2, 2],
       [1, 0, 2],
       [1, 0, 0]])
>>> c
array([[[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]],

       [[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]],

       [[ nan,  nan,  nan],
        [ nan,  nan,  nan],
        [ nan,  nan,  nan]]])

So, in the above example, I would want to end up with c[0,0,2] = 1.

I know I probably do this with some nested loops, but ideally I want this done in a more efficient/vectorized way.

Answer 1

You can use fancy indexing as this, assuming the largest value in b is always less than c.shape[2]:

n1, n2 = a.shape
c[np.arange(n1)[:,None], np.arange(n2), b] = a

c
#array([[[ nan,  nan,   1.],
#        [ nan,  nan,   1.],
#        [ nan,  nan,   1.]],

#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [ nan,  nan,   1.]],

#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [  1.,  nan,  nan]]])

---

Here we use integer arrays for all dimensions to trigger advanced indexing, and the three arrays are broadcasted against each other as follows (here we use numpy.broacast_arrays to visualize this):

i, j, k = np.broadcast_arrays(np.arange(3)[:,None], np.arange(3), b)

print("first dimension index: ")
print(i)
print("second dimension index: ")
print(j)
print("third dimension index: ")
print(k)

first dimension index: 
[[0 0 0]
 [1 1 1]
 [2 2 2]]
second dimension index: 
[[0 1 2]
 [0 1 2]
 [0 1 2]]
third dimension index: 
[[2 2 2]
 [1 0 2]
 [1 0 0]]

Now the advanced indexing goes as (0, 0, 2), (0, 1, 2), (0, 2, 2) ... ,i.e. pick one value from each array at the same positions to form an index for an element:

---

Some testing cases:

c[0,0,2]
#1.0

c[0,1,2]
#1.0

c[2,1,0]
#1.0

Answer 2

Ok, so this feels like a total hack, but does the trick:

a = np.ones((3,3))
b = np.array([[2, 2, 2],
              [1, 0, 2],
              [1, 0, 0]])
c = np.empty((3,3,3))*np.nan

z_coords = np.arange(3)

c[z_coords[None, None, :] == b[..., None]] = a.ravel()

What I do is create an boolean indexing array that is true for the indices we want to assign, and then assign these.

array([[[ nan,  nan,   1.],
        [ nan,  nan,   1.],
        [ nan,  nan,   1.]],

       [[ nan,   1.,  nan],
        [  1.,  nan,  nan],
        [ nan,  nan,   1.]],

       [[ nan,   1.,  nan],
        [  1.,  nan,  nan],
        [  1.,  nan,  nan]]])