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

@@ -146,6 +146,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
@@ -311,7 +314,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
@@ -373,16 +378,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)
@@ -598,18 +604,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,11 @@
 import numpy as np
 import treecorr
 
+from sp_validation.correlation import (
+    measurement_matches,
+    process_gg,
+    write_measurement_metadata,
+)
 from sp_validation.statistics import jackknife_patch_centers
 
 
@@ -20,6 +25,7 @@ def calculate_2pcf_version(
         ver,
         npatch=None,
         compute_tomography=False,
+        read_cached=False,
         **treecorr_config,
     ):
         """
@@ -42,6 +48,10 @@ def calculate_2pcf_version(
             compute_tomography (bool, optional): Whether to compute tomographic
             correlations. Defaults to False.
 
+            read_cached (bool, optional): Reuse columns-only dumps for plotting,
+            including patched runs. They carry the published means and variances,
+            but no patch results for dense or derived jackknife covariance.
+
             **treecorr_config: Additional TreeCorr configuration parameters that
             will override the instance's default `treecorr_config`. For example,
             `min_sep=1`.
@@ -52,9 +62,12 @@ def calculate_2pcf_version(
             key is ``"tomo_bin_all_tomo_bin_all"``.
 
         Notes:
-            - The non-tomographic pair is written to the columns-only TreeCorr
-              dump ``xi_{basename}.txt``. If that file already exists, the pair is
-              read back from it instead of being recomputed.
+            - The non-tomographic pair is written to a columns-only TreeCorr
+              dump with a configuration/mean-source JSON sidecar. Unpatched
+              runs can reuse it; patched measurements remeasure to retain
+              covariance. Plotting can opt into reading the saved columns.
+            - Full-sample means use ``bin_slop`` from the configured means pass.
+              The patched covariance pass omits it and uses TreeCorr's default.
             - Seeded patch centres are computed once from the full catalogue and
               shared by every tomographic bin pair.
         """
@@ -81,7 +94,11 @@ def calculate_2pcf_version(
         for bin1, bin2 in tomo_bin_pairs:
             if (bin1, bin2) == ("all", "all"):
                 out_fname = self._xi_txt_path(ver, treecorr_config, npatch)
-                if os.path.exists(out_fname):
+                if (
+                    (int(npatch) == 1 or read_cached)
+                    and not self.force_run
+                    and measurement_matches(out_fname, treecorr_config)
+                ):
                     self.print_done(f"Skipping 2PCF calculation, {out_fname} exists")
                     gg = treecorr.GGCorrelation(treecorr_config)
                     gg.read(out_fname)
@@ -94,16 +111,14 @@ def calculate_2pcf_version(
             patch_centers = self._patch_centers(cols, npatch)
 
             for bin1, bin2 in to_compute:
-                gg = treecorr.GGCorrelation(treecorr_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)
+                gg = process_gg(treecorr_config, cat_gal1, cat_gal2)
 
                 if (bin1, bin2) == ("all", "all"):
                     # Columns only. The covariance matrix lives in the SACC part;
@@ -114,6 +129,10 @@ def calculate_2pcf_version(
                         self._xi_txt_path(ver, treecorr_config, npatch),
                         write_patch_results=False,
                         write_cov=False,
+                        precision=17,
+                    )
+                    write_measurement_metadata(
+                        self._xi_txt_path(ver, treecorr_config, npatch), treecorr_config
                     )
 
                 ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = gg
@@ -183,6 +202,7 @@ def calculate_2pcf(
         self,
         npatch=None,
         compute_tomography=False,
+        read_cached=False,
         **treecorr_config,
     ):
         """
@@ -199,6 +219,9 @@ def calculate_2pcf(
             compute_tomography (bool, optional): Whether to compute tomographic
             correlations. Defaults to False.
 
+            read_cached (bool, optional): Reuse columns-only dumps for plotting;
+            defaults to False so patched measurements retain resampling state.
+
             **treecorr_config: Additional TreeCorr configuration parameters passed
             through to each per-version call.
 
@@ -212,6 +235,7 @@ def calculate_2pcf(
                 ver,
                 npatch=npatch,
                 compute_tomography=compute_tomography,
+                read_cached=read_cached,
                 **treecorr_config,
             )
 
@@ -230,7 +254,12 @@ def calculate_aperture_mass_dispersion(
         theta_map = np.geomspace(theta_min * 5, theta_max / 2, nbins_map)
         self._map2["theta_map"] = theta_map
 
-        treecorr_config = self._binning(theta_min, theta_max, nbins)
+        treecorr_config = self._binning(
+            theta_min,
+            theta_max,
+            nbins,
+            var_method="jackknife" if int(npatch) > 1 else "shot",
+        )
 
         for ver in self.versions:
             if compute_tomography:
@@ -252,16 +281,14 @@ def calculate_aperture_mass_dispersion(
             patch_centers = self._patch_centers(cols, npatch)
 
             for bin1, bin2 in tomo_bin_pairs:
-                gg = treecorr.GGCorrelation(treecorr_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)
+                gg = process_gg(treecorr_config, cat_gal1, cat_gal2)
 
                 mapsq, mapsq_im, mxsq, mxsq_im, varmapsq = gg.calculateMapSq(
                     R=theta_map,
@@ -292,7 +319,7 @@ def plot_2pcf(
         Writes ``xi_pm_tomography_{tomography}.png`` and
         ``xi_pm_theta_tomography_{tomography}.png`` under the output directory.
         """
-        self.calculate_2pcf(compute_tomography=tomography)
+        self.calculate_2pcf(compute_tomography=tomography, read_cached=True)
 
         for times_theta in (False, True):
             prefix = r"$\theta\,$" if times_theta else ""
@@ -333,7 +360,7 @@ def plot_ratio_xi_sys_xi(
                 "xi_psf_sys; construct CosmologyValidation with "
                 "compute_tomography=True"
             )
-        self.calculate_2pcf(compute_tomography=tomography)
+        self.calculate_2pcf(compute_tomography=tomography, read_cached=True)
 
         y_label = r"$\xi^{{\rm PSF, sys}}_{0} / \xi_{0}$"
         self.plot_2pcf_tomography(