Numpy - Searching in a 4D matrix (AKA messed-up meshgrids)

Question

I am sorry if a similar question has been already posted in some way, but I could not find it anywhere so far. My problem is the following:

Suppose I have a 4D numpy matrix like this one

M= array([[[[0.  , 0.  , 0.  ],
            [0.  , 0.  , 0.01]],

           [[0.  , 0.01, 0.  ],
            [0.  , 0.01, 0.01]]],

          [[[0.01, 0.  , 0.  ],
            [0.01, 0.  , 0.01]],

           [[0.01, 0.01, 0.  ],
            [0.01, 0.01, 0.01]]]])

Which can be seen as a 3D meshgrid, where each point in space is a triplet of values (rows / axis=3 of the matrix). I have another 2D np array, corresponding to a set of points (in this case 2):

Points= array([[0.01, 0.01, 0.], [0., 0., 0.]])

I would like to look into M and find the coordinates, or indices, corresponding to those points. Something like this

coordinates= array([[1,1,0], [0,0,0]])

Unfortunately I have to avoid for loops as much as possible. I am looking for an equivalent of numpy.where() for such cases.

Thanks!

Sayandip Dutta · Accepted Answer · 2021-02-15T22:17:44.250

I don't see a strictly numpy solution because of this. However, you can achieve this without loops:

>>> np.vstack([*map(lambda x: np.argwhere(np.equal(M, x).all(-1)), Points)])

array([[1, 1, 0],
       [0, 0, 0]], dtype=int64)

However, loop is preferable in this case, both in terms of speed and readability

>>> np.vstack([np.argwhere(np.equal(M, p).all(-1)) for p in Points])
array([[1, 1, 0],
       [0, 0, 0]], dtype=int64)

Or,

>>> np.hstack([np.equal(M, p).all(-1).nonzero() for p in Points]).T
array([[1, 1, 0],
       [0, 0, 0]], dtype=int64)

Or,

>>> np.array([np.equal(M, p).all(-1).nonzero() for p in Points]).squeeze()
array([[1, 1, 0],
       [0, 0, 0]], dtype=int64)

Timings:

>>> %timeit np.vstack([*map(lambda x: np.argwhere(np.equal(M, x).all(-1)), Points)])
42.8 µs ± 8.44 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit np.vstack([np.argwhere(np.equal(M, p).all(-1)) for p in Points])
33.3 µs ± 840 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit np.hstack([np.equal(M, p).all(-1).nonzero() for p in Points]).T
23.1 µs ± 786 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit np.array([np.equal(M, p).all(-1).nonzero() for p in Points]).squeeze()
15.9 µs ± 62.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Numpy - Searching in a 4D matrix (AKA messed-up meshgrids)

1 Answers1