RFI Visibility Calculation Components

class tabascal.components.rfi_vis.RiemannVis

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

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

Return pure, JIT-compatible function

build_set_params()

Build parameter sampling function (optional)

setup(config)

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisFFI

build_forward()

Return pure, JIT-compatible function

build_set_params()

Build parameter sampling function (optional)

setup(config)

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisVariable

build_constants()

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

Return pure, JIT-compatible function

build_set_params()

Build parameter sampling function (optional)

setup(config)

All validation and error-prone operations here

class tabascal.components.rfi_vis.RiemannVisVariableFFI

build_forward()

Return pure, JIT-compatible function

build_set_params()

Build parameter sampling function (optional)

setup(config)

All validation and error-prone operations here