fix: stale SMW cache and in-place beam smoothing on a live PICSLike instance - #48
Conversation
Two silent-correctness bugs that only bite when a PICSLike instance is used for more than one evaluation pass. Neither raised; both just returned wrong chi-squared. The SMW cache holds (projected1, projected2, term1), all derived from the basis's noise factorisation, and was invalidated only in setup_maps(). A second setup_computation_basis() on a live instance left the cache keyed off the previous basis's V N^-1 while evaluation used the new basis's K_chol. PICSLike now drops the cache when the basis is rebuilt. The invalidation stays in PICSLike: the cache is PICSLike's, not Core's. Fixing that exposed a second, independent bug underneath. SpectraManager. set_cls() shallow-copied the input dict, so the stored spectra were the caller's own arrays, and BeamManager.apply_smoothing() multiplies them in place. Every evaluation therefore re-smoothed the parameter grid's theory spectra: the same parameter point evaluated three times gave 619, 863, 1227. The array branch of set_cls() already copied; the dict branch now does too, which also honours the documented promise that an injected cls_data= array is never modified. This one is pre-existing on master and independent of the basis — it is fixed here because the acceptance test for the cache bug ( evaluate, rebuild, evaluate) cannot pass while it stands. Both fixes are covered by tests that fail without them.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #48 +/- ##
=======================================
Coverage 92.63% 92.64%
=======================================
Files 40 40
Lines 5092 5096 +4
=======================================
+ Hits 4717 4721 +4
Misses 375 375
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
setup_maps() and setup_computation_basis() are the two invalidation sites, and run() calls both under `if self.rank == 0`. Worker ranks populate the cache in compute() and receive their basis through _broadcast_variables(), so nothing ever cleared theirs: a second pipeline pass left every worker pairing a projection from the previous basis with the new one, while rank 0 was correct. The broadcast is the workers' equivalent of the two rank-0 setup calls — it hands them exactly the basis and maps the cached projections are derived from — so it is the right place to invalidate. Not reachable in a single-pass run(), and not reproducible on a box without mpi4py (cosmocore._mpi falls back to a serial stub). The test drives _broadcast_variables() directly, so it covers the invariant at one rank and CI exercises it for real under mpirun -n 2.
There was a problem hiding this comment.
Pull request overview
Fixes two silent correctness issues that affect repeated evaluations on long-lived likelihood instances in CosmoForge: (1) PICSLike’s harmonic (SMW) cache persisting across computation-basis rebuilds, and (2) theory spectra being mutated in-place via beam smoothing due to shallow copying.
Changes:
- Clear PICSLike’s SMW cached projections on
setup_computation_basis()and on MPI broadcast to prevent stale basis-dependent reuse. - Deep-copy spectra dict inputs in
SpectraManager.set_cls()to prevent beam smoothing from mutating caller-owned arrays across evaluations. - Add regression tests covering live-instance reuse, basis rebuild correctness, and non-mutation of injected spectra; document fixes in the changelog.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/cosmoforge.picslike/picslike/picslike.py | Invalidate SMW cache on basis rebuild and on MPI broadcast to avoid stale harmonic projections. |
| src/cosmoforge.cosmocore/cosmocore/spectra_io.py | Copy spectra dict values to prevent in-place beam smoothing from mutating caller-provided arrays. |
| src/cosmoforge.picslike/tests/test_live_instance_reuse.py | New tests ensuring repeat evaluation stability and basis rebuild matches a from-scratch reference. |
| src/cosmoforge.picslike/tests/test_mpi.py | New MPI test asserting broadcast clears stale SMW cache state. |
| src/cosmoforge.cosmocore/tests/test_cls_kwarg.py | New test ensuring beam smoothing does not mutate the caller’s spectra dict. |
| docs/source/changelog.rst | Changelog entries describing both fixed silent-correctness issues. |
| # Copy the values, not just the dict: apply_smoothing multiplies the | ||
| # stored spectra in place, so sharing the arrays would reach back | ||
| # into the caller's — re-smoothing a parameter grid's theory spectra | ||
| # on every evaluation, or an injected ``cls_data`` on every use. | ||
| self._cls_dict = {label: arr.copy() for label, arr in cls_data.items()} |
There was a problem hiding this comment.
Good catch — taken, in 73a67c0.
Confirmed the failure mode before changing anything: with lmax=8 and a 29-element input, _cls_dict held 29 entries against a 9-row _cls_matrix, and apply_smoothing died at beam.py:409 with operands could not be broadcast together with shapes (29,) (9,). It was loud rather than silent, and it predates this PR — but you're right that the dict branch is the odd one out: the array branch two lines down already does cls_data[:n_ell].copy(). The dict branch now truncates the same way, which also stops the new copy from duplicating a full CAMB tail nobody reads.
Handing a full CAMB run to a small-lmax analysis is the normal case, so this was a real crash. Regression test added (test_cls_longer_than_lmax_is_truncated).
| def setup_computation_basis(self, *args, **kwargs): | ||
| """ | ||
| Build the computation basis, dropping any cached SMW data first. | ||
|
|
||
| Thin wrapper over :meth:`cosmocore.Core.setup_computation_basis` — see | ||
| there for the parameters. The cached ``(projected1, projected2, term1)`` | ||
| are derived from the basis's noise factorisation, so they belong to a | ||
| single basis lifetime: a rebuild on a live instance must not leave the | ||
| next evaluation pairing a stale projection with the new kernel. | ||
| """ | ||
| self._smw_data_cache = None | ||
| return super().setup_computation_basis(*args, **kwargs) |
There was a problem hiding this comment.
Considered, but declining — happy to be overruled by a human reviewer.
Mirroring the base signature means duplicating 10 parameters to add one line of cache invalidation, and that mirror rots silently. Core.setup_computation_basis is not a stable signature: git log -L on those lines shows it churning repeatedly — compress/delta_m added, compression renamed to basis, method="auto" added, per-field thresholds added. Every one of those would have needed a matching edit here, and a missed one means PICSLike rejects a kwarg that Core accepts — a TypeError on a call that should work.
On the three costs you name:
- Type checking — there is no mypy or pyright in this repo's CI, so there is nothing to weaken today.
- Docs — the docstring delegates explicitly (
Thin wrapper over Core.setup_computation_basis — see there for the parameters), and Sphinx builds with zero new warnings against master. - Introspection — this is the real cost, and it is what
*args, **kwargsbuys the forwarding correctness with. It is the standard idiom for a pass-through override whose only job is a side effect beforesuper().
The alternative that gets both — pinning __signature__ from Core — is more magic than this override deserves.
If the preference is to mirror it anyway, say so and it is a five-minute change.
The dict branch stored the caller's arrays at full length while the matrix was built at lmax+1. apply_smoothing pairs the stored spectra against ell-indexed factors of that length and writes them back into the matrix, so a full CAMB run handed to a small-lmax analysis died on a broadcast error. The array branch already truncated; the dict branch now does too.
…ke-smw-cache-invalidation # Conflicts: # docs/source/changelog.rst
Closes #17 (PICSLike SMW data cache not invalidated on basis rebuild).
Three silent-correctness bugs that only bite when a
PICSLikeinstance is used for more than one evaluation pass. None of them raises; all just return wrong chi-squared. The issue was found by static review and never reproduced — so the first thing this branch does is reproduce it.1. The reported bug: stale SMW cache on basis rebuild
The compressed-SMW path caches
(projected1, projected2, term1), all derived from the basis's noise factorisation, and invalidated it only insetup_maps().setup_computation_basis()is inherited fromCore, which knows nothing about the cache, so a rebuild on a live instance left the cache keyed off the previous basis'sV N⁻¹while evaluation used the new basis'sK_chol.Measured on the red test: chi² came out as -4.3e7 where the correct value is 481. Negative — impossible for a chi-squared — and still no exception.
PICSLike.setup_computation_basis()now drops the cache and delegates tosuper(). The invalidation stays in PICSLike: the cache is PICSLike's, notCore's. No basis-versioned cache scheme, per the issue's out-of-scope note.2. Underneath it: theory spectra smoothed in place
Fixing the cache moved chi² from
-4.3e7to669, still against a true value of481. The bases were byte-identical and the maps were identical; the theory spectra had changed.SpectraManager.set_cls()shallow-copied the input dict (cls_data.copy()), so the stored spectra were the caller's own arrays — andBeamManager.apply_smoothing()multiplies them in place. Every evaluation therefore re-smoothed the parameter grid's stored theory spectra. The same parameter point, evaluated three times:The array branch of
set_cls()already copied its input; the dict branch now does too. This also honours the documented promise that an injectedcls_data=array is never modified by the run that consumes it (ADR-0017).This bug is pre-existing on master and independent of the basis. It is fixed here because the acceptance criterion for the cache bug — evaluate a point, rebuild the basis, evaluate again, and match a from-scratch computation — cannot pass while it stands. Happy to split it out if you would rather review it separately.
3. The same cache hole on the MPI worker ranks
setup_maps()andsetup_computation_basis()are the two invalidation sites, andrun()calls both underif self.rank == 0. Worker ranks populate the cache incompute()and receive their basis through_broadcast_variables()— so nothing ever cleared theirs. A second pipeline pass left every worker pairing a projection from the previous basis with the new one, while rank 0 was correct. Fixing only the rank-0 path would have left this standing.The broadcast is the workers' equivalent of the two rank-0 setup calls — it hands them exactly the basis and maps the cached projections are derived from — so it is the right place to invalidate.
Caveat, stated plainly: this one is not reproducible on a box without mpi4py (
cosmocore._mpifalls back to a serial stub, so every local run is size 1). The test drives_broadcast_variables()directly, so it covers the invariant at one rank, and CI exercises it for real undermpirun -n 2.Why the suite was green
None of these affect a single-pass
run(), where each grid point is evaluated exactly once against one basis. They need a reused instance — a second basis, a second pass, or a repeated point.Tests
picslike/tests/test_live_instance_reuse.py— repeat evaluation is stable; basis rebuild matches a from-scratch computation against the new basis.cosmocore/tests/test_cls_kwarg.py— beam smoothing does not mutate the caller's spectra.picslike/tests/test_mpi.py— the broadcast drops a stale cache.Each fails without its fix; I verified them independently (disabling the cache fix alone re-reds only the rebuild test).
Full suites green: cosmocore 535 passed, picslike 102 passed / 2 skipped, qube 172 passed / 10 skipped. Sphinx build produces zero new warnings against master (62 on both).