How to logically combine integer indices in numpy?

Question

Does anyone know how to combine integer indices in numpy? Specifically, I've got the results of a few np.wheres and I would like to extract the elements that are common between them.

For context, I am trying to populate a large 3d array with the number of elements that are between boundary values of each cell, i.e. I have records of individual events including their time, latitude and longitude. I want to grid this into a 3D frequency matrix, where the dimensions are time, lat and lon.

I could loop round the array elements doing an np.where(timeCondition & latCondition & lonCondition), population with the length of the where result, but I figured this would be very inefficient as you would have to repeat a lot of the wheres.

What would be better is to just have a list of wheres for each of the cells in each dimension, and then loop through the logically combining them?

Answer 1

as @ali_m said, use bitwise and should be much faster, but to answer your question:

• call ravel_multi_index() to convert the multi-dim index into 1-dim index.
• call intersect1d() to get the index that in both condition.
• call unravel_index() to convert the 1-dim index back to multi-dim index.

Here is the code:

import numpy as np

a = np.random.rand(10, 20, 30)

idx1 = np.where(a>0.2)
idx2 = np.where(a<0.4)

ridx1 = np.ravel_multi_index(idx1, a.shape)
ridx2 = np.ravel_multi_index(idx2, a.shape)
ridx = np.intersect1d(ridx1, ridx2)
idx = np.unravel_index(ridx, a.shape)

np.allclose(a[idx], a[(a>0.2) & (a<0.4)])

or you can use ridx directly:

a.ravel()[ridx]