[FIXED] Updating a numpy array using an index list with multiple occurrences of the same index

Issue

I have a numpy array of size [n_rows, n_cols, n_channels]. In my code I have a loop where the array is constantly updated and clipped:

def update(arr, row_idx, col_idx, ch_idx):
    arr[row_idx, col_idx, ch_idx] += 1
    arr[arr > 10] = 10

arr = np.array(n_rows, n_cols, n_channels)
while True:
    update(arr, 0, 1, 2)

In order to optimize my code, I can use a cache with index list and update the array once every N iterations:

def update(arr, rows_list, cols_list, ch_list):
        arr[rows_list, cols_list, ch_list] += 1
        arr[arr > 10] = 10

arr = np.array(n_rows, n_cols, n_channels)
cache_length = 3
rows_list, cols_list, ch_list = [], [], []
while True:
    rows_list.append(something1)
    cols_list.append(something2)
    ch_list.append(something3)
    if len(row_list) == cache_length:
        update(arr, rows_list, cols_list, ch_list)
        rows_list, cols_list, ch_list = [], [], []

This saves time, but it could happen that the cache contains the same array index multiple times, e.g.:

# arr[0, 0, 6] should be updated twice
update(arr, [0, 0, 2], [3, 3, 5], [6, 6, 6])

How can I change my code to make this optimization work?

Solution

You can aggregate using numpy.unique:

def update(arr, row_idx, col_idx, ch_idx):
    idx, cnt = np.unique([row_idx, col_idx, ch_idx],
                         return_counts=True, axis=1)
    arr[tuple(idx)] += cnt
    arr[arr > 10] = 10

And you can potentially optimize further by only clipping the updated values (and not the full array):

def update(arr, row_idx, col_idx, ch_idx):
    idx, cnt = np.unique([row_idx, col_idx, ch_idx],
                         return_counts=True, axis=1)
    idx = tuple(idx)
    arr[idx] = np.clip(arr[idx]+cnt, -np.inf, 10)

Example:

arr = np.zeros((2, 3, 4), dtype='int')
update(arr, [0, 0, 1], [1, 1, 2], [3, 3, 3])

# arr
array([[[0, 0, 0, 0],
        [0, 0, 0, 2],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0],
        [0, 0, 0, 1]]])

Answered By - mozway