Randomly sample points from coordinates that fall within a square

Question

I have an array of individual ids and a separate array with their corresponding spatial coordinates on a 20x20 plane. I can easily select individuals from specific parts of the plane like so:

from numpy.random import randint
from numpy.random import default_rng

rng = default_rng()
inds = rng.choice(50, size=50, replace=False) #individual tags
locs = np.random.randint(0, 20, size=(50, 2)) #corresponding locations

groups = {
   'left' : inds[np.logical_and(locs[:, 0] < 10, locs[:, 1] < 20)],
   'right' : inds[np.logical_and(locs[:, 0] > 10, locs[:, 1] > 0)]
}

which gives me a "left" and "right" group. However, I'd like to randomly drop a square of size 4x4 onto the plane and sample only individuals that fall within that square. Any help would be greatly appreciated!

Answer 1

You can use numpy.clip with the rectangle coordinates and then check whether the clipped coordinates are the same as the original ones (if so, clipping didn't do anything which means the coordinates where already inside the rectangle):

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import numpy as np

rng = np.random.default_rng(seed=0)
locs = rng.integers(0, 20, size=(250, 2))

rect_size = np.array([4, 4])
rect_bot_left = rng.integers(0, 20-rect_size, size=2)
rect_top_right = rect_bot_left + rect_size

inside = (locs == np.clip(locs, rect_bot_left, rect_top_right)).all(axis=1)

fig, ax = plt.subplots(figsize=(6, 6))
ax.scatter(*locs.T, c=inside)
ax.add_patch(Rectangle(rect_bot_left, *rect_size, edgecolor='red', facecolor='none'))
plt.show()