argmin.py

import time
import numba as nb
import numpy as np


def timeit(method):
    "Decorator to time a function"

    def timed(*args, **kw):
        ts = time.time()
        result = method(*args, **kw)
        te = time.time()
        print("{0}  {1:.2f} ms".format(method.__name__, (te - ts) * 1000))
        return result

    return timed


INTMAX = np.iinfo("int64").max
NTHREADS = nb.get_num_threads()

@nb.njit(parallel=True) # type: ignore
def _pfmin():
    N = 2 ** 36
    xmins = [[0]]
    xmins.clear()
    minvals = [INTMAX]
    minvals.clear()

    for i in nb.prange(NTHREADS):
        vmin = INTMAX
        # Initialize xmin with one value, so numba can infer it's type
        xmin = [0]
        xmin.clear()
        for j in range(i * N//NTHREADS , (i+1) * N//NTHREADS):
            v = abs(j - 123)
            if v == vmin:
                xmin.append(j)
            if v < vmin:
                vmin = v
                xmin.clear()
                xmin.append(j)
        
        xmins.append(xmin)
        minvals.append(vmin)
    
    return np.array(minvals), xmins


@timeit
@nb.njit
def pfmin():
    vmins, xmins = _pfmin()
    idxs,  = np.where(vmins == vmins.min())
    argmins = []
    for idx in idxs:
        argmins.extend(xmins[idx])

    return argmins


@nb.njit(parallel=False) # type: ignore
def _fmin():
    N = 2 ** 36
    xmins = [[0]]
    xmins.clear()
    minvals = [INTMAX]
    minvals.clear()

    for i in nb.prange(NTHREADS):
        vmin = INTMAX
        # Initialize xmin with one value, so numba can infer it's type
        xmin = [0]
        xmin.clear()
        for j in range(i * N//NTHREADS , (i+1) * N//NTHREADS):
            v = abs(j - 123)
            if v == vmin:
                xmin.append(j)
            if v < vmin:
                vmin = v
                xmin.clear()
                xmin.append(j)
        
        xmins.append(xmin)
        minvals.append(vmin)
    
    return np.array(minvals), xmins

@timeit
@nb.njit
def fmin():
    vmins, xmins = _fmin()
    idxs,  = np.where(vmins == vmins.min())
    argmins = []
    for idx in idxs:
        argmins.extend(xmins[idx])

    return argmins

if __name__ == "__main__":
  pfmin()
  fmin()