Non-uniform grid indexing benchmark

non_uniform_grid_indexing.py

from torch.utils import benchmark
import numpy as np


h0 = 0.1        # 1st round grid h
h1 = 0.01       # 2nd round grid h
Bq = 100        # query batch size
results = []

rng = [0., 10.]

# edge-based grid
x0 = np.linspace(rng[0], rng[1], num=int((rng[1] - rng[0])/h0))
x1 = np.linspace(rng[0], rng[1], num=int((rng[1] - rng[0])/h0))

idx0 = np.arange(len(x0))
idx1 = np.arange(len(x1))

# (get finer grid) randomly select indices and split it
# batch operation isn't supported
np.random.seed(41)
sel0 = np.random.choice(len(x0))
sel1 = np.random.choice(len(x1))

x0 = np.hstack([x0[:sel0], np.linspace(x0[sel0-1], x0[sel0] - h1, num=int(h0/h1) - 2), x0[sel0:]])
x1 = np.hstack([x1[:sel1], np.linspace(x1[sel1-1], x1[sel1] - h1, num=int(h0/h1) - 2), x1[sel1:]])
x = np.stack([x0, x1])

# (xy2ind) query index of (x0, x1) in batch
qx0 = np.random.random(Bq) * (rng[1] - rng[0]) + rng[0]
qx1 = np.random.random(Bq) * (rng[1] - rng[0]) + rng[0]
qx = np.stack([qx0, qx1], axis=0)


def xy2ind_0():
    """ indexing based on minimum distance """
    min_dist = np.argmin((np.expand_dims(x, axis=-2) - qx[..., None]) ** 2, axis=-1)
    a = min_dist[1] + len(x[0]) * min_dist[0]

    return a


def xy2ind_1():
    """ indexing based on binary search """
    s0 = np.searchsorted(x0, qx0)
    s1 = np.searchsorted(x1, qx1)

    # inner-bin calibration
    a0 = s0 - ((x0[s0] - qx0) ** 2 > (x0[s0 - 1] - qx0) ** 2)
    a1 = s1 - ((x1[s1] - qx1) ** 2 > (x1[s1 - 1] - qx1) ** 2)

    a = a1 + len(x0) * a0

    return a


assert np.all(xy2ind_0() == xy2ind_1()), 'different result!'


results.append(benchmark.Timer(
    stmt='xy2ind()',
    globals={'xy2ind': xy2ind_0},
    description='method0',
).blocked_autorange())

results.append(benchmark.Timer(
    stmt='xy2ind()',
    globals={'xy2ind': xy2ind_1},
    description='method1',
).blocked_autorange())

compare = benchmark.Compare(results)
compare.trim_significant_figures()
compare.colorize()
compare.print()

Results are:

[----------------- -----------------]
| method0 | method1
1 threads: ---------------------------
xy2ind() | 60 | 10
Times are in microseconds (us).