From 980559279a581eee39ab622f5e2a462726a18078 Mon Sep 17 00:00:00 2001
From: Matthew Kolopanis <matthew.kolopanis@gmail.com>
Date: Mon, 6 Nov 2023 11:35:33 -0700
Subject: [PATCH] simplify bl to ants calc a little

---
 pyuvdata/utils.pyx | 55 +++++++++++++++++++++++-----------------------
 1 file changed, 28 insertions(+), 27 deletions(-)

diff --git a/pyuvdata/utils.pyx b/pyuvdata/utils.pyx
index 572d1a6138..84315624a8 100644
--- a/pyuvdata/utils.pyx
+++ b/pyuvdata/utils.pyx
@@ -49,7 +49,7 @@ cdef inline int_or_float max(int_or_float a, int_or_float b):
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef int_or_float arraymin(int_or_float[::1] array) nogil:
+cdef int_or_float arraymin(int_or_float[::1] array) noexcept nogil:
     cdef int_or_float minval = array[0]
     cdef Py_ssize_t i
     for i in range(array.shape[0]):
@@ -59,7 +59,7 @@ cdef int_or_float arraymin(int_or_float[::1] array) nogil:
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef int_or_float arraymax(int_or_float[::1] array) nogil:
+cdef int_or_float arraymax(int_or_float[::1] array) noexcept nogil:
     cdef int_or_float maxval = array[0]
     cdef Py_ssize_t i
     for i in range(array.shape[0]):
@@ -70,28 +70,23 @@ cdef int_or_float arraymax(int_or_float[::1] array) nogil:
 @cython.boundscheck(False)
 @cython.wraparound(False)
 cdef inline void _bl_to_ant_256(
-    numpy.uint64_t[::1] _bl,
-    numpy.uint64_t[:, ::1] _ants,
-    long nbls,
+    numpy.uint64_t _bl,
+    numpy.uint64_t[:] ants,
 ):
-  cdef Py_ssize_t i
+  ants[1] = _bl % 256
+  ants[0] = (_bl - (ants[1])) // 256
 
-  for i in range(nbls):
-    _ants[1, i] = (_bl[i]) % 256
-    _ants[0, i] = (_bl[i] - (_ants[1, i])) // 256
   return
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef inline void _bl_to_ant_2048(
-    numpy.uint64_t[::1] _bl,
-    numpy.uint64_t[:, ::1] _ants,
-    int nbls
+cdef inline  void _bl_to_ant_2048(
+    numpy.uint64_t _bl,
+    numpy.uint64_t[:] ants,
 ):
-  cdef Py_ssize_t i
-  for i in range(nbls):
-    _ants[1, i] = (_bl[i] - 2 ** 16) % 2048
-    _ants[0, i] = (_bl[i] - 2 ** 16 - (_ants[1, i])) // 2048
+  ants[1] = (_bl - 2 ** 16) % 2048
+  ants[0] = (_bl - 2 ** 16 - (ants[1])) // 2048
+
   return
 
 # defining these constants helps cython not cast the large
@@ -101,15 +96,15 @@ cdef numpy.uint64_t large_mod = 2147483648
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
+@cython.cdivision(True)
 cdef inline void _bl_to_ant_2147483648(
-    numpy.uint64_t[::1] _bl,
-    numpy.uint64_t[:, ::1] _ants,
-    int nbls
+    numpy.uint64_t _bl,
+    numpy.uint64_t[:] ants,
 ):
-  cdef Py_ssize_t i
-  for i in range(nbls):
-    _ants[1, i] = (_bl[i] - bl_large) % large_mod
-    _ants[0, i] = (_bl[i] - bl_large - (_ants[1, i])) // large_mod
+
+  ants[1] = (_bl - bl_large) % large_mod
+  ants[0] = (_bl - bl_large - (ants[1])) // large_mod
+
   return
 
 
@@ -124,13 +119,19 @@ cpdef numpy.ndarray[dtype=numpy.uint64_t, ndim=2] baseline_to_antnums(
   cdef numpy.npy_intp * dims = [2, <numpy.npy_intp> nbls]
   cdef numpy.ndarray[ndim=2, dtype=numpy.uint64_t] ants = numpy.PyArray_EMPTY(ndim, dims, numpy.NPY_UINT64, 0)
   cdef numpy.uint64_t[:, ::1] _ants = ants
+  cdef Py_ssize_t cnt
+
 
   if  _min >= (2 ** 16 + 2 ** 22):
-    _bl_to_ant_2147483648(_bl, _ants, nbls)
+    bl_fn = _bl_to_ant_2147483648
   elif _min >= 2 ** 16:
-    _bl_to_ant_2048(_bl, _ants, nbls)
+   bl_fn =  _bl_to_ant_2048
   else:
-    _bl_to_ant_256(_bl, _ants,  nbls)
+   bl_fn = _bl_to_ant_256
+
+  for cnt in range(0, nbls):
+    bl_fn(_bl[cnt], _ants[:, cnt] )
+
   return ants
 
 @cython.boundscheck(False)