move values of 3D array knowing new coordinates with mask

Question

I would like to extend a problem I already discussed on Stackoverflow . It was dealing with 2D numpy arrays and I would like to do the same with a 3-dimensional array.

I would like to "move" elements of a 2D array to new coordinates which are stored in 2 other arrays. I'm looking to automate this, because in reality my arrays are large (400x200x100). Some values wont find his coordinates and wont be used, Some of these coordinates are masked, which I have indicated in the example below by using the value 0. If the coordinate is masked, the elements in the array I want to reshuffle won't be used.

import numpy as np

#My new coordinates in X and Y directions   

mx = np.array([[[ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.]],

       [[ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.],
       [ 1.,  2.,  3.,  4.,  0.]]])

my = np.array([[[ 0.,  2.,  2.,  2.,  2.],
       [ 0.,  3.,  3.,  3.,  3.],
       [ 0.,  4.,  4.,  4.,  4.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.]],

       [[ 0.,  2.,  2.,  2.,  2.],
       [ 0.,  3.,  3.,  3.,  3.],
       [ 0.,  4.,  4.,  4.,  4.],
       [ 0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.]]])


IRtest = np.array([[[-0.07383495, -0.08606554, -0.08480594, -0.08099556, -0.08218414],
       [-0.07866761, -0.08373   , -0.08253587, -0.08106102, -0.08220205],
       [-0.07727436, -0.08271511, -0.0807254 , -0.07832416, -0.08021686],
       [-0.07612349, -0.08190446, -0.07996929, -0.07842754, -0.08024891],
       [-0.07488144, -0.08150557, -0.08038229, -0.07895656, -0.07997815]],

       [[-0.07383495, -0.08606554, -0.08480594, -0.08099556, -0.08218414],
       [-0.07866761, -0.08373   , -0.08253587, -0.08106102, -0.08220205],
       [-0.07727436, -0.08271511, -0.0807254 , -0.07832416, -0.08021686],
       [-0.07612349, -0.08190446, -0.07996929, -0.07842754, -0.08024891],
       [-0.07488144, -0.08150557, -0.08038229, -0.07895656, -0.07997815]]])

So the array expected looks like :

array_expected = np.array([[[-0.08271511, -0.0807254 , -0.07832416, -0.08021686, 0],
       [-0.08190446, -0.07996929, -0.07842754, -0.08024891, 0],
       [-0.08150557, -0.08038229, -0.07895656, -0.07997815, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]],

       [[-0.08271511, -0.0807254 , -0.07832416, -0.08021686, 0],
       [-0.08190446, -0.07996929, -0.07842754, -0.08024891, 0],
       [-0.08150557, -0.08038229, -0.07895656, -0.07997815, 0],
       [0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0]]])

I try with the part of code i got with my last post.

b = np.zeros_like(IRtest)

for i in range(IRtest.shape[1]):
    for j in range(IRtest.shape[2]):
        for k in range(IRtest.shape[0]):
            b[k, j, i] = IRtest[k,my[k,j,i],mx[k,j,i]]*(mx[k,j,i]!=-1)*(my[k,j,i]!=-1)       

b

But the result isn t the same i have expected :

  array([[[-0.08606554, -0.0807254 , -0.07832416, -0.08021686, -0.07727436],
        [-0.08606554, -0.07996929, -0.07842754, -0.08024891, -0.07612349],
        [-0.08606554, -0.08038229, -0.07895656, -0.07997815, -0.07488144],
        [-0.08606554, -0.08480594, -0.08099556, -0.08218414, -0.07383495],
        [-0.08606554, -0.08480594, -0.08099556, -0.08218414, -0.07383495]],

       [[-0.08606554, -0.0807254 , -0.07832416, -0.08021686, -0.07727436],
        [-0.08606554, -0.07996929, -0.07842754, -0.08024891, -0.07612349],
        [-0.08606554, -0.08038229, -0.07895656, -0.07997815, -0.07488144],
        [-0.08606554, -0.08480594, -0.08099556, -0.08218414, -0.07383495],
        [-0.08606554, -0.08480594, -0.08099556, -0.08218414, -0.07383495]]])

Answer 1

Your map matrices are wrong, to get the result you want they need to be like these (since, when you put values into b, you are checking if m[k, j, i] != -1 and you want the last columns to be 0, not the first ones)

mx = np.array([[[ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.]],

       [[ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
       [ 1.,  2.,  3.,  4.,  -1.],
      [ 1.,  2.,  3.,  4.,  -1.]]])

my = np.array([[[ 2.,  2.,  2.,  2.,  -1.],
       [ 3.,  3.,  3.,  3.,  -1.],
       [ 4.,  4.,  4.,  4.,  -1.],
       [ -1.,  -1.,  -1.,  -1.,  -1.],
       [ -1.,  -1.,  -1.,  -1.,  -1.]],

       [[ 2.,  2.,  2.,  2.,  -1.],
       [ 3.,  3.,  3.,  3.,  -1.],
       [ 4.,  4.,  4.,  4.,  -1.],
       [ -1.,  -1.,  -1.,  -1.,  -1.],
       [ -1.,  -1.,  -1.,  -1.,  -1.]]])

Also in your loops, it would be better to switch the dimensions in the first and second loops so they become

for i in range(IRtest.shape[2]):
    for j in range(IRtest.shape[1]):
        for k in range(IRtest.shape[0]):

This does not matter for the case you give here since the matrices are square, but the real problem you mention the matrices are not square so it becomes an issue there.

Answer 2

To summarize the answer to your previous question regarding 2dim, you can simply use fancy-indexing like (after fixing the dtypes of your arrays, of course):

b = IRtest[my,mx] * ~(mask_my | mask_mx)

---

Now, in order to apply the same technique to the 3dim case, you need to create a "neutral" index array, to apply to the first axis. This is where np.indices comes useful:

mz = np.indices(IRtest.shape)[0]  # take [0] because we need to be neutral w.r.t. axis=0

Now apply the fancy indexing:

b = IRtest[mz, my, mx]

To apply the mask, it is easiest to add a dimension to the mask array in order to make it 3dim, then let numpy's broadcasting do its magic. (I'm assuming we're using the same masks from the 2dim question, with dtype=bool).

unified_mask = mask_my | mask_mx
b *= ~unified_mask[np.newaxis,...]

As in the 2dim case, you could alrenatively use fancy-assignment (again, axis=0 is being broadcast):

b[unified_mask[np.newaxis,...]] = 0.