Fused Radio Interferometer Measurement Equation

Radio Interferometer Measurement Equation

The Radio Interferometer Measurement Equation (RIME) describes the response of an interferometer to a sky model. As described in A full-sky Jones formalism, a RIME could be written as follows:

\[V_{pq} = G_{p} \left( \sum_{s} E_{ps} L_{p} K_{ps} B_{s} K_{qs}^H L_{q}^H E_{qs}^H \right) G_{q}^H\]

where for antenna \(p\) and \(q\), and source \(s\):

  • \(G_{p}\) represents direction-independent effects.

  • \(E_{ps}\) represents direction-dependent effects.

  • \(L_{p}\) represents the feed rotation.

  • \(K_{ps}\) represents the phase delay term.

  • \(B_{s}\) represents the brightness matrix.

The RIME is more formally described in the following four papers:

The Fused RIME

The RIME poses a number of implementation challenges which focus on flexibility, speed and ease of use.

Firstly, the RIME can be composed of many terms representing various physical effects. It is useful for scientist to be able to specify many different terms in the above Equation, for example.

Secondly, the computational complexity of the RIME O(S x V) where S is the number of source and V is the number of visibilities. This is computionationally expensive relative to degridding strategies.

Thirdly, it should be as easy as possible to define the RIME, but not at the cost of the previous two constraints.

The Fused RIME therefore implements a “RIME Compiler” using Numba for speed, which compiles a RIME Specification defined by a number of Terms into a single, optimal unit of execution.

A Simple Example

In the following example, we will define a simple RIME using the Fused RIME API to define terms for computing:

  1. the Phase Delay.

  2. the Brightness Matrix.

The RIME Specification

The specification for this RIME is as follows:

rime_spec = RimeSpecification("(Kpq, Bpq): [I,Q,U,V] -> [XX,XY,YX,YY]",
                              terms={"K": Phase},
                              transformers=[LMTransformer])

(Kpq, Bpq) specifies the onion including the Phase Delay and Brightness more formally defined here, while the the pq in both terms signifies that they are calculated per-baseline. [I,Q,U,V] -> [XX,XY,YX,YY] defines the stokes to correlation conversion within the RIME and also identifies whether the RIME is handling linear or circular feeds. terms={"K": Phase} indicates that the K term is implemented as a custom Phase term, described in the next section. Finally, LMTransformer is a Transformer that precomputes lm coordinates for use by all terms.

Custom Phase Term

Within the RIME, each term is sampled at an individual source, row and channel.

Therefore each term must provide a sampling function that will provide the necessary data for multiplication within the RIME. Consider the following Phase Term:

from africanus.experimental.rime.fused.terms.core import Term

class Phase(Term):
    def sampler(self):
        def phase_sample(state, s, r, t, f1, f2, a1, a2, c):
            p = state.real_phase[s, r] * state.chan_freq[c]
            return np.cos(p) + np.sin(p)*1j

        return phase_sample

This may look simple: we compute the complex phase by multiplying the real phase at each source and row by the channel frequency and return the complex exponential of this value.

However, questions remain: What is the state object and how do we know that the real_phase and chan_freq are members? To answer this, we must define (and understand) a second method defined on the Phase term, called init_fields.

import numba
from africanus.experimental.rime.fused.terms.core import Term

class Phase(Term)
    def init_fields(self, typingctx, 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,
                                           uvw.dtype,
                                           chan_freq.dtype)

        # Define the type of new fields on the state object
        # in this case a 2D Numba array with dtype numba_type
        fields = [("real_phase", numba_type[:, :])]

        def real_phase(lm, uvw, chan_freq):
            """Compute the real_phase upfront, instead of in
            the sampling function"""
            real_phase = np.empty((lm.shape[0], uvw.shape[0]), numba_type)

            for s in range(lm.shape[0]):
                l, m = lm[s]
                n = 1.0 - l**2 - m**2
                n = np.sqrt(0.0 if n <= 0.0 else n) - 1.0

                for r in range(uvw.shape[0]):
                    u, v, w = uvw[r]
                    real_phase[s, r] = -2.0*np.pi*(l*u + m*v + n*w)/3e8

            return real_phase

        # Return the new field definition and
        # the function for creating it
        return fields, real_phase

init_fields serves multiple purposes:

  1. It requests input for the Phase term. The above definition of init_fields signifies that the Phase term desires the lm, uvw and chan_freq arrays. Additionally, these arrays will be stored on the state object provided to the sampling function.

  2. It supports reasoning about Numba types in a manner similar to numba.generated_jit(). The lm, uvw and chan_freq arguments contain the Numba types of the variables supplied to the RIME, while the typingctx argument contains a Numba Typing Context which can be useful for reasoning about these types. For example typingctx.unify_types(lm.dtype, uvw.dtype, chan_freq.dtype) returns a type with sufficient precision, given the input types, similar to numpy.result_type().

  3. It allows the user to define new fields, as well as a function for defining those fields on the state object. The above definition of init_fields returns a list of (name, type) tuples defining the new field names and their types, while real_phase defines the creation of this new field.

    This is useful for optimising the sampling function by pre-computing values. For example, it is wasteful to compute the real phase for each source, row and channel.

Returning to our definition of the Phase Term sampling function, we can see that it uses the new field real_phase defined in init_fields, as well as the chan_freq array requested in init_fields to compute a complex exponential.

class Phase(Term):
    def sampler(self):
        def phase_sample(state, s, r, t, f1, f2, a1, a2, c):
            p = state.real_phase[s, r] * state.chan_freq[c]
            return np.cos(p) + np.sin(p)*1j

        return phase_sample

Transformers

Using Term.init_fields(), we can precompute data for use in sampling functions, within a single Term. However, sometimes we wish to precompute data for use by multiple Terms. This can be achieved through the use of Transformers. A good example of data that it is useful to precompute for multiple Terms are lm coordinates, which are in turn, derived from phase_dir and radec which are the phase centre of an observation and the position of a source, respectively. In the following code snippet, LMTransformer.init_fields

from africanus.experimental.rime.fused.transformers import Transformer

class LMTransformer(Transformer):
    # Must specify list of outputs produced by this transformer on the
    # OUTPUTS class attribute
    OUTPUTS = ["lm"]

    def init_fields(self, typingctx, radec, phase_dir):
        # Type and provide method for initialising the lm output
        dt = typingctx.unify_types(radec.dtype, phase_dir.dtype)
        fields = [("lm", dt[:, :])]

        def lm(radec, phase_dir):
            lm = np.empty_like(radec)
            pc_ra = phase_dir[0]
            pc_dec = phase_dir[1]

            sin_pc_dec = np.sin(pc_dec)
            cos_pc_dec = np.cos(pc_dec)

            for s in range(radec.shape[0]):
                da = radec[s, 0] - pc_ra
                sin_ra_delta = np.sin(da)
                cos_ra_delta = np.cos(da)

                sin_dec = np.sin(radec[s, 1])
                cos_dec = np.cos(radec[s, 1])

                lm[s, 0] = cos_dec*sin_ra_delta
                lm[s, 1] = sin_dec*cos_pc_dec - cos_dec*sin_pc_dec*cos_ra_delta

            return lm

        return fields, lm

The lm array will be available on the state object and as a valid input for Term.init_fields().

Invoking the RIME

We then invoke the RIME by passing in the RimeSpecification, as well as a dataset containing the required arguments:

from africanus.experimental.rime.fused.core import rime
import numpy as np

dataset = {
    "radec": np.random.random((10, 2))*1e-5,
    "phase_dir": np.random.random((2,))*1e-5,
    "uvw": np.random.random((100, 3))*1e5,
    "chan_freq:" np.linspace(.856e9, 2*.856e9, 16),
    ...,
    "stokes": np.random.random((10, 4)),
    # other required data
}

rime_spec = RimeSpecification("(Kpq, Bpq)",
                              terms={"K": Phase},
                              transformers=LMTransformer)
model_visibilities = rime(rime_spec, dataset)

Dask Support

Dask wrappers are provided for the africanus.experimental.rime.fused.core.rime() function. In order to support this, both Term and Transformer classes need to supply a dask_schema function which is used to define the schema for each supplied argument, which in turn is supplied to a dask.array.blockwise() call.

The schema should be a tuple of dimension string names. In particular, the rime function assigns special meaning to source, row, chan and corr – These names are are associated with individual sources (fields) and Measurement Set rows, channels and correlations, respectively. Dask Array chunking is supported along these dimensions in the sense that the rime will be computed for each chunk along these dimensions.

Note

Chunks in dimensions other than source, row, chan and corr will be contracted into a single array within the rime function. It is recommended that other dimensions contain a single chunk, or contain small quantities of data relative to the special dimensions.

Therefore, Phase.dask_schema could be implemented as follows:

class Phase(Term):
    def dask_schema(self, lm, uvw, chan_freq):
        assert lm.ndim == 2
        assert uvw.ndim == 2
        assert chan_freq.ndim == 1

        return {
            "lm": ("source", "lm-component"),
            "uvw": ("row", "uvw-component"),
            "chan_freq": ("chan",),
        }

The dask_schema for a Transformer is slightly different as, in addition a schema for the inputs, it must also provide an array_like variable describing the number of dimensions and data type of the output arrays. The array_like variables are in turn passed into Term.dask_schema. Thus, LMTransformer.dask_schema could be implemented as follows;

class LMTransformer(Transformer):
    OUTPUTS = ["lm"]

    def dask_schema(self, phase_dir, radec):
        dt = np.result_type(phase_dir.dtype, radec.dtype)
        return ({
            "phase_dir": ("radec-component",),
            "radec": ("source", "radec-component",),
        },
        {
            "lm": np.empty((0,0), dtype=dt)
        })

Then, in a paradigm very similar to the non-dask case, we create a RimeSpecification and supply it, along with a dictionary or dataset of dask arrays, to the rime() function. This will produce a dask array representing the model visibilities.

from africanus.experimental.rime.fused.dask import rime
import dask.array as da
import numpy as np

dataset = {
    "radec": da.random.random((10, 2), chunks=(2, 2))*1e-5,
    "phase_dir": da.random.random((2,), chunks=(2,))*1e-5,
    "uvw": da.random.random((100, 3), chunks=(10, 3))*1e5,
    "chan_freq:" da.linspace(.856e9, 2*.856e9, 16, chunks=(4,)),
    ...,
    "stokes": da.random.random((10, 4), chunks=(2, 4)),
    # other required data
}

rime_spec = RimeSpecification("(Kpq, Bpq)",
                              terms={"K": Phase},
                              transformers=LMTransformer)
model_visibilities = rime(rime_spec, dataset)
model_visibilities.compute()

API

class africanus.experimental.rime.fused.terms.core.Term[source]

Base class for Terms which describe parts of the Fused RIME. Implementors of a RIME Term should inherit from it.

A Term is an object that defines how a term in the RIME should be sampled to produces the Jones Terms that make up the RIME. It therefore defines a sampling function, which in turn depends on arbitrary inputs for performing the sampling.

A high degree of flexibility and leeway is afforded when implementing a Term. It might be implemented by merely indexing an array of Jones Matrices, or by implementing some computational model describing the Jones Terms.

class Phase(Term):
    def __init__(self, configuration):
        super().__init__(configuration)
init_fields(self, typing_ctx, arg1, ..., argn, kwarg1=None, ..., kwargn=None)

Requests inputs to the RIME term, ensuring that they are stored on a state object supplied to the sampling function and allows for new fields to be initialised and stored on the state object.

Requested inputs arg1...argn are required to be passed to the Fused RIME by the caller and are supplied to init_fields as Numba types. kwarg1...kwargn are optional – if omitted by the caller, their default types (and values) will be supplied.

init_fields should return a (fields, function) tuple. fields should be a list of the form [(name, numba_type)], while function should be a function of the form fn(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 Simple Example.

Parameters
  • typingctx – A Numba typing context.

  • arg1...argn – Required RIME inputs for this Term.

  • kwarg1...kwargn – Optional RIME inputs for this Term. Types here should be simple: ints, floats, complex numbers and strings are ideal.

Return type

tuple

Returns

A (fields, function) tuple.

Warning

The function returned by init_fields must be compileable in Numba’s nopython mode.

sampler(self)

Return a sampling function of the following form:

def sampler(self):
    def sample(state, s, r, t, f1, f2, a1, a2, c):
        ...

return sample
Parameters
  • state – A state object containing the inputs requested by all Term objects in the RIME, as well as any fields created by Term.init_fields.

  • s – Source index.

  • r – Row index.

  • t – Time index.

  • f1 – Feed 1 index.

  • f2 – Feed 2 index.

  • a1 – Antenna 1 index.

  • a2 – Antenna2 index.

  • c – Channel index.

Return type

scalar or a tuple

Returns

a scalar or a tuple of two scalars or a tuple of four scalars.

Warning

The sampling function returned by sampler must be compileable in Numba’s nopython mode.

dask_schema(self, arg1, ..., argn, kwargs1=None, ..., kwargn=None)
Parameters
  • arg1...argn – Required RIME inputs for this Transformer.

  • kwarg1...kwargn – Optional RIME inputs for this Transformer. Types here should be simple: ints, floats, complex numbers and strings are ideal.

Return type

dict

Returns

A dictionary of the form {name: schema} defining the blockwise() dimension schema of each supplied argument and keyword argument.

class africanus.experimental.rime.fused.transformers.core.Transformer[source]

Base class for precomputing data for consumption by Term’s.

OUTPUTS

This class attributes should contain names of the outputs produced by the Transformer class. This should correspond to the fields produced by Transformer.init_fields().

init_fields(self, typing_ctx, 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. Functionally, this method behaves exactly the same as the init_fields() method, the difference being that the outputs are available to all Terms.

Return type

tuple

Returns

A (fields, function) tuple.

Warning

The function returned by init_fields must be compileable in Numba’s nopython mode.

dask_schema(self, arg1, ..., argn, kwargs1=None, ..., kwargn=None)
Return type

tuple

Returns

A (inputs, outputs) tuple.

inputs should be a dictionary of the form {name: schema} where schema is a dimension schema suitable for use in dask.array.blockwise(). A suitable schema for visibility data would be (row, chan, corr), while a uvw coordinate schema could be (row, uvw-component).

outputs should be a dictionary of the form {name: array_like}, where array_like is an object with dtype and ndim attributes. A suitable array_like for lm data could be np.empty((0,0), dtype=np.float64).

Predefined Terms

class africanus.experimental.rime.fused.terms.phase.Phase(configuration)[source]

Phase Delay Term

Attributes
configuration
class africanus.experimental.rime.fused.terms.brightness.Brightness(configuration, stokes, corrs)[source]

Brightness Matrix Term

Attributes
configuration
class africanus.experimental.rime.fused.terms.gaussian.Gaussian(configuration)[source]

Gaussian Amplitude Term

Attributes
configuration
class africanus.experimental.rime.fused.terms.feed_rotation.FeedRotation(configuration, feed_type, corrs)[source]

Feed Rotation Term

Attributes
configuration
class africanus.experimental.rime.fused.terms.cube_dde.BeamCubeDDE(configuration, corrs)[source]

Voxel Beam Cube Term

Attributes
configuration

Predefined Transformers

class africanus.experimental.rime.fused.transformers.lm.LMTransformer[source]
class africanus.experimental.rime.fused.transformers.parangle.ParallacticTransformer(process_pool)[source]

Numpy

Dask