Rename x_index to x_inverse on state objects

ratt-ru · Oct 31, 2024 · 04f759b · 04f759b
1 parent d10c748
commit 04f759b
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 30 deletions.
diff --git a/africanus/experimental/rime/fused/arguments.py b/africanus/experimental/rime/fused/arguments.py
@@ -45,13 +45,13 @@ class ArgumentDependencies:
     REQUIRED_ARGS = ("time", "antenna1", "antenna2", "feed1", "feed2")
     KEY_ARGS = (
         "utime",
-        "time_index",
+        "time_inverse",
         "uantenna",
-        "antenna1_index",
-        "antenna2_index",
+        "antenna1_inverse",
+        "antenna2_inverse",
         "ufeed",
-        "feed1_index",
-        "feed2_index",
+        "feed1_inverse",
+        "feed2_inverse",
     )
 
     def __init__(self, arg_names, terms, transformers):

diff --git a/africanus/experimental/rime/fused/core.py b/africanus/experimental/rime/fused/core.py
@@ -94,7 +94,7 @@ def impl(*args):
 
             for s in range(nsrc):
                 for r in range(nrow):
-                    t = state.time_index[r]
+                    t = state.time_inverse[r]
                     a1 = state.antenna1[r]
                     a2 = state.antenna2[r]
                     f1 = state.feed1[r]

diff --git a/africanus/experimental/rime/fused/intrinsics.py b/africanus/experimental/rime/fused/intrinsics.py
@@ -215,13 +215,13 @@ def _add(x, y):
 class IntrinsicFactory:
     KEY_ARGS = (
         "utime",
-        "time_index",
+        "time_inverse",
         "uantenna",
-        "antenna1_index",
-        "antenna2_index",
+        "antenna1_inverse",
+        "antenna2_inverse",
         "ufeed",
-        "feed1_index",
-        "feed2_index",
+        "feed1_inverse",
+        "feed2_inverse",
     )
 
     def __init__(self, arg_dependencies):
@@ -315,13 +315,13 @@ def pack_index(typingctx, args):
 
             key_types = {
                 "utime": arg_info["time"][0],
-                "time_index": types.int64[:],
+                "time_inverse": types.int64[:],
                 "uantenna": arg_info["antenna1"][0],
-                "antenna1_index": types.int64[:],
-                "antenna2_index": types.int64[:],
+                "antenna1_inverse": types.int64[:],
+                "antenna2_inverse": types.int64[:],
                 "ufeed": arg_info["feed1"][0],
-                "feed1_index": types.int64[:],
-                "feed2_index": types.int64[:],
+                "feed1_inverse": types.int64[:],
+                "feed2_inverse": types.int64[:],
             }
 
             if tuple(key_types.keys()) != argdeps.KEY_ARGS:
@@ -368,23 +368,23 @@ def codegen(context, builder, signature, args):
                 fn_sig = types.Tuple(list(key_types.values()))(*fn_arg_types)
 
                 def _indices(time, antenna1, antenna2, feed1, feed2):
-                    utime, _, time_index, _ = _unique_internal(time)
+                    utime, _, time_inverse, _ = _unique_internal(time)
                     uants = np.unique(np.concatenate((antenna1, antenna2)))
                     ufeeds = np.unique(np.concatenate((feed1, feed2)))
-                    antenna1_index = np.searchsorted(uants, antenna1)
-                    antenna2_index = np.searchsorted(uants, antenna2)
-                    feed1_index = np.searchsorted(ufeeds, feed1)
-                    feed2_index = np.searchsorted(ufeeds, feed2)
+                    antenna1_inverse = np.searchsorted(uants, antenna1)
+                    antenna2_inverse = np.searchsorted(uants, antenna2)
+                    feed1_inverse = np.searchsorted(ufeeds, feed1)
+                    feed2_inverse = np.searchsorted(ufeeds, feed2)
 
                     return (
                         utime,
-                        time_index,
+                        time_inverse,
                         uants,
-                        antenna1_index,
-                        antenna2_index,
+                        antenna1_inverse,
+                        antenna2_inverse,
                         ufeeds,
-                        feed1_index,
-                        feed2_index,
+                        feed1_inverse,
+                        feed2_inverse,
                     )
 
                 index = context.compile_internal(builder, _indices, fn_sig, fn_args)

diff --git a/africanus/experimental/rime/fused/terms/cube_dde.py b/africanus/experimental/rime/fused/terms/cube_dde.py
@@ -171,8 +171,8 @@ def sampler(self):
         zero_vis = zero_vis_factory(ncorr)
 
         def cube_dde(state, s, r, t, f1, f2, a1, a2, c):
-            a = state.antenna1_index[r] if left else state.antenna2_index[r]
-            feed = state.feed1_index[r] if left else state.feed2_index[r]
+            a = state.antenna1_inverse[r] if left else state.antenna2_inverse[r]
+            feed = state.feed1_inverse[r] if left else state.feed2_inverse[r]
             sin_pa = state.beam_parangle[t, feed, a, 0]
             cos_pa = state.beam_parangle[t, feed, a, 1]
 

diff --git a/africanus/experimental/rime/fused/terms/feed_rotation.py b/africanus/experimental/rime/fused/terms/feed_rotation.py
@@ -43,8 +43,8 @@ def sampler(self):
         linear = self.feed_type == "linear"
 
         def feed_rotation(state, s, r, t, f1, f2, a1, a2, c):
-            a = state.antenna1_index[r] if left else state.antenna2_index[r]
-            f = state.feed1_index[r] if left else state.feed2_index[r]
+            a = state.antenna1_inverse[r] if left else state.antenna2_inverse[r]
+            f = state.feed1_inverse[r] if left else state.feed2_inverse[r]
             sin_a = state.feed_parangle[t, f, a, 0, 0]
             cos_a = state.feed_parangle[t, f, a, 0, 1]
             sin_b = state.feed_parangle[t, f, a, 1, 0]