Add return_cc option to run_pipeline

CHANGELOG.md

@@ -4,6 +4,17 @@ All notable changes to `codameter` will be documented in this file.
 The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+### Added
+
+- **`run_pipeline(..., return_cc=True)`** — optionally return the per-epoch
+  stretching correlation coefficient alongside `(dvv, valid)`, for
+  coherence-based error models (`uq_measurement.weaver_stretching_error`).
+  Works for fixed and moving references with the stretching estimator; NaN
+  otherwise. The default two-tuple return and all gating behavior are
+  unchanged (CC-gating remains fixed-reference-only).
+
 ## 0.3.0 — 2026-07-27
 
 ### Added

src/codameter/deviations.py

@@ -31,6 +31,7 @@
 import numpy as np
 
 from .synthetic_demo import (
+    METHODS,
     YEAR_D,
     C,
     Synth,
@@ -83,27 +84,51 @@
 # ---------------------------------------------------------------------------
 # Run one pipeline configuration on a shared set of daily CCFs.
 # ---------------------------------------------------------------------------
-def _moving_reference(name, ccfs, t, *, band, fs, window, ref_days=45, **kw):
+def _moving_reference(
+    name, ccfs, t, *, band, fs, window, ref_days=45, collect_cc=False, **kw
+):
     """Generic trailing-reference measurement for *any* estimator.
 
     A moving reference re-baselines each epoch against the previous
     ``ref_days`` — the deviation that erases slow trends (best_practices rule 7).
+
+    With ``collect_cc=True``, also returns the per-epoch correlation
+    coefficient for estimators that produce one (stretching); NaN otherwise.
     """
     ndays = ccfs.shape[0]
     out = np.full(ndays, np.nan)
+    cc_out = np.full(ndays, np.nan)
     for d in range(ref_days, ndays):
         ref = ccfs[d - ref_days : d].mean(axis=0)
-        val = measure(name, ccfs[d], ref, t, band=band, fs=fs, window=window, **kw)
-        out[d] = np.atleast_1d(val)[0]
-    return out
+        if collect_cc:
+            res = METHODS[name](ccfs[d], ref, t, band=band, fs=fs, window=window, **kw)
+            if isinstance(res, tuple):
+                out[d] = np.atleast_1d(res[0])[0]
+                cc_out[d] = np.atleast_1d(res[1])[0]
+            else:
+                out[d] = np.atleast_1d(res)[0]
+        else:
+            val = measure(name, ccfs[d], ref, t, band=band, fs=fs, window=window, **kw)
+            out[d] = np.atleast_1d(val)[0]
+    return (out, cc_out) if collect_cc else out
 
 
-def run_pipeline(ccfs, t, fs, cfg, *, eps_max=0.05):
+def run_pipeline(ccfs, t, fs, cfg, *, eps_max=0.05, return_cc=False):
     """Recover dv/v(t) under one processing configuration ``cfg``.
 
     Returns ``(dvv, valid)``: the per-day series and a boolean mask of epochs the
     pipeline actually produced (moving/inversion references have a warm-up gap;
     CC-gating drops low-coherence epochs).
+
+    With ``return_cc=True``, returns ``(dvv, valid, cc)`` where ``cc`` is the
+    per-epoch stretching correlation coefficient — the input to coherence-based
+    error models such as :func:`codameter.uq_measurement.weaver_stretching_error`.
+    ``cc`` is NaN wherever the configuration does not produce one (non-stretching
+    estimators, the inversion reference, and warm-up epochs).
+
+    CC-gating (``cfg["gate"]``) applies to the fixed reference only, as it
+    always has; the moving-reference CC is returned for error modelling but
+    does not change ``valid``.
     """
     name = cfg["estimator"]
     band, window, k, ref = cfg["band"], cfg["window"], cfg["stack"], cfg["reference"]
@@ -122,19 +147,34 @@ def run_pipeline(ccfs, t, fs, cfg, *, eps_max=0.05):
                 name, stacked, reference, t, band=band, fs=fs, window=window, **extra
             )
     elif ref == "moving":
-        dvv = _moving_reference(
-            name, stacked, t, band=band, fs=fs, window=window, **extra
-        )
+        if name == "stretching (TS)":
+            dvv, cc = _moving_reference(
+                name,
+                stacked,
+                t,
+                band=band,
+                fs=fs,
+                window=window,
+                collect_cc=True,
+                **extra,
+            )
+        else:
+            dvv = _moving_reference(
+                name, stacked, t, band=band, fs=fs, window=window, **extra
+            )
     elif ref == "inversion":  # Brenguier et al. (2014) joint inversion (stretching)
         dvv = measure_inversion(ccfs, t, band=band, fs=fs, window=window)
     else:
         raise ValueError(ref)
 
     dvv = np.asarray(dvv, float)
     valid = np.isfinite(dvv)
-    if cfg.get("gate") and cc is not None:
+    if cfg.get("gate") and ref == "fixed" and cc is not None:
         keep = cc > 0.6
         valid &= keep
+    if return_cc:
+        cc_arr = np.full(dvv.shape, np.nan) if cc is None else np.asarray(cc, float)
+        return dvv, valid, cc_arr
     return dvv, valid
 
 

tests/test_deviations.py

@@ -4,10 +4,12 @@
 
 import numpy as np
 import pytest
-
 from codameter import deviations as D
 from codameter.synthetic_demo import (
-    Synth, _days, daily_ccfs, volcano_truth,
+    Synth,
+    _days,
+    daily_ccfs,
+    volcano_truth,
 )
 
 
@@ -56,3 +58,37 @@ def test_multiverse_sobol_sums_sensible():
         assert -1e-9 <= v <= 1.0 + 1e-9
     # The pipeline spread is non-trivial (the whole point).
     assert np.nanstd(mv["rms"]) > 0
+
+
+class TestReturnCC:
+    def test_default_still_two_tuple(self, small_dataset):
+        s, days, truth, ccfs = small_dataset
+        out = D.run_pipeline(ccfs, s.t, s.fs, D.BASELINE)
+        assert len(out) == 2
+
+    def test_fixed_stretching_returns_cc(self, small_dataset):
+        s, days, truth, ccfs = small_dataset
+        dvv, valid, cc = D.run_pipeline(ccfs, s.t, s.fs, D.BASELINE, return_cc=True)
+        assert cc.shape == dvv.shape
+        # On a clean synthetic the coherence should be high wherever valid.
+        assert np.all(cc[valid] > 0.6)
+        # dvv/valid identical to the two-tuple call (return_cc is read-only).
+        dvv2, valid2 = D.run_pipeline(ccfs, s.t, s.fs, D.BASELINE)
+        np.testing.assert_array_equal(dvv, dvv2)
+        np.testing.assert_array_equal(valid, valid2)
+
+    def test_moving_stretching_returns_cc_after_warmup(self, small_dataset):
+        s, days, truth, ccfs = small_dataset
+        cfg = dict(D.BASELINE, reference="moving")
+        dvv, valid, cc = D.run_pipeline(ccfs, s.t, s.fs, cfg, return_cc=True)
+        assert np.isnan(cc[:10]).all()  # warm-up gap
+        assert np.isfinite(cc[valid]).all()
+        # Gating stays fixed-reference-only: valid must match the legacy call.
+        dvv2, valid2 = D.run_pipeline(ccfs, s.t, s.fs, cfg)
+        np.testing.assert_array_equal(valid, valid2)
+
+    def test_non_stretching_cc_is_nan(self, small_dataset):
+        s, days, truth, ccfs = small_dataset
+        cfg = dict(D.BASELINE, estimator="MWCS", gate=False)
+        dvv, valid, cc = D.run_pipeline(ccfs, s.t, s.fs, cfg, return_cc=True)
+        assert np.isnan(cc).all()