Separate correlation means from the jackknife patch layout

papers/cosmo_val/config/config.yaml

@@ -24,6 +24,8 @@ cosmo_val:
   theta_min: 1.0
   theta_max: 250.0
   nbins: 20
+  # Applies to published means; patched covariance uses TreeCorr's default bin_slop.
+  b_target: 0.01
   theta_min_plot: 0.8
   theta_max_plot: 260.0
   ylim_alpha: [-0.01, 0.05]

src/sp_validation/correlation.py

@@ -0,0 +1,123 @@
+"""Layout-independent full-sample means with patched covariance measurements."""
+
+import json
+from pathlib import Path
+
+import treecorr
+
+
+def _measurement_configs(config):
+    """Return the explicit means settings and default-tolerance patched settings."""
+    means = dict(config)
+    patched = dict(config)
+    patched.pop("bin_slop", None)
+    return {"means": means, "patched": patched}
+
+
+def measurement_matches(path, config):
+    """Whether a cache records the current means and patched configurations."""
+    metadata = Path(str(path) + ".json")
+    if not Path(path).exists() or not metadata.exists():
+        return False
+    recorded = json.loads(metadata.read_text())
+    recorded_configs = recorded.get("configs")
+    if recorded.get("means") != "unpatched" or not isinstance(recorded_configs, dict):
+        return False
+    expected_configs = _measurement_configs(config)
+    if config.get("min_top") is not None:
+        # With the root depth pinned, threads affect reduction round-off,
+        # not the approximation. A plotting job can reuse a multi-core product.
+        for settings in (*expected_configs.values(), *recorded_configs.values()):
+            if isinstance(settings, dict):
+                settings.pop("num_threads", None)
+    return recorded_configs == expected_configs
+
+
+def write_measurement_metadata(path, config):
+    """Record both TreeCorr configurations beside a cached product."""
+    Path(str(path) + ".json").write_text(
+        json.dumps({"means": "unpatched", "configs": _measurement_configs(config)})
+        + "\n"
+    )
+
+
+def _unpatched_catalog(cat):
+    """The same positions, shears and weights, with no spatial partition."""
+    if cat.ra is not None:
+        positions = dict(ra=cat.ra, dec=cat.dec, ra_units="rad", dec_units="rad")
+        if cat.r is not None:
+            positions["r"] = cat.r
+    else:
+        positions = dict(x=cat.x, y=cat.y)
+        if cat.z is not None:
+            positions["z"] = cat.z
+    return treecorr.Catalog(
+        **positions, g1=cat.g1, g2=cat.g2, w=cat.w, wpos=cat.wpos, npatch=1
+    )
+
+
+def measure_with_patches(measure, catalogs, config, means=None):
+    """Return means and a patched measurement with distinct TreeCorr settings.
+
+    ``measure(catalogs, config)`` must return a fresh measurement of the
+    supplied catalogue mapping. ``config`` describes the means pass, including
+    its explicit ``bin_slop``. The patched pass omits ``bin_slop`` so TreeCorr
+    uses its standard default. A cached ``means`` can be reused for repeated
+    covariance draws of the same catalogue. With no patches, one means-config
+    measurement supplies both products.
+    """
+    configs = _measurement_configs(config)
+    if all(cat.npatch == 1 for cat in catalogs.values()):
+        patched = measure(catalogs, configs["means"])
+        return (patched if means is None else means), patched
+
+    patched = measure(catalogs, configs["patched"])
+    if means is None:
+        unpatched = {}
+        # Preserve aliases, including an auto-correlation passed as a cross-pair.
+        copies = {}
+        for key, cat in catalogs.items():
+            if id(cat) not in copies:
+                copies[id(cat)] = _unpatched_catalog(cat)
+            unpatched[key] = copies[id(cat)]
+        means = measure(unpatched, configs["means"])
+    return means, patched
+
+
+def process_gg(config, cat1, cat2=None):
+    """Measure GG means with ``config`` and patched resampling at TreeCorr's default.
+
+    The shared two-pass mechanism omits ``bin_slop`` only from the patched
+    configuration. Covariance estimates, including joint and derived-statistic
+    jackknives, retain TreeCorr's per-patch results. Published pair counts,
+    weights, separations and complex correlations come from the means pass.
+    """
+    catalogs = {"cat1": cat1}
+    if cat2 is not None:
+        catalogs["cat2"] = cat2
+
+    def measure(cats, measurement_config):
+        cfg = dict(measurement_config)
+        if all(cat.npatch == 1 for cat in cats.values()):
+            cfg["var_method"] = "shot"
+        gg = treecorr.GGCorrelation(cfg)
+        gg.process(cats["cat1"], cat2=cats.get("cat2"))
+        return gg
+
+    means, gg = measure_with_patches(measure, catalogs, config)
+    if means is not gg:
+        # TreeCorr estimates covariance lazily from fields that the means pass
+        # replaces below, so freeze the patched-pass estimates first.
+        _ = gg.varxip, gg.varxim, gg.cov
+        for name in (
+            "xip",
+            "xim",
+            "xip_im",
+            "xim_im",
+            "meanr",
+            "meanlogr",
+            "npairs",
+            "weight",
+        ):
+            getattr(gg, name)[:] = getattr(means, name)
+    return gg

src/sp_validation/cosmo_val/core.py

@@ -147,6 +147,9 @@ class CosmologyValidation(
         Maximum angular separation in arcminutes for correlation function binning.
     nbins : int, default 20
         Number of angular bins for TreeCorr real-space correlation functions.
+    b_target : float, default 0.01
+        Maximum logarithmic-bin tolerance for full-sample means, with bin_slop
+        capped at 1 per grid. Patched covariance passes use TreeCorr's default.
     var_method : {'jackknife', 'sample', 'bootstrap', 'marked_bootstrap'}, default 'jackknife'
         TreeCorr variance estimation method.
     npatch : int, default 20
@@ -312,7 +315,9 @@ def __init__(
         cosmo_params=None,
         compute_tomography=False,
         force_run=False,
+        b_target=0.01,
     ):
+        self.b_target = b_target
         self.rho_tau_method = rho_tau_method
         self.cov_estimate_method = cov_estimate_method
         self.compute_cov_rho = compute_cov_rho
@@ -374,16 +379,17 @@ def __init__(
             "nbins": nbins,
             "var_method": var_method,
             "cross_patch_weight": "match" if var_method == "jackknife" else "simple",
-            # min_top sets the depth of TreeCorr's root cells, hence which pairs
-            # bin_slop approximates. Left unset, TreeCorr derives it from its
-            # thread count (max(3, ceil(log2 n))) and ξ± depends on the machine.
-            # 6 is what TreeCorr derives on candide's 48- and 64-CPU nodes.
+            # min_top fixes TreeCorr's root-cell depth, which controls the pairs
+            # bin_slop approximates. Its default depends on thread count and would
+            # make ξ± depend on the machine; 6 is used on candide's 48/64-CPU nodes.
             "min_top": 6,
             # The CPUs this process may use; TreeCorr's own default is the
             # node's count, whatever share of it the job holds.
             "num_threads": len(os.sched_getaffinity(0)),
         }
 
+        self.treecorr_config = self._binning()
+
         self.catalog_config_path = Path(catalog_config)
         with self.catalog_config_path.open("r") as file:
             self.cc = cc = yaml.load(file, Loader=yaml.FullLoader)
@@ -634,18 +640,25 @@ def basename(
         )
 
     def _binning(self, min_sep=None, max_sep=None, nbins=None, **extra):
-        """treecorr_config with min_sep/max_sep/nbins overridden.
+        """Means-pass TreeCorr config with min_sep/max_sep/nbins overridden.
 
         None falls back to the instance's treecorr_config value for that key;
-        any further keys in `extra` override on top.
+        any further keys in `extra` override on top. The shared two-pass
+        measurement removes ``bin_slop`` from the patched config, leaving
+        TreeCorr's default for covariance and resampling products.
         """
-        return {
+        config = {
             **self.treecorr_config,
             "min_sep": min_sep or self.treecorr_config["min_sep"],
             "max_sep": max_sep or self.treecorr_config["max_sep"],
             "nbins": nbins or self.treecorr_config["nbins"],
             **extra,
         }
+        bin_size = np.log(config["max_sep"] / config["min_sep"]) / config["nbins"]
+        # This is the means tolerance; measure_with_patches drops it for the
+        # patched covariance pass so TreeCorr uses its standard default there.
+        config["bin_slop"] = extra.get("bin_slop", min(1, self.b_target / bin_size))
+        return config
 
     def _read_shear_cols(self, ver, *keys):
         """Read shear-catalog columns by their config-key names.

src/sp_validation/cosmo_val/real_space.py

@@ -11,6 +11,7 @@
 import numpy as np
 import treecorr
 
+from sp_validation.correlation import process_gg
 from sp_validation.statistics import jackknife_patch_centers
 
 from .. import sacc_io
@@ -56,6 +57,10 @@ def calculate_2pcf_version(
             **treecorr_config: Overrides of the instance's ``treecorr_config``,
                 e.g. ``min_sep=1``.
 
+        Full-sample means use an unpatched pass with the configured ``bin_slop``.
+        Patched covariance uses TreeCorr's default tolerance.
+        No unsealed text cache is written; figures reuse the sealed SACC part.
+
         Returns:
             dict: ``"tomo_bin_{b1}_tomo_bin_{b2}"`` → ξ± shaped like a TreeCorr
             ``GGCorrelation``: the part's
@@ -98,16 +103,16 @@ def calculate_2pcf_version(
             patch_centers = self._patch_centers(cols, npatch)
 
             for bin1, bin2 in to_compute:
-                gg = treecorr.GGCorrelation(gg_config)
-
                 cat_gal1 = self._bin_catalog(cols, bin1, npatch, patch_centers)
                 cat_gal2 = (
                     self._bin_catalog(cols, bin2, npatch, patch_centers)
                     if bin1 != bin2
                     else None
                 )
 
-                gg.process(cat_gal1, cat2=cat_gal2)
+                # Publish unpatched means; retain the default-tolerance patched
+                # covariance before sealing the part under its catalogue's blind.
+                gg = process_gg(gg_config, cat_gal1, cat_gal2)
 
                 if (bin1, bin2) != ("all", "all"):
                     ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = gg