fix: stale SMW cache and in-place beam smoothing on a live PICSLike instance

docs/source/changelog.rst

@@ -117,6 +117,24 @@ array instead of a file path, and nothing is written to disk unless asked.
   compression and PICSLike ``do_cross``) now raise a ``RuntimeError`` naming
   the contract when called after ``release_pixel_projector()``, instead of an
   opaque ``ValueError`` from numpy.
+* Evaluating the same parameter point twice on one ``PICSLike`` instance no
+  longer drifts. Beam smoothing multiplied the theory spectra in place, so each
+  evaluation re-smoothed the parameter grid's own arrays and chi-squared grew
+  with every pass. Spectra handed to ``set_cls`` are now copied, which also
+  means an injected ``cls_data=`` array is no longer modified by the run that
+  consumes it.
+* A C_ℓ longer than the analysis ``lmax`` is truncated on the way in rather
+  than failing to broadcast inside beam smoothing. Handing a full CAMB run to a
+  small-``lmax`` analysis is the normal case; only the dict form of ``cls_data``
+  was affected, the array form already truncated.
+* Rebuilding the computation basis on a live ``PICSLike`` instance no longer
+  produces wrong chi-squared. The compressed-SMW path cached data projected
+  through the *previous* basis and kept using it against the new one — silently,
+  with no error. ``setup_computation_basis()`` now drops that cache, so the
+  documented ``setup_computation_basis → setup_maps → compute`` order is no
+  longer the only safe one. The same staleness hit the MPI worker ranks on a
+  second pipeline pass, where the cache outlived the broadcast that replaced the
+  basis it came from; rank 0 stayed correct, so the disagreement was invisible.
 
 Version 1.0.1 (2026-05-21)
 --------------------------

src/cosmoforge.cosmocore/cosmocore/spectra_io.py

@@ -273,7 +273,18 @@ def set_cls(
         effective_lmax = lmax if lmax is not None else self.fields[0].lmax
         n_ell = effective_lmax + 1
         if isinstance(cls_data, dict):
-            self._cls_dict = cls_data.copy()
+            # 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.
+            # Truncate like the array branch below: apply_smoothing pairs these
+            # against ℓ-indexed factors of length n_ell and writes them back into
+            # _cls_matrix, so an over-long C_ℓ (a full CAMB run against a small
+            # analysis lmax) has to lose its tail here or fail to broadcast there.
+            self._cls_dict = {
+                label: np.asarray(arr, dtype=np.float64)[:n_ell].copy()
+                for label, arr in cls_data.items()
+            }
             # Build matrix from dictionary
             self._cls_matrix = np.zeros((n_ell, self.n_spectra))
             for idx, label in enumerate(self._spectra_labels):

src/cosmoforge.cosmocore/tests/test_cls_kwarg.py

@@ -51,3 +51,61 @@ def test_injected_cls_data_matches_file_path():
         assert ref_cls.keys() == inj_cls.keys()
         for label in ref_cls:
             np.testing.assert_array_equal(inj_cls[label], ref_cls[label])
+
+
+def test_beam_smoothing_does_not_mutate_the_caller_spectra():
+    """``set_cls`` takes ownership of its input: the caller's arrays stay put.
+
+    Beam smoothing multiplies the stored spectra in place. A shallow copy of the
+    input dict would let that reach back into the caller's arrays — re-smoothing
+    them on every subsequent ``set_cls`` with the same source.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        params = _make_params(tmpdir, nside=4, lmax=8)
+
+        core = ConcreteCore(params)
+        core.setup_fields()
+
+        caller_cls = readcl(params.inputclfile, params, lmax=8)
+        pristine = {label: arr.copy() for label, arr in caller_cls.items()}
+
+        core.collection.set_cls(caller_cls, lmax=8)
+        core.setup_beams(lmax=8)
+        core.collection.beam_manager.apply_smoothing(
+            core.collection.spectra_manager, lmax=8
+        )
+
+        for label in pristine:
+            np.testing.assert_array_equal(caller_cls[label], pristine[label])
+
+
+def test_cls_longer_than_lmax_is_truncated():
+    """A C_ℓ longer than the analysis ``lmax`` keeps only the part that is used.
+
+    Handing in a full CAMB run and analysing a few multipoles of it is the
+    normal case. The stored spectra have to be truncated to ``lmax + 1`` like
+    the array branch already does, or beam smoothing — which pairs them with
+    ℓ-indexed factors of that length — cannot broadcast.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        params = _make_params(tmpdir, nside=4, lmax=8)
+
+        core = ConcreteCore(params)
+        core.setup_fields()
+
+        cls = readcl(params.inputclfile, params, lmax=8)
+        overlong = {
+            label: np.concatenate([arr, np.ones(20)]) for label, arr in cls.items()
+        }
+
+        core.collection.set_cls(overlong, lmax=8)
+        core.setup_beams(lmax=8)
+        core.collection.beam_manager.apply_smoothing(
+            core.collection.spectra_manager, lmax=8
+        )
+
+        stored = _cls_dict(core)
+        matrix = core.collection.spectra_manager._cls_matrix
+        for label in stored:
+            assert stored[label].shape == (9,)
+        assert matrix.shape[0] == 9

src/cosmoforge.picslike/picslike/picslike.py

@@ -323,6 +323,19 @@ def prepare_covariance_matrix(self):
         self.inv_cov = matrix_inverse_symm(self.total_cov)
         self.log("Computed inverse of primary covariance matrix", level=4)
 
+    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)
+
     def setup_maps(self):
         """
         Read observational map data from FITS files.
@@ -419,6 +432,12 @@ def setup_parameter_grid(self) -> None:
 
     def _broadcast_variables(self):
         """Broadcast essential data from master to all MPI worker processes."""
+        # Workers never run setup_maps / setup_computation_basis (both are
+        # rank-0 only in run()), so this is their only invalidation point: the
+        # basis and maps they are about to receive replace exactly what the
+        # cached projections were derived from.
+        self._smw_data_cache = None
+
         # Python objects (small, serialization is fine)
         self.params = self.comm.bcast(self.params if self.rank == 0 else None, root=0)
         self.collection: FieldCollection = self.comm.bcast(

src/cosmoforge.picslike/tests/test_live_instance_reuse.py

@@ -0,0 +1,78 @@
+"""Reusing a live PICSLike instance must not corrupt chi-squared.
+
+Two pieces of per-evaluation state used to survive longer than they should:
+
+* the harmonic (SMW) cache of ``(projected1, projected2, term1)``, which is
+  derived from the basis's noise factorisation and so belongs to a single basis
+  lifetime — a rebuild must not leave the next chi-squared evaluation pairing a
+  stale projection with the new basis's kernel;
+* the theory spectra themselves, which beam smoothing multiplied in place,
+  re-smoothing the parameter grid's arrays on every evaluation.
+
+Both were silent: no exception, just wrong numbers.
+"""
+
+import numpy as np
+import pytest
+
+from picslike import PICSLike
+
+#: Noise rescaling between basis #1 and basis #2. Same shapes on both sides, so
+#: a stale cache is silently wrong rather than loudly broken.
+NOISE_SCALE = 3.0
+
+
+@pytest.fixture
+def cfg(sandboxed_config):
+    return sandboxed_config("tests/data/nside4/TQU/fast_config.yaml")
+
+
+def _prepared(config_path):
+    """A PICSLike instance set up to the point just before the basis is built."""
+    like = PICSLike(config_path)
+    like.setup_parameter_grid()
+    like.setup_fields()
+    like.setup_geometry()
+    like.setup_covariance_matrices()
+    like.setup_cls(lmax=like.lmax_signal)
+    like.setup_beams(lmax=like.lmax_signal)
+    return like
+
+
+def test_repeated_evaluation_is_stable(cfg):
+    """The same parameter point evaluates to the same chi-squared every time."""
+    like = _prepared(cfg)
+    like.setup_computation_basis(method="harmonic")
+    like.setup_maps()
+
+    point = like.parameter_grid.grid_points[0]
+    first, _ = like._compute_likelihood_point(point)
+    second, _ = like._compute_likelihood_point(point)
+
+    np.testing.assert_allclose(second, first, rtol=1e-12)
+
+
+def test_basis_rebuild_invalidates_smw_cache(cfg):
+    """A second ``setup_computation_basis`` re-derives the cached SMW data."""
+    like = _prepared(cfg)
+    like.setup_computation_basis(method="harmonic")
+    like.setup_maps()
+
+    point = like.parameter_grid.grid_points[0]
+    like._compute_likelihood_point(point)  # populates the cache from basis #1
+
+    # Rebuild on the live instance against a different noise model. The basis
+    # consumed (and nulled) noise_cov1, so the noise has to be re-materialised.
+    like.setup_covariance_matrices()
+    like.noise_cov1 *= NOISE_SCALE
+    like.setup_computation_basis(method="harmonic")
+    chi2, log_like = like._compute_likelihood_point(point)
+
+    ref = _prepared(cfg)
+    ref.noise_cov1 *= NOISE_SCALE
+    ref.setup_computation_basis(method="harmonic")
+    ref.setup_maps()
+    chi2_ref, log_like_ref = ref._compute_likelihood_point(point)
+
+    np.testing.assert_allclose(chi2, chi2_ref, rtol=1e-10)
+    np.testing.assert_allclose(log_like, log_like_ref, rtol=1e-10)

src/cosmoforge.picslike/tests/test_mpi.py

@@ -28,3 +28,21 @@ def test_picslike_pipeline_under_mpi(comm, fast_config_path):
         assert result is not None
         assert np.all(np.isfinite(result.log_likelihood_values))
         assert np.all(np.isfinite(result.chi_squared_values))
+
+
+def test_broadcast_drops_stale_smw_cache(fast_config_path):
+    """The broadcast is the worker ranks' only SMW-cache invalidation point.
+
+    ``setup_maps`` and ``setup_computation_basis`` — the two invalidation sites
+    on rank 0 — are both rank-0 only inside ``run()``. Workers populate the
+    cache during ``compute()`` and receive their basis and maps here, so a
+    broadcast that left the cache standing would pair a projection from the
+    previous basis with the new one.
+    """
+    pl = PICSLike(fast_config_path)
+    pl.run()
+
+    pl._smw_data_cache = ("stale", "stale", "stale")
+    pl._broadcast_variables()
+
+    assert pl._smw_data_cache is None