[FIXED] Numpy/Torch insert smallest value in case of collision

Issue

I have an empty numpy array, a list of indices, and list of values associated with the indices. The issue is that there may be duplicates in the indices. In all these "collision" cases, I'd like the smallest value to be picked. Just wondering what is the best way to go about it. Eg:

array = [0,0,0,0,0,0,0]
indices = [0, 0, 2, 3, 2, 4]
values = [1.0, 3.0, 3.5, 1.5, 2.5, 8.0]

Result:

out = [1.0, 0, 2.5, 1.5, 8.0, 0.0, 0.0]

Solution

You can always implement something manually like:

import numpy as np


def index_reduce(arr, indices, out, reducer=min):
    touched = np.zeros_like(out, dtype=np.bool_)
    for i, x in enumerate(indices):
        if not touched[x]:
            out[x] = arr[i]
            touched[x] = True
        else:
            out[x] = reducer(out[x], arr[i])
    return out

which essentially loops through the indices and assign the values of arr to out if not already touched (keeping track of this with the touched array) and reducing the output with the specified reducer.

NOTE: The reducer function needs to be such that the final result can only depend on the current and previous value.

---

The usage of this would be:

indices = [0, 0, 2, 3, 2, 4]
values = [1.0, 3.0, 3.5, 1.5, 2.5, 8.0]
array = np.zeros(7)


index_reduce(values, indices, array)
# array([1. , 0. , 2.5, 1.5, 8. , 0. , 0. ])

---

If performances are of concern, you can also accelerate the above code with Numba with a simple decoration provided that also the values and indices inputs are NumPy arrays:

import numba as nb

index_reduce_nb = nb.njit(index_reduce)


indices = np.array([0, 0, 2, 3, 2, 4])
values = np.array([1.0, 3.0, 3.5, 1.5, 2.5, 8.0])
array = np.zeros(7)


index_reduce_nb(values, indices, array)
# array([1. , 0. , 2.5, 1.5, 8. , 0. , 0. ])

---

Benchmarks

The above solutions can be compared to a Torch-based solution (reworked from @Shai's answer):

import torch


def index_reduce_torch(arr, indices, out, reduce_="amin"):
    arr = torch.from_numpy(arr)
    indices = torch.from_numpy(indices)
    out = torch.from_numpy(out)
    return out.index_reduce_(dim=0, index=indices, source=arr, reduce=reduce_, include_self=False).numpy()

or, with additional skipping of Torch gradients:

index_reduce_torch_ng = torch.no_grad()(index_reduce_torch)
index_reduce_torch_ng.__name__ = "index_reduce_torch_ng"

and a Pandas-based solution (reworked from @bpfrd's answer):

import pandas as pd


def index_reduce_pd(arr, indices, out, reducer=min):
    df = pd.DataFrame(data=zip(indices, arr))
    df1 = df.groupby(0, as_index=False).agg(reducer)
    out[df1[0]] = df1[1]
    return out

using the following code:

funcs = index_reduce, index_reduce_nb, index_reduce_pd, index_reduce_torch, index_reduce_torch_ng

timings = {}
for i in range(4, 18):
    n = 2 ** i
    print(f"n = {n}, i = {i}")
    extrema = 0, 2 * n
    indices = np.random.randint(*extrema, n)
    values = np.random.random(n)
    out = np.zeros(extrema[1] + 1)
    timings[n] = []
    base = funcs[0](values, indices, out)
    for func in funcs:
        res = func(values, indices, out)
        is_good = np.allclose(base, res)
        timed = %timeit -r 16 -n 16 -q -o func(values, indices, out)
        timing = timed.best * 1e6
        timings[n].append(timing if is_good else None)
        print(f"{func.__name__:>24}  {is_good}  {timing:10.3f} µs")

to produce with the additional lines:

import matplotlib.pyplot as plt


df = pd.DataFrame(data=timings, index=[func.__name__ for func in funcs]).transpose()
df.plot(marker='o', xlabel='Input size / #', ylabel='Best timing / µs', figsize=(6, 4))
df.plot(marker='o', xlabel='Input size / #', ylabel='Best timing / µs', ylim=[0, 500], figsize=(6, 4))
fig = plt.gcf()
fig.patch.set_facecolor('white')

these plots:

(the second is a zoomed-in version of the first).

These indicate that the Numba accelerated solution could be the fastest, closely followed by the Torch-based solution while the Pandas approach could be the slowest, even slower than the explicit solution without acceleration.

Answered By - norok2