Say I have two large numpy arrays a1 and a2 (each with 10000 numbers). I want to find the indices in each array that meet the condition of f(x1, x2) > 0. To be clear, for each number in a1 (or a2), if there's any number in a2 (or a1) that meets the condition, then it is a valid number. For now I am using one loop over a1 and only using numpy vectorization for calculations on a2:
import numpy as np
a1 = np.random.rand(10000)
a2 = np.random.rand(10000)
def f(x1, x2):
return np.exp(x2 - x1) / 1.2 - 1
valid_i1_set = set()
valid_i2_set = set()
for i1 in range(len(a1)):
valid_i2 = np.nonzero(f(a1[i1], a2) > 0)[0].tolist()
valid_i2_set.update(valid_i2)
if len(valid_i2) > 0:
valid_i1_set.add(i1)
print(sorted(valid_i1_set))
print(sorted(valid_i2_set))
I imagine there should be a faster way to do this without a for loop. Making a product array of 10000x10000 numbers also doesn't feel right. Any suggestions?
A possible solution:
i1 = np.flatnonzero(a1 < a2.max() - np.log(1.2))
i2 = np.flatnonzero(a2 > a1.min() + np.log(1.2))
We avoid computing the full 10 000 × 10 000 matrix by noticing that the condition f(x1, x2) > 0 is equivalent to x2 > x1 + log(1.2). This means that for any value in a1 to satisfy the condition with at least one value in a2, it must be less than max(a2) - log(1.2), and similarly, for any value in a2 to satisfy the condition with some value in a1, it must be greater than min(a1) + log(1.2).
We compute log(1.2) with log, find the thresholds via max and min, and finally extract the qualifying indices in one go with np.flatnonzero.