forming subarrays based upon a condition in numpy

Question

I have an array of values like

import numpy as np 
dats = np.array([[r1,x1,y1],[r2,x2,y2],...,[rn,xn,yn]])

and I need to find the subarray which has elements such that the coordinates xi,yi lie within a circle of radius R around another point xp,yp. Here is a solution:

def dats_within_radius(R,xp,yp): 
    temp = np.array([])
    for loop in range(dats.shape[0]):
        r_loop = dats[loop][0]
        x_loop = dats[loop][1]
        y_loop = dats[loop][2]
        if (x_loop-x)**2+(y_loop-y)**2 <= R**2:
            temp = np.append(temp, [r_loop,x_loop,y_loop])
    temp = temp.reshape((int(temp.size/3),3))
    return temp

There must be a much better way, right? Can anyone give me a recommendation?

Answer 1

Here's a vectorized one-liner making use of broadcasting -

def dats_within_radius_vectorized(R,xp,yp): 
    return dats[((dats - [R,xp,yp])**2).sum(1) <= R**2]

For performance, I would suggest using np.einsum , like so -

def dats_within_radius_vectorized2(dats, R,xp,yp): 
    subs = dats - [R,xp,yp]
    return dats[np.einsum('ij,ij->i',subs, subs) <= R**2]

---

Runtime test -

In [82]: # Inputs
    ...: dats = np.random.randint(0,9,(10000,3))
    ...: R,xp,yp = 5,2,1
    ...: 

In [83]: %timeit dats_within_radius(dats,R,xp,yp) # Original soln
    ...: %timeit dats_within_radius_loop_compr(dats,R,xp,yp) # @BloodyD's soln
    ...: %timeit dats_within_radius_vectorized(dats,R,xp,yp)
    ...: %timeit dats_within_radius_vectorized2(dats,R,xp,yp)
    ...: 
10 loops, best of 3: 117 ms per loop
10 loops, best of 3: 38.8 ms per loop
1000 loops, best of 3: 517 µs per loop
1000 loops, best of 3: 410 µs per loop

Answer 2

I would try something like

def dats_within_radius(R,x,y): 
    return [[r_loop, X, Y] for r_loop,X,Y in dats if (X-x)**2+(Y-y)**2 <= R**2]

Here you use the fact, that in python you can 1. iterate over list (or numpy array) elements directly and 2. in python you can expand values from iterables directly to variables. And finally List Comprehensions are much faster than loops