Read the MS SIGMA column per baseline instead of collapsing it to a scalar

[chrisfinlay]

Progresses #121. Unblocks #110.

read_ms did SIGMA.mean(), so the per-baseline noise that uvfits2ms estimates was computed, written to the MS, and then thrown away. Every run with data.noise: null was fit under a uniform noise.

Why it matters

On EDA2 the per-baseline SIGMA spans a factor of ~30, so a scalar under-weights the quietest baselines by up to ~200x in the likelihood.

It is worse for anything that fits gains. The per-antenna noise correlates with the per-antenna gain — measured sigma_a ~ amplitude_a^0.76, R = 0.96. A uniform-noise likelihood cannot tell a loud antenna from a noisy one, so the fitted gain absorbs the noise structure. That is a bias in the calibration solution, not merely lost efficiency.

What changed

New tabascal/noise.py owns per-baseline noise:

• Median over time, not mean. SIGMA is per row and constant in time per baseline, so a median keeps a few corrupted rows from dragging a baseline's estimate.
• Dead baselines take the median of the rest. A non-positive or non-finite estimate carries no information, and a zero would divide the likelihood by nothing. They are flagged out of it anyway.
• data.noise still overrides, and now also accepts an .npz carrying sigma_bl (per baseline) or s_ant (per antenna, combined as sqrt(s_p^2 + s_q^2)/sqrt(2) — normalised so uniform antennas reproduce themselves).

Consumers split by what they actually need. The likelihood and reduced_chi2 take the per-baseline array. The time-integration sampling heuristic and the truth-metric normalisation take a new noise_scalar, because both reduce over every baseline and have no baseline axis left to align with.

The subtle part

reduced_chi2 broadcasts the noise onto the data before masking:

noise = jnp.broadcast_to(broadcast_to_vis(noise, true.shape), true.shape)
rchi2 = jnp.sum((jnp.abs(pred[~flags] - true[~flags]) / noise[~flags]) ** 2) / norm

x[~flags] flattens, so applying a per-baseline noise afterwards would silently recycle values across baselines. Verified by mutation: reverting the ordering fails 2 tests — and only shows up under a ragged flag mask, which is why the test builds one (one sample dropped from baseline 0, baseline 3 dropped entirely).

One test of mine was initially unsound and I fixed it: with the same residual on every baseline, sum((r / noise_i)^2) is invariant under permuting the noise, so a shuffled pairing scores identically and proves nothing. The residual has to vary per baseline for the check to bite.

Under sharding, noise was cast with float(); a per-baseline array cannot survive that, so it is made global and replicated instead.

No reference re-recording

The simulated benchmark MS has a uniform SIGMA — one unique value, 0.6496226 — so per_baseline_sigma reproduces the old scalar exactly. All 10 pipeline chi2 references pass unchanged, on the nose. The variation this fixes lives in real EDA2 data.

Testing

852 tests pass, including 29 new ones in tests/test_noise.py. Docs updated for data.noise.

Not in this PR

#121 also asks for a model-free channel-differencing noise estimator, for data whose SIGMA is absent or untrustworthy. That is a separate input path rather than part of unbreaking the existing one, and it is easier to review on its own — the issue stays open for it.

[read-the-docs-community]

Documentation build overview

📚 tabascal | 🛠️ Build #34091226 | 📁 Comparing ed6c1fb against latest (2de5545)

  🔍 Preview build  

2 files changed

± config.html
± _modules/tabascal/ms.html