RFI Visibility Calculation Components

class tabascal.components.rfi_vis.PolyInterpVis[source]

RFI visibilities from the data grid, through tabascal.coarse_rfi_vis.coarse_rfi_vis().

Reads the signal, the phase and the delay polynomial on the data grid – from ComplexRFIVarAntCoarse and FixedOrbitCoarse – and rebuilds the fine samples of each cell inside the visibility calculation instead of reading them from fine-grid state. The signal is interpolated by the polynomial through the 2 * rfi.poly_interp_stencil + 1 nearest cells on each axis; the phase is rebuilt from its Taylor series across the cell.

This component does nothing but build the tables at setup and call the one function. That function is the whole of what a compiled kernel replaces; its module docstring is the specification, and the weight tables are inputs to it, so a different interpolant is a different table through the same kernel.

build_constants()[source]

Return arrays that do not change during the forward pass.

Returns a dict of array_name -> array_value. These will be stored in constants as “_c/<ClassName>/array_name” by Model.__init__.

build_forward()[source]

Return pure, JIT-compatible function

build_set_params()[source]

Build parameter sampling function (optional)

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.PolyInterpVisFFI[source]

PolyInterpVis through the compiled ri_kernels operator.

The same tables, the same inputs and the same result as PolyInterpVis; the one function that component calls is replaced by ri_kernels.jax_api.RFIInterpVisOp, whose CPU and GPU kernels carry the primal, the JVP and the transpose. The operator wants the antenna axis first, so the three data-grid arrays are transposed on the way in – data grid sized, so cheap. Needs an ri_kernels build that has the operator; a release without it is refused at setup rather than at the first forward.

build_forward()[source]

Return pure, JIT-compatible function

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.PolyInterpVisHybrid[source]

Slow baselines use quadrature; fast baselines use an analytic cell integral.

rfi.poly_analytic.quadrature_limit is the crossover in samples per cell. A null limit uses the measured VJP crossover for the working precision. Both groups reuse the compact antenna maps and scatter their results into MS order, so shared antennas accumulate both cotangents. The fast group’s time table holds monomial coefficients and its dt slot holds the cell duration; it never constructs a fine time grid.

The analytic phase is quadratic with a perturbative cubic correction, and carries the full interpolated amplitude polynomial. Splitting bounds local curvature and cubic phase independently of winding; segments, terms and cubic_terms control those expansions and their accuracy.

class tabascal.components.rfi_vis.PolyInterpVisHybridFFI[source]

PolyInterpVisHybrid through compiled quadrature and analytic ops.

Each nonempty group has an operator on its compact antenna axis. Analytic calls receive monomial coefficients and a scalar cell duration in the quadrature time-table and offset slots. The shared FFI forward handles gathers, antenna-first layout, baseline scatter and the source-shard sum. The analytic operator differentiates amplitude only; phase and delay are fixed trajectory inputs. Requires ri_kernels from the interp-analytic branch.

build_forward()

Return pure, JIT-compatible function

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.PolyInterpVisVariable[source]

PolyInterpVis with the fine sampling set per baseline group.

rfi.poly_time_sampling chooses at most two groups from the per-baseline fringe-rate requirements. Each group has its own odd quadrature count and interpolation tables, built at fine_offsets for that count. One group therefore evaluates the same grid as the non-variable route.

The split minimises the work of materialising antenna samples, including antennas used by both groups. Only the antennas a group uses enter its call, with baseline endpoints remapped to that compact axis. The gathers’ transposes add the cotangents from both calls at a shared antenna; an identity antenna map bypasses the gather altogether.

build_forward()[source]

Return pure, JIT-compatible function

class tabascal.components.rfi_vis.PolyInterpVisVariableFFI[source]

PolyInterpVisVariable through the operator, one call per group.

Each operator is constructed for the group’s compact antenna axis and local baseline endpoints. It receives independent time tables and only that group’s antenna inputs, transposed to the antenna-first layout the operator expects. Shared antennas are gathered into both calls, so their cotangents accumulate back onto the original data-grid signal. No kernel change is needed: the groups differ only in input shapes and tables.

build_forward()[source]

Return pure, JIT-compatible function

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVis[source]

Riemann-sum RFI visibilities in pure JAX, scanned over the baseline axis.

The reference implementation of the same integral as RiemannVisFFI, and the one that kernel is validated against in value, forward mode and reverse mode. The baseline axis is walked in blocks of rfi.baseline_block_size under checkpoint (see tabascal.interferometry.calculate_rfi_vis_blocked()) so that the fine grid it integrates is bounded by the block rather than by the whole array: what the forward pass leaves behind for reverse mode is the result and a transposed copy of its per-antenna inputs, not the (n_bl, n_rfi, n_freq_fine, n_time_fine) intermediate the reduction is built from.

It trades recomputation for memory rather than aiming at speed. The block size does not change the result – baselines are independent – only how much of the fine grid is live at once, and how many scan steps that takes. A null block size is every baseline in a single step: the fine grid is still recomputed rather than stored, so the tape stays small, but it is formed whole. Measured on one GH200 that peaks where the unscanned kernel did, and across four it is well under it, the tape being per-device memory that the collective does not divide.

build_constants()[source]

Return arrays that do not change during the forward pass.

Returns a dict of array_name -> array_value. These will be stored in constants as “_c/<ClassName>/array_name” by Model.__init__.

build_forward()[source]

Return pure, JIT-compatible function

build_set_params()[source]

Build parameter sampling function (optional)

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisFFI[source]
build_forward()[source]

Return pure, JIT-compatible function

build_set_params()[source]

Build parameter sampling function (optional)

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisVariable[source]
build_constants()[source]

Return arrays that do not change during the forward pass.

Returns a dict of array_name -> array_value. These will be stored in constants as “_c/<ClassName>/array_name” by Model.__init__.

build_forward()[source]

Return pure, JIT-compatible function

build_set_params()[source]

Build parameter sampling function (optional)

setup(config)[source]

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisVariableFFI[source]
build_forward()[source]

Return pure, JIT-compatible function

build_set_params()[source]

Build parameter sampling function (optional)

setup(config)[source]

All validation and error-prone operations here