diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b994a074..a7f8350d 100644 --- a/.github/workflows/ci.yml +++ b/.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] diff --git a/HISTORY.rst b/HISTORY.rst index 0fc54441..cc3c2771 100644 --- a/HISTORY.rst +++ b/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`) diff --git a/africanus/experimental/rime/fused/arguments.py b/africanus/experimental/rime/fused/arguments.py index b7cfc52e..36fa0b95 100644 --- 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 index 4cf39a2a..362197ad 100644 --- 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 index b34c6016..a2fce378 100644 --- 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 index 97924ce0..aaa87e8c 100644 --- a/africanus/experimental/rime/fused/terms/cube_dde.py +++ b/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)): diff --git a/africanus/experimental/rime/fused/terms/feed_rotation.py b/africanus/experimental/rime/fused/terms/feed_rotation.py index 6a5d52e7..f97c8247 100644 --- 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] diff --git a/docs/experimental.rst b/docs/experimental.rst index 8b5f50d6..f0d40bfc 100644 --- a/docs/experimental.rst +++ b/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,7 +501,7 @@ 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`` @@ -453,6 +509,7 @@ API :ref:`Simple Example `. :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 `_ 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 ++++++++++++++++