The following basically does what you want. Notice that it does actually contain a loop but not of the type one should avoid. The loop has only 4 iterations so it is not a performance problem. What one should avoid with numpy/pytorch is looping over all entries in a large array.
img = torch.zeros([9,9,3])
points = np.stack(np.indices([9,9])).reshape(2,-1)
blue = (0,0,1)
red = (1,0,0)
orange = (1,1/2,0)
green = (0,1,0)
angles = [0,45,225,280]
origions = np.array([(4,4),(4,4),(4,4),(4,4)])
colors = [blue, red, orange, green]
def angle_from(p):
return np.rad2deg(np.arctan2(*(points-p.reshape(2,1)))) % 360
def set_color(angle, origin, color):
angles = angle_from(np.array(origin))
mask = angles - angle == 0
img.view(-1,3)[np.where(mask),:] = torch.tensor(color, dtype=torch.float)
for angle, origin, color in zip(angles,origions,colors):
set_color(angle, origin, color)
rad = np.deg2rad(angle)
x = (origin[0],origin[1]+10*np.cos(rad))
y = (origin[0],origin[1]+10*np.sin(rad))
plt.plot(x, y, c='white')
angle = 280
rad = np.deg2rad(angle)
plt.scatter(*points)
plt.imshow(img, origin='lower')

Seems almost perfect except that you cheated with the green area. As you can see in the picture the centers of the green squares are not actually on the ray with the angle you claim it is. I plotted the centers and rays with the origin and angle you chose so one can see that more easily. I suspect that is often going to be the case and you want a way to choose the green spots. My approach there was to pick not one but two rays and show the squares with their centers being between the rays.
def set_color_between_rays(angles, origins, color):
angles1 = angle_from(np.array(origins[0]))
angles2 = angle_from(np.array(origins[1]))
mask = ((angles1 - angles[0]) >= 0) & ((angles2 - angles[1]) <= 0)
img.view(-1,3)[np.where(mask),:] = torch.tensor(color, dtype=torch.float)
angle = 295
rad = np.deg2rad(angle)
origins = np.array([[3.6,3.6],[4.3,4.3]])
plt.plot((origins[0,0],origins[0,0]+10*np.cos(rad)),(origins[0,1],origins[0,1]+10*np.sin(rad)),c='green')
plt.plot((origins[1,0],origins[1,0]+10*np.cos(rad)),(origins[1,1],origins[1,1]+10*np.sin(rad)),c='green')
set_color_between_rays([angle, angle], origins, green)
plt.scatter(*points)
plt.imshow(img, origin='lower')
