try to get around np.abs not being implemented in numba yet

Author: mkolopanis

src/py21cmwedge/uvgridder.py

@@ -31,7 +31,6 @@ def norm(x, y, out, scale_factor):
         out /= out_norm
 
 
-# 2-d adaptation of https://stackoverflow.com/a/70614173
 @njit(numba.types.UniTuple(numba.float64[:, :], 2)(numba.float64[:], numba.float64[:]))
 def meshgrid(x, y):
     xx = np.empty(shape=(x.size, y.size), dtype=x.dtype)
@@ -43,6 +42,16 @@ def meshgrid(x, y):
     return xx, yy
 
 
+@njit(numba.int32(numba.float64, numba.float64[:]))
+def get_nearest(u, u_bins):
+    dists = u - u_bins
+    for cnt in range(dists.size):
+        if dists[cnt] < 0:
+            dists[cnt] = abs(dists[cnt])
+
+    return dists.argmin()
+
+
 class UVGridder(object):
     """Base uvgridder object."""
 
@@ -392,8 +401,8 @@ def _weights_nearest(
               on self.wavelength_scale slope
         """
         for freq_cnt, (_u, _v) in enumerate(zip(u, v)):
-            u_index = np.abs(_u - x[:]).argmin()
-            v_index = np.abs(_v - x[:]).argmin()
+            u_index = get_nearest(_u, x)
+            v_index = get_nearest(_v, x)
 
             # v,u indexing because y is the outer dimension in memory
             uvf_cube[freq_cnt, v_index, u_index] += nbls