Interpolation Tables and Path Fits

Setup-time tables for the data-grid RFI route.

Everything here runs once, on the host, in float64 numpy, and produces the constant inputs of tabascal.coarse_rfi_vis.coarse_rfi_vis():

• interp_tables(): for each data cell on an axis, the weights that turn the 2h + 1 coarse samples nearest to it into its fine samples, and where that stencil starts. These are the polynomial (Lagrange) weights. The kernel takes the tables as data, so any other linear interpolant – the conditional mean of a Gaussian process, a windowed sinc – is the same kernel fed a different table.

• fit_path(): the polynomial in time through each cell’s propagated delay samples, from whose coefficients the kernel rebuilds the fine phase.

• fine_offsets(): where a cell’s fine samples sit relative to its own data-grid sample.

• poly_time_groups(): the baseline partition and compact antenna maps that decide which independent time tables each variable call needs.

class tabascal.poly_interp.DeviceGroup(positions: NDArray, a1: NDArray, a2: NDArray, n_real: int, n_ghost_antennas: int = 0)

One device’s share of a group, shaped identically on every device.

a1/a2 index the group’s antenna axis extended by n_ghost_antennas dark antennas, the same count on every device. n_real output rows are real baselines; ghosts scatter to a spare row that is discarded before returning the visibility array. positions says where the real rows belong in the device’s own block of it.

property n_padded: int

Rows the operator is asked for, real and ghost, the same on every device.

class tabascal.poly_interp.PolyTimeGroup(baseline_indices: NDArray, n_g: int, antennas: NDArray, a1: NDArray, a2: NDArray)

Host-side baseline membership and compact antenna indices for one call.

tabascal.poly_interp.analytic_sampling_cut(half_width: int, segments: int, terms: int, cubic_terms: int = 3) → int

A conservative operation-count crossover, in quadrature samples.

For product degree P=4h and J cubic terms, moments reach D=P+3(J-1), and the linear recurrence reaches M=D+2(K-1). Per piece it takes M upward and M+64 downward stages, K(D+1) curvature and J(P+1) cubic contractions, and (2h+1)^2 amplitude products. Counting each as one sample is deliberately conservative: a quadrature sample also interpolates and exponentiates. This is an arithmetic diagnostic, not a runtime model. The hybrid components use measured precision-dependent VJP crossovers instead.

tabascal.poly_interp.device_groups(group: PolyTimeGroup, n_dev: int, n_bl_total: int, offset: int = 0) → tuple[DeviceGroup, ...]

Split a group by which device owns each baseline, padding to one shape.

With block = ceil(n_bl_total / n_dev), device d owns the contiguous range [d*block, (d+1)*block) of the visibility array in the order the data already has, so nothing is reordered: the observed visibilities, the flags and the noise are simply sharded along the axis they already have, and every array downstream keeps that split all the way to the likelihood. A device therefore never needs anyone else’s visibility rows. Likelihood scalars and reverse-mode gradients of shared per-antenna inputs still need reductions. For indivisible totals the last block extends past the real array; the caller trims those rows and uses replicated placement at component boundaries.

Which of a group’s baselines land on a device is then whatever the data ordering puts there, so the counts differ between devices. shard_map runs one program, so they are padded up to a common count with dark ghost baselines: distinct pairs against extra antennas carrying no signal, which collide with no real baseline, stay distinct from each other, and contribute neither visibility nor gradient – the same device padded_rfi_count() uses on the source axis.

positions are local to the owning device’s block, so a group writes into the device’s own rows, and are padded to the same length as a1 with block – one past the block’s last row. The caller gives its local visibility array that one extra row, lets the ghosts land in it and drops it, so the scatter has a fixed shape and a ghost can never overwrite a real baseline. offset is where this group’s share starts within that block, since several groups share it.

tabascal.poly_interp.fine_offsets(n_int: int, spacing: float) → NDArray

Offsets of a cell’s fine samples from its data-grid sample, (n_int,).

TabConfig lays a fine axis out with n_int samples per cell, in cell order, spacing / n_int apart, with the cell’s own data-grid sample at index n_int // 2 of its block (the crop in tabascal.fft_gp.latent_to_signal_init()), so the offsets are (v - n_int // 2) * spacing / n_int. They are formed here from the count and the spacing rather than read off the config’s fine grid: that grid is built from a unit grid in the run’s precision, and in single precision its samples jitter by 1e-4 of a cell, which is metres of path at orbital speed. tests/test_poly_interp.py holds this to the layout that function makes.

tabascal.poly_interp.fit_path(path_fine: NDArray, offsets: NDArray, order: int) → NDArray

Least-squares polynomial through each cell’s fine path samples.

Parameters:

• path_fine (Array (..., n_time, n_int)) – The path – or delay, or any quantity – at the fine samples of each cell, in whatever unit the caller keeps it in.

• offsets (Array (n_int,)) – Offsets of the fine samples from the cell centre, in seconds.

• order (int) – Degree of the polynomial. Reduced to n_int - 1 when a cell has fewer samples than a higher degree needs; at one sample per cell that is the path at the centre and nothing else.

Returns:

Coefficients L_k with path(t_c + d) ~ sum_k L_k d^k / k!: the quantity and its first order time derivatives at each cell centre, in its unit per second^k.

Return type:

Array (…, n_time, order + 1)

tabascal.poly_interp.interp_tables(n_cells: int, half_width: int, offsets: NDArray) → tuple[NDArray, NDArray]

Per-cell interpolation weights and stencil starts for one axis.

Parameters:

• n_cells (int) – Number of data cells on the axis.

• half_width (int) – Stencil half-width h: 2h + 1 cells enter each cell’s interpolant, which is the polynomial of degree 2h through them. Reduced to what the axis can hold when it has fewer cells than that.

• offsets (Array (n_int,)) – Positions of a cell’s fine samples relative to its centre, in units of the cell spacing (see fine_offsets()).

Returns:

• weights (Array (n_cells, n_stencil, n_int))

• start (Array (n_cells,) of int) – fine[c, v] = sum_k weights[c, k, v] * coarse[start[c] + k].

Notes

The stencil of cell c is the 2h + 1 cells nearest to it: centred on c away from the edges, and at the edges shifted inwards, where the polynomial through those same nearest cells is evaluated off-centre rather than a shorter one fitted. Away from the edges every row of weights is the same table; a kernel may exploit that, the reference does not.

tabascal.poly_interp.lagrange_basis(nodes: NDArray, x: NDArray) → NDArray

Lagrange basis through nodes evaluated at x, as (n_nodes, n_x).

Row k is the polynomial of degree n_nodes - 1 that is 1 at nodes[k] and 0 at every other node, so values @ basis is the interpolating polynomial through (nodes, values) sampled at x.

tabascal.poly_interp.make_poly_time_group(requirements, a1, a2, indices) → PolyTimeGroup

Materialise a chosen partition with the same compact maps on every route.

tabascal.poly_interp.monomial_tables(n_cells: int, half_width: int) → tuple[NDArray, NDArray]

Lagrange coefficients in x = 2*tau/T, including the shifted edge stencils.

These occupy the time-table slot of fine_signal: the identical stencil contraction now yields polynomial coefficients rather than sampled values. No Vandermonde fit is needed; multiplying each basis’s linear factors gives its monomials directly on the host in double precision.

tabascal.poly_interp.poly_sample_counts(requirements: NDArray) → NDArray

Round fringe-rate requirements up to positive odd quadrature counts.

The estimate bounds midpoint quadrature error. With fine_offsets an even count instead shifts the grid half a sample to the left, introducing a first-order phase error that rounding up alone does not control. Odd counts keep that convention and sample the midpoints of equal sub-intervals.

tabascal.poly_interp.poly_time_groups(requirements: NDArray, a1: NDArray, a2: NDArray, *, max_groups: int = 2, split_at: int | None = None) → tuple[PolyTimeGroup, ...]

Choose one or two groups by the work of materialising antenna samples.

Each group’s count is the largest rounded requirement of its baselines. Search every threshold between distinct requirements, scoring a partition by sum(n_g * len(antennas_g)). The source, channel and cell dimensions multiply every candidate equally, so they need not enter the score. An antenna shared by the two groups is materialised twice and counted twice. This is work, not a runtime prediction for either implementation.

split_at restricts the search to requirements at or below that threshold versus requirements above it. An empty side or a split that does not strictly improve on one group falls back to one group; ties favour fewer groups. Equal two-group scores choose the lowest threshold deterministically.

Sorting once and accumulating prefix/suffix antenna incidence gives all candidate antenna counts without constructing a baseline mask per split. Only the winning partition is materialised into static group records.

tabascal.poly_interp.split_group_over_devices(group: PolyTimeGroup, n_dev: int) → tuple[PolyTimeGroup, ...]

Divide one group’s baselines evenly over devices, keeping its antennas.

Each device computes a share of the group’s baselines and none of anyone else’s, so nothing has to be summed across devices – unlike sharding the source axis, where every device computes every baseline for a few sources and the partial visibilities must be added back together.

The antenna set is deliberately not recompacted per device, which is what separates this from make_poly_time_group(). Every device keeps the whole group’s antennas, so the per-antenna signal is identical on all of them and enters the map replicated; recompacting would give each device a different antenna count, and shard_map runs one program with one set of shapes. The signal is small – 0.24 GB at 512 stations against 1.256 GB for a single visibility array – so replicating it costs far less than the visibilities it lets us divide.

a1 and a2 stay indices into the group’s compact antenna axis, and baseline_indices stay indices into the global visibility array, so a device knows where its own results belong.

Raises when the count does not divide: shard_map needs one shape for every device, and a silent remainder would drop baselines.