NumPy 2D array: selecting indices in a circle

Question

For some rectangular we can select all indices in a 2D array very efficiently:

arr[y:y+height, x:x+width]

...where (x, y) is the upper-left corner of the rectangle and height and width the height (number of rows) and width (number of columns) of the rectangular selection.

Now, let's say we want to select all indices in a 2D array located in a certain circle given center coordinates (cx, cy) and radius r. Is there a numpy function to achieve this efficiently?

Currently I am pre-computing the indices manually by having a Python loop that adds indices into a buffer (list). Thus, this is pretty inefficent for large 2D arrays, since I need to queue up every integer lying in some circle.

# buffer for x & y indices
indices_x = list()
indices_y = list()

# lower and upper index range
x_lower, x_upper = int(max(cx-r, 0)), int(min(cx+r, arr.shape[1]-1))
y_lower, y_upper = int(max(cy-r, 0)), int(min(cy+r, arr.shape[0]-1))
range_x = range(x_lower, x_upper)
range_y = range(y_lower, y_upper)

# loop over all indices
for y, x in product(range_y, range_x):
    # check if point lies within radius r
    if (x-cx)**2 + (y-cy)**2 < r**2:
        indices_y.append(y)
        indices_x.append(x)

# circle indexing
arr[(indices_y, indices_x)]

As mentioned, this procedure gets quite inefficient for larger arrays / circles. Any ideas for speeding things up?

If there is a better way to index a circle, does this also apply for "arbitrary" 2D shapes? For example, could I somehow pass a function that expresses membership of points for an arbitrary shape to get the corresponding numpy indices of an array?

Answer 1

You could define a mask that contains the circle. Below, I have demonstrated it for a circle, but you could write any arbitrary function in the mask assignment. The field mask has the dimensions of arr and has the value True if the condition on the righthand side is satisfied, and False otherwise. This mask can be used in combination with the indexing operator to assign to only a selection of indices, as the line arr[mask] = 123. demonstrates.

import numpy as np
import matplotlib.pyplot as plt

x = np.arange(0, 32)
y = np.arange(0, 32)
arr = np.zeros((y.size, x.size))

cx = 12.
cy = 16.
r = 5.

# The two lines below could be merged, but I stored the mask
# for code clarity.
mask = (x[np.newaxis,:]-cx)**2 + (y[:,np.newaxis]-cy)**2 < r**2
arr[mask] = 123.

# This plot shows that only within the circle the value is set to 123.
plt.figure(figsize=(6, 6))
plt.pcolormesh(x, y, arr)
plt.colorbar()
plt.show()