The Data-Grid RFI Visibility (kernel boundary)

RFI visibilities from the data grid: the function a compiled kernel replaces.

coarse_rfi_vis() is the whole of the boundary. Everything comes in on the data grid, the data-grid visibilities go out, and the fine grid – the n_int_freq x n_int_time samples inside each data cell that the fine-grid route (trajectory:FixedOrbit, rfi_signal:ComplexRFIVarAnt, rfi_vis:RiemannVisFFI) carries as model state – exists only inside it, one time cell at a time. A kernel implementing this function, with its JVP and VJP with respect to rfi_A, drops into tabascal.components.rfi_vis.PolyInterpVis exactly where ri_kernels drops into RiemannVisFFI. Nothing else in the route needs to change: the components that produce the inputs run on the host at setup or are a plain inverse FFT, and are the same for a compiled kernel as for this reference.

analytic_rfi_vis() is the pure-JAX alternative for fast baselines. It uses the same coarse inputs and frequency quadrature, but integrates the amplitude polynomial against a quadratic phase, with a cubic correction. Its time table holds monomial coefficients rather than fine-sample weights, so the work has no Nyquist floor and does not grow with fringe winding.

Inputs

Per source r, antenna a, channel f and time cell t. Only rfi_A carries a gradient; every other input is a constant of the run.

• rfi_A (n_rfi, n_ant, n_freq, n_time), complex: the RFI signal on the data grid.

• rfi_phase (n_rfi, n_ant, n_freq, n_time): the phase at the channel and cell centre, reduced to one turn.

• rfi_delay_poly_us (n_rfi, n_ant, n_time, n_path): the geometric delay tau (us) and its time derivatives tau_k (us/s^k) at the cell centre, relative to the array mean (see below). The delay is -(range + w) / c, so the phase is 2 pi f tau: the fine-grid route’s rfi_delay_us convention.

• w_freq (n_freq, n_sf, n_int_freq) and start_freq (n_freq,): interpolation weights across each channel, and the first channel of each channel’s stencil.

• w_time (n_time, n_st, n_int_time) and start_time (n_time,): the same across each cell.

• dnu_mhz (n_int_freq,) and dt (n_int_time,): the fine offsets from the channel centre (MHz) and the cell centre (s).

• freqs_mhz (n_freq,): channel centres (MHz). MHz times microseconds is cycles, so the phase is formed without a scaling constant.

• a1, a2 (n_bl,): the two antennas of each baseline.

Output: vis_rfi (n_bl, n_freq, n_time), complex.

The computation

For one cell (f, t) and one of its fine samples (u, v), with c the speed of light:

signal   A[r, a](u, v) = sum_k sum_l  w_freq[f, k, u] w_time[t, l, v]
                                      rfi_A[r, a, start_freq[f] + k, start_time[t] + l]
delay    dtau[r, a](v) = sum_{k >= 1} rfi_delay_poly_us[r, a, t, k] dt[v]^k / k!
phase    phi[r, a](u, v) = rfi_phase[r, a, f, t]
                           + 2 pi ( (freqs_mhz[f] + dnu_mhz[u]) dtau[r, a](v)
                                    + dnu_mhz[u] rfi_delay_poly_us[r, a, t, 0] )
sample   S[r, a](u, v) = A[r, a](u, v) exp(i phi[r, a](u, v))
result   vis_rfi[b, f, t] = mean_{u, v} sum_r S[r, a1[b]](u, v) conj(S[r, a2[b]](u, v))

The last line is the integrand of the fine-grid Riemann sum (tabascal.interferometry.calculate_rfi_vis_fine(), averaged over each cell as calculate_rfi_vis_blocked() does); the three lines before it are what replace reading A and phi from the fine-grid state. The signal interpolates the 2h + 1 nearest coarse samples on each axis; the phase is exact across the channel (linear in frequency) and a Taylor series across the cell.

The weights are data. tabascal.poly_interp.interp_tables() fills them with the polynomial through the stencil; the conditional mean of a Gaussian process prior, or any other linear interpolant, is a different table and the same kernel.

Precision

The phase is arranged so that a single-precision kernel never forms a large number and then reduces it. The unreduced phase 2 pi freqs L / c is of order a million turns, which float32 cannot hold to a fraction of a turn; rfi_phase carries it reduced, computed in float64 on the host. The kernel adds to it only the change across the cell and across the channel. Do not rebuild rfi_phase from rfi_delay_poly_us[..., 0] inside a kernel.

Derivatives

rfi_A is the only differentiated input; rfi_phase and rfi_delay_poly_us come from a fixed orbit and the rest are tables. The result is bilinear in the fine samples S, and S is linear in rfi_A (the interpolation is linear and the phase factor is a constant), so both derivatives are structural:

• JVP: with dS the tangent of rfi_A pushed through the same interpolation and phase factor, dvis[b] = mean sum_r ( dS[r, a1[b]] conj(S[r, a2[b]]) + S[r, a1[b]] conj(dS[r, a2[b]]) ).

• VJP: the cotangent of vis_rfi is scattered to the fine samples of each baseline’s two antennas, each weighted by the other antenna’s sample – the fine-grid kernel’s own transpose – then multiplied by the conjugate phase factor and pushed back through the weight tables, which is a stencil-sized scatter-add onto the data grid, and nothing of the fine grid survives it.

Sign and conjugation conventions are JAX’s for complex inputs; the reference for both is JAX’s derivative of this function, and tests/test_coarse_rfi_vis.py holds that to finite differences. A kernel is validated the way ri_kernels is: value, JVP and VJP against this function.

Memory

The reference forms one time cell at a time, (n_bl, n_rfi, n_freq, n_int_freq, n_int_time) complex, under jax.checkpoint so the reverse pass recomputes the cell rather than keeping it. That is what makes it usable at the sizes the fine-grid route runs at, not what makes it fast: it gathers per baseline and recomputes each antenna’s samples for every baseline it is on. A kernel would stage each antenna’s fine samples once per cell.

tabascal.coarse_rfi_vis.analytic_rfi_vis(rfi_A: Array, rfi_phase: Array, rfi_delay: Array, w_freq: Array, start_freq: Array, g_time: Array, start_time: Array, dnu_mhz: Array, int_time: Array, freqs_mhz: Array, a1: Array, a2: Array, *, segments: int = 2, terms: int = 6, cubic_terms: int = 3) → Array

Integrate the amplitude polynomial against the quadratic delay phase.

g_time[t,l,m] is the Lagrange basis in powers of x = 2*tau/T. The frequency contraction is unchanged, so finite channel integration retains exactly the reference’s frequency offsets. Time integration is analytic: multiply the antenna polynomials by convolution, then contract against phase moments. The quadratic phase is integrated analytically and residual cubic phase is expanded on each piece. Three terms suffice at the measured cubic coefficients; zero deliberately drops cubic phase for comparison. Derivatives above order three are omitted.

Equal pieces keep curvature small without imposing a Nyquist sample count. Translation of both the amplitude and phase is exact, including the constant phase of each piece. The working arrays carry polynomial degree, not fringe winding. Only the data-grid amplitude is differentiated.

tabascal.coarse_rfi_vis.cell_vis(S: Array, a1: Array, a2: Array) → Array

The cell’s visibilities from its fine samples, (n_bl, n_freq).

S is (n_rfi, n_ant, n_freq, n_int_freq, n_int_time): the summed product over sources, averaged over the cell’s fine samples.

tabascal.coarse_rfi_vis.coarse_rfi_vis(rfi_A: Array, rfi_phase: Array, rfi_delay: Array, w_freq: Array, start_freq: Array, w_time: Array, start_time: Array, dnu_mhz: Array, dt: Array, freqs_mhz: Array, a1: Array, a2: Array) → Array

The data-grid RFI visibilities, (n_bl, n_freq, n_time).

The reference implementation of the function described in the module docstring, and the boundary a compiled kernel replaces. One time cell per step of a lax.map, each cell under jax.checkpoint.

tabascal.coarse_rfi_vis.fine_phase(rfi_phase: Array, rfi_delay: Array, freqs_mhz: Array, dnu_mhz: Array, dt: Array) → Array

The fine phase of one time cell, for every source, antenna and channel.

rfi_phase (n_rfi, n_ant, n_freq) and rfi_delay (n_rfi, n_ant, n_path) are the cell’s own slices; the delay in microseconds and its derivatives in microseconds per second^k, the frequencies in MHz.

Returns (n_rfi, n_ant, n_freq, n_int_freq, n_int_time).

tabascal.coarse_rfi_vis.fine_signal(rfi_A: Array, w_freq: Array, start_freq: Array, w_time: Array, start_time: Array) → Array

The fine samples of one time cell, for every source, antenna and channel.

w_time (n_st, n_int_time) and start_time (a scalar) are the cell’s own row of the tables; the frequency tables are whole, since every channel of the cell is interpolated at once.

Returns (n_rfi, n_ant, n_freq, n_int_freq, n_int_time).

tabascal.coarse_rfi_vis.linear_phase_moments(a: Array, degree: int) → Array

The moments mean_{[-1,1]} x**m exp(i*a*x), through degree.

Integration by parts divides by the large winding, so upward recurrence is stable only while m <= |a|. Above that point we run the same identity downwards from a zero tail, 64 orders beyond the last requested moment. The unwanted solution then contracts by |a|/m at every step. This also supplies the zero-frequency limit without a division by zero or a cancellation-prone Taylor series at moderate winding.

tabascal.coarse_rfi_vis.quadratic_phase_moments(a: Array, b: Array, degree: int, terms: int = 16) → Array

Normalised moments on [-1, 1] of exp(i*(a*x+b*x*x)).

Outside or beyond the neighbourhood of the stationary point, expand the curvature about linear-phase moments. The caller splits the cell first: with |b| <= 1 per piece, 16 terms leave less than 2e-13 absolute remainder. Near the stationary point the Fresnel seed and its upward recurrence are stable, except at small b where division by b is itself ill-conditioned. There the same convergent series supplies the continuous limit instead.