Astronomical Visibility Components

class tabascal.components.ast_vis.DiscreteSkyVis

Visibilities of a discrete sky by direct DFT, with the full w term.

Reads ast_radec (n_src, 2) radians, ast_I (n_src, n_freq) Jy and ast_shape (n_src, 3) radians from the state — see FixedDiscreteSky, which must be listed before this component — and ACCUMULATES into vis_ast using the visibility equation

V(u,v,w) = sum_k I_k G_k(u,v) exp(-2i pi (u l_k + v m_k + w (n_k - 1)) / lambda)

For a discrete sky the direct sum is exact, gridless, differentiable, and unaffected by field of view or baseline length; “discrete” is the set of sources, not their size, so ImageSkyVis is reserved for a sky carried as an image.

I_k enters undivided: there is no 1 / n. The RIME integrand carries B / n because dOmega = dl dm / n, but a source of integrated flux S is B = S delta_Omega and delta_Omega = n delta(l) delta(m), so the Jacobian cancels and a source contributes its catalogue flux exactly, in every direction. (u, v, w) is the ANTENNA2 - ANTENNA1 baseline the equation above is written for, which is ast.uvw_sign times the UVW column. Which baseline that column holds is a property of whatever wrote the data, so the sign is a config key; the default, (-1, -1, -1), is the convention tab-sim writes — see setup.

G_k is the uv-plane envelope of an elliptical Gaussian source,

G(u,v) = exp(-pi^2 / (4 ln 2) * (a^2 u’^2 + b^2 v’^2))

for FWHM a (major) and b (minor) in radians, with (u', v') the baseline in wavelengths rotated into the source frame:

u' = u sin(phi) + v cos(phi)        along the major axis
v' = u cos(phi) - v sin(phi)        along the minor axis

phi is the position angle in the radio convention, measured from north (the m axis) through east (the l axis), so phi = 0 puts the major axis north-south and a north-south baseline is the one that resolves the source out. A zero FWHM gives G = 1 exactly, so points and Gaussians are the same code path.

It accumulates rather than assigns, so it composes with the astronomical GP: with both listed, vis_ast is the GP plus the fixed sources. Model zeroes vis_ast before the forward chain runs, so the two may be listed in either order.

The source axis is walked in blocks of ast.source_block_size with jax.lax.scan(), and the block body is rematerialised: the delay array is (n_bl, n_time, n_src), which for a real catalogue is the largest array in the model, and blocking replaces n_src in that shape with the block size at the cost of recomputing each block in the backward pass.

Sources more than 90 degrees from the phase centre are modelled, not rejected: only n - 1 enters the phase and radec_to_lmn() computes it exactly over the whole sphere, so nothing here breaks down at n <= 0. Such a source is almost always a catalogue mistake rather than a real one, so FixedDiscreteSky warns about it at setup and leaves the decision to the caller.

A fixed sky exists to make a per-antenna gain identifiable (see issue #124); the flux scale it fixes the gain against is only physical if the data are calibrated to Jy.

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

Build the forward computation function

setup(config)

Initialize component with configuration

class tabascal.components.ast_vis.GPVisAst

build_block_signal()

The latent-to-signal transform of one block of baselines.

Split out so a subclass can replace the transform and nothing else: the priors, the initialisation and the blocking above are the same whichever way the modes are turned into a signal.

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

The baseline axis is walked in blocks of ast.baseline_block_size rather than vmapped whole. latent_to_signal pads the latent block up to the padded k-grid, shifts, inverse-transforms and crops back, so a vmap over every baseline holds (n_bl, n_freq_pad, n_time_pad) three times over – at the default padding each axis is about twice the data one, so about four times the elements of the visibilities, and the crop throws all of that away. The scan replaces n_bl in that shape with the block, which is the whole of the term this component contributes to peak memory.

There is deliberately no checkpoint on the body, unlike the RFI components’ scans. The chain from the latent parameters to vis_ast is affine – an elementwise sigma * base + mu, then pad, shift, ifftn and crop, every one of them linear – so reverse mode is its transpose and stores no primal intermediates to begin with. Issue #153 records the measurement; a remat here would be a no-op today and a silent recomputation if the chain ever stopped being linear.

The affine transform runs inside the body, on the block, so that neither it nor the padded grid is ever formed for every baseline at once.

Splitting the baseline axis across devices was measured alongside this and is deliberately not here: on the astronomical benchmark the scan takes the value-and-gradient peak from 1008 MB to 163.5 MB on its own, and adding the split takes it back up to 174.5 MB, because the gather the visibilities need on the way out makes the backward pass scatter a cotangent the forward has just collected. See issue #209 for the change that would remove the padded grid altogether.

build_set_params()

Build parameter sampling function (optional)

setup(config)

All validation and error-prone operations here

validate_and_test()

Call this before using in JIT context

class tabascal.components.ast_vis.GPVisAstDFT

GPVisAst with the padded grid never formed.

The prior, the initialisation and the blocking over baselines are the parent’s; only the step from the surviving Fourier modes to the visibilities differs. latent_to_signal() pads those modes back up to the padded grid, inverse-transforms the whole of it and crops – at the default padding, about four times the elements of the visibilities formed and thrown away, per block, in the forward pass and again in reverse. The chain is linear and separable, so the same values come out of one small matrix per axis (latent_to_signal_dft()), and the largest thing formed is the size of the visibilities.

The modes of this component are set by the array’s geometry rather than by a choice, so the matrices are fixed at setup and shared by every baseline.

Select it where the parent would go:

model:
  components:
    - ast_vis:GPVisAstDFT

build_block_signal()

The latent-to-signal transform of one block of baselines.

Split out so a subclass can replace the transform and nothing else: the priors, the initialisation and the blocking above are the same whichever way the modes are turned into a signal.

tabascal.components.ast_vis.radec_to_lmn(ra, dec, ra0, dec0)

Direction cosines of sources at (ra, dec) about a phase centre, in radians.

Returns (l, m, n, n - 1). n - 1 is returned alongside n because it is the quantity the w term actually needs and the two cannot both be computed accurately from one expression.

n is the exact spherical form sin(d) sin(d0) + cos(d) cos(d0) cos(da) rather than sqrt(1 - l^2 - m^2). The square root is the cosine of the angular distance only on the near hemisphere: it is unsigned, so it folds a source more than 90 degrees from the phase centre back onto the near side instead of giving it the negative n it has.

n - 1 uses the haversine identity n - 1 = -2 h with h = sin^2((d - d0)/2) + cos(d) cos(d0) sin^2(da/2), the haversine of the angular distance. h runs over [0, 1] across the whole sphere — 0 at the phase centre, 1 at the antipode — so n - 1 runs over [-2, 0] and n < 0 exactly when h > 1/2. Subtracting a nearby n from 1 cancels catastrophically: at a 40 arcsec offset 1 - n ~ 2e-8, which in single precision is below the spacing of the floats either expression lands on, so the difference comes out as exactly zero and the w term disappears. The haversine form never forms the difference, so it keeps full relative accuracy at any offset.