From 982c8e06febfbdef279372dcb61157703c0b9f73 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Tue, 29 Sep 2026 16:33:40 +0200 Subject: [PATCH 1/3] calculate_2pcf: jackknife patches from a seeded k-means, no patch-centre file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TreeCorr seeds its k-means afresh on every unseeded call, so the jackknife patches, and with them the ξ± jackknife covariance, differed between runs that did not share a {ver}_patches_npatch=N.dat. The k-means is now seeded with a fixed generator, which makes the patches a pure function of the catalogue's positions and weights on any thread count; the centre file and its reuse logic go. The aperture-mass ξ± takes the same seed. Co-Authored-By: Claude Opus 5.5 --- src/sp_validation/cosmo_val/real_space.py | 15 ++---- src/sp_validation/tests/test_cosmo_val.py | 61 ++++++++++++----------- 2 files changed, 38 insertions(+), 38 deletions(-) diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 89dad2a9..ec7ac4b3 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -43,8 +43,9 @@ def calculate_2pcf(self, ver, npatch=None, **treecorr_config): Notes: - If the output file for the given configuration already exists, the calculation is skipped, and the results are loaded from the file. - - If a patch file for the given configuration does not exist, it is - created during the process. + - Jackknife patches come from a k-means seeded with a fixed + generator, so they are a pure function of the catalogue's + positions and weights, the same on every run and thread count. - The ``.txt`` TreeCorr dump is the only raw byproduct written here. """ @@ -73,9 +74,6 @@ def calculate_2pcf(self, ver, npatch=None, **treecorr_config): g1, g2 = self._calibrated_g(ver) w = self._read_shear_cols(ver, "w_col") - # Use patch file if it exists - patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - cat_gal = treecorr.Catalog( ra=self.results[ver].dat_shear["RA"], dec=self.results[ver].dat_shear["Dec"], @@ -85,13 +83,9 @@ def calculate_2pcf(self, ver, npatch=None, **treecorr_config): ra_units=self.treecorr_config["ra_units"], dec_units=self.treecorr_config["dec_units"], npatch=npatch, - patch_centers=patch_file if os.path.exists(patch_file) else None, + rng=np.random.default_rng(0), ) - # If no patch file exists, save the current patches - if not os.path.exists(patch_file): - cat_gal.write_patch_centers(patch_file) - # Process the catalog & write the correlation functions gg.process(cat_gal) # Columns only. The covariance matrix lives in the SACC part; a @@ -356,6 +350,7 @@ def calculate_aperture_mass_dispersion( ra_units=self.treecorr_config["ra_units"], dec_units=self.treecorr_config["dec_units"], npatch=npatch, + rng=np.random.default_rng(0), ) gg.process(cat_gal) diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index 032eb2ff..401115d1 100644 --- a/src/sp_validation/tests/test_cosmo_val.py +++ b/src/sp_validation/tests/test_cosmo_val.py @@ -501,39 +501,44 @@ def test_a_patched_xi_dump_reads_back(self, tmp_path): getattr(read, column), getattr(measured, column), rtol=1e-4 ) - def test_calculate_2pcf_does_not_depend_on_thread_count(self, tmp_path): - """calculate_2pcf's ξ± is the same on 4 and on 48 TreeCorr threads. + def test_calculate_2pcf_is_reproducible_across_runs_and_threads(self, tmp_path): + """Two fresh output trees on 4 and 16 threads measure the same ξ±. - Production binning (default bin_slop/angle_slop), both runs on the - jackknife patches the first one writes, each from a fresh Catalog; they - must agree to far below the jackknife σ. + calculate_2pcf draws its jackknife patches by seeded k-means, so each + run splits the catalogue alike: the per-patch-pair counts, ξ± and its + jackknife variance agree. The footprint is wide enough that unseeded + draws land on different patches. """ import treecorr - params, version = self._write_synthetic_catalogs( - tmp_path, n_gal=4000, coherent_shear=True - ) - cv = CosmologyValidation( - versions=[version], - npatch=8, - theta_min=15.0, - theta_max=70.0, - nbins=6, - **params, - ) - - xi = {} - for n_threads in (4, 48): - # calculate_2pcf reads back an existing text dump instead of measuring. - for dump in Path(params["output_dir"]).glob(f"{version}_xi_*.txt"): - dump.unlink() - gg = cv.calculate_2pcf(version, num_threads=n_threads) + xi, var, counts = {}, {}, {} + for tree, n_threads in (("a", 4), ("b", 16)): + (tmp_path / tree).mkdir() + params, version = self._write_synthetic_catalogs( + tmp_path / tree, + n_gal=4000, + ra_range=(0.0, 60.0), + dec_range=(-10.0, 30.0), + coherent_shear=True, + ) + gg = CosmologyValidation( + versions=[version], + npatch=10, + theta_min=15.0, + theta_max=70.0, + nbins=6, + **params, + ).calculate_2pcf(version, num_threads=n_threads) assert treecorr.get_omp_threads() == n_threads # the count took effect - xi[n_threads] = np.concatenate([gg.xip, gg.xim]) - sigma = np.sqrt(np.concatenate([gg.varxip, gg.varxim])) - - shift = np.max(np.abs(xi[48] - xi[4]) / sigma) - assert shift < 1e-6, f"ξ± moves by {shift:.3g}σ between 4 and 48 threads" + xi[tree] = np.concatenate([gg.xip, gg.xim]) + var[tree] = np.concatenate([gg.varxip, gg.varxim]) + counts[tree] = {k: r.npairs for k, r in gg.results.items()} + + assert counts["a"].keys() == counts["b"].keys() + for k in counts["a"]: + np.testing.assert_array_equal(counts["a"][k], counts["b"][k]) + assert np.max(np.abs(xi["a"] - xi["b"]) / np.sqrt(var["a"])) < 1e-6 + np.testing.assert_allclose(var["a"], var["b"], rtol=1e-6) def test_calculate_scale_dependent_leakage_runs_on_synthetic_catalog( self, tmp_path From 2f70c9357ec0492a49bd9261dbc7803950a88542 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Wed, 30 Sep 2026 23:59:15 +0200 Subject: [PATCH 2/3] rho_tau: seed the per-patch k-means as calculate_2pcf does Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01UzeJqdtGgoeWizre32L5fD --- src/sp_validation/rho_tau.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 62ae8bf8..bb8d2953 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -372,7 +372,7 @@ def get_jackknife_cov( field = rho_stat_handler.catalogs.catalogs_dict[ f"psf_{version}{i}" ].getNField(max_top=int.bit_length(npatch) - 1, coords="spherical") - patch, centers = field.run_kmeans(npatch) + patch, centers = field.run_kmeans(npatch, rng=np.random.default_rng(0)) # Update the patch centers of the catalogs for key, cat in rho_stat_handler.catalogs.catalogs_dict.items(): From 8c55e0a6840eebbbd16dfbd2475021f1b0c126e7 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Thu, 1 Oct 2026 05:27:45 +0200 Subject: [PATCH 3/3] Jackknife patch centres from a k-means on a pinned-depth tree statistics.jackknife_patch_centers runs the seeded k-means on an NField with min_top=6, so the initial centres no longer depend on the OpenMP thread count TreeCorr reads when it sizes the tree. calculate_2pcf and calculate_aperture_mass_dispersion build their shear catalogue through RealSpaceMixin._shear_catalog, which passes those centres as patch_centers; rho_tau.get_jackknife_cov takes its per-patch centres from the same helper. The reproducibility test now runs 100 patches and simulates 4- and 16-CPU machines, and asserts identical patch labels, per-patch-pair counts, xi+- and jackknife variance. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01UzeJqdtGgoeWizre32L5fD --- src/sp_validation/cosmo_val/real_space.py | 60 +++++++++++------------ src/sp_validation/rho_tau.py | 9 ++-- src/sp_validation/statistics.py | 39 ++++++++++++++- src/sp_validation/tests/test_cosmo_val.py | 42 +++++++++++----- 4 files changed, 100 insertions(+), 50 deletions(-) diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index ec7ac4b3..fce6d55c 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -14,8 +14,30 @@ import treecorr from cs_util import plots as cs_plots +from sp_validation.statistics import jackknife_patch_centers + class RealSpaceMixin: + def _shear_catalog(self, ver, npatch): + """Calibrated shear catalogue of ``ver`` with seeded jackknife patches. + + Call inside ``self.results[ver].temporarily_read_data()``. + """ + positions = { + "ra": self.results[ver].dat_shear["RA"], + "dec": self.results[ver].dat_shear["Dec"], + "w": self._read_shear_cols(ver, "w_col"), + "ra_units": self.treecorr_config["ra_units"], + "dec_units": self.treecorr_config["dec_units"], + } + centers = None + if int(npatch) > 1: + centers = jackknife_patch_centers( + treecorr.Catalog(**positions), int(npatch) + ) + g1, g2 = self._calibrated_g(ver) + return treecorr.Catalog(**positions, g1=g1, g2=g2, patch_centers=centers) + def calculate_2pcf(self, ver, npatch=None, **treecorr_config): """ Calculate the two-point correlation function (2PCF) ξ± for a given catalog @@ -43,9 +65,10 @@ def calculate_2pcf(self, ver, npatch=None, **treecorr_config): Notes: - If the output file for the given configuration already exists, the calculation is skipped, and the results are loaded from the file. - - Jackknife patches come from a k-means seeded with a fixed - generator, so they are a pure function of the catalogue's - positions and weights, the same on every run and thread count. + - Jackknife patches come from a seeded k-means on a fixed-depth tree + (``statistics.jackknife_patch_centers``), so they are a pure + function of the catalogue's positions and weights, the same on + every run and thread count. - The ``.txt`` TreeCorr dump is the only raw byproduct written here. """ @@ -71,20 +94,7 @@ def calculate_2pcf(self, ver, npatch=None, **treecorr_config): else: # Load data and create a catalog with self.results[ver].temporarily_read_data(): - g1, g2 = self._calibrated_g(ver) - w = self._read_shear_cols(ver, "w_col") - - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=g2, - w=w, - ra_units=self.treecorr_config["ra_units"], - dec_units=self.treecorr_config["dec_units"], - npatch=npatch, - rng=np.random.default_rng(0), - ) + cat_gal = self._shear_catalog(ver, npatch) # Process the catalog & write the correlation functions gg.process(cat_gal) @@ -340,24 +350,10 @@ def calculate_aperture_mass_dispersion( gg.read(out_fname) else: with self.results[ver].temporarily_read_data(): - g1, g2 = self._calibrated_g(ver) - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=g2, - w=self._read_shear_cols(ver, "w_col"), - ra_units=self.treecorr_config["ra_units"], - dec_units=self.treecorr_config["dec_units"], - npatch=npatch, - rng=np.random.default_rng(0), - ) - + cat_gal = self._shear_catalog(ver, npatch) gg.process(cat_gal) gg.write(out_fname) del cat_gal - del g1 - del g2 mapsq, mapsq_im, mxsq, mxsq_im, varmapsq = gg.calculateMapSq( R=theta_map, diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index bb8d2953..62f40959 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -9,6 +9,7 @@ # SquareRootScale now lives in sp_validation.plots; re-exported here so that # `from sp_validation.rho_tau import SquareRootScale` keeps working. from sp_validation.plots import SquareRootScale # noqa: F401 +from sp_validation.statistics import jackknife_patch_centers def _extract_xip(correlations): @@ -369,10 +370,10 @@ def get_jackknife_cov( print(f"Computing the patch centers for patch {i + 1}/{ncov}") npatch = rho_stat_handler.catalogs._params["patch_number"] - field = rho_stat_handler.catalogs.catalogs_dict[ - f"psf_{version}{i}" - ].getNField(max_top=int.bit_length(npatch) - 1, coords="spherical") - patch, centers = field.run_kmeans(npatch, rng=np.random.default_rng(0)) + centers = jackknife_patch_centers( + rho_stat_handler.catalogs.catalogs_dict[f"psf_{version}{i}"], + npatch, + ) # Update the patch centers of the catalogs for key, cat in rho_stat_handler.catalogs.catalogs_dict.items(): diff --git a/src/sp_validation/statistics.py b/src/sp_validation/statistics.py index b8c6d1d5..e8a3fd49 100644 --- a/src/sp_validation/statistics.py +++ b/src/sp_validation/statistics.py @@ -3,13 +3,50 @@ :Name: statistics.py :Description: Cosmology-independent statistical helpers (jackknife resampling, - chi2/PTE, covariance<->correlation, OneCovariance reshaping). + jackknife patch centres, chi2/PTE, covariance<->correlation, + OneCovariance reshaping). Extracted verbatim from the former basic.py. """ import numpy as np from scipy import stats +#: Depth of the ball-tree layers that seed the jackknife k-means. +PATCH_MIN_TOP = 6 + + +def jackknife_patch_centers(cat, npatch, seed=0): + """Seeded k-means jackknife patch centres for a TreeCorr catalogue. + + Seeding the k-means does not by itself fix the patches. TreeCorr starts the + k-means from the top layers of a ball tree whose depth, unless ``min_top`` + is given, grows with the OpenMP thread count (``Field._determine_top``), + and ``Catalog(npatch=..., rng=...)`` offers no way to set it. Pinning that + depth makes the centres a function of the catalogue's positions, weights + and ``seed`` alone, on every machine. + + Parameters + ---------- + cat : treecorr.Catalog + Catalogue with spherical (RA, Dec) positions. + npatch : int + Number of patches. + seed : int, optional + Seed for the k-means initialisation. + + Returns + ------- + numpy.ndarray + Patch centres, to pass as ``patch_centers`` to ``treecorr.Catalog``. + """ + field = cat.getNField( + min_top=PATCH_MIN_TOP, + max_top=int.bit_length(npatch) - 1, + coords="spherical", + ) + _, centers = field.run_kmeans(npatch, rng=np.random.default_rng(seed)) + return centers + def jackknif_weighted_average2( data, diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index 401115d1..d2885071 100644 --- a/src/sp_validation/tests/test_cosmo_val.py +++ b/src/sp_validation/tests/test_cosmo_val.py @@ -501,18 +501,34 @@ def test_a_patched_xi_dump_reads_back(self, tmp_path): getattr(read, column), getattr(measured, column), rtol=1e-4 ) - def test_calculate_2pcf_is_reproducible_across_runs_and_threads(self, tmp_path): - """Two fresh output trees on 4 and 16 threads measure the same ξ±. - - calculate_2pcf draws its jackknife patches by seeded k-means, so each - run splits the catalogue alike: the per-patch-pair counts, ξ± and its - jackknife variance agree. The footprint is wide enough that unseeded - draws land on different patches. + def test_calculate_2pcf_is_reproducible_across_machines( + self, tmp_path, monkeypatch + ): + """Two fresh output trees on 4- and 16-CPU machines measure the same ξ±. + + TreeCorr's default k-means tree depth follows the machine's OpenMP + thread count (3 levels on 4 CPUs, 4 on 16). With 100 patches (up to 6 + levels, and not a power of two) that changes the k-means start, so a + seed alone would split the catalogue differently. calculate_2pcf pins + the depth: patch labels, per-patch-pair counts, ξ± and its jackknife + variance all agree. """ import treecorr + import treecorr.field + + patches = [] + process = treecorr.GGCorrelation.process + + def recording_process(gg, cat, *args, **kwargs): + patches.append(np.array(cat.patch)) + return process(gg, cat, *args, **kwargs) + + monkeypatch.setattr(treecorr.GGCorrelation, "process", recording_process) xi, var, counts = {}, {}, {} - for tree, n_threads in (("a", 4), ("b", 16)): + for tree, n_cpu in (("a", 4), ("b", 16)): + # The thread count TreeCorr's default tree depth reads. + monkeypatch.setattr(treecorr.field, "get_omp_threads", lambda n=n_cpu: n) (tmp_path / tree).mkdir() params, version = self._write_synthetic_catalogs( tmp_path / tree, @@ -523,22 +539,22 @@ def test_calculate_2pcf_is_reproducible_across_runs_and_threads(self, tmp_path): ) gg = CosmologyValidation( versions=[version], - npatch=10, + npatch=100, theta_min=15.0, theta_max=70.0, nbins=6, **params, - ).calculate_2pcf(version, num_threads=n_threads) - assert treecorr.get_omp_threads() == n_threads # the count took effect + ).calculate_2pcf(version) xi[tree] = np.concatenate([gg.xip, gg.xim]) var[tree] = np.concatenate([gg.varxip, gg.varxim]) counts[tree] = {k: r.npairs for k, r in gg.results.items()} + np.testing.assert_array_equal(patches[0], patches[1]) assert counts["a"].keys() == counts["b"].keys() for k in counts["a"]: np.testing.assert_array_equal(counts["a"][k], counts["b"][k]) - assert np.max(np.abs(xi["a"] - xi["b"]) / np.sqrt(var["a"])) < 1e-6 - np.testing.assert_allclose(var["a"], var["b"], rtol=1e-6) + np.testing.assert_allclose(xi["a"], xi["b"], rtol=0, atol=1e-12) + np.testing.assert_allclose(var["a"], var["b"], rtol=1e-10) def test_calculate_scale_dependent_leakage_runs_on_synthetic_catalog( self, tmp_path