Merge branch 'master' into optimise-beam-cube

.github/workflows/ci.yml

@@ -36,7 +36,7 @@ jobs:
           python -m pip install -U
           pip
           setuptools
-          pipenv
+          virtualenv
 
       - name: Checkout source
         uses: actions/checkout@v4
@@ -64,11 +64,18 @@ jobs:
           echo "NUMBA_NRT_STATS=1" >> .env
           echo "NUMBA_CAPTURED_ERRORS='new_style'" >> .env
 
+      - name: Create base virtual environment
+        run: python -m virtualenv -p ${{ matrix.python-version }} /tmp/base
+
       - name: Install base codex-africanus
-        run: pipenv install .[testing]
+        run: |
+          source /tmp/base/bin/activate
+          pip install .[testing]
 
       - name: Run base test suite
-        run: pipenv run py.test -s -vvv africanus/
+        run: |
+          source /tmp/base/bin/activate
+          py.test -s -vvv africanus/
 
       - name: List the measures directory
         run: curl ftp://ftp.astron.nl/outgoing/Measures/ > measures_dir.txt
@@ -89,17 +96,18 @@ jobs:
           curl ftp://ftp.astron.nl/outgoing/Measures/WSRT_Measures.ztar | tar xvzf - -C ~/measures
           echo "measures.directory: ~/measures" > ~/.casarc
 
-      - name: Install complete codex-africanus
-        run: >
-          pipenv install
-          .[complete]
-          git+https://gitlab.mpcdf.mpg.de/ift/nifty_gridder.git#egg=nifty-gridder
+      - name: Create complete virtual environment
+        run: python -m virtualenv -p ${{ matrix.python-version }} /tmp/complete
 
-      - name: Log installed package versions
-        run: pipenv graph
+      - name: Install complete codex-africanus
+        run: |
+          source /tmp/complete/bin/activate
+          pip install .[complete] git+https://gitlab.mpcdf.mpg.de/ift/nifty_gridder.git#egg=nifty-gridder
 
       - name: Run complete test suite
-        run: pipenv run py.test -s -vvv africanus/
+        run: |
+          source /tmp/complete/bin/activate
+          py.test -s -vvv africanus/
 
   deploy:
     needs: [test]

HISTORY.rst

@@ -4,6 +4,9 @@ History
 
 0.3.8 (2024-09-29)
 ------------------
+* Support an `init_state` argument into both `Term.init_fields`
+  and `Transformer.init_fields` (:pr:`319`)
+* Use virtualenv to setup github CI test environments (:pr:`321`)
 * Update to NumPy 2.0.0 (:pr:`317`)
 * Update to python-casacore 3.6.1 (:pr:`317`)
 * Test on Python 3.12 (:pr:`318`)

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):

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]

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)

africanus/experimental/rime/fused/terms/cube_dde.py

@@ -362,8 +362,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]
-            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]
             result = zero_vis(state.beam.dtype.type(0))
 
             for co in numba.literal_unroll(range(ncorr)):

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]

docs/experimental.rst

@@ -138,7 +138,7 @@ defined on the `Phase` term, called `init_fields`.
     from africanus.experimental.rime.fused.terms.core import Term
 
     class Phase(Term)
-        def init_fields(self, typingctx, lm, uvw, chan_freq):
+        def init_fields(self, typingctx, init_state, lm, uvw, chan_freq):
             # Given the numba types of the lm, uvw and chan_freq
             # arrays, derive a unified output numba type
             numba_type = typingctx.unify_types(lm.dtype,
@@ -241,7 +241,7 @@ In the following code snippet, ``LMTransformer.init_fields``
         # OUTPUTS class attribute
         OUTPUTS = ["lm"]
 
-        def init_fields(self, typingctx, radec, phase_dir):
+        def init_fields(self, typingctx, init_state, radec, phase_dir):
             # Type and provide method for initialising the lm output
             dt = typingctx.unify_types(radec.dtype, phase_dir.dtype)
             fields = [("lm", dt[:, :])]
@@ -272,6 +272,61 @@ In the following code snippet, ``LMTransformer.init_fields``
 The ``lm`` array will be available on the ``state`` object and as a valid input
 for :meth:`Term.init_fields`.
 
+Indexing arrays
++++++++++++++++
+
+The ``init_state`` and ``state`` objects contains NumPy arrays storing
+Measurement Set v2.0 indexing information.
+
+.. code-block:: python
+
+    class State:
+        utime             # Unique times
+        uantenna          # Unique antenna indices
+        ufeed             # Unique feed indices
+        time_inverse      # Maps the time at a row into utime
+        antenna1_inverse  # Maps the antenna1 index at a row into uantenna
+        antenna2_inverse  # Maps the antenna2 index at a row into uantenna
+        feed1_inverse     # Maps the feed1 index at a row into ufeed
+        feed2_inverse     # Maps the feed2 index at a row into ufeed
+        ...
+
+These arrays are useful in cases where the developer wishes to avoid
+recomputing values multiple times for each row in the sampling function.
+Instead they can be pre-computed for unique times, antennas and feeds
+in :meth:`Term.init_fields` and then looked up in :meth:`Term.sampler`.
+
+.. code-block:: python
+
+    class MyTerm(Term):
+        def init_fields(self, typingctx, init_state, ...):
+            fields = [("precomputed", numba.float64[:, :, :])]
+
+            def precompute(init_state, ...):
+                ntime = init_state.utime.shape[0]
+                nfeed = init_state.ufeed.shape[0]
+                nant = init_state.uantenna.shape[0]
+                precomputed = np.empty((ntime, nfeed, nant), np.float64)
+
+                for t in range(ntime):
+                    for f in range(nfeed):
+                        for a in range(nant):
+                            precomputed[t, f, a] = ...
+
+                return precomputed
+
+            return fields, precompute
+
+        def sampler(self, state, s, r, t, f1, f2, a1, a2, c):
+            left = self.configuration == "left"
+
+            def sample_precomputed(state, s, r, t, f1, f2, a1, a2, c):
+                f = state.feed1_inverse[r] if left else state.feed2_inverse[r]
+                a = state.antenna1_inverse[r] if left else state.antenna2_inverse[r]
+                return state.precomputed[t, f, a]
+
+            return sample_precomputed
+
 
 Invoking the RIME
 +++++++++++++++++
@@ -429,7 +484,8 @@ API
             def __init__(self, configuration):
                 super().__init__(configuration)
 
-    .. py:method:: Term.init_fields(self, typing_ctx, arg1, ..., argn, \
+    .. py:method:: Term.init_fields(self, typing_ctx, init_state, \
+                                    arg1, ..., argn, \
                                     kwarg1=None, ..., kwargn=None)
 
         Requests inputs to the RIME term, ensuring that they are
@@ -445,14 +501,15 @@ API
         ``init_fields`` should return a :code:`(fields, function)` tuple.
         ``fields`` should be a list of the form :code:`[(name, numba_type)]`, while
         ``function`` should be a function of the form
-        :code:`fn(arg1, ..., argn, kwarg1=None, .., kwargn=None)`
+        :code:`fn(init_state, arg1, ..., argn, kwarg1=None, .., kwargn=None)`
         and should return the variables of the type defined
         in ``fields``. Note that it's signature therefore matches
         that of ``init_fields`` from after the ``typingctx``
         argument. See the
         :ref:`Simple Example <experimental-fused-rime-example-anchor>`.
 
         :param typingctx: A Numba typing context.
+        :param init_state: State object holding index information.
         :param arg1...argn: Required RIME inputs for this Term.
         :param kwarg1...kwargn: Optional RIME inputs for this Term. \
             Types here should be simple: ints, floats, complex numbers
@@ -528,8 +585,9 @@ API
         This should correspond to the fields produced by
         :meth:`Transformer.init_fields`.
 
-    .. py:method:: Transformer.init_fields(self, typing_ctx, arg1, ..., argn, \
-                                    kwarg1=None, ..., kwargn=None)
+    .. py:method:: Transformer.init_fields(self, typing_ctx, init_state, \
+                                           arg1, ..., argn, \
+                                           kwarg1=None, ..., kwargn=None)
 
         Requests inputs to the Transformer, and specifies new fields and
         the function for creating them on the ``state`` object.
@@ -547,8 +605,9 @@ API
             in Numba's
             `nopython <https://numba.pydata.org/numba-doc/latest/user/jit.html#nopython_>`_ mode.
 
-    .. py:method:: dask_schema(self, arg1, ..., argn, \
-                        kwargs1=None, ..., kwargn=None)
+    .. py:method:: dask_schema(self, init_state, \
+                               arg1, ..., argn, \
+                               kwargs1=None, ..., kwargn=None)
 
 
 
@@ -571,7 +630,6 @@ API
 
 
 
-
 Predefined Terms
 ++++++++++++++++