find 2d elements in a 3d array which are similar to 2d elements in another 3d array

Question

I have two 3D arrays and want to identify 2D elements in one array, which have one or more similar counterparts in the other array.

This works in Python 3:

import numpy as np
import random

np.random.seed(123)
A = np.round(np.random.rand(25000,2,2),2)
B = np.round(np.random.rand(25000,2,2),2)

a_index = np.zeros(A.shape[0])

for a in range(A.shape[0]):
    for b in range(B.shape[0]):
        if np.allclose(A[a,:,:].reshape(-1, A.shape[1]), B[b,:,:].reshape(-1, B.shape[1]),
                       rtol=1e-04, atol=1e-06):
            a_index[a] = 1
            break

np.nonzero(a_index)[0]

But of course this approach is awfully slow. Please tell me, that there is a more efficient way (and what it is). THX.

Answer 1

You are trying to do an all-nearest-neighbor type query. This is something that has special O(n log n) algorithms, I'm not aware of a python implementation. However you can use regular nearest-neighbor which is also O(n log n) just a bit slower. For example scipy.spatial.KDTree or cKDTree .

import numpy as np
import random
np.random.seed(123)
A = np.round(np.random.rand(25000,2,2),2)
B = np.round(np.random.rand(25000,2,2),2)

import scipy.spatial
tree = scipy.spatial.cKDTree(A.reshape(25000, 4))
results = tree.query_ball_point(B.reshape(25000, 4), r=1e-04, p=1)

print [r for r in results if r != []]
# [[14252], [1972], [7108], [13369], [23171]]

query_ball_point() is not an exact equivalent to allclose() but it is close enough , especially if you don't care about the rtol parameter to allclose() . You also get a choice of metric ( p=1 for city block, or p=2 for Euclidean).

PS Consider using query_ball_tree() for very large data sets. Both A and B have to be indexed in that case.

PS I'm not sure what effect the 2d-ness of the elements should have; the sample code I gave treats them as 1d and that is identical at least when using city block metric.

Answer 2

From the docs of np.allclose , we have :

If the following equation is element-wise True, then allclose returns True.

absolute(a - b) <= (atol + rtol * absolute(b))

Using that criteria, we can have a vectorized implementation using broadcasting , customized for the stated problem, like so -

# Setup parameters
rtol,atol = 1e-04, 1e-06

# Use np.allclose criteria to detect true/false across all pairwise elements
mask = np.abs(A[:,None,] - B) <= (atol + rtol * np.abs(B))

# Use the problem context to get final output
out = np.nonzero(mask.all(axis=(2,3)).any(1))[0]