[FIXED] Assigning values to 2d array with list of coordinates and time series without for-loop

Issue

I have data with four attributes. Two are coordinates xs and ys; the other two are timestamps ts and polarity ps. I need to set the value of an image at each position (x,y) with the newest polarity. The xs and ys are not unique.

I can do it with for-loop as shown below and I need a vectorized way:

import numpy as np

color_p = (0, 0, 255)
color_n = (255, 0, 0)
H, W = 128, 128
n_p = 100

xs, ys = np.random.randint(0, W, n_p), np.random.randint(0, H, n_p)
ts = np.random.rand(n_p)
ts.sort()
ps = np.random.randint(0, 2, n_p)

img = np.zeros((H, W, 3), dtype=np.uint8)
for i in range(n_p):
    x, y, p = xs[i], ys[i], ps[i]
    img[y, x] = color_p if p > 0 else color_n

Solution

The issue here is that if you have duplicate (x,y) coordinates, the later ones will overwrite the earlier values. So, what you can do is find the unique coordinates in reverse order and use those to set the values.

points = np.vstack((xs, ys))

# unique indices are taken from the end of the array, so they need to be
# recalculated to be based off the front of the array
unique_indices = np.unique(points[:,::-1], axis=1, return_index=True)[1]
unique_indices = n_p - unique_indices - 1
x_unique, y_unique = points[:,unique_indices]
img2 = np.zeros((H, W, 3), dtype=np.uint8)

cond = ps[unique_indices] > 0
img2[y_unique[cond], x_unique[cond]] = color_p
img2[y_unique[~cond], x_unique[~cond]] = color_n

I tested this code with n_p = 1000 to ensure there were duplicates and verified the results were correct using (img == img2).all().

Answered By - jared