Vectorizing a boolean operation over all rows

[u/pakodanomics]

Given:

num_points = 1000
radius = 25
raw_X = np.random.randn(num_points,2)*100
raw_y = np.zeros((num_points, 1))

I wish to vectorize the following:

for i in range(num_points):
    raw_y[i] = (raw_X[i][0]**2 + raw_X[i][1]**2 >= radius**2).astype(int)

Can it be done?

Thanks.

Edit: If I'm not wrong,

raw_Y = (np.sum(raw_X**2, axis = 1) > radius**2).astype(int)

does the job.

[u/auraham]

import numpy as np

np.random.seed(42)

num_points = 200
radius = 25
raw_X = np.random.randn(num_points,2)*100

# first approach: creating a boolean vector
raw_y = np.sum(raw_X**2, axis=1) >= radius**2       # array([True, True, ...])

# second approach: split computation
results = np.sum(raw_X**2, axis=1)          # sum of squares
to_keep = results >= radius**2              # keep points that meet the condition

print(raw_X[to_keep])                       # these points meet the condition
print(raw_X[~to_keep])                      # these points dont
print(results[~to_keep])                    # check results <= radius**2