@numba.jit(nopython=True, fastmath=True, boundscheck=False, parallel=False)
def abs_argmax2(re, im):
    N = re.shape[0]
    i = 0
    max_idx = 0
    current_max = 0
    current_abs2 = 0

    # main loop
    while i < N:
        current_abs2 = re[i]**2 + im[i]**2
        if current_abs2 > current_max:
            max_idx = i
            current_max = current_abs2
        i += 1
    return max_idx