From 878f14b1fc002c2930e42965cb16f441c893b11b Mon Sep 17 00:00:00 2001
From: Simon Perkins <simon.perkins@gmail.com>
Date: Tue, 5 Nov 2024 16:25:13 +0200
Subject: [PATCH] Unroll correlation loops in init_fields

---
 .../experimental/rime/fused/terms/cube_dde.py | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/africanus/experimental/rime/fused/terms/cube_dde.py b/africanus/experimental/rime/fused/terms/cube_dde.py
index 7ce24acf..7ad0fd11 100644
--- a/africanus/experimental/rime/fused/terms/cube_dde.py
+++ b/africanus/experimental/rime/fused/terms/cube_dde.py
@@ -235,35 +235,35 @@ def beam(
                                 md = vm - gm0
 
                                 # Zero the accumulators
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     absc_sum[co] = 0
                                     corr_sum[co] = 0
 
                                 # Lower cube
                                 weight = (1.0 - ld) * (1.0 - md) * nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl0, gm0, gc0, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = ld * (1.0 - md) * nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl1, gm0, gc0, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = (1.0 - ld) * md * nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl0, gm1, gc0, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = ld * md * nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl1, gm1, gc0, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
@@ -271,34 +271,34 @@ def beam(
                                 # Upper cube
                                 weight = (1.0 - ld) * (1.0 - md) * inv_nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl0, gm0, gc1, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = ld * (1.0 - md) * inv_nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl1, gm0, gc1, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = (1.0 - ld) * md * inv_nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl0, gm1, gc1, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 weight = ld * md * inv_nud
 
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     value = beam[gl1, gm1, gc1, co]
                                     absc_sum[co] += weight * np.abs(value)
                                     corr_sum[co] += weight * value
 
                                 # Assign interpolated values
-                                for co in range(ncorr):
+                                for co in numba.literal_unroll(range(ncorr)):
                                     div = np.abs(corr_sum[co])
                                     value = corr_sum[co] * absc_sum[co]