Fix two gain-handling bugs when writing results to the MS - #134
Fix two gain-handling bugs when writing results to the MS#134chrisfinlay wants to merge 4 commits into
Conversation
|
The What failed: one perf check of eight — The metric: Why it is not the change in this PR.
Seven of eight are within ±0.36 s. The failing one is +13.7 s, and it was the first check to run, consistent with a cold-start or scratch-filesystem stall. Not a time-limit kill either: #131 passed Unrelated, but it cost diagnosis time: the reframe failure line reads while the report table shows |
|
cscs-ci run default |
Both were latent under UnitaryGains, which every test and shipped config uses,
and both surface as soon as a run fits a non-unit gain.
The baseline gain was built from ANTENNA1 twice:
a1 = xds_ms.ANTENNA1.data[:n_bl]
a2 = xds_ms.ANTENNA1.data[:n_bl] # ANTENNA2
so gains_bl was |g_p|^2 -- real, positive, blind to the second antenna, all
phase discarded, and wrong on every baseline whose antennas differ.
Residuals were formed as vis_obs - (vis_ast + vis_rfi) using the un-gained model
from the results zarr, while the forward model is gains_bl * (vis_ast + vis_rfi).
The correct expressions were present but commented out, and gains_bl was
computed and then never used at all.
Residuals are now formed in the data frame, vis_obs - gains_bl * model, rather
than the calibrated frame vis_obs/gains_bl - model: dividing by the gain inflates
the noise on low-gain baselines and distorts noise-referenced residual metrics.
CORRECTED_DATA carries the calibrated data, vis_obs / gains_bl. Moving every
column into one calibrated frame with matching weights is #123.
All of this reduces exactly to the previous behaviour when the gains are unity,
so existing results are unchanged and no references move.
The arithmetic is factored into baseline_gains, data_frame_residuals and
read_antenna_pairs so it can be tested without an MS fixture -- read_antenna_pairs
exists specifically so that reading one column twice is a testable mistake rather
than a two-line typo. Verified by mutation: reverting the ANTENNA1/2 bug fails 2
tests, reverting the residual frame fails 3.
The legacy 3-d ast_vis layout carries no baseline axis, so its per-baseline gain
cannot be reconstructed. It now raises rather than silently writing columns in
the wrong frame, unless the stored gains are unity -- the only case it ever
handled correctly.
Closes #122
Both from codex review. The unity guard for the legacy 3-d layout tested the MEAN of the stored gains, so samples of 0.9 and 1.1 average to exactly 1 and slipped through -- waving past precisely the case the guard exists to catch. Now every sample is tested. The per-baseline gain was formed after averaging the sample axis, giving E[g_p] conj(E[g_q]) instead of E[g_p conj(g_q)]. These differ whenever the two antennas' gains vary together across samples: for gains rising 1 -> 2 on both, E[g^2] = 2.5 against E[g]^2 = 2.25. Not reachable today -- both writers of this zarr use Predictive with x[None] and batch_ndims=1, i.e. exactly one sample, and run_opt carries a comment forbidding multi-sample under sharding -- but free to get right, so baseline_gains grew an ant_axis and the product is formed per sample via mean_baseline_gains. The model visibilities are still averaged independently of the gains, so with more than one sample the residuals would be E[g]E[model] rather than E[g*model]. That is a deeper restructure than this fix warrants, so it warns instead. Verified by mutation: reverting the guard to the mean fails 1 test, and making mean_baseline_gains reduce before multiplying fails 1.
Codex review, second pass: a warning does not stop downstream use of incorrect residual columns, so form them correctly instead. The residuals now subtract E[g*model], computed per sample and reduced after, rather than E[g]*E[model] built from independently averaged factors. These differ whenever the gains and the model covary across samples, which posterior draws from a joint fit generally do: for a gain rising 1 -> 2 alongside a model rising 1 -> 3, E[gm] = 3.5 against E[g]E[m] = 3.0. data_frame_residuals now takes models that are already gained, since only the caller still holds the sample axis. gained_model_mean does that reduction and exists as a named function so the order is pinnable -- verified by mutation: reducing before multiplying fails a test. Not doing the other review point, reconstructing gains for the legacy 3-d layout. The claim is that ast_vis is (sample, bl, time) there so n_bl is recoverable, but no artefact in this repository supports it: every write.py in the history, back to the oldest commit that has the file, writes 4-d ["sample", "bl", "freq", "time"]. No writer here has ever produced a 3-d ast_vis, so the axis order is an inference and the branch is unreachable dead code. Writing speculative reconstruction against a layout with no fixture and no producer is worse than an error that says to re-run; the guard stays.
15275be to
5e9e2cc
Compare
| mistake rather than a two-line typo. | ||
| """ | ||
|
|
||
| a1 = xds_ms.ANTENNA1.data[:n_bl].compute() |
There was a problem hiding this comment.
This seems potentially dangerous if the row axis is stacked with all times for one baseline first.
There was a problem hiding this comment.
Checked before changing this, and you are right to flag it.
Taking the first n_bl rows assumes the MS is ordered time-major — all baselines of one timestep before the next. A baseline-major store would hand back n_bl rows of the same pair, which reproduces exactly the |g_p|^2 bug this PR exists to fix, silently.
The assumption is not new here — it is systemic in the reader, which does reshape(n_time, n_bl) in four places (TIME, UVW, the data column, FLAG). So the fix is not to work around it in one function but to check it where the pairs are read:
if len(set(zip(a1.tolist(), a2.tolist()))) != n_bl:
raise ValueError(
f"The first {n_bl} rows do not hold {n_bl} distinct antenna pairs, so "
"the MS is not ordered time-major. tabascal reads visibilities as "
"(n_time, n_bl); sort the MS by TIME before running."
)A baseline-major MS now stops with that message instead of producing wrong columns. Verified the repo's example MS satisfies the assumption (first 28 rows: 1 unique TIME, 28 distinct pairs), and added two tests for the violating cases.
Worth noting the check only protects this function. The reader's other four reshape(n_time, n_bl) sites would still mis-shape a baseline-major MS — happening to be self-consistent, so no error, just wrong data throughout. Validating the row order once at read time would be the complete fix; happy to file that separately if you want it.
Delete the ndim == 3 path and its unity guard. It predates multi-frequency data: every write.py in this repository's history, back to the oldest commit holding the file, writes 4-d ["sample", "bl", "freq", "time"], so nothing here has ever produced a 3-d ast_vis. An unknown ndim now raises directly. That also removes the branch's gains_bl special-casing, so the 4-d path is the only path and the function loses a level of indentation. Cut the docstrings back. The UnitaryGains rationale belongs in the PR, not in every docstring that touches a gain; write.py is down to 36 comment/docstring lines in 217. Review comment on read_antenna_pairs: taking the first n_bl rows assumes the MS is time-major, and a baseline-major store would return n_bl rows of the SAME pair -- silently reproducing the |g_p|^2 bug this PR fixes. The assumption is systemic (the reader does reshape(n_time, n_bl) in four places), so it is now checked where the pairs are read: n_bl rows must hold n_bl distinct pairs, or the run stops and says to sort by TIME.
1 participant
Closes #122.
Two bugs in
tabascal/write.py, both latent underUnitaryGains— which every test and every shipped config uses. That is why neither the test suite nor CI has ever seen them, and why they surface the moment a run fits a non-unit gain.1. The baseline gain was built from
ANTENNA1twiceSo
gains_blwasg_p * conj(g_p)=|g_p|^2: real, positive, blind to the second antenna, with all phase information discarded, and wrong on every baseline whose two antennas differ — i.e. all of them.2. Residuals were formed in the wrong frame
Residuals used the un-gained model visibilities from the results zarr:
while the forward model is
gains_bl * (vis_ast + vis_rfi). With a non-unit gain that subtracts a model in the wrong frame and the residual columns are meaningless.The correct expressions were already present, commented out — and
gains_blwas computed and then never used at all.The fix
Residuals are formed in the data frame,
vis_obs - gains_bl * model, rather than the calibrated framevis_obs / gains_bl - model: dividing by the gain inflates the noise on low-gain baselines and distorts any noise-referenced residual metric.CORRECTED_DATAnow carries the calibrated data,vis_obs / gains_bl.Moving every column into a single calibrated frame, with the
WEIGHT_SPECTRUMthat belongs to it, is #123 — deliberately not done here.All of this reduces exactly to the current behaviour when the gains are unity, so existing
UnitaryGainsresults are unchanged and no references need re-recording.Making it testable
There are no tests for
write_results_msand testing it directly needs an MS plus a results zarr. Rather than skip the test, the arithmetic is factored into three small pure functions —baseline_gains,data_frame_residualsandread_antenna_pairs.read_antenna_pairsexists specifically for bug 1: reading the same column twice was a two-line typo with no seam to test at. Behind a named function it becomes a testable mistake.Every test uses a non-unit, non-uniform gain with distinct amplitude and phase per antenna. A unity gain cannot distinguish
g_p conj(g_q)from|g_p|^2, nor a gained model from a raw one — which is the whole reason these survived.Verified by mutation
Reverting each bug and re-running:
a2back toANTENNA1The suite catches both, rather than merely passing alongside the fix.
The legacy 3-d layout
The
ndim == 3ast_visbranch carries no baseline axis, so the per-baseline gain cannot be reconstructed there at all. It now raises a clearNotImplementedErrorrather than silently writing columns in the wrong frame — unless the stored gains are unity, the only case it has ever handled correctly. The results writer has produced the 4-d layout since #93, so this is only reachable for an older zarr.Testing
754 tests pass, including 13 new ones in
tests/test_write.py.