diff --git a/papers/cosmo_val/config/config.yaml b/papers/cosmo_val/config/config.yaml index f469cc44..9a31e0b8 100644 --- a/papers/cosmo_val/config/config.yaml +++ b/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] diff --git a/src/sp_validation/correlation.py b/src/sp_validation/correlation.py new file mode 100644 index 00000000..131a81bc --- /dev/null +++ b/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 diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index 4fbd2a89..bd9137ae 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/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. diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 7dea3542..a55c1106 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/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,8 +111,6 @@ 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) @@ -103,7 +118,7 @@ def calculate_2pcf_version( 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,8 +281,6 @@ 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) @@ -261,7 +288,7 @@ def calculate_aperture_mass_dispersion( 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( diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index e744fdf8..454c0891 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -6,10 +6,16 @@ import numpy as np from astropy.io import fits +from astropy.table import Table from shear_psf_leakage.rho_tau_cov import CovTauTh from shear_psf_leakage.rho_tau_stat import RhoStat, TauStat from sp_validation import grammar, io +from sp_validation.correlation import ( + measure_with_patches, + measurement_matches, + write_measurement_metadata, +) # SquareRootScale lives in sp_validation.plots; re-exported here so that # `from sp_validation.rho_tau import SquareRootScale` keeps working. @@ -22,6 +28,58 @@ def _extract_xip(correlations): return np.array([corr.xip for corr in correlations]).flatten() +def _compute_stats_with_patches( + handler, kind, catalog_id, filename, save_cov=False, means=None +): + """Publish unpatched ρ/τ means with the patched variances and covariance. + + The handlers expose catalogue mappings and FITS tables, rather than their + internal GG objects. The shared measurement mechanism uses the configured + bin_slop for means and TreeCorr's default for the patched covariance pass. + Cached means avoid repeating the full-sample pass across covariance draws. + """ + catalogs = handler.catalogs.catalogs_dict + config = handler._treecorr_config + selected = { + key: cat for key, cat in catalogs.items() if key.endswith(f"_{catalog_id}") + } + attribute = f"{kind}_stats" + compute = getattr(handler, f"compute_{kind}_stats") + + def measure(cats, measurement_config): + patched = any(cat.npatch > 1 for cat in cats.values()) + handler.catalogs.catalogs_dict = cats + handler._treecorr_config = { + **measurement_config, + "var_method": measurement_config.get("var_method", "shot") + if patched + else "shot", + } + try: + compute( + catalog_id, + filename, + save_cov=save_cov and patched, + func=_extract_xip if save_cov and patched else None, + var_method="jackknife" if patched else "shot", + ) + return getattr(handler, attribute).copy() + finally: + handler.catalogs.catalogs_dict = catalogs + handler._treecorr_config = config + + means, patched = measure_with_patches(measure, selected, config, means=means) + table = Table(patched) + for name in table.colnames: + if not name.startswith("var"): + table[name] = means[name] + setattr(handler, attribute, table) + getattr(handler, f"save_{kind}_stats")(filename) + getattr(handler, f"load_{kind}_stats")(filename) + write_measurement_metadata(Path(handler.catalogs._output) / filename, config) + return means + + class _CatalogueLoader: """Supply FITS paths to the file-based rho/tau readers. @@ -287,7 +345,7 @@ def get_rho_tau( ) with _CatalogueLoader(config[version], params) as load: - rho_stats_exists = rho_path.exists() + rho_stats_exists = measurement_matches(rho_path, treecorr_config) cov_exists = True if not cov_rho else cov_rho_path.exists() need_compute = (not rho_stats_exists) or (not cov_exists) or force_run @@ -297,12 +355,8 @@ def get_rho_tau( load("psf"), catalog_id=catalog_id, mask=mask_star ) - rho_stat_handler.compute_rho_stats( - catalog_id, - rho_path.name, - save_cov=cov_rho, - func=_extract_xip if cov_rho else None, - var_method="jackknife" if cov_rho else None, + _compute_stats_with_patches( + rho_stat_handler, "rho", catalog_id, rho_path.name, save_cov=cov_rho ) rho_stat_handler.load_rho_stats(rho_path.name) else: @@ -319,7 +373,7 @@ def get_rho_tau( verbose=True, ) - if tau_path.exists() and not force_run: + if measurement_matches(tau_path, treecorr_config) and not force_run: print( f"Skipping tau statistics computation, file {tau_path} already exists." ) @@ -338,7 +392,7 @@ def get_rho_tau( load("shear"), cat_type="gal", catalog_id=version, mask=mask_gal ) - tau_stat_handler.compute_tau_stats(version, tau_path.name, var_method=None) + _compute_stats_with_patches(tau_stat_handler, "tau", version, tau_path.name) print(f"Time to compute rho and tau statistics: {time.time() - start_time:.2f} s") return rho_stat_handler, tau_stat_handler @@ -425,7 +479,13 @@ def get_jackknife_cov( tau_filename = f"tau_stats_{base_tau}.fits" tau_cov_path = Path(outdir) / f"cov_tau_{base_tau}_jk.npy" - if tau_cov_path.exists() and not force_run: + if ( + tau_cov_path.exists() + and (Path(outdir) / f"cov_rho_{base_rho}_jk.npy").exists() + and measurement_matches(Path(outdir) / rho_filename, treecorr_config) + and measurement_matches(Path(outdir) / tau_filename, treecorr_config) + and not force_run + ): print(f"Skipping covariance computation, file {tau_cov_path} already exists.") rho_stat_handler = RhoStat( output=outdir, treecorr_config=treecorr_config, verbose=False @@ -447,6 +507,7 @@ def get_jackknife_cov( return rho_stat_handler, tau_stat_handler params = get_params_rho_tau(config[version]) + params["patch_number"] = npatch rho_stat_handler = RhoStat( output=outdir, treecorr_config=treecorr_config, verbose=False @@ -475,6 +536,7 @@ def chunks(i): for kind in ("rho", "tau") ] + rho_means = tau_means = None with _CatalogueLoader(config[version], params) as load: for i in range(ncov): if not all(os.path.exists(chunk) for chunk in chunks(i)): @@ -521,21 +583,21 @@ def chunks(i): ) # Compute and save rho stats - rho_stat_handler.compute_rho_stats( + rho_means = _compute_stats_with_patches( + rho_stat_handler, + "rho", catalog_id(i), rho_filename, save_cov=True, - func=_extract_xip, - var_method="jackknife", + means=rho_means, ) - - # function to extract the tau_+ - tau_stat_handler.compute_tau_stats( + tau_means = _compute_stats_with_patches( + tau_stat_handler, + "tau", catalog_id(i), tau_filename, save_cov=True, - func=_extract_xip, - var_method="jackknife", + means=tau_means, ) # Update the keys in the dictionaries diff --git a/src/sp_validation/tests/test_correlation.py b/src/sp_validation/tests/test_correlation.py new file mode 100644 index 00000000..070490d7 --- /dev/null +++ b/src/sp_validation/tests/test_correlation.py @@ -0,0 +1,140 @@ +"""Two-pass measurement contracts, including tolerance and cache settings.""" + +import json + +import numpy as np +import treecorr + +from sp_validation.correlation import ( + measure_with_patches, + measurement_matches, + process_gg, + write_measurement_metadata, +) + + +def test_unpatched_measurement_is_one_pass(): + cat = treecorr.Catalog(x=[0, 1, 2], y=[0, 0, 1], g1=[0.1] * 3, g2=[0.2] * 3) + config = {"bin_slop": 0.04} + calls = [] + + def measure(catalogs, measurement_config): + calls.append((catalogs, measurement_config)) + return object() + + means, covariance = measure_with_patches(measure, {"a": cat}, config) + assert len(calls) == 1 + assert means is covariance + assert calls[0][1] == config + + +def test_covariance_draws_reuse_means_and_preserve_catalog_aliases(): + cat = treecorr.Catalog( + x=[0, 1, 2, 3], + y=[0, 0, 1, 1], + g1=[0.1] * 4, + g2=[0.2] * 4, + npatch=2, + rng=np.random.default_rng(1), + ) + config = {"bin_slop": 0.05} + calls = [] + + def measure(catalogs, measurement_config): + assert catalogs["a"] is catalogs["b"] + calls.append((catalogs["a"].npatch, dict(measurement_config))) + return object() + + catalogs = {"a": cat, "b": cat} + means, _ = measure_with_patches(measure, catalogs, config) + reused, _ = measure_with_patches(measure, catalogs, config, means=means) + assert [npatch for npatch, _ in calls] == [2, 1, 2] + assert "bin_slop" not in calls[0][1] + assert calls[1][1]["bin_slop"] == config["bin_slop"] + assert "bin_slop" not in calls[2][1] + assert reused is means + + +def test_patched_gg_uses_default_slop_and_means_use_policy(monkeypatch): + rng = np.random.default_rng(2) + cat = treecorr.Catalog( + x=rng.random(100), + y=rng.random(100), + g1=rng.normal(size=100), + g2=rng.normal(size=100), + npatch=2, + rng=np.random.default_rng(1), + ) + config = dict( + min_sep=0.01, + max_sep=1, + nbins=8, + min_top=3, + num_threads=1, + bin_slop=0.04, + var_method="jackknife", + ) + original = treecorr.GGCorrelation + default_config = {key: value for key, value in config.items() if key != "bin_slop"} + default_slop = original(default_config).bin_slop + observed = [] + + def recording_correlation(measurement_config): + gg = original(measurement_config) + observed.append((dict(measurement_config), gg.bin_slop)) + return gg + + monkeypatch.setattr(treecorr, "GGCorrelation", recording_correlation) + process_gg(config, cat) + + assert len(observed) == 2 + assert observed[0][1] == default_slop + assert observed[1][1] == config["bin_slop"] + assert observed[0][1] != observed[1][1] + + +def test_measurement_metadata_records_both_tolerance_settings(tmp_path): + path = tmp_path / "measurement.txt" + path.write_text("cached columns\n") + config = dict( + min_sep=0.01, + max_sep=1, + nbins=8, + min_top=6, + num_threads=8, + bin_slop=0.04, + ) + + write_measurement_metadata(path, config) + metadata_path = tmp_path / "measurement.txt.json" + metadata = json.loads(metadata_path.read_text()) + assert metadata["configs"]["means"]["bin_slop"] == config["bin_slop"] + assert "bin_slop" not in metadata["configs"]["patched"] + assert measurement_matches(path, {**config, "num_threads": 12}) + + metadata_path.write_text(json.dumps({"means": "unpatched", "config": config})) + assert not measurement_matches(path, config) + + +def test_unpatched_gg_matches_direct_measurement(): + rng = np.random.default_rng(1) + cat = treecorr.Catalog( + x=rng.random(500), + y=rng.random(500), + g1=rng.normal(size=500), + g2=rng.normal(size=500), + ) + config = dict( + min_sep=0.01, + max_sep=1, + nbins=8, + min_top=6, + num_threads=1, + bin_slop=0.1, + var_method="shot", + ) + direct = treecorr.GGCorrelation(config) + direct.process(cat) + measured = process_gg(config, cat) + for name in ("xip", "xim", "meanr", "meanlogr", "weight", "npairs", "cov"): + np.testing.assert_array_equal(getattr(measured, name), getattr(direct, name)) diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index 7029cdcd..b056f531 100644 --- a/src/sp_validation/tests/test_cosmo_val.py +++ b/src/sp_validation/tests/test_cosmo_val.py @@ -469,7 +469,7 @@ def test_xi_part_carries_the_covariance_the_measurement_estimated( params, version = self._write_synthetic_catalogs(tmp_path) part = tmp_path / "part.sacc" - gg = run_2pcf.run_2pcf( + kwargs = dict( ver=version, min_sep=5.0, max_sep=100.0, @@ -479,25 +479,24 @@ def test_xi_part_carries_the_covariance_the_measurement_estimated( output_dir=params["output_dir"], sacc_out=str(part), ) - - cov = sacc_io.load(str(part), allow_unblinded=True).covariance - if npatch > 1: - assert isinstance(cov, sacc.covariance.FullCovariance) - np.testing.assert_array_equal(cov.dense, gg.cov) - else: - assert isinstance(cov, sacc.covariance.DiagonalCovariance) - np.testing.assert_array_equal( - cov.diag, np.concatenate([gg.varxip, gg.varxim]) - ) - - def test_a_patched_xi_dump_reads_back(self, tmp_path): - """calculate_2pcf_version reads back the dump a patched measurement wrote. - - The ξ± figure rules re-enter calculate_2pcf on the reporting grid, which - has patches, and are handed the dump rule xi wrote (to its precision). - """ + # The second producer call encounters its own columns-only text cache. + for _ in range(2): + gg = run_2pcf.run_2pcf(**kwargs) + cov = sacc_io.load(str(part), allow_unblinded=True).covariance + if npatch > 1: + assert gg.results + assert isinstance(cov, sacc.covariance.FullCovariance) + np.testing.assert_array_equal(cov.dense, gg.cov) + else: + assert isinstance(cov, sacc.covariance.DiagonalCovariance) + np.testing.assert_array_equal( + cov.diag, np.concatenate([gg.varxip, gg.varxim]) + ) + + def test_a_patched_xi_dump_remeasures_covariance(self, tmp_path): + """Columns-only dumps cannot supply the dense jackknife covariance.""" params, version = self._write_synthetic_catalogs(tmp_path) - binning = dict(npatch=4, min_sep=5.0, max_sep=100.0, nbins=6) + binning = dict(npatch=4, min_sep=5.0, max_sep=100.0, nbins=6, num_threads=1) pair = "tomo_bin_all_tomo_bin_all" measured = CosmologyValidation( versions=[version], **params @@ -507,10 +506,33 @@ def test_a_patched_xi_dump_reads_back(self, tmp_path): read = CosmologyValidation(versions=[version], **params).calculate_2pcf_version( version, **binning )[pair] - for column in ("meanr", "npairs", "xip", "xim", "varxip", "varxim"): - np.testing.assert_allclose( - getattr(read, column), getattr(measured, column), rtol=1e-4 + for column in ("meanr", "npairs", "xip", "xim", "varxip", "varxim", "cov"): + np.testing.assert_array_equal( + getattr(read, column), getattr(measured, column) + ) + assert read.results + + def test_xi_plotting_reads_the_published_columns_without_remeasurement( + self, tmp_path, monkeypatch + ): + params, version = self._write_synthetic_catalogs(tmp_path) + cv = CosmologyValidation( + [version], npatch=4, theta_min=5, theta_max=100, nbins=6, **params + ) + pair = "tomo_bin_all_tomo_bin_all" + measured = cv.calculate_2pcf_version(version, num_threads=1)[pair] + + def no_catalogue_read(*args): + raise AssertionError("plotting must reuse its declared measurement input") + + monkeypatch.setattr(cv, "_shear_columns", no_catalogue_read) + cached = cv.calculate_2pcf_version(version, read_cached=True)[pair] + for column in ("xip", "xim", "meanr", "meanlogr", "weight", "npairs"): + np.testing.assert_array_equal( + getattr(cached, column), getattr(measured, column) ) + np.testing.assert_allclose(cached.varxip, measured.varxip, rtol=1e-15) + assert not cached.results def test_calculate_2pcf_is_reproducible_across_machines( self, tmp_path, monkeypatch @@ -523,7 +545,8 @@ def test_calculate_2pcf_is_reproducible_across_machines( process = treecorr.GGCorrelation.process def recording_process(gg, cat, *args, **kwargs): - patches.append(np.array(cat.patch)) + if cat.npatch > 1: + patches.append(np.array(cat.patch)) return process(gg, cat, *args, **kwargs) monkeypatch.setattr(treecorr.GGCorrelation, "process", recording_process) @@ -607,6 +630,199 @@ def recording_bin_catalog(cols, bin_id, npatch, patch_centers=None): cross_pair_catalogs[1][2], expected_centers, rtol=0, atol=1e-14 ) + @pytest.mark.parametrize("b_target", [0.01, 0.02]) + def test_bin_slop_follows_the_grid(self, tmp_path, b_target): + params, version = self._write_synthetic_catalogs(tmp_path) + cv = CosmologyValidation( + [version], + theta_min=1, + theta_max=250, + nbins=20, + b_target=b_target, + **params, + ) + expected = b_target / (np.log(250) / 20) + assert cv.treecorr_config["bin_slop"] == pytest.approx(expected) + assert cv._binning()["bin_slop"] == pytest.approx(expected) + assert cv._binning(0.08, 300, 1000)["bin_slop"] == 1 + assert cv._binning(bin_slop=0)["bin_slop"] == 0 + assert "angle_slop" not in cv.treecorr_config + + @pytest.mark.parametrize("tomography", [False, True]) + def test_published_means_are_layout_independent( + self, tmp_path, monkeypatch, tomography + ): + """Both auto- and cross-bin means ignore the covariance patch layout.""" + import treecorr + + from sp_validation.statistics import jackknife_patch_centers + + params, version = self._write_synthetic_catalogs( + tmp_path, + n_gal=4000, + with_tomography=tomography, + coherent_shear=True, + ) + cv = CosmologyValidation( + [version], npatch=8, theta_min=5, theta_max=200, nbins=8, **params + ) + # Small trees exercise cell aggregation without a large catalogue. + cv.treecorr_config.update(min_top=2, num_threads=1) + cols = cv._shear_columns(version, tomography) + cat = cv._bin_catalog(cols, "all", 1) + centers = jackknife_patch_centers(cat, 8) + angle = np.deg2rad(0.35) + rotation = np.array( + [ + [np.cos(angle), -np.sin(angle), 0], + [np.sin(angle), np.cos(angle), 0], + [0, 0, 1], + ] + ) + layouts = [centers, centers @ rotation.T] + assert not np.array_equal( + cv._bin_catalog(cols, "all", 8, layouts[0]).patch, + cv._bin_catalog(cols, "all", 8, layouts[1]).patch, + ) + measured = [] + for centers in layouts: + monkeypatch.setattr(cv, "_patch_centers", lambda *args, c=centers: c) + measured.append( + cv.calculate_2pcf_version(version, compute_tomography=tomography) + ) + for pair in measured[0]: + a, b = measured[0][pair], measured[1][pair] + for column in ( + "xip", + "xim", + "xip_im", + "xim_im", + "meanr", + "meanlogr", + "weight", + "npairs", + ): + np.testing.assert_array_equal(getattr(a, column), getattr(b, column)) + assert not np.array_equal(a.cov, b.cov) + # Covariance is still exactly the direct patched measurement. + bin1, bin2 = ( + (1, 2) + if pair == "tomo_bin_1_tomo_bin_2" + else ( + (1, 1) + if pair == "tomo_bin_1_tomo_bin_1" + else (2, 2) + if pair == "tomo_bin_2_tomo_bin_2" + else ("all", "all") + ) + ) + covariance_config = cv._binning() + covariance_config.pop("bin_slop") + direct = treecorr.GGCorrelation(covariance_config) + direct.process( + cv._bin_catalog(cols, bin1, 8, layouts[0]), + cat2=( + cv._bin_catalog(cols, bin2, 8, layouts[0]) if bin1 != bin2 else None + ), + ) + np.testing.assert_array_equal(a.cov, direct.cov) + np.testing.assert_array_equal(a.estimate_cov("jackknife"), direct.cov) + + def test_rho_tau_means_are_layout_independent(self, tmp_path, monkeypatch): + from sp_validation import rho_tau + + params, version = self._write_synthetic_catalogs( + tmp_path, + n_gal=1200, + n_star=1000, + with_psf=True, + ) + cv = CosmologyValidation( + [version], npatch=8, theta_min=5, theta_max=200, nbins=8, **params + ) + cv.treecorr_config["num_threads"] = 1 + cv.cc[version]["patch_number"] = 8 + from shear_psf_leakage.rho_tau_stat import Catalogs + + from sp_validation.statistics import jackknife_patch_centers + + cols = cv._shear_columns(version, False) + centers = jackknife_patch_centers(cv._bin_catalog(cols, "all", 1), 8) + angle = np.deg2rad(0.35) + rotation = np.array( + [ + [np.cos(angle), -np.sin(angle), 0], + [np.sin(angle), np.cos(angle), 0], + [0, 0, 1], + ] + ) + build = Catalogs.build_catalog + outputs = [] + for i, layout in enumerate((centers, centers @ rotation.T)): + + def with_layout(self, *args, layout=layout, **kwargs): + kwargs["patch_centers"] = layout + return build(self, *args, **kwargs) + + monkeypatch.setattr(Catalogs, "build_catalog", with_layout) + out = tmp_path / f"rho_tau_{i}" + out.mkdir() + rho, tau = rho_tau.get_jackknife_cov( + cv.cc, + version, + cv.treecorr_config, + str(out), + "test", + "test", + npatch=8, + ncov=1, + ) + outputs.append( + ( + rho.rho_stats.copy(), + tau.tau_stats.copy(), + np.load(out / "cov_tau_test_jk.npy"), + ) + ) + for i in (0, 1): + for column in outputs[0][i].dtype.names: + if not column.startswith("var"): + np.testing.assert_array_equal( + outputs[0][i][column], outputs[1][i][column] + ) + assert not np.array_equal(outputs[0][2], outputs[1][2]) + assert not np.array_equal( + outputs[0][1]["vartau_0_p"], outputs[1][1]["vartau_0_p"] + ) + + def test_aperture_mass_means_are_layout_independent(self, tmp_path, monkeypatch): + from sp_validation.statistics import jackknife_patch_centers + + params, version = self._write_synthetic_catalogs(tmp_path, n_gal=3000) + cv = CosmologyValidation([version], npatch=8, **params) + cv.treecorr_config.update(min_top=2, num_threads=1) + cols = cv._shear_columns(version, False) + cat = cv._bin_catalog(cols, "all", 1) + centers = jackknife_patch_centers(cat, 8) + angle = np.deg2rad(0.35) + rotation = np.array( + [ + [np.cos(angle), -np.sin(angle), 0], + [np.sin(angle), np.cos(angle), 0], + [0, 0, 1], + ] + ) + measured = [] + for layout in (centers, centers @ rotation.T): + monkeypatch.setattr(cv, "_patch_centers", lambda *args, c=layout: c) + cv.calculate_aperture_mass_dispersion( + theta_min=5, theta_max=200, nbins=30, npatch=8 + ) + measured.append(cv.map2[version]["tomo_bin_all_tomo_bin_all"]) + for key in ("mapsq", "mapsq_im", "mxsq", "mxsq_im"): + np.testing.assert_array_equal(measured[0][key], measured[1][key]) + assert not np.array_equal(measured[0]["varmapsq"], measured[1]["varmapsq"]) + def test_treecorr_runs_on_the_cpus_the_process_holds(self, tmp_path): """By default TreeCorr takes the process's CPU affinity, not the node's count.""" import treecorr diff --git a/src/sp_validation/tests/test_rho_tau_jackknife.py b/src/sp_validation/tests/test_rho_tau_jackknife.py index bc64537a..3ea8348e 100644 --- a/src/sp_validation/tests/test_rho_tau_jackknife.py +++ b/src/sp_validation/tests/test_rho_tau_jackknife.py @@ -59,12 +59,17 @@ def test_jackknife_draws_use_distinct_shared_layouts(tmp_path, monkeypatch): ) } layouts, catalogs, chunks = [], [], {"rho": [], "tau": []} + unpatched = [] compute = rho_tau.RhoStat.compute_rho_stats save = np.save def record_compute(self, catalog_id, *args, **kwargs): cats = self.catalogs.catalogs_dict psf = cats[f"psf_{catalog_id}"] + if psf.npatch == 1: + assert not kwargs["save_cov"] + unpatched.append(catalog_id) + return compute(self, catalog_id, *args, **kwargs) layouts.append(psf.patch_centers.copy()) catalogs.append(psf) for prefix in ("psf_error", "psf_size_error", "gal"): @@ -100,6 +105,7 @@ def record_save(path, arr, *args, **kwargs): npatch=7, ncov=3, ) + assert len(unpatched) == 1 assert len(layouts) == 3 assert all( not np.array_equal(layouts[i], layouts[j]) for i in range(3) for j in range(i) diff --git a/workflow/common.py b/workflow/common.py index 111da5f5..fb99cd4c 100644 --- a/workflow/common.py +++ b/workflow/common.py @@ -417,6 +417,7 @@ def cv_cov_tau(version, fiducial=None): "compute_cov_rho", "n_cov", "n_sim_cov", + "b_target", "theta_min", "theta_max", "nbins", diff --git a/workflow/rules/twopoint.smk b/workflow/rules/twopoint.smk index 93abebe4..3c7b2890 100644 --- a/workflow/rules/twopoint.smk +++ b/workflow/rules/twopoint.smk @@ -34,6 +34,7 @@ rule xi: cat_config=CAT_CONFIG, output_dir=str(COSMO_VAL), grid=lambda w: grid_of(XI_GRIDS, w), + b_target=config.get("cosmo_val", {}).get("b_target", 0.01), resources: # The fine integration grid needs more memory and wall time than the # ~20-bin reporting one; scale on nbins rather than splitting the rule. @@ -59,6 +60,7 @@ rule rho_tau_stats: npatch="{npatch}", cat_config=CAT_CONFIG, output_dir=str(COSMO_VAL), + b_target=config.get("cosmo_val", {}).get("b_target", 0.01), resources: mem_mb=64000, disk_mb=20000, diff --git a/workflow/scripts/run_2pcf.py b/workflow/scripts/run_2pcf.py index 1fcd9a0c..911778cb 100644 --- a/workflow/scripts/run_2pcf.py +++ b/workflow/scripts/run_2pcf.py @@ -45,6 +45,7 @@ def run_2pcf( output_dir, sacc_out=None, grid="reporting", + b_target=0.01, ): """Measure ξ±(θ) for ``ver`` and write its reporting SACC part. @@ -68,6 +69,7 @@ def run_2pcf( output_dir=output_dir, # so the SACC provenance metadata stamps the npatch actually measured npatch=npatch, + b_target=b_target, ) gg = cv.calculate_2pcf_version( ver, @@ -113,6 +115,7 @@ def _from_snakemake(smk): cat_config=p["cat_config"], output_dir=p["output_dir"], grid=p.get("grid", "reporting"), + b_target=p.get("b_target", 0.01), # The SACC part goes exactly where the rule declares it; the .txt # byproduct still lands under the resolved output dir. sacc_out=smk.output["sacc"], @@ -145,6 +148,7 @@ def _from_cli(argv=None): ap.add_argument( "--grid", default="reporting", help="SACC grid tag for the measured points" ) + ap.add_argument("--b-target", type=float, default=0.01) a = ap.parse_args(argv) run_2pcf( ver=a.ver, @@ -155,6 +159,7 @@ def _from_cli(argv=None): cat_config=a.cat_config, output_dir=a.out, grid=a.grid, + b_target=a.b_target, ) diff --git a/workflow/scripts/run_rho_tau.py b/workflow/scripts/run_rho_tau.py index d04da3f0..6e638ddd 100644 --- a/workflow/scripts/run_rho_tau.py +++ b/workflow/scripts/run_rho_tau.py @@ -16,6 +16,7 @@ theta_max=float(params["max_sep"]), nbins=int(params["nbins"]), npatch=int(params["npatch"]), + b_target=params.get("b_target", 0.01), catalog_config=params["cat_config"], output_dir=params["output_dir"], )