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:
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:
the Phase Delay.
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:
It requests input for the Phase term. The above definition of
init_fields
signifies that the Phase term desires thelm
,uvw
andchan_freq
arrays. Additionally, these arrays will be stored on thestate
object provided to the sampling function.It supports reasoning about Numba types in a manner similar to
numba.generated_jit()
. Thelm
,uvw
andchan_freq
arguments contain the Numba types of the variables supplied to the RIME, while thetypingctx
argument contains a Numba Typing Context which can be useful for reasoning about these types. For exampletypingctx.unify_types(lm.dtype, uvw.dtype, chan_freq.dtype)
returns a type with sufficient precision, given the input types, similar tonumpy.result_type()
.It allows the user to define new fields, as well as a function for defining those fields on the
state
object. The above definition ofinit_fields
returns a list of(name, type)
tuples defining the new field names and their types, whilereal_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 thestate
object.Requested inputs
arg1...argn
are required to be passed to the Fused RIME by the caller and are supplied toinit_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)]
, whilefunction
should be a function of the formfn(arg1, ..., argn, kwarg1=None, .., kwargn=None)
and should return the variables of the type defined infields
. Note that it’s signature therefore matches that ofinit_fields
from after thetypingctx
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
- Returns
A
(fields, function)
tuple.
Warning
The
function
returned byinit_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 byTerm.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
- Returns
A dictionary of the form
{name: schema}
defining theblockwise()
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 theinit_fields()
method, the difference being that the outputs are available to all Terms.- Return type
- Returns
A
(fields, function)
tuple.
Warning
The
function
returned byinit_fields
must be compileable in Numba’s nopython mode.
- dask_schema(self, arg1, ..., argn, kwargs1=None, ..., kwargn=None)¶
- Return type
- Returns
A
(inputs, outputs)
tuple.inputs
should be a dictionary of the form{name: schema}
whereschema
is a dimension schema suitable for use indask.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}
, wherearray_like
is an object withdtype
andndim
attributes. A suitable array_like for lm data could benp.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