diff --git a/papers/cosmo_val/config/config.yaml b/papers/cosmo_val/config/config.yaml index 99f553c9..084509ce 100644 --- a/papers/cosmo_val/config/config.yaml +++ b/papers/cosmo_val/config/config.yaml @@ -51,7 +51,19 @@ cosmo_val: noise_bias_method: analytic cell_seed: 8192 path_onecovariance: "/home/guerrini/OneCovariance/" - rho_tau_method: lsq + # PSF-leakage fit: 'eiv' is the joint errors-in-variables fit of the rho/tau + # jackknife measurement below; 'lsq'/'emcee' fit the published tau statistics + # with the cov_estimate_method covariance. + rho_tau_method: eiv + psf_leakage: + # k-means patches shared by stars and galaxies in the joint measurement + npatch_fine: 800 + # jackknife patch counts the fit reports (Ward unions of the fine patches) + patch_counts: [200, 400] + # the patch count whose fit gives the samples used downstream + fiducial_npatch: 400 + # beta(T_PSF) = beta0 + beta1 (T_PSF / - 1) + beta_psf_size: false # Fiducial scale cut (arcmin) for B-mode PTE extraction. Paper II: [12, 83]. fiducial_scale_cut: [12, 83] diff --git a/pyproject.toml b/pyproject.toml index 10eb70ac..ffecce33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,11 +81,8 @@ dependencies = [ "sacc>=2.4,<3", "scipy", "seaborn", - # Intentionally tracks the mutable `develop` branch (NOT SHA-pinned): - # shear_psf_leakage is actively co-developed CosmoStat code we want to stay - # current with, so freezing it would only create bump-churn. (Contrast the - # getdist feature-branch below, which is an external fork we pin for repro.) - "shear_psf_leakage @ git+https://github.com/CosmoStat/shear_psf_leakage.git@develop", + # Centred galaxy-position alpha template (shear_psf_leakage PR #59). + "shear_psf_leakage @ git+https://github.com/CosmoStat/shear_psf_leakage.git@d0e10d83ecb1978ed56599f0b46326e3d000680e", "skyproj", # UNIONS-WL fork of DESC Smokescreen. Tracks the fork's main; uv.lock # freezes the exact commit, so updates land via diff --git a/src/sp_validation/correlation.py b/src/sp_validation/correlation.py index 131a81bc..80de8370 100644 --- a/src/sp_validation/correlation.py +++ b/src/sp_validation/correlation.py @@ -1,11 +1,30 @@ """Layout-independent full-sample means with patched covariance measurements.""" +import hashlib import json from pathlib import Path +import numpy as np import treecorr +def selection_digest(*masks): + """Digest of row selections; ``None`` stands for every row. + + Recorded beside a cached measurement so a product made from different rows, + e.g. another star flag or tomographic bin, is never reused. + """ + digest = hashlib.sha1() + for mask in masks: + if mask is None: + digest.update(b"all;") + else: + mask = np.asarray(mask, dtype=bool) + digest.update(f"{mask.size};".encode()) + digest.update(np.packbits(mask).tobytes()) + return digest.hexdigest() + + def _measurement_configs(config): """Return the explicit means settings and default-tolerance patched settings.""" means = dict(config) @@ -14,15 +33,19 @@ def _measurement_configs(config): return {"means": means, "patched": patched} -def measurement_matches(path, config): - """Whether a cache records the current means and patched configurations.""" +def measurement_matches(path, config, selection=None, definition=None): + """Whether a cache records the current configurations and row selection.""" metadata = Path(str(path) + ".json") if not Path(path).exists() or not metadata.exists(): return False recorded = json.loads(metadata.read_text()) + if recorded.get("definition") != definition: + return False recorded_configs = recorded.get("configs") if recorded.get("means") != "unpatched" or not isinstance(recorded_configs, dict): return False + if recorded.get("selection") != selection: + 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, @@ -33,12 +56,14 @@ def measurement_matches(path, config): 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 write_measurement_metadata(path, config, selection=None, definition=None): + """Record both TreeCorr configurations and the row selection beside a product.""" + record = {"means": "unpatched", "configs": _measurement_configs(config)} + if definition is not None: + record["definition"] = definition + if selection is not None: + record["selection"] = selection + Path(str(path) + ".json").write_text(json.dumps(record) + "\n") def _unpatched_catalog(cat): diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index a0beaf73..1c623fc9 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -22,7 +22,7 @@ from .catalog_characterization import CatalogCharacterizationMixin from .cosebis import CosebisMixin from .pseudo_cl import PseudoClMixin -from .psf_systematics import PSFSystematicsMixin +from .psf_systematics import PSF_LEAKAGE_DEFAULTS, PSFSystematicsMixin from .pure_eb import PureEBMixin from .real_space import RealSpaceMixin @@ -128,8 +128,18 @@ class CosmologyValidation( and analysis settings for each version. output_dir : str, optional Output directory. If None, the catalog config's paths.output. - rho_tau_method : {'lsq', 'mcmc'}, default 'lsq' - Fitting method for PSF leakage systematics parameters. + rho_tau_method : {'eiv', 'lsq', 'emcee', 'none'}, default 'eiv' + Fitting method for the PSF-leakage parameters (alpha, beta, eta). + 'eiv' is the joint errors-in-variables fit of the rho/tau jackknife + measurement configured by ``psf_leakage``; 'lsq' and 'emcee' fit the + published tau statistics with the ``cov_estimate_method`` covariance. + psf_leakage : dict, optional + Settings of the PSF-leakage measurement and 'eiv' fit, over + ``PSF_LEAKAGE_DEFAULTS``: ``npatch_fine`` (k-means patches shared by + stars and galaxies), ``patch_counts`` (jackknife patch counts the fit + reports, unions of the fine patches), ``fiducial_npatch`` (the count + whose fit gives the samples used downstream), ``beta_psf_size`` (fit + beta(T_PSF) = beta0 + beta1 (T_PSF / - 1)). cov_estimate_method : {'th', 'jk', 'sim'}, default 'th' Covariance estimation: 'th' for semi-analytic theory, 'jk' for jackknife, 'sim' for a covariance measured on simulations. @@ -284,7 +294,8 @@ def __init__( versions, catalog_config="./cat_config.yaml", output_dir=None, - rho_tau_method="lsq", + rho_tau_method="eiv", + psf_leakage=None, cov_estimate_method="th", compute_cov_rho=True, n_cov=100, @@ -318,7 +329,20 @@ def __init__( b_target=0.01, ): self.b_target = b_target - self.rho_tau_method = rho_tau_method + self.set_psf_parameter_sampling_method(rho_tau_method) + unknown = set(psf_leakage or {}) - set(PSF_LEAKAGE_DEFAULTS) + if unknown: + raise ValueError(f"Unknown psf_leakage settings: {sorted(unknown)}") + self.psf_leakage = {**PSF_LEAKAGE_DEFAULTS, **(psf_leakage or {})} + npatch_fine = self.psf_leakage["npatch_fine"] + patch_counts = self.psf_leakage["patch_counts"] + if not all(1 < n < npatch_fine for n in patch_counts): + raise ValueError( + f"psf_leakage patch_counts {patch_counts} must each lie strictly " + f"between 1 and npatch_fine={npatch_fine}" + ) + if self.psf_leakage["fiducial_npatch"] not in patch_counts: + raise ValueError("psf_leakage fiducial_npatch must be one of patch_counts") self.cov_estimate_method = cov_estimate_method self.compute_cov_rho = compute_cov_rho self.n_cov = n_cov diff --git a/src/sp_validation/cosmo_val/diagnostic_plots.py b/src/sp_validation/cosmo_val/diagnostic_plots.py index 6634c03a..495d840e 100644 --- a/src/sp_validation/cosmo_val/diagnostic_plots.py +++ b/src/sp_validation/cosmo_val/diagnostic_plots.py @@ -93,20 +93,39 @@ def plot_objectwise_leakage(coefficients, styles, output): plt.close(fig) +def leakage_contour_axes(nparam): + """getdist names, labels and reference markers of the leakage parameters. + + Three parameters are (alpha, beta, eta); four, with the PSF-size slope, + (alpha, beta0, eta, beta1). + """ + if nparam == 3: + labels = [r"\alpha", r"\beta", r"\eta"] + elif nparam == 4: + labels = [r"\alpha", r"\beta_0", r"\eta", r"\beta_1"] + else: + raise ValueError(f"expected 3 or 4 leakage parameters, got {nparam}") + names = [f"x{i}" for i in range(nparam)] + markers = dict(zip(names, (0, 1, 1, 0))) + return names, labels, markers + + def plot_rho_tau_fits(products, styles, contours, samples, summary): """Render posterior contours and propagated PSF errors from cached arrays.""" from shear_psf_leakage import plots as psf_plots versions = list(products) + sample_list = [products[v]["flat_samples"] for v in versions] + names, labels, markers = leakage_contour_axes(sample_list[0].shape[1]) psf_plots.plot_contours( - [products[v]["flat_samples"] for v in versions], - names=["x0", "x1", "x2"], - labels=[r"\alpha", r"\beta", r"\eta"], + sample_list, + names=names, + labels=labels, savefig=str(contours), legend_labels=versions, legend_loc="upper right", contour_colors=[styles[v]["colour"] for v in versions], - markers={"x0": 0, "x1": 1, "x2": 1}, + markers=markers, show=False, close=True, ) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 4e67c4aa..735475d6 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -5,7 +5,9 @@ object-wise) for catalogue validation. """ +import json import os +import zlib from pathlib import Path import matplotlib.pyplot as plt @@ -19,15 +21,31 @@ from shear_psf_leakage.rho_tau_stat import PSFErrorFit from uncertainties import ufloat -from .. import blinding, sacc_io +from .. import blinding, leakage_fit, sacc_io from ..io import open_entry from ..masks import catalogue_selection from ..rho_tau import ( + get_params_rho_tau, get_rho_tau_w_cov, get_samples, + get_samples_eiv, + leakage_cache_key, + measure_leakage_jackknife, + read_leakage_jackknife, ) +from .diagnostic_plots import leakage_contour_axes from .sacc_writers import rho_tau_to_sacc +# Settings of the joint rho/tau leakage measurement and its 'eiv' fit; the +# CosmologyValidation ``psf_leakage`` argument overrides them key by key. +PSF_LEAKAGE_DEFAULTS = { + "npatch_fine": 800, + "patch_counts": [200, 400], + "fiducial_npatch": 400, + "beta_psf_size": False, +} +RHO_TAU_METHODS = ("eiv", "lsq", "emcee", "none") + # TODO: Reorganise the order of functions so it is more readable class PSFSystematicsMixin: @@ -404,6 +422,107 @@ def _get_star_mask(self, ver): ) return star_mask + def leakage_jackknife(self, ver, tomo_bin_id="all", measure=False): + """Joint rho/tau jackknife measurement of one version and bin. + + Read from ``rho_tau_stats/leakage_eiv_.npz``, which must have + been measured with the current ``psf_leakage`` settings and star and + galaxy selections. Only ``measure=True`` (the workflow rule + ``cv_leakage_jackknife``) measures it, by + `rho_tau.measure_leakage_jackknife`, when no matching cache exists or + ``force_run`` is set; otherwise a missing or stale cache raises. + """ + cache = self.__dict__.setdefault("_leakage_jackknife", {}) + if (ver, tomo_bin_id) not in cache: + path = self.leakage_jackknife_path(ver, tomo_bin_id) + settings = self.psf_leakage + mask_star = self._get_star_mask(ver) + mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) + key = leakage_cache_key( + self.treecorr_config, + settings["npatch_fine"], + settings["patch_counts"], + settings["beta_psf_size"], + mask_star=mask_star, + mask_gal=mask_gal, + params=get_params_rho_tau(self.cc[ver]), + ) + reuse = not (measure and self.force_run) + products = read_leakage_jackknife(path, key) if reuse else None + if products is None and not measure: + state = "is stale for" if path.exists() else "is missing for" + raise FileNotFoundError( + f"The joint leakage measurement {path} {state} the current " + "psf_leakage settings and selections. Measure it with the " + "Snakemake rule cv_leakage_jackknife, or interactively with " + f"cv.leakage_jackknife({ver!r}, {tomo_bin_id!r}, measure=True)." + ) + if products is None: + products = measure_leakage_jackknife( + self.cc, + ver, + self.treecorr_config, + path, + key, + mask_star=mask_star, + mask_gal=mask_gal, + ) + cache[(ver, tomo_bin_id)] = products + return cache[(ver, tomo_bin_id)] + + def leakage_jackknife_path(self, ver, tomo_bin_id="all", kind="leakage_eiv"): + """Path of a joint leakage product (``kind`` leakage_eiv or leakage_eiv_fit).""" + base = self.basename( + ver, tomo_bin_a=tomo_bin_id, npatch=self.psf_leakage["npatch_fine"] + ) + suffix = "json" if kind == "leakage_eiv_fit" else "npz" + return Path(self._output_path("rho_tau_stats", f"{kind}_{base}.{suffix}")) + + def _leakage_eiv_fit(self, ver, tomo_bin_id, nsamples): + """EIV fits of one version and bin, and samples of the fiducial fit. + + The fits at every ``patch_counts`` entry are also written as JSON + beside the measurement. The sample seed derives from the bin id, so + different bins draw independent samples. + """ + cache = self.__dict__.setdefault("_leakage_eiv_fits", {}) + if (ver, tomo_bin_id) not in cache: + settings = self.psf_leakage + samples, fits, quantiles = get_samples_eiv( + self.leakage_jackknife(ver, tomo_bin_id), + settings["patch_counts"], + settings["fiducial_npatch"], + beta_psf_size=settings["beta_psf_size"], + nsamples=nsamples, + seed=zlib.crc32(str(tomo_bin_id).encode()), + ) + summary = { + str(n): { + key: fit[key] + for key in ( + "parameters", + "sigma", + "parameter_covariance", + "chi2", + "dof", + "pte", + "hartlap", + "bfgs_success", + "gradient_norm", + "profile_hessian_positive", + "profile_hessian_eigenvalues", + "residual", + "residual_sigma", + ) + } + for n, fit in fits.items() + } + self.leakage_jackknife_path(ver, tomo_bin_id, "leakage_eiv_fit").write_text( + json.dumps({"settings": settings, "fits": summary}, indent=2) + "\n" + ) + cache[(ver, tomo_bin_id)] = samples, fits, quantiles + return cache[(ver, tomo_bin_id)] + def set_params_rho_tau(self, ver, params, params_psf): params = {**params} @@ -485,8 +604,11 @@ def set_params_leakage_object(self, ver): return params_in def set_psf_parameter_sampling_method(self, rho_tau_method): - if rho_tau_method not in ["emcee", "lsq"]: - raise ValueError("Invalid PSF parameter sampling method.") + if rho_tau_method not in RHO_TAU_METHODS: + raise ValueError( + f"Invalid PSF parameter sampling method {rho_tau_method!r}; " + f"choose one of {RHO_TAU_METHODS}." + ) self.rho_tau_method = rho_tau_method def set_psf_parameter_nsamples(self, nsamples): @@ -496,15 +618,14 @@ def set_psf_parameter_nwalkers(self, nwalkers): self.psf_error_nwalkers = nwalkers def get_samples(self, version, params, tomo_bin_id, track_result=False): - # Hartlap debiasing of the inverse covariance: the number of jackknife - # patches or of simulations the covariance was estimated from. - n_realisations = {"jk": params["patch_number"], "sim": self.n_sim_cov}.get( - self.cov_estimate_method - ) - - base_rho = self.basename(version) - base_tau = self.basename(version, tomo_bin_a=tomo_bin_id) + """PSF-leakage parameter samples of one version and tomographic bin. + With ``rho_tau_method`` 'eiv' these are draws from the Gaussian of the + ``fiducial_npatch`` joint EIV fit, (alpha, beta, eta) or, with + ``beta_psf_size``, (alpha, beta0, eta, beta1); the tracked result + holds the fits at every patch count. Otherwise the 'lsq' or 'emcee' + samples of the published tau statistics. + """ # Set the number of samples and walkers and fallback to defaults if not attributed. n_samples = ( self.psf_error_nsamples if hasattr(self, "psf_error_nsamples") else 10_000 @@ -513,16 +634,28 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): self.psf_error_nwalkers if hasattr(self, "psf_error_nwalkers") else 124 ) - flat_samples, result, q = get_samples( - self.psf_fitter, - base_rho, - base_tau, - cov_type=self.cov_estimate_method, - apply_debias=n_realisations, - sampler=self.rho_tau_method, - nsamples=n_samples, - nwalkers=n_walkers, - ) + if self.rho_tau_method == "eiv": + flat_samples, result, q = self._leakage_eiv_fit( + version, tomo_bin_id, n_samples + ) + else: + # Hartlap debiasing of the inverse covariance: the number of + # jackknife patches or of simulations the covariance was + # estimated from. + n_realisations = { + "jk": params["patch_number"], + "sim": self.n_sim_cov, + }.get(self.cov_estimate_method) + flat_samples, result, q = get_samples( + self.psf_fitter, + self.basename(version), + self.basename(version, tomo_bin_a=tomo_bin_id), + cov_type=self.cov_estimate_method, + apply_debias=n_realisations, + sampler=self.rho_tau_method, + nsamples=n_samples, + nwalkers=n_walkers, + ) if track_result: if version not in self.rho_tau_fits.keys(): @@ -540,14 +673,37 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): return flat_samples def get_xi_psf_sys_samples(self, ver, params, tomo_bin_a, tomo_bin_b): - base_rho = self.basename(ver) - self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") - nbins = self.psf_fitter.rho_stat_handler._treecorr_config["nbins"] + """xi_sys+ and xi_sys- samples, one per pair of parameter samples. + With 'eiv', xi_sys = a^T M b over the rho (and, with the PSF-size + slope, u de) correlations of the joint leakage measurement of the + full star sample (`leakage_fit.systematic_correlation_matrix`); + otherwise from the published rho statistics. + """ # Get the samples for the given tomographic bins flat_samples_a = self.get_samples(ver, params, tomo_bin_a, track_result=False) flat_samples_b = self.get_samples(ver, params, tomo_bin_b, track_result=False) + if self.rho_tau_method == "eiv": + products = self.leakage_jackknife(ver, tomo_bin_a) + nbins = len(products["theta"]) + rho_xip = products["vector"][: 6 * nbins].reshape(6, nbins) + slope = self.psf_leakage["beta_psf_size"] + return tuple( + leakage_fit.xi_psf_sys( + flat_samples_a, + flat_samples_b, + leakage_fit.systematic_correlation_matrix( + rho, products[f"slope_{part}"] if slope else None + ), + ) + for rho, part in ((rho_xip, "xip"), (products["rho_xim"], "xim")) + ) + + base_rho = self.basename(ver) + self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + nbins = self.psf_fitter.rho_stat_handler._treecorr_config["nbins"] + xi_psf_sys_samples_plus = np.array( [ self.psf_fitter.compute_xi_psf_sys(sample_a, sample_b, p_or_m="p") @@ -1090,15 +1246,16 @@ def plot_rho_tau_fits( sample_list = [ self.rho_tau_fits[ver]["flat_samples"]["tomo_bin_all"] for ver in versions ] + names, labels, markers = leakage_contour_axes(sample_list[0].shape[1]) psfleak_plots.plot_contours( sample_list, - names=["x0", "x1", "x2"], - labels=[r"\alpha", r"\beta", r"\eta"], + names=names, + labels=labels, savefig=savefig, legend_labels=versions, legend_loc="upper right", contour_colors=colors, - markers={"x0": 0, "x1": 1, "x2": 1}, + markers=markers, show=show, close=close, ) diff --git a/src/sp_validation/dashboard/psf.py b/src/sp_validation/dashboard/psf.py index 736b3eaf..764aa314 100644 --- a/src/sp_validation/dashboard/psf.py +++ b/src/sp_validation/dashboard/psf.py @@ -26,14 +26,34 @@ 2: ((0, 2), (1, 1), (2, 4)), 5: ((0, 5), (1, 4), (2, 3)), } -PARAMS = (r"\alpha", r"\beta", r"\eta") +# Labels of (α, β, η) and, with the PSF-size slope, (α, β₀, η, β₁). +PARAMS = { + 3: (r"\alpha", r"\beta", r"\eta"), + 4: (r"\alpha", r"\beta_0", r"\eta", r"\beta_1"), +} + + +# τ0, τ2, τ5 rows of the PSF-size slope template ⟨u δe, {p, δe, q}⟩. +SLOPE_ROW = {0: 0, 2: 1, 5: 2} def tau_model_samples(rt, k, samples): - """τ_{k,+}(θ) for each coefficient draw, shape ``(n_samples, n_theta)``.""" + """τ_{k,+}(θ) for each coefficient draw, shape ``(n_samples, n_theta)``. + + Draws of (α, β₀, η, β₁) from the PSF-size-slope fit add β₁ times the + fit's ``tau_slope_template`` (⟨u δe, p⟩, ⟨u δe, δe⟩, ⟨u δe, q⟩ for τ0, + τ2, τ5). + """ model = 0 for param, rho_k in TAU_MODEL[k]: model = model + samples[:, [param]] * rt["rho"][rho_k]["p"][None, :] + if samples.shape[1] == 4: + if "tau_slope_template" not in rt["fit"]: + raise ValueError("PSF-size slope samples need the fit's tau_slope_template") + template = np.asarray(rt["fit"]["tau_slope_template"])[SLOPE_ROW[k]] + model = model + samples[:, [3]] * template[None, :] + elif samples.shape[1] != 3: + raise ValueError(f"expected 3 or 4 leakage parameters, got {samples.shape[1]}") return model @@ -101,10 +121,11 @@ def _contour_levels(h, fractions=(0.68, 0.95)): return sorted(flat[np.searchsorted(cum, f)] for f in fractions) -def coefficient_draws(rt, out, markers=(0, 1, 1)): +def coefficient_draws(rt, out, markers=(0, 1, 1, 0)): """Cached coefficient draws; axes omit 0.3% of each marginal tail.""" samples = rt["fit"]["flat_samples"] n = samples.shape[1] + labels = PARAMS[n] with S.figure_style(): fig, axes = plt.subplots(n, n, figsize=(3.6, 3.6)) lims = [np.percentile(samples[:, i], [0.3, 99.7]) for i in range(n)] @@ -130,7 +151,7 @@ def coefficient_draws(rt, out, markers=(0, 1, 1)): ) lo, med, hi = np.percentile(x, [16, 50, 84]) ax.set_title( - rf"${PARAMS[i]} = {med:.3g}^{{+{hi - med:.2g}}}" + rf"${labels[i]} = {med:.3g}^{{+{hi - med:.2g}}}" rf"_{{-{med - lo:.2g}}}$", fontsize=7.5, ) @@ -170,11 +191,11 @@ def coefficient_draws(rt, out, markers=(0, 1, 1)): ax.set_xlim(lims[j]) ax.locator_params(nbins=3) if i == n - 1: - ax.set_xlabel(rf"${PARAMS[j]}$") + ax.set_xlabel(rf"${labels[j]}$") else: ax.set_xticklabels([]) if j == 0 and i > 0: - ax.set_ylabel(rf"${PARAMS[i]}$") + ax.set_ylabel(rf"${labels[i]}$") elif j > 0 and i != j: ax.set_yticklabels([]) fig.get_layout_engine().set(hspace=0.02, wspace=0.02) diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py new file mode 100644 index 00000000..4a945633 --- /dev/null +++ b/src/sp_validation/leakage_fit.py @@ -0,0 +1,453 @@ +"""Numerical tools for joint errors-in-variables PSF-leakage fits. + +This module has no catalogue I/O. It builds the rho-to-tau design and replaces its alpha predictor when supplied, +profiles noisy templates from a joint delete-one covariance, and provides +TreeCorr jackknife and patch-grouping helpers. +""" + +import numpy as np +from scipy.cluster.hierarchy import fcluster, linkage +from scipy.linalg import cho_factor, cho_solve +from scipy.optimize import minimize +from scipy.stats import chi2 + + +def rho_tau_design(rho): + """Return the linear rho-to-tau design for rho0..rho5. + + ``rho`` has shape ``(6, nbin)``. The output rows are ordered as tau0, + tau2, tau5, each with ``nbin`` angular bins, and columns are alpha, beta, + eta. This is the field definition used by ``shear_psf_leakage``. + """ + rho = np.asarray(rho, dtype=float) + if rho.ndim != 2 or rho.shape[0] != 6: + raise ValueError("rho must have shape (6, nbin)") + return np.concatenate( + [ + np.stack((rho[0], rho[2], rho[5]), axis=-1), + np.stack((rho[2], rho[1], rho[4]), axis=-1), + np.stack((rho[5], rho[4], rho[3]), axis=-1), + ], + axis=0, + ) + + +def ward_group_map(centres, ngroups): + """Group fine patch centres into ``ngroups`` Ward-linkage unions. + + The centre coordinates are used as supplied (TreeCorr spherical patch + centres are three-dimensional unit vectors). The result maps every fine + patch index to a contiguous coarse group index. + """ + centres = np.asarray(centres, dtype=float) + if centres.ndim != 2: + raise ValueError("centres must be a two-dimensional array of patch centres") + if centres.shape[0] < 3: + raise ValueError( + f"Ward grouping needs at least 3 fine patch centres, got {centres.shape[0]}" + ) + if not 1 < ngroups < len(centres): + raise ValueError("ngroups must be between 2 and the number of fine patches") + groups = ( + fcluster(linkage(centres, method="ward"), ngroups, criterion="maxclust") - 1 + ) + if len(np.unique(groups)) != ngroups: + raise ValueError( + f"Ward grouping produced {len(np.unique(groups))} groups, expected {ngroups}" + ) + return groups.astype(int, copy=False) + + +def treecorr_jackknife_rows(correlations, groups=None, check_rows=True): + """Build exact TreeCorr '+' match-jackknife rows in O(number of pairs). + + When ``groups`` maps fine patch IDs to coarse unions, all fine pairs within + one coarse union are treated as internal pairs and deleted together. + Returns ``(rows, max_checked_error)``: the rows of all correlations side by + side, and the largest difference from TreeCorr's own delete-one result over + the checked fine rows. + """ + blocks = [] + max_error = 0.0 + group_map = None if groups is None else np.asarray(groups, dtype=int) + + for corr in correlations: + fine_n = corr.npatch1 + if fine_n != corr.npatch2 or fine_n < 2: + raise ValueError( + "each correlation must have matching patches and npatch > 1" + ) + if group_map is None: + n_patch = fine_n + else: + if group_map.shape != (fine_n,): + raise ValueError("groups must contain one index per fine patch") + unique = np.unique(group_map) + if not np.array_equal(unique, np.arange(len(unique))): + raise ValueError("group indices must be contiguous and start at zero") + n_patch = len(unique) + if n_patch < 2: + raise ValueError("grouped jackknife requires at least two groups") + + alpha = n_patch / (2 + np.sqrt(2) * (n_patch - 1)) + keys = np.asarray(list(corr.results), dtype=int) + if keys.ndim != 2 or keys.shape[1] != 2: + raise ValueError("TreeCorr correlation has no patch-pair results") + i, j = keys.T + if group_map is not None: + i, j = group_map[i], group_map[j] + pair_values = np.asarray([corr.results[tuple(key)]._xi[0] for key in keys]) + pair_weights = np.asarray([corr.results[tuple(key)].weight for key in keys]) + lost = np.zeros((n_patch, len(corr.xip)), dtype=float) + lost_weight = np.zeros_like(lost) + cross = i != j + diagonal = ~cross + for values, output in ( + (pair_values, lost), + (pair_weights, lost_weight), + ): + np.add.at(output, i[cross], alpha * values[cross]) + np.add.at(output, j[cross], alpha * values[cross]) + np.add.at(output, i[diagonal], values[diagonal]) + + denominator = pair_weights.sum(axis=0)[None, :] - lost_weight + rows = (pair_values.sum(axis=0)[None, :] - lost) / denominator + point = pair_values.sum(axis=0) / pair_weights.sum(axis=0) + if not np.allclose(point, corr.xip, rtol=1e-9, atol=1e-22): + raise AssertionError("pair totals do not reproduce TreeCorr's full result") + + if check_rows and group_map is None: + check = corr.copy() + pair_iter = corr._jackknife_pairs("match") + for patch in sorted({0, n_patch // 2, n_patch - 1}): + check._calculate_xi_from_pairs( + check._keep_ok(pair_iter[patch]), corr_only=True + ) + error = float(np.max(np.abs(check.xip - rows[patch]))) + max_error = max(max_error, error) + if not np.allclose(check.xip, rows[patch], rtol=1e-8, atol=1e-22): + raise AssertionError( + f"fast jackknife row {patch} differs from TreeCorr by {error}" + ) + + blocks.append(rows) + + if not blocks: + raise ValueError("at least one TreeCorr correlation is required") + return np.concatenate(blocks, axis=1), max_error + + +def residual_covariance_model(x_rows, y_rows, cross_term=True): + """The covariance V(p) of ``r = y - x p`` for a noisy design ``x``. + + ``x_rows`` (shape ``(npatch, ndata, nparam)``) and ``y_rows`` (shape + ``(npatch, ndata)``) are delete-one jackknife rows of ``x`` and ``y``. + With the jackknife covariances C_xx, C_xy, C_yy built from them, + V(p) = C_yy - (B C_xy + C_yx B^T) + B C_xx B^T, B the linear map + x -> x p. ``cross_term=False`` zeroes C_xy. Returns two functions of + ``p``: V(p), shape ``(ndata, ndata)``, and dV/dp, shape + ``(nparam, ndata, ndata)``. + """ + n_patch = len(x_rows) + jk_factor = (n_patch - 1) / n_patch + dx = x_rows - x_rows.mean(axis=0) + dy = y_rows - y_rows.mean(axis=0) + cov_yy = jk_factor * dy.T @ dy + cov_xy = jk_factor * np.einsum("nai,nb->iab", dx, dy) + if not cross_term: + cov_xy.fill(0.0) + linear = -cov_xy - cov_xy.transpose(0, 2, 1) + cov_xx = jk_factor * np.einsum("nai,nbj->ijab", dx, dx) + + def covariance(parameters): + return ( + cov_yy + + np.einsum("i,iab->ab", parameters, linear) + + np.einsum("i,j,ijab->ab", parameters, parameters, cov_xx) + ) + + def derivative(parameters): + return ( + linear + + np.einsum("j,ijab->iab", parameters, cov_xx) + + np.einsum("j,jiab->iab", parameters, cov_xx) + ) + + return covariance, derivative + + +def _profile_fit(x, y, x_rows, y_rows, scales, cross_term, start): + """Fit ``y = x p`` with noisy ``x`` by profiling its latent true values. + + Minimises r^T V(p)^-1 r (no log-determinant term), V from + `residual_covariance_model`, by BFGS from several starts. The best trial + is the lowest objective among those BFGS reports converged, or among all + trials if none did (``bfgs_success`` False). + """ + n_patch, n_data, n_param = x_rows.shape + if y_rows.shape != (n_patch, n_data): + raise ValueError("y_rows must have shape (npatch, ndata)") + if x.shape != (n_data, n_param): + raise ValueError("x must have shape (ndata, nparam)") + if scales.shape != (n_param,) or np.any(scales <= 0): + raise ValueError("scales must contain one positive value per parameter") + hartlap = (n_patch - n_data - 2) / (n_patch - 1) + if hartlap <= 0: + raise ValueError("Hartlap factor is non-positive for this patch/data count") + + residual_covariance, covariance_derivative = residual_covariance_model( + x_rows, y_rows, cross_term + ) + + def objective(scaled_parameters): + parameters = scaled_parameters * scales + residual = y - x @ parameters + factor = cho_factor( + residual_covariance(parameters), lower=True, check_finite=False + ) + weighted_residual = cho_solve(factor, residual, check_finite=False) + value = float(residual @ weighted_residual) + gradient = ( + -2 * x.T @ weighted_residual + - np.einsum( + "a,iab,b->i", + weighted_residual, + covariance_derivative(parameters), + weighted_residual, + ) + ) * scales + return value, gradient + + # Generalised least squares with the template noise ignored. + factor = cho_factor(residual_covariance(np.zeros(n_param)), lower=True) + fixed_start = np.linalg.lstsq( + x.T @ cho_solve(factor, x), x.T @ cho_solve(factor, y), rcond=None + )[0] + starts = [fixed_start, np.zeros(n_param)] + if n_param >= 3: + starts.extend( + [ + np.pad([0.02, 1.7, 3.0], (0, n_param - 3)), + np.pad([-0.01, 1.0, -2.0], (0, n_param - 3)), + ] + ) + if start is not None: + start = np.asarray(start, dtype=float) + if start.shape != (n_param,): + raise ValueError("start must contain one value per fit parameter") + starts.append(start) + rng = np.random.default_rng(1028) + starts.extend(fixed_start + rng.normal(size=n_param) * scales for _ in range(3)) + + fitted = [ + minimize( + objective, + initial / scales, + jac=True, + method="BFGS", + options={"gtol": 1e-7, "maxiter": 500}, + ) + for initial in starts + ] + trials = [ + { + "parameters": (fit.x * scales).tolist(), + "chi2_raw": float(fit.fun), + "success": bool(fit.success), + "gradient_norm": float(np.linalg.norm(fit.jac)), + } + for fit in fitted + ] + best = min([fit for fit in fitted if fit.success] or fitted, key=lambda f: f.fun) + parameters = best.x * scales + covariance = residual_covariance(parameters) + residual = y - x @ parameters + chi2_raw = float(residual @ np.linalg.solve(covariance, residual)) + chi2_value = hartlap * chi2_raw + + step = 1e-4 + hessian = np.column_stack( + [ + ( + objective(best.x + np.eye(n_param)[index] * step)[1] + - objective(best.x - np.eye(n_param)[index] * step)[1] + ) + / (2 * step) + for index in range(n_param) + ] + ) + hessian = (hessian + hessian.T) / 2 + try: + np.linalg.cholesky(hessian) + positive = True + except np.linalg.LinAlgError: + positive = False + if positive: + parameter_covariance = ( + 2 * np.linalg.inv(hessian) * scales[:, None] * scales[None, :] / hartlap + ) + sigma = np.sqrt(np.diag(parameter_covariance)).tolist() + parameter_covariance = parameter_covariance.tolist() + else: + parameter_covariance = None + sigma = [None] * n_param + dof = n_data - n_param + residual_sigma = np.sqrt(np.diag(covariance)) + return { + "parameters": parameters.tolist(), + "sigma": sigma, + "parameter_covariance": parameter_covariance, + "chi2": chi2_value, + "chi2_raw": chi2_raw, + "dof": dof, + "pte": float(chi2.sf(chi2_value, dof)), + "hartlap": hartlap, + "residual": residual.tolist(), + "prediction": (x @ parameters).tolist(), + "residual_sigma": residual_sigma.tolist(), + "residual_sigma_normalized": (residual / residual_sigma).tolist(), + "profile_hessian_eigenvalues": np.linalg.eigvalsh(hessian).tolist(), + "profile_hessian_positive": positive, + "bfgs_success": bool(best.success), + "gradient_norm": float(np.linalg.norm(best.jac)), + "trials": trials, + "objective": "joint noisy-template Gaussian profile; no residual log determinant", + "cross_term": bool(cross_term), + } + + +def fit_joint_eiv( + vector, + jk_rows, + *, + nbin=None, + alpha_vector=None, + alpha_rows=None, + template_vector=None, + template_rows=None, + scales=None, + start=None, + cross_term=True, +): + """Fit the noisy-predictor model of tau0/tau2/tau5 with joint covariance. + + ``vector`` is ordered ``[rho0..rho5, tau0, tau2, tau5]`` and ``jk_rows`` + contains the same delete-one measurements, one row per patch. Optional + alpha_vector/alpha_rows replace the alpha predictor and its matched rows, + preserving the joint covariance with all other predictors and tau. Optional + template columns add measured predictors such as ``u * delta_e``; their + point values and rows have shapes ``(ndata, ntemplate)`` and + ``(npatch, ndata, ntemplate)``. + + The objective is the profile chi-squared with the latent true templates + profiled out, minimised by multi-start BFGS; it has no log-determinant + term. The reported chi-squared is scaled by the Hartlap factor + (N - d - 2) / (N - 1), N the number of rows and d the number of tau data; + the parameter covariance is twice the inverse numerical profile Hessian, + divided by the same factor. ``cross_term=False`` zeroes the rho-tau block + of the covariance. + + Returns a dict with the ``parameters``, their ``sigma`` and + ``parameter_covariance`` (None where the profile Hessian is not positive + definite, ``profile_hessian_positive`` False), ``chi2``, ``dof``, + ``pte``, ``hartlap``, the residuals, and the convergence diagnostics + ``bfgs_success`` and ``gradient_norm`` of the chosen trial. + """ + vector = np.asarray(vector, dtype=float) + jk_rows = np.asarray(jk_rows, dtype=float) + if vector.ndim != 1 or jk_rows.ndim != 2: + raise ValueError("vector and jk_rows must be one- and two-dimensional") + if nbin is None: + if vector.size % 9: + raise ValueError("vector length must be divisible by 9") + nbin = vector.size // 9 + nbin = int(nbin) + n_rho = 6 * nbin + n_data = 3 * nbin + if vector.size != 9 * nbin or jk_rows.shape[1] != vector.size: + raise ValueError("vector/jk_rows must contain six rho and three tau blocks") + + x = rho_tau_design(vector[:n_rho].reshape(6, nbin)) + y = vector[n_rho:] + x_rows = np.stack( + [rho_tau_design(row[:n_rho].reshape(6, nbin)) for row in jk_rows], axis=0 + ) + y_rows = jk_rows[:, n_rho:] + if alpha_vector is not None: + alpha_vector = np.asarray(alpha_vector, dtype=float) + alpha_rows = np.asarray(alpha_rows, dtype=float) + if alpha_vector.shape != (n_data,) or alpha_rows.shape != ( + len(jk_rows), + n_data, + ): + raise ValueError("alpha vector/rows must match the tau data and patches") + x[:, 0] = alpha_vector + x_rows[:, :, 0] = alpha_rows + elif alpha_rows is not None: + raise ValueError("alpha_rows requires alpha_vector") + if template_vector is not None: + template_vector = np.asarray(template_vector, dtype=float) + template_rows = np.asarray(template_rows, dtype=float) + if template_vector.ndim == 1: + template_vector = template_vector[:, None] + if template_rows.ndim == 2: + template_rows = template_rows[:, :, None] + if template_vector.shape[0] != n_data: + raise ValueError("template_vector must have one row per tau datum") + if template_rows.shape != ( + len(jk_rows), + n_data, + template_vector.shape[1], + ): + raise ValueError("template_rows must have shape (npatch, ndata, ntemplate)") + x = np.concatenate((x, template_vector), axis=1) + x_rows = np.concatenate((x_rows, template_rows), axis=2) + elif template_rows is not None: + raise ValueError("template_rows requires template_vector") + + n_parameters = x.shape[1] + if scales is None: + scales = [0.01, 1.0, 2.0] + [1.0] * (n_parameters - 3) + scales = np.asarray(scales, dtype=float) + if scales.shape != (n_parameters,): + raise ValueError("scales must contain one value per fit parameter") + return _profile_fit(x, y, x_rows, y_rows, scales, cross_term, start) + + +def systematic_correlation_matrix(rho, slope=None): + """Return the matrix M of field correlations behind xi_sys. + + The PSF-error model is g_sys = alpha p + beta de + eta q (+ beta1 u de), + with p the PSF ellipticity, de the star-minus-PSF ellipticity, q the size + residual field, and u = T_PSF / - 1. ``rho`` has shape + ``(6, nbin)`` (rho0..rho5, one of xi+ or xi-); ``slope``, if given, has + shape ``(4, nbin)``: the correlations of u de with p, de, q and itself. + The result has shape ``(k, k, nbin)``, k = 3 or 4, symmetric in its first + two axes, so that xi_sys = a^T M b for parameter vectors a, b. + """ + rho = np.asarray(rho, dtype=float) + if rho.ndim != 2 or rho.shape[0] != 6: + raise ValueError("rho must have shape (6, nbin)") + pairs = {(0, 0): 0, (1, 1): 1, (0, 1): 2, (2, 2): 3, (1, 2): 4, (0, 2): 5} + size = 3 if slope is None else 4 + matrix = np.zeros((size, size, rho.shape[1])) + for (i, j), index in pairs.items(): + matrix[i, j] = matrix[j, i] = rho[index] + if slope is not None: + slope = np.asarray(slope, dtype=float) + if slope.shape != (4, rho.shape[1]): + raise ValueError("slope must have shape (4, nbin)") + for i in range(4): + matrix[3, i] = matrix[i, 3] = slope[i] + return matrix + + +def xi_psf_sys(parameters_a, parameters_b, matrix): + """Return xi_sys = a^T M b for parameter samples a and b. + + ``parameters_a`` and ``parameters_b`` have shape ``(nsample, k)`` or + ``(k,)`` and ``matrix`` comes from `systematic_correlation_matrix`. + """ + a = np.atleast_2d(parameters_a) + b = np.atleast_2d(parameters_b) + return np.einsum("si,ijn,sj->sn", a, matrix, b) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 2da5ccb3..ae7535b5 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -1,19 +1,23 @@ import gc +import hashlib +import json import os import time from pathlib import Path from tempfile import TemporaryDirectory import numpy as np +import treecorr 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 import grammar, io, leakage_fit from sp_validation.correlation import ( measure_with_patches, measurement_matches, + selection_digest, write_measurement_metadata, ) @@ -22,6 +26,36 @@ from sp_validation.plots import SquareRootScale # noqa: F401 from sp_validation.statistics import jackknife_patch_centers +TAU_FIELD_DEFINITION = "selected-weighted-mean-v1" +ALPHA_TEMPLATE_DEFINITION = "centred-galaxy-psf-v1" + + +def tau_field_definition(params): + """Selected galaxy field, including column names and component responses.""" + return { + "centring": TAU_FIELD_DEFINITION, + **{ + key: params.get(key) + for key in ( + "ra_col", + "dec_col", + "w_col", + "e1_col", + "e2_col", + "R11", + "R22", + "ra_PSF_col", + "dec_PSF_col", + "e1_PSF_col", + "e2_PSF_col", + "e1_star_col", + "e2_star_col", + "PSF_size", + "star_size", + ) + }, + } + def _extract_xip(correlations): """Return flattened array of xip values from a list of correlations.""" @@ -29,7 +63,7 @@ def _extract_xip(correlations): def _compute_stats_with_patches( - handler, kind, catalog_id, filename, save_cov=False, means=None + handler, kind, catalog_id, filename, save_cov=False, means=None, selection=None ): """Publish unpatched ρ/τ means with the patched variances and covariance. @@ -76,7 +110,14 @@ def measure(cats, measurement_config): 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) + write_measurement_metadata( + Path(handler.catalogs._output) / filename, + config, + selection=selection, + definition=tau_field_definition(handler.catalogs._params) + if kind == "tau" + else None, + ) return means @@ -122,7 +163,15 @@ def _columns(self, block): "star_flag", ) elif block == "shear": - keys = ("ra_col", "dec_col", "w_col", "e1_col", "e2_col") + keys = ( + "ra_col", + "dec_col", + "w_col", + "e1_col", + "e2_col", + "e1_gal_PSF_col", + "e2_gal_PSF_col", + ) else: raise KeyError(f"unknown rho/tau catalogue block {block!r}") @@ -201,6 +250,8 @@ def get_params_rho_tau(cat): params["w_col"] = cat["shear"]["w_col"] params["e1_col"] = cat["shear"]["e1_col"] params["e2_col"] = cat["shear"]["e2_col"] + params["e1_gal_PSF_col"] = cat["shear"].get("e1_PSF_col") + params["e2_gal_PSF_col"] = cat["shear"].get("e2_PSF_col") params["tomo_bin_col"] = cat["shear"].get("tomo_bin_col") params["R11"] = cat["shear"].get("R11") params["R22"] = cat["shear"].get("R22") @@ -345,7 +396,11 @@ def get_rho_tau( ) with _CatalogueLoader(config[version], params) as load: - rho_stats_exists = measurement_matches(rho_path, treecorr_config) + rho_selection = selection_digest(mask_star) + tau_selection = selection_digest(mask_star, mask_gal) + rho_stats_exists = measurement_matches( + rho_path, treecorr_config, selection=rho_selection + ) 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 @@ -356,7 +411,12 @@ def get_rho_tau( ) _compute_stats_with_patches( - rho_stat_handler, "rho", catalog_id, rho_path.name, save_cov=cov_rho + rho_stat_handler, + "rho", + catalog_id, + rho_path.name, + save_cov=cov_rho, + selection=rho_selection, ) rho_stat_handler.load_rho_stats(rho_path.name) else: @@ -373,7 +433,15 @@ def get_rho_tau( verbose=True, ) - if measurement_matches(tau_path, treecorr_config) and not force_run: + if ( + measurement_matches( + tau_path, + treecorr_config, + selection=tau_selection, + definition=tau_field_definition(params), + ) + and not force_run + ): print( f"Skipping tau statistics computation, file {tau_path} already exists." ) @@ -392,12 +460,333 @@ def get_rho_tau( load("shear"), cat_type="gal", catalog_id=version, mask=mask_gal ) - _compute_stats_with_patches(tau_stat_handler, "tau", version, tau_path.name) + _compute_stats_with_patches( + tau_stat_handler, + "tau", + version, + tau_path.name, + selection=tau_selection, + ) print(f"Time to compute rho and tau statistics: {time.time() - start_time:.2f} s") return rho_stat_handler, tau_stat_handler +# The correlations of the joint leakage measurement, as (first, second) +# catalogue types of shear_psf_leakage's Catalogs: p = "psf", +# de = "psf_error", q = "psf_size_error", galaxies = "gal", and the PSF-size +# slope field u de = "psf_error_size_slope". +LEAKAGE_RHO_PAIRS = ( + ("psf", "psf"), + ("psf_error", "psf_error"), + ("psf_error", "psf"), + ("psf_size_error", "psf_size_error"), + ("psf_error", "psf_size_error"), + ("psf", "psf_size_error"), +) +LEAKAGE_TAU_PAIRS = (("gal", "psf"), ("gal", "psf_error"), ("gal", "psf_size_error")) +LEAKAGE_SLOPE_PAIRS = ( + ("psf_error_size_slope", "psf"), + ("psf_error_size_slope", "psf_error"), + ("psf_error_size_slope", "psf_size_error"), +) +LEAKAGE_SLOPE_AUTO = ("psf_error_size_slope", "psf_error_size_slope") + + +def leakage_cache_key( + treecorr_config, + npatch_fine, + patch_counts, + beta_psf_size, + mask_star=None, + mask_gal=None, + *, + params, +): + """The settings and row selections that identify a joint leakage measurement. + + ``mask_star``/``mask_gal`` are the star and galaxy selections the + measurement uses, recorded as `selection_digest`s. TreeCorr settings the + measurement overrides (``var_method``, ``cross_patch_weight``) or that + only change round-off (``num_threads``) are left out. + """ + config = { + key: value + for key, value in treecorr_config.items() + if key not in ("num_threads", "var_method", "cross_patch_weight") + } + return { + "field_definition": tau_field_definition(params), + "alpha_template": { + "mode": ALPHA_TEMPLATE_DEFINITION, + "e1": params.get("e1_gal_PSF_col"), + "e2": params.get("e2_gal_PSF_col"), + }, + "treecorr_config": config, + "npatch_fine": int(npatch_fine), + "patch_counts": sorted(int(n) for n in patch_counts), + "beta_psf_size": bool(beta_psf_size), + "selection": { + "stars": selection_digest(mask_star), + "stars_galaxies": selection_digest(mask_star, mask_gal), + }, + } + + +def read_leakage_jackknife(path, key): + """Products of a cached joint leakage measurement made with ``key``, or None.""" + path = Path(path) + metadata_path = path.with_suffix(".json") + if not (path.exists() and metadata_path.exists()): + return None + if json.loads(metadata_path.read_text()).get("key") != key: + return None + with np.load(path) as data: + return dict(data) + + +def measure_leakage_jackknife( + config, + version, + treecorr_config, + path, + key, + mask_star=None, + mask_gal=None, +): + """Measure the rho/tau vector and its joint delete-one jackknife rows. + + ``key`` (from `leakage_cache_key` with these masks) holds the settings: + TreeCorr configuration, ``npatch_fine``, ``patch_counts``, + ``beta_psf_size``, and the selections of ``mask_star``/``mask_gal``. + Stars and galaxies are built into TreeCorr catalogues exactly as + `get_rho_tau` builds them (``shear_psf_leakage`` field definitions; + galaxy ellipticities are centred on the selected weighted mean), sharing one + set of ``npatch_fine`` k-means patch centres found on the selected stars. One + pass of GG correlations, with 'match' cross-patch weights, measures + rho0..rho5 and tau0, tau2, tau5 and, with ``beta_psf_size``, the + correlations of u de (u = T_PSF / - 1 over the selected stars) + with p, de, q and itself. Delete-one jackknife rows come from the + per-patch-pair results at ``npatch_fine`` and, for every count in + ``patch_counts``, at Ward-linkage unions of the fine patches. + + Writes the NPZ products and JSON metadata. + Alpha point values and matched rows are saved separately as alpha_vector + and alpha_rows_; they replace the alpha design column only. + Returns the products as a dict: ``theta``, ``vector`` (rho0..rho5, tau0, tau2, tau5 + xi+, flattened), ``rho_xim``, ``rows_`` (one row per patch, the + vector's order followed, with the slope, by the three u de templates), + ``centres``, ``groups_``, and with the slope ``template_vector``, + ``slope_xip`` and ``slope_xim`` (u de with p, de, q, u de). + """ + path = Path(path) + metadata_path = path.with_suffix(".json") + npatch_fine = key["npatch_fine"] + patch_counts = key["patch_counts"] + beta_psf_size = key["beta_psf_size"] + start_time = time.time() + params = get_params_rho_tau(config[version]) + params["patch_number"] = npatch_fine + measurement_config = { + **treecorr_config, + "var_method": "shot", + "cross_patch_weight": "match", + } + rho_handler = RhoStat(output=str(path.parent), treecorr_config=measurement_config) + tau_handler = TauStat( + catalogs=rho_handler.catalogs, + output=str(path.parent), + treecorr_config=measurement_config, + ) + catalogs = rho_handler.catalogs + catalogs.set_params(params, str(path.parent)) + with _CatalogueLoader(config[version], params) as load: + rho_handler.build_cat_to_compute_rho( + load("psf"), catalog_id="leakage", mask=mask_star + ) + tau_handler.build_cat_to_compute_tau( + load("shear"), cat_type="gal", catalog_id="leakage", mask=mask_gal + ) + tau_handler.build_cat_to_compute_tau( + load("shear"), cat_type="gal_psf", catalog_id="leakage", mask=mask_gal + ) + cats = { + name.removesuffix("_leakage"): cat + for name, cat in catalogs.catalogs_dict.items() + } + centres = cats["psf"].patch_centers + mean_psf_size = None + if beta_psf_size: + stars = catalogs.read_shear_cat(path_gal=None, path_psf=load("psf")) + size = np.asarray(stars[params["PSF_size"]], dtype=float) + if mask_star is not None: + size = size[mask_star] + del stars + mean_psf_size = float(size.mean()) + u = size / mean_psf_size - 1 + de = cats["psf_error"] + if u.size != de.ntot: + raise ValueError( + f"PSF-size template has {u.size} rows but the de catalogue " + f"{de.ntot}; the two must be the same selected stars" + ) + cats["psf_error_size_slope"] = treecorr.Catalog( + ra=de.ra, + dec=de.dec, + ra_units="rad", + dec_units="rad", + g1=u * de.g1, + g2=u * de.g2, + w=de.w, + patch_centers=centres, + ) + del size, u + + groups = {n: leakage_fit.ward_group_map(centres, n) for n in patch_counts} + counts = [npatch_fine, *patch_counts] + pairs = list(LEAKAGE_RHO_PAIRS + LEAKAGE_TAU_PAIRS) + if beta_psf_size: + pairs += [*LEAKAGE_SLOPE_PAIRS, LEAKAGE_SLOPE_AUTO] + xip, xim, rows = [], [], {n: [] for n in counts} + theta = None + max_check_error = 0.0 + for first, second in pairs: + print(f"Joint leakage correlation <{first}, {second}>", flush=True) + corr = treecorr.GGCorrelation(measurement_config) + corr.process(cats[first], cats[second]) + xip.append(corr.xip.copy()) + xim.append(corr.xim.copy()) + if (first, second) == LEAKAGE_TAU_PAIRS[0]: + theta = corr.meanr.copy() + if (first, second) != LEAKAGE_SLOPE_AUTO: + fine, error = leakage_fit.treecorr_jackknife_rows([corr]) + max_check_error = max(max_check_error, error) + rows[npatch_fine].append(fine) + for n in patch_counts: + grouped, _ = leakage_fit.treecorr_jackknife_rows( + [corr], groups=groups[n], check_rows=False + ) + rows[n].append(grouped) + del corr + gc.collect() + + n_joint = len(LEAKAGE_RHO_PAIRS) + len(LEAKAGE_TAU_PAIRS) + products = { + "theta": theta, + "vector": np.concatenate(xip[:n_joint]), + "rho_xim": np.array(xim[: len(LEAKAGE_RHO_PAIRS)]), + "centres": centres, + **{f"rows_{n}": np.concatenate(rows[n], axis=1) for n in counts}, + **{f"groups_{n}": groups[n] for n in patch_counts}, + } + if beta_psf_size: + products["template_vector"] = np.concatenate(xip[n_joint : n_joint + 3]) + products["slope_xip"] = np.array(xip[n_joint:]) + products["slope_xim"] = np.array(xim[n_joint:]) + + # Release the shear tree before building the galaxy-PSF tree. + cats["gal"].clear_cache() + alpha_correlations = tau_handler.compute_alpha_template("leakage") + products["alpha_vector"] = np.concatenate([c.xip for c in alpha_correlations]) + for n in counts: + alpha_rows, error = leakage_fit.treecorr_jackknife_rows( + alpha_correlations, groups=None if n == npatch_fine else groups[n] + ) + products[f"alpha_rows_{n}"] = alpha_rows + max_check_error = max(max_check_error, error) + del alpha_correlations + + path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(path, **products) + metadata = { + "key": key, + "version": version, + "stars": int(cats["psf"].nobj), + "galaxies": int(cats["gal"].nobj), + "mean_psf_size": mean_psf_size, + "order": [f"<{a}, {b}>" for a, b in pairs], + "jackknife_row_check_max_abs_error": max_check_error, + "treecorr_version": treecorr.__version__, + "seconds": time.time() - start_time, + } + metadata_path.write_text(json.dumps(metadata, indent=2) + "\n") + print(f"Joint leakage measurement written to {path} in {metadata['seconds']:.0f} s") + return products + + +def fit_leakage(products, npatch, beta_psf_size=False, cross_term=True): + """Joint EIV fit of `measure_leakage_jackknife` products at ``npatch``. + + Returns the `leakage_fit.fit_joint_eiv` result dict, whose + ``parameters`` are (alpha, beta, eta), or (alpha, beta0, eta, beta1) + with the PSF-size slope. + """ + rows = np.asarray(products[f"rows_{npatch}"]) + n_joint = len(products["vector"]) + template = {} + if beta_psf_size: + template = dict( + template_vector=products["template_vector"], + template_rows=rows[:, n_joint:], + ) + return leakage_fit.fit_joint_eiv( + products["vector"], + rows[:, :n_joint], + alpha_vector=products["alpha_vector"], + alpha_rows=products[f"alpha_rows_{npatch}"], + cross_term=cross_term, + **template, + ) + + +def get_samples_eiv( + products, patch_counts, fiducial_npatch, beta_psf_size=False, nsamples=10000, seed=0 +): + """Joint EIV fits at each patch count and Gaussian samples at the fiducial. + + Returns ``(samples, fits, quantiles)``: ``nsamples`` draws (generator + ``seed``) from N(mu, Sigma_p) of the ``fiducial_npatch`` fit, the fit + results keyed by patch count, and the 16/50/84 per cent sample quantiles + per parameter. Raises if the fiducial fit's profile Hessian is not + positive definite, since it then has no parameter covariance. + """ + fits = { + int(n): fit_leakage(products, n, beta_psf_size=beta_psf_size) + for n in patch_counts + } + for n, result in fits.items(): + errors = ", ".join( + f"{mu:.4g} +/- {sigma:.2g}" if sigma is not None else f"{mu:.4g} +/- ?" + for mu, sigma in zip(result["parameters"], result["sigma"]) + ) + flags = [ + message + for ok, message in ( + (result["bfgs_success"], "BFGS did not converge"), + (result["profile_hessian_positive"], "Hessian not positive definite"), + ) + if not ok + ] + print( + f"EIV leakage fit, {n} patches: chi2 = {result['chi2']:.2f}/" + f"{result['dof']} (PTE {result['pte']:.3g}), parameters {errors}" + + (f" [{'; '.join(flags)}]" if flags else "") + ) + fiducial = fits[int(fiducial_npatch)] + if not fiducial["profile_hessian_positive"]: + raise ValueError( + f"The {fiducial_npatch}-patch EIV leakage fit has a profile Hessian " + "that is not positive definite (eigenvalues " + f"{fiducial['profile_hessian_eigenvalues']}); its parameters are not " + "constrained, so no samples can be drawn." + ) + samples = np.random.default_rng(seed).multivariate_normal( + fiducial["parameters"], fiducial["parameter_covariance"], size=nsamples + ) + quantiles = np.percentile(samples, [16, 50, 84], axis=0) + return samples, fits, quantiles + + def get_theory_cov( config, version, @@ -415,7 +804,8 @@ def get_theory_cov( ``CovTauTh`` takes the survey area and effective galaxy density from the galaxy catalogue after ``mask_gal``. The masks select rows as in - ``get_rho_tau``. + ``get_rho_tau``; their `selection_digest` is recorded beside the cached + covariance, which is reused only for the same selection. """ params = get_params_rho_tau(config[version]) @@ -423,8 +813,14 @@ def get_theory_cov( info = config[version] target_cov = Path(outdir) / f"cov_tau_{base}_th.npy" + metadata = Path(str(target_cov) + ".json") + selection = selection_digest(mask_star, mask_gal) - if target_cov.exists(): + if ( + target_cov.exists() + and metadata.exists() + and json.loads(metadata.read_text()).get("selection") == selection + ): print(f"Skipping covariance computation, file {target_cov} already exists.") return @@ -451,23 +847,26 @@ def get_theory_cov( print(f"--- Covariance matrix assembled in {time.time() - start_time:.2f}s ---") target_cov.parent.mkdir(parents=True, exist_ok=True) np.save(target_cov, cov) + metadata.write_text(json.dumps({"selection": selection}) + "\n") print("Saved covariance matrix of version: ", version) del cov_tau_th gc.collect() return -def rho_cache_valid(rho_path, cov_rho_path, treecorr_config, force_run=False): +def rho_cache_valid( + rho_path, cov_rho_path, treecorr_config, force_run=False, selection=None +): """Whether a cached ρ measurement and its jackknife covariance can be reused. - The ρ FITS must record the current means and patched configurations - (:func:`~sp_validation.correlation.measurement_matches`); the PSF-only + The ρ FITS must record the current means, patched configurations and star + selection (:func:`~sp_validation.correlation.measurement_matches`); the PSF-only result is then shared by every source bin of a tomographic run. """ return ( not force_run and Path(cov_rho_path).exists() - and measurement_matches(rho_path, treecorr_config) + and measurement_matches(rho_path, treecorr_config, selection=selection) ) @@ -491,12 +890,21 @@ def get_jackknife_cov( rho_filename = f"rho_stats_{base_rho}.fits" tau_filename = f"tau_stats_{base_tau}.fits" tau_cov_path = Path(outdir) / f"cov_tau_{base_tau}_jk.npy" + rho_selection = selection_digest(mask_star) + tau_selection = selection_digest(mask_star, mask_gal) 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 measurement_matches( + Path(outdir) / rho_filename, treecorr_config, selection=rho_selection + ) + and measurement_matches( + Path(outdir) / tau_filename, + treecorr_config, + selection=tau_selection, + definition=tau_field_definition(get_params_rho_tau(config[version])), + ) and not force_run ): print(f"Skipping covariance computation, file {tau_cov_path} already exists.") @@ -539,7 +947,9 @@ def get_jackknife_cov( rho_path = Path(outdir) / rho_filename cov_rho_path = Path(outdir) / f"cov_rho_{base_rho}_jk.npy" - reuse_rho = rho_cache_valid(rho_path, cov_rho_path, treecorr_config, force_run) + reuse_rho = rho_cache_valid( + rho_path, cov_rho_path, treecorr_config, force_run, selection=rho_selection + ) if reuse_rho: rho_stat_handler.load_rho_stats(rho_filename) cached_rho_cov = np.load(cov_rho_path) @@ -613,6 +1023,7 @@ def chunks(i): rho_filename, save_cov=True, means=rho_means, + selection=rho_selection, ) tau_means = _compute_stats_with_patches( tau_stat_handler, @@ -621,6 +1032,7 @@ def chunks(i): tau_filename, save_cov=True, means=tau_means, + selection=tau_selection, ) # Update the keys in the dictionaries @@ -655,6 +1067,29 @@ def chunks(i): return rho_stat_handler, tau_stat_handler +def _sample_cache_name(psf_fitter, base_tau, sampler, cov_type): + """Name of a sample cache, keyed on the tau statistics' row selection. + + The selection is the `selection_digest` recorded beside + ``tau_stats_.fits`` (stars and galaxies), so samples of a + different selection are never reused. + """ + name = f"{base_tau}_sampler_{sampler}_cov_tau_type_{cov_type}" + metadata = ( + Path(psf_fitter.tau_stat_handler.catalogs._output) + / f"tau_stats_{base_tau}.fits.json" + ) + if not metadata.exists(): + return name + record = json.loads(metadata.read_text()) + selection = record.get("selection") + definition = hashlib.sha256( + json.dumps(record.get("definition"), sort_keys=True).encode() + ).hexdigest()[:12] + name = f"{name}_field={definition}" + return name if selection is None else f"{name}_selection={selection[:12]}" + + def get_samples( psf_fitter, base_rho, @@ -743,7 +1178,7 @@ def get_samples_emcee( psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") # Check if the path exists (use base for cache key to account for different TreeCorr configs) - base_sample = f"{base_tau}_sampler_emcee_cov_tau_type_{cov_type}" + base_sample = _sample_cache_name(psf_fitter, base_tau, "emcee", cov_type) sample_file_path = psf_fitter.get_sample_path(base_sample) if os.path.exists(sample_file_path): print(f"Skipping sampling; {sample_file_path} exists.") @@ -795,7 +1230,7 @@ def get_samples_lsq( psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") - base_sample = f"{base_tau}_sampler_lsq_cov_tau_type_{cov_type}" + base_sample = _sample_cache_name(psf_fitter, base_tau, "lsq", cov_type) # Check if the path exists (use base for cache key to account for different TreeCorr configs) sample_file_path = psf_fitter.get_sample_path(base_sample) if os.path.exists(sample_file_path): diff --git a/src/sp_validation/tests/regression/test_lsq_sample_cache_ignores_cov_type.py b/src/sp_validation/tests/regression/test_lsq_sample_cache_ignores_cov_type.py index 88e778df..22e10834 100644 --- a/src/sp_validation/tests/regression/test_lsq_sample_cache_ignores_cov_type.py +++ b/src/sp_validation/tests/regression/test_lsq_sample_cache_ignores_cov_type.py @@ -78,3 +78,36 @@ def small_sample(**kwargs): f"jk after th returned posterior std {std_jk}; expected ~{SIGMA_JK}, " f"not cached th width ~{SIGMA_TH}" ) + + +def test_lsq_samples_follow_the_recorded_selection( + tmp_path, monkeypatch, restore_random_state +): + """Samples cached for one row selection are not reused for another. + + The tau statistics record their star/galaxy selection beside the FITS + file. Remeasuring them on other rows (here tau shifts from 0 to 1) must + give samples of the new tau, not the cached draws of the old selection. + """ + _write_fixture(tmp_path) + metadata = tmp_path / "tau_stats_fit.fits.json" + metadata.write_text('{"selection": "aaaaaaaaaaaaaaaa"}\n') + rho = RhoStat(output=str(tmp_path), treecorr_config=CFG) + tau = TauStat(output=str(tmp_path), treecorr_config=CFG) + fitter = PSFErrorFit(rho, tau, str(tmp_path)) + original = fitter.get_least_squares_params_samples + + def small_sample(**kwargs): + kwargs.update(n_samples=256, verbose=False) + return original(**kwargs) + + monkeypatch.setattr(fitter, "get_least_squares_params_samples", small_sample) + before, _, _ = _run(fitter, "th") + assert np.allclose(before.mean(axis=0), 0, atol=1e-3) + + Table({"theta": [10.0], **{f"tau_{i}_p": [1.0] for i in (0, 2, 5)}}).write( + tmp_path / "tau_stats_fit.fits", overwrite=True + ) + metadata.write_text('{"selection": "bbbbbbbbbbbbbbbb"}\n') + after, _, _ = _run(fitter, "th") + assert np.allclose(after.mean(axis=0), 1, atol=1e-2), after.mean(axis=0) diff --git a/src/sp_validation/tests/regression/test_tomo_tau_centering_before_bin_mask.py b/src/sp_validation/tests/regression/test_tomo_tau_centering_before_bin_mask.py index b789438c..25350e8d 100644 --- a/src/sp_validation/tests/regression/test_tomo_tau_centering_before_bin_mask.py +++ b/src/sp_validation/tests/regression/test_tomo_tau_centering_before_bin_mask.py @@ -84,11 +84,6 @@ def _stars(n=2000, seed=1): return st -@pytest.mark.xfail( - strict=True, - raises=AssertionError, - reason="CosmoStat/shear_psf_leakage#52: centring precedes bin masking", -) def test_masked_galaxy_catalog_is_centred_on_bin_mean(): """A bin's TreeCorr galaxy catalogue must have zero weighted mean e1. @@ -107,11 +102,6 @@ def test_masked_galaxy_catalog_is_centred_on_bin_mean(): ) -@pytest.mark.xfail( - strict=True, - raises=AssertionError, - reason="CosmoStat/shear_psf_leakage#52: full-catalogue mean biases bin tau", -) def test_tomographic_tau0_plus_vanishes_for_bin_constant_ellipticity(tmp_path): """TauStat must give tau_0+ = 0 for a bin with constant galaxy e1. diff --git a/src/sp_validation/tests/test_correlation.py b/src/sp_validation/tests/test_correlation.py index 070490d7..98e56a94 100644 --- a/src/sp_validation/tests/test_correlation.py +++ b/src/sp_validation/tests/test_correlation.py @@ -9,6 +9,7 @@ measure_with_patches, measurement_matches, process_gg, + selection_digest, write_measurement_metadata, ) @@ -116,6 +117,21 @@ def test_measurement_metadata_records_both_tolerance_settings(tmp_path): assert not measurement_matches(path, config) +def test_measurement_cache_is_keyed_by_row_selection(tmp_path): + path = tmp_path / "measurement.txt" + path.write_text("cached") + config = dict(min_sep=1, max_sep=10, nbins=4, bin_slop=0.1) + unmasked = np.ones(10, dtype=bool) + masked = unmasked.copy() + masked[3] = False + + write_measurement_metadata(path, config, selection=selection_digest(unmasked)) + assert measurement_matches(path, config, selection=selection_digest(unmasked)) + assert not measurement_matches(path, config, selection=selection_digest(masked)) + assert not measurement_matches(path, config) + assert selection_digest(None) != selection_digest(unmasked) + + def test_unpatched_gg_matches_direct_measurement(): rng = np.random.default_rng(1) cat = treecorr.Catalog( diff --git a/src/sp_validation/tests/test_dashboard_psf.py b/src/sp_validation/tests/test_dashboard_psf.py index 50912349..76ccbf38 100644 --- a/src/sp_validation/tests/test_dashboard_psf.py +++ b/src/sp_validation/tests/test_dashboard_psf.py @@ -117,3 +117,33 @@ def test_corrected_tau_loader_needs_no_assembled_or_fit(tmp_path, monkeypatch): monkeypatch.setattr(P, "declared_blind", lambda c, v: "toy") monkeypatch.setattr(P, "load_tau", lambda s, c, b: (s, c, b)) assert P.load_corrected_tau(tmp_path, "V_leak_corr", "config") == (part, cov, "toy") + + +def test_slope_samples_use_the_tau_slope_template(rt, tmp_path, monkeypatch): + """(α, β₀, η, β₁) draws add β₁ ⟨u δe, ·⟩ to τ; without the template, raise.""" + samples = np.array([[0.1, 1.0, 2.0, 3.0]]) + template = np.arange(9.0).reshape(3, 3) * 1e-6 + with pytest.raises(ValueError, match="tau_slope_template"): + psf.tau_model_samples(rt, 2, samples) + rt["fit"]["tau_slope_template"] = template + rho = 1e-5 * (0.1 + 1.0 + 2.0) + for k, row in ((0, 0), (2, 1), (5, 2)): + np.testing.assert_allclose( + psf.tau_model_samples(rt, k, samples)[0], rho + 3.0 * template[row] + ) + rt["fit"]["flat_samples"] = np.random.default_rng(1).normal(size=(64, 4)) + saved = [] + monkeypatch.setattr(psf.S, "save", lambda fig, *args: saved.append(fig)) + psf.coefficient_draws(rt, tmp_path / "draws.png") + assert "$\\beta_1$" in [ax.get_xlabel() for ax in saved[0].axes] + plt.close(saved[0]) + + +def test_contour_axes_follow_the_parameter_count(): + from sp_validation.cosmo_val.diagnostic_plots import leakage_contour_axes + + names, labels, markers = leakage_contour_axes(4) + assert names == ["x0", "x1", "x2", "x3"] + assert labels[3] == r"\beta_1" + assert markers == {"x0": 0, "x1": 1, "x2": 1, "x3": 0} + assert leakage_contour_axes(3)[1] == [r"\alpha", r"\beta", r"\eta"] diff --git a/src/sp_validation/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py new file mode 100644 index 00000000..bfb6b3fb --- /dev/null +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -0,0 +1,416 @@ +"""Synthetic and TreeCorr checks for the joint PSF-leakage estimator.""" + +import numpy as np +import pytest +import treecorr +from scipy.optimize import minimize +from shear_psf_leakage.rho_tau_stat import PSFErrorFit, RhoStat, TauStat + +from sp_validation.leakage_fit import ( + fit_joint_eiv, + residual_covariance_model, + rho_tau_design, + systematic_correlation_matrix, + treecorr_jackknife_rows, + ward_group_map, + xi_psf_sys, +) +from sp_validation.rho_tau import get_samples_eiv + +TRUE_PARAMETERS = np.array([0.013, 1.28, 0.96, 1.9]) + + +def _synthetic_joint(seed, nbin=5, npatch=200, slope=False, noise=True): + """Point vector and delete-one rows for tau = B(rho) p (+ beta1 * template). + + Each row perturbs rho (and the template) and propagates the perturbation + into tau, plus independent tau noise, so the rows carry a real rho-tau + cross-covariance. Delete-one rows scatter by sigma / sqrt(N - 1) for a + point-vector scatter sigma, so the point vector is perturbed by the + same draws scaled up by sqrt(N - 1). With ``noise=False`` the point vector + is exact. + """ + rng = np.random.default_rng(seed) + parameters = TRUE_PARAMETERS[: 4 if slope else 3] + rho = rng.normal(size=(6, nbin)) + template = rng.normal(size=3 * nbin) if slope else None + + def draw(scale): + d_rho = rng.normal(scale=0.08 * scale, size=6 * nbin) + d_template = rng.normal(scale=0.05 * scale, size=3 * nbin) + d_tau = rho_tau_design(d_rho.reshape(6, nbin)) @ parameters[:3] + d_tau += rng.normal(scale=0.04 * scale, size=3 * nbin) + if slope: + d_tau += parameters[3] * d_template + return d_rho, d_template, d_tau + + def tau(rho_values, template_values): + value = rho_tau_design(rho_values.reshape(6, nbin)) @ parameters[:3] + return value + (parameters[3] * template_values if slope else 0) + + rows, template_rows = [], [] + for _ in range(npatch): + d_rho, d_template, d_tau = draw(1 / np.sqrt(npatch - 1)) + row_rho = rho.ravel() + d_rho + row_template = (template + d_template) if slope else None + rows.append(np.concatenate((row_rho, tau(rho.ravel(), template) + d_tau))) + template_rows.append(row_template) + d_rho, d_template, d_tau = draw(1.0 if noise else 0.0) + vector = np.concatenate((rho.ravel() + d_rho, tau(rho.ravel(), template) + d_tau)) + template_vector = (template + d_template) if slope else None + return ( + vector, + np.array(rows), + template_vector, + np.array(template_rows) if slope else None, + parameters, + ) + + +@pytest.mark.parametrize("slope", [False, True]) +def test_joint_eiv_is_exact_without_noise(slope): + """An exact point vector gives the true parameters and chi-squared zero.""" + vector, rows, template, template_rows, expected = _synthetic_joint( + seed=204, slope=slope, noise=False + ) + result = fit_joint_eiv( + vector, rows, template_vector=template, template_rows=template_rows + ) + np.testing.assert_allclose(result["parameters"], expected, rtol=1e-7, atol=1e-9) + assert result["chi2"] < 1e-12 + assert result["dof"] == 15 - len(expected) + assert result["hartlap"] == pytest.approx((200 - 15 - 2) / 199) + assert np.all(np.asarray(result["sigma"]) > 0) + + +@pytest.mark.parametrize("slope", [False, True]) +def test_joint_eiv_recovers_parameters_within_errors(slope): + """With noise, the fit lands within its errors and chi-squared is sane.""" + pulls, chi2s = [], [] + for seed in range(5): + vector, rows, template, template_rows, expected = _synthetic_joint( + seed=seed, slope=slope + ) + result = fit_joint_eiv( + vector, rows, template_vector=template, template_rows=template_rows + ) + pulls.append((np.array(result["parameters"]) - expected) / result["sigma"]) + chi2s.append(result["chi2"]) + assert result["profile_hessian_positive"] + pulls = np.array(pulls) + assert np.all(np.abs(pulls) < 4), pulls + # Five draws of a chi-squared with 11-12 degrees of freedom. + assert 3 < np.mean(chi2s) < 25, chi2s + + +def test_joint_eiv_cross_covariance_changes_the_fit(): + """The rho-tau block matters when the delete-one rows correlate.""" + vector, rows, _, _, _ = _synthetic_joint(seed=41) + full = fit_joint_eiv(vector, rows, cross_term=True) + no_cross = fit_joint_eiv(vector, rows, cross_term=False) + # Without the cross term the propagated rho noise is counted twice in V. + assert full["chi2"] > 1.3 * no_cross["chi2"] + assert full["chi2"] - no_cross["chi2"] > 2.0 + + +def test_treecorr_rows_match_native_match_jackknife_and_grouped_catalogues(): + """Fine and unioned rows agree with TreeCorr's native design matrix.""" + rng = np.random.default_rng(1234) + n_fine, n_coarse = 12, 4 + centres = rng.normal(size=(n_fine, 2)) + groups = ward_group_map(centres, n_coarse) + config = dict( + min_sep=0.01, + max_sep=4, + nbins=10, + bin_slop=0.0, + angle_slop=0.0, + var_method="shot", + num_threads=1, + cross_patch_weight="match", + ) + + def catalogue(size): + return treecorr.Catalog( + x=rng.normal(size=size), + y=rng.normal(size=size), + g1=rng.normal(size=size), + g2=rng.normal(size=size), + w=rng.uniform(0.2, 3, size=size), + patch_centers=centres, + ) + + def coarse(catalogue): + return treecorr.Catalog( + x=catalogue.x, + y=catalogue.y, + g1=catalogue.g1, + g2=catalogue.g2, + w=catalogue.w, + patch=groups[catalogue.patch], + npatch=n_coarse, + ) + + first, second = catalogue(1000), catalogue(800) + for left, right in ((first, second), (first, first)): + fine_corr = treecorr.GGCorrelation(config) + fine_corr.process(left, right) + rows, _ = treecorr_jackknife_rows([fine_corr]) + native, _ = treecorr.build_multi_cov_design_matrix( + [fine_corr], + "jackknife", + func=lambda correlations: correlations[0].xip, + cross_patch_weight="match", + ) + np.testing.assert_allclose(rows, native, rtol=1e-11, atol=1e-14) + centered = rows - rows.mean(axis=0) + covariance = (len(rows) - 1) / len(rows) * centered.T @ centered + np.testing.assert_allclose( + covariance, + treecorr.estimate_multi_cov( + [fine_corr], + "jackknife", + func=lambda correlations: correlations[0].xip, + cross_patch_weight="match", + ), + rtol=1e-10, + atol=1e-14, + ) + + coarse_corr = treecorr.GGCorrelation(config) + coarse_corr.process(coarse(left), coarse(right)) + grouped, _ = treecorr_jackknife_rows( + [fine_corr], groups=groups, check_rows=False + ) + native_grouped, _ = treecorr.build_multi_cov_design_matrix( + [coarse_corr], + "jackknife", + func=lambda correlations: correlations[0].xip, + cross_patch_weight="match", + ) + np.testing.assert_allclose( + fine_corr.xip, coarse_corr.xip, rtol=1e-11, atol=1e-14 + ) + np.testing.assert_allclose(grouped, native_grouped, rtol=1e-11, atol=1e-14) + + +def test_xi_psf_sys_matches_term_by_term_expansion(): + """a^T M b reproduces the production rho-term expansion, plus the slope.""" + rng = np.random.default_rng(3) + rho = rng.normal(size=(6, 4)) + slope = rng.normal(size=(4, 4)) + a, b = rng.normal(size=4), rng.normal(size=4) + (al_a, be_a, et_a, b1_a), (al_b, be_b, et_b, b1_b) = a, b + expected = ( + al_a * al_b * rho[0] + + be_a * be_b * rho[1] + + (al_a * be_b + be_a * al_b) * rho[2] + + et_a * et_b * rho[3] + + (be_a * et_b + et_a * be_b) * rho[4] + + (al_a * et_b + et_a * al_b) * rho[5] + ) + constant = xi_psf_sys(a[:3], b[:3], systematic_correlation_matrix(rho)) + np.testing.assert_allclose(constant[0], expected, rtol=1e-12) + + expected += ( + b1_a * b1_b * slope[3] + + (al_a * b1_b + b1_a * al_b) * slope[0] + + (be_a * b1_b + b1_a * be_b) * slope[1] + + (et_a * b1_b + b1_a * et_b) * slope[2] + ) + with_slope = xi_psf_sys(a, b, systematic_correlation_matrix(rho, slope)) + np.testing.assert_allclose(with_slope[0], expected, rtol=1e-12) + + +def test_ward_groups_are_contiguous_unions(): + centres = np.random.default_rng(5).normal(size=(40, 3)) + groups = ward_group_map(centres, 7) + assert groups.shape == (40,) + np.testing.assert_array_equal(np.unique(groups), np.arange(7)) + + +def _linear_map(parameters, nbin): + """L with r = L [rho, tau, template] to first order: tau - B(p) rho - b1 t.""" + alpha_beta_eta = parameters[:3] + linear = np.zeros((3 * nbin, (9 + 3 * (len(parameters) > 3)) * nbin)) + for block, columns in enumerate(((0, 2, 5), (2, 1, 4), (5, 4, 3))): + rows = slice(block * nbin, (block + 1) * nbin) + for value, rho in zip(alpha_beta_eta, columns): + linear[rows, rho * nbin : (rho + 1) * nbin] -= value * np.eye(nbin) + linear[:, 6 * nbin : 9 * nbin] = np.eye(3 * nbin) + if len(parameters) > 3: + linear[:, 9 * nbin :] = -parameters[3] * np.eye(3 * nbin) + return linear + + +def _dense_covariance(rows): + centred = rows - rows.mean(axis=0) + return (len(rows) - 1) / len(rows) * centred.T @ centred + + +def _design_rows(rows, template_rows, nbin): + return np.stack( + [ + np.column_stack((rho_tau_design(row[: 6 * nbin].reshape(6, nbin)), t)) + for row, t in zip(rows, template_rows) + ] + ) + + +def test_residual_covariance_is_the_propagated_joint_jackknife_covariance(): + """V(p) = L C L^T for the explicit joint covariance of [rho, tau, template].""" + rng = np.random.default_rng(0) + nbin, npatch = 4, 50 + rows = rng.normal(size=(npatch, 9 * nbin)) + template_rows = rng.normal(size=(npatch, 3 * nbin)) + covariance, derivative = residual_covariance_model( + _design_rows(rows, template_rows, nbin), rows[:, 6 * nbin :] + ) + dense = _dense_covariance(np.concatenate((rows, template_rows), axis=1)) + parameters = np.array([0.02, 1.3, -0.7, 0.5]) + linear = _linear_map(parameters, nbin) + np.testing.assert_allclose( + covariance(parameters), linear @ dense @ linear.T, rtol=1e-12, atol=1e-12 + ) + step = 1e-6 + for index in range(4): + shift = np.eye(4)[index] * step + numeric = (covariance(parameters + shift) - covariance(parameters - shift)) / ( + 2 * step + ) + np.testing.assert_allclose( + derivative(parameters)[index], numeric, rtol=1e-6, atol=1e-8 + ) + + +def test_noisy_fit_minimises_the_dense_joint_chi2(): + """The fit's chi2 and minimum are those of the dense r^T (L C L^T)^-1 r.""" + nbin = 5 + vector, rows, template, template_rows, _ = _synthetic_joint(seed=3, slope=True) + result = fit_joint_eiv( + vector, rows, template_vector=template, template_rows=template_rows + ) + dense = _dense_covariance(np.concatenate((rows, template_rows), axis=1)) + x = np.column_stack((rho_tau_design(vector[: 6 * nbin].reshape(6, nbin)), template)) + + def dense_chi2(parameters): + linear = _linear_map(parameters, nbin) + residual = vector[6 * nbin :] - x @ parameters + return residual @ np.linalg.solve(linear @ dense @ linear.T, residual) + + fitted = np.array(result["parameters"]) + assert result["chi2_raw"] == pytest.approx(dense_chi2(fitted), rel=1e-10) + direct = minimize( + dense_chi2, + fitted * 1.05, + method="Nelder-Mead", + options={"xatol": 1e-10, "fatol": 1e-12, "maxiter": 20000}, + ) + np.testing.assert_allclose(direct.x, fitted, rtol=1e-4, atol=1e-6) + assert result["chi2_raw"] <= direct.fun * (1 + 1e-8) + assert result["profile_hessian_positive"] and result["bfgs_success"] + + +def test_unconstrained_parameter_is_flagged_not_given_an_error(): + """A template column of zeros leaves eta unconstrained: no sigma, flagged.""" + nbin = 5 + vector, rows, _, _, _ = _synthetic_joint(seed=8) + for target in (vector, *rows): + target[3 * nbin : 6 * nbin] = 0.0 # rho3, rho4, rho5: the eta column + result = fit_joint_eiv(vector, rows) + assert not result["profile_hessian_positive"] + assert result["sigma"] == [None, None, None] + assert result["parameter_covariance"] is None + with pytest.raises(ValueError, match="not positive definite"): + get_samples_eiv( + { + "vector": vector, + "rows_200": rows, + "alpha_vector": rho_tau_design(vector[:30].reshape(6, 5))[:, 0], + "alpha_rows_200": np.stack( + [rho_tau_design(r[:30].reshape(6, 5))[:, 0] for r in rows] + ), + }, + [200], + 200, + ) + + +def _production_fitter(tmp_path, rho): + config = dict(min_sep=1.0, max_sep=30.0, nbins=rho.shape[1], sep_units="arcmin") + fitter = PSFErrorFit( + RhoStat(output=str(tmp_path), treecorr_config=config), + TauStat(output=str(tmp_path), treecorr_config=config), + str(tmp_path), + ) + stats = {"theta": np.geomspace(1, 30, rho.shape[1])} + stats.update({f"rho_{i}_p": rho[i] for i in range(6)}) + stats.update({f"rho_{i}_m": -2 * rho[i] for i in range(6)}) + fitter.rho_stat_handler.rho_stats = stats + return fitter + + +def test_design_and_xi_sys_match_shear_psf_leakage(tmp_path): + """The rho-to-tau design and constant-beta xi_sys are PSFErrorFit's.""" + rng = np.random.default_rng(11) + rho = rng.normal(size=(6, 7)) + fitter = _production_fitter(tmp_path, rho) + a, b = rng.normal(size=3), rng.normal(size=3) + np.testing.assert_allclose( + rho_tau_design(rho) @ a, + fitter.model((*a, 0.0, 0.0), use_eta=True, use_fourth_moment=False), + rtol=1e-13, + ) + for sign, part in ((1, "p"), (-2, "m")): + np.testing.assert_allclose( + xi_psf_sys(a, b, systematic_correlation_matrix(sign * rho))[0], + fitter.compute_xi_psf_sys(a, b, p_or_m=part), + rtol=1e-12, + ) + + +@pytest.mark.parametrize("slope", [False, True]) +def test_alpha_replaces_only_its_column_and_matched_rows(monkeypatch, slope): + """The full joint rows reach the objective with alpha replaced, beta/eta intact.""" + from sp_validation import leakage_fit + + vector, rows, template, template_rows, _ = _synthetic_joint(14, slope=slope) + nbin = len(vector) // 9 + n_rho = 6 * nbin + rng = np.random.default_rng(22) + alpha = rng.normal(size=3 * nbin) + alpha_rows = rng.normal(size=(len(rows), 3 * nbin)) + original_vector, original_rows = vector.copy(), rows.copy() + expected_x = rho_tau_design(vector[:n_rho].reshape(6, nbin)) + expected_rows = np.stack([rho_tau_design(r[:n_rho].reshape(6, nbin)) for r in rows]) + expected_x[:, 0] = alpha + expected_rows[:, :, 0] = alpha_rows + if slope: + expected_x = np.column_stack((expected_x, template)) + expected_rows = np.concatenate( + (expected_rows, template_rows[:, :, None]), axis=2 + ) + + def inspect(x, y, xr, yr, scales, cross_term, start): + np.testing.assert_array_equal(x, expected_x) + np.testing.assert_array_equal(xr, expected_rows) + np.testing.assert_array_equal(y, vector[n_rho:]) + np.testing.assert_array_equal(yr, rows[:, n_rho:]) + assert x.shape[1] == (4 if slope else 3) + assert cross_term + return "inspected" + + monkeypatch.setattr(leakage_fit, "_profile_fit", inspect) + assert ( + fit_joint_eiv( + vector, + rows, + alpha_vector=alpha, + alpha_rows=alpha_rows, + template_vector=template, + template_rows=template_rows, + ) + == "inspected" + ) + np.testing.assert_array_equal(vector, original_vector) + np.testing.assert_array_equal(rows, original_rows) diff --git a/src/sp_validation/tests/test_leakage_measurement.py b/src/sp_validation/tests/test_leakage_measurement.py new file mode 100644 index 00000000..bac80e40 --- /dev/null +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -0,0 +1,418 @@ +"""The joint leakage measurement and its 'eiv' fit through the cosmo_val mixin.""" + +import json +import os + +import numpy as np +import pytest +from astropy.io import fits + +from sp_validation import rho_tau +from sp_validation.cosmo_val.psf_systematics import PSFSystematicsMixin + +VERSION = "SP_synthetic" + + +def _write_catalogues(tmp_path, n_star=4000, n_gal=6000): + rng = np.random.default_rng(7) + ra, dec = rng.uniform(10, 12, n_star), rng.uniform(0, 2, n_star) + e_psf = rng.normal(0, 0.02, (2, n_star)) + e_star = e_psf + rng.normal(0, 0.01, (2, n_star)) + t_psf = rng.uniform(0.4, 0.6, n_star) + t_star = t_psf * (1 + rng.normal(0, 0.02, n_star)) + flags = (rng.random(n_star) < 0.05).astype(np.int16) + psf = tmp_path / "psf.fits" + fits.BinTableHDU.from_columns( + [ + fits.Column("RA", "D", array=ra), + fits.Column("DEC", "D", array=dec), + fits.Column("HSM_G1_PSF", "D", array=e_psf[0]), + fits.Column("HSM_G2_PSF", "D", array=e_psf[1]), + fits.Column("HSM_G1_STAR", "D", array=e_star[0]), + fits.Column("HSM_G2_STAR", "D", array=e_star[1]), + fits.Column("HSM_T_PSF", "D", array=t_psf), + fits.Column("HSM_T_STAR", "D", array=t_star), + fits.Column("HSM_FLAG_PSF", "I", array=flags), + fits.Column("HSM_FLAG_STAR", "I", array=np.zeros(n_star, np.int16)), + ] + ).writeto(psf) + shear = tmp_path / "shear.fits" + fits.BinTableHDU.from_columns( + [ + fits.Column("RA", "D", array=rng.uniform(10, 12, n_gal)), + fits.Column("Dec", "D", array=rng.uniform(0, 2, n_gal)), + fits.Column("e1_PSF", "D", array=rng.normal(0, 0.02, n_gal)), + fits.Column("e2_PSF", "D", array=rng.normal(0, 0.02, n_gal)), + fits.Column("e1", "D", array=rng.normal(0, 0.3, n_gal)), + fits.Column("e2", "D", array=rng.normal(0, 0.3, n_gal)), + fits.Column("w_des", "D", array=rng.uniform(0.5, 2, n_gal)), + ] + ).writeto(shear) + return psf, shear + + +class _Validation(PSFSystematicsMixin): + """Minimal host with the configuration the leakage methods read.""" + + def __init__(self, tmp_path, beta_psf_size=False): + psf, shear = _write_catalogues(tmp_path) + (tmp_path / "rho_tau_stats").mkdir() + self.output = tmp_path + self.cc = { + "paths": {"output": str(tmp_path)}, + VERSION: { + "patch_number": 30, + "psf": { + "path": str(psf), + "ra_col": "RA", + "dec_col": "DEC", + "e1_PSF_col": "HSM_G1_PSF", + "e2_PSF_col": "HSM_G2_PSF", + "e1_star_col": "HSM_G1_STAR", + "e2_star_col": "HSM_G2_STAR", + "PSF_size": "HSM_T_PSF", + "star_size": "HSM_T_STAR", + "PSF_flag": "HSM_FLAG_PSF", + "star_flag": "HSM_FLAG_STAR", + }, + "shear": { + "path": str(shear), + "ra_col": "RA", + "dec_col": "Dec", + "w_col": "w_des", + "e1_PSF_col": "e1_PSF", + "e2_PSF_col": "e2_PSF", + "e1_col": "e1", + "e2_col": "e2", + }, + }, + } + self.versions = [VERSION] + self.treecorr_config = { + "ra_units": "degrees", + "dec_units": "degrees", + "sep_units": "arcmin", + "min_sep": 5.0, + "max_sep": 60.0, + "nbins": 3, + "bin_slop": 0.0, + "angle_slop": 0.0, + "num_threads": 1, + } + self.psf_leakage = { + "npatch_fine": 30, + "patch_counts": [15, 20], + "fiducial_npatch": 20, + "beta_psf_size": beta_psf_size, + } + self.rho_tau_method = "eiv" + self.force_run = False + + def basename(self, version, tomo_bin_a="all", npatch=None, **kwargs): + return f"{version}_{tomo_bin_a}_npatch={npatch}" + + def _output_path(self, *parts): + return os.path.join(self.output, *parts) + + +def _masks(val): + """The star and galaxy selections of the leakage measurement.""" + return val._get_star_mask(VERSION), val._get_galaxy_mask(VERSION, "all") + + +def _fresh(val, **psf_leakage): + """A new host on the same outputs, without in-memory products.""" + fresh = _Validation.__new__(_Validation) + fresh.__dict__.update( + {k: v for k, v in val.__dict__.items() if not k.startswith("_leakage")} + ) + fresh.psf_leakage = {**val.psf_leakage, **psf_leakage} + return fresh + + +def test_joint_vector_equals_production_rho_tau(tmp_path): + """With exact pair counting the joint vector is production's rho/tau.""" + val = _Validation(tmp_path) + products = val.leakage_jackknife(VERSION, measure=True) + mask_star, mask_gal = _masks(val) + (tmp_path / "production").mkdir() + rho_handler, tau_handler = rho_tau.get_rho_tau( + val.cc, + VERSION, + val.treecorr_config, + str(tmp_path / "production"), + "base", + "base", + mask_star=mask_star, + mask_gal=mask_gal, + ) + expected = np.concatenate( + [rho_handler.rho_stats[f"rho_{i}_p"] for i in range(6)] + + [tau_handler.tau_stats[f"tau_{i}_p"] for i in (0, 2, 5)] + ) + np.testing.assert_allclose(products["vector"], expected, rtol=1e-10, atol=1e-16) + np.testing.assert_allclose( + products["rho_xim"], + [rho_handler.rho_stats[f"rho_{i}_m"] for i in range(6)], + rtol=1e-10, + atol=1e-16, + ) + assert products["rows_30"].shape == (30, 27) + assert products["rows_15"].shape == (15, 27) + np.testing.assert_array_equal(np.unique(products["groups_20"]), np.arange(20)) + + +@pytest.mark.parametrize("beta_psf_size", [False, True]) +def test_eiv_samples_and_xi_sys_through_the_mixin(tmp_path, monkeypatch, beta_psf_size): + val = _Validation(tmp_path, beta_psf_size=beta_psf_size) + val.psf_error_nsamples = 50 + val.leakage_jackknife(VERSION, measure=True) + samples = val.get_samples(VERSION, {}, "all") + k = 4 if beta_psf_size else 3 + assert samples.shape == (50, k) + fits_json = json.loads( + val.leakage_jackknife_path(VERSION, "all", "leakage_eiv_fit").read_text() + ) + assert set(fits_json["fits"]) == {"15", "20"} + assert fits_json["fits"]["20"]["dof"] == 9 - k + + plus, minus = val.get_xi_psf_sys_samples(VERSION, {}, "all", "all") + assert plus.shape == minus.shape == (50, 3) + products = val.leakage_jackknife(VERSION) + rho = products["vector"][:18].reshape(6, 3) + alpha, beta, eta = samples[0, :3] + constant = ( + alpha**2 * rho[0] + + beta**2 * rho[1] + + 2 * alpha * beta * rho[2] + + eta**2 * rho[3] + + 2 * beta * eta * rho[4] + + 2 * alpha * eta * rho[5] + ) + if beta_psf_size: + beta1 = samples[0, 3] + slope = products["slope_xip"] + constant += beta1**2 * slope[3] + 2 * beta1 * ( + alpha * slope[0] + beta * slope[1] + eta * slope[2] + ) + np.testing.assert_allclose(plus[0], constant, rtol=1e-10) + + # A fresh host reads the cached measurement instead of measuring. + def fail(*args, **kwargs): + raise AssertionError("cached measurement was not reused") + + monkeypatch.setattr( + "sp_validation.cosmo_val.psf_systematics.measure_leakage_jackknife", fail + ) + np.testing.assert_array_equal( + _fresh(val).leakage_jackknife(VERSION)["vector"], products["vector"] + ) + + +def test_reading_without_a_measurement_names_the_rule(tmp_path): + """Only measure=True measures; every other caller needs the cache.""" + val = _Validation(tmp_path) + with pytest.raises(FileNotFoundError, match="is missing.*cv_leakage_jackknife"): + val.leakage_jackknife(VERSION) + with pytest.raises(FileNotFoundError, match="cv_leakage_jackknife"): + val.get_samples(VERSION, {}, "all") + assert not val.leakage_jackknife_path(VERSION).exists() + + +@pytest.mark.parametrize("change", ["beta_psf_size", "selection"]) +def test_cache_is_stale_after_a_change(tmp_path, change): + """PSF-size slope and star selection each invalidate the cache.""" + val = _Validation(tmp_path) + val.leakage_jackknife(VERSION, measure=True) + if change == "beta_psf_size": + fresh = _fresh(val, beta_psf_size=True) + else: + fresh = _fresh(val) + star_mask = PSFSystematicsMixin._get_star_mask + + def fewer_stars(self, ver): + # Drop a strip of stars, so the selection changes. + ra = fits.getdata(self.cc[ver]["psf"]["path"])["RA"] + return star_mask(self, ver) & (ra > 10.2) + + fresh._get_star_mask = fewer_stars.__get__(fresh) + with pytest.raises(FileNotFoundError, match="is stale"): + fresh.leakage_jackknife(VERSION) + assert _fresh(val).leakage_jackknife(VERSION) is not None + + +@pytest.mark.parametrize( + "settings, message", + [ + ({"npatch_fine": 300}, "strictly between 1 and npatch_fine"), + ({"patch_counts": [1, 400]}, "strictly between 1 and npatch_fine"), + ({"fiducial_npatch": 300}, "fiducial_npatch must be one of"), + ], +) +def test_psf_leakage_settings_are_validated(settings, message): + from sp_validation.cosmo_val import CosmologyValidation + + with pytest.raises(ValueError, match=message): + CosmologyValidation(versions=[], psf_leakage=settings) + + +def test_theory_covariance_follows_the_selection(tmp_path, monkeypatch): + """cov_tau_*_th.npy is rebuilt when the star or galaxy selection changes.""" + val = _Validation(tmp_path) + built = [] + + class FakeCovTauTh: + def __init__(self, mask_star=None, mask_gal=None, **kwargs): + built.append((mask_star.sum(), mask_gal.sum())) + + def build_cov(self, **kwargs): + return np.eye(2) * len(built) + + monkeypatch.setattr(rho_tau, "CovTauTh", FakeCovTauTh) + mask_star, mask_gal = _masks(val) + outdir = tmp_path / "rho_tau_stats" + + def covariance(star, gal): + rho_tau.get_theory_cov( + val.cc, + VERSION, + val.treecorr_config, + str(outdir), + "base", + mask_star=star, + mask_gal=gal, + ) + return np.load(outdir / "cov_tau_base_th.npy")[0, 0] + + assert covariance(mask_star, mask_gal) == 1 + assert covariance(mask_star, mask_gal) == 1 + fewer = mask_gal.copy() + fewer[np.flatnonzero(fewer)[0]] = False + assert covariance(mask_star, fewer) == 2 + assert len(built) == 2 + + +def test_field_and_alpha_definitions_invalidate_leakage_cache(tmp_path, monkeypatch): + """Old field definitions, responses and galaxy PSF mappings cannot reuse rows.""" + val = _Validation(tmp_path) + products = val.leakage_jackknife(VERSION, measure=True) + assert products["alpha_vector"].shape == (9,) + assert products["alpha_rows_15"].shape == (15, 9) + path = val.leakage_jackknife_path(VERSION) + mask_star, mask_gal = _masks(val) + params = rho_tau.get_params_rho_tau(val.cc[VERSION]) + + def key(): + return rho_tau.leakage_cache_key( + val.treecorr_config, + 30, + [15, 20], + False, + mask_star, + mask_gal, + params=params, + ) + + assert rho_tau.read_leakage_jackknife(path, key()) is not None + for field, value in (("R11", 0.8), ("e1_gal_PSF_col", "other_psf")): + previous = params[field] + params[field] = value + assert rho_tau.read_leakage_jackknife(path, key()) is None + params[field] = previous + monkeypatch.setattr(rho_tau, "TAU_FIELD_DEFINITION", "different-centring") + assert rho_tau.read_leakage_jackknife(path, key()) is None + monkeypatch.setattr(rho_tau, "TAU_FIELD_DEFINITION", "selected-weighted-mean-v1") + monkeypatch.setattr(rho_tau, "ALPHA_TEMPLATE_DEFINITION", "different-alpha") + assert rho_tau.read_leakage_jackknife(path, key()) is None + + +def test_ordinary_tau_rejects_old_field_definition(tmp_path, monkeypatch): + """An existing FITS and matching masks do not preserve obsolete tau centring.""" + val = _Validation(tmp_path) + outdir = tmp_path / "rho_tau_stats" + mask_star, mask_gal = _masks(val) + args = (val.cc, VERSION, val.treecorr_config, str(outdir), "base", "base") + rho_tau.get_rho_tau(*args, mask_star=mask_star, mask_gal=mask_gal) + path = outdir / "tau_stats_base.fits.json" + record = json.loads(path.read_text()) + assert record.pop("definition")["centring"] == rho_tau.TAU_FIELD_DEFINITION + path.write_text(json.dumps(record)) + calls = [] + compute = rho_tau._compute_stats_with_patches + + def track(handler, kind, *args, **kwargs): + calls.append(kind) + return compute(handler, kind, *args, **kwargs) + + monkeypatch.setattr(rho_tau, "_compute_stats_with_patches", track) + rho_tau.get_rho_tau(*args, mask_star=mask_star, mask_gal=mask_gal) + assert calls == ["tau"] + calls.clear() + rho_tau.get_rho_tau(*args, mask_star=mask_star, mask_gal=mask_gal) + assert calls == [] + + +def test_measured_alpha_closes_selected_response_calibrated_injection(tmp_path): + """The consumer measures the selected, response-calibrated injection derivative.""" + val = _Validation(tmp_path) + val.cc[VERSION]["shear"].update(R11=0.7, R22=1.2) + path = val.cc[VERSION]["shear"]["path"] + data = fits.getdata(path) + mask = np.asarray(data["RA"] > 10.4) + val._get_galaxy_mask = lambda *args: mask + clean = val.leakage_jackknife(VERSION, measure=True) + data["e1"] += 0.1 * data["e1_PSF"] + data["e2"] += 0.1 * data["e2_PSF"] + fits.BinTableHDU(data).writeto(path, overwrite=True) + leaked_host = _fresh(val) + leaked_host.force_run = True + leaked = leaked_host.leakage_jackknife(VERSION, measure=True) + np.testing.assert_allclose(clean["centres"], leaked["centres"], rtol=0, atol=1e-14) + np.testing.assert_allclose( + clean["vector"][:18], leaked["vector"][:18], rtol=1e-12, atol=1e-18 + ) + np.testing.assert_allclose( + leaked["vector"][18:] - clean["vector"][18:], + 0.1 * clean["alpha_vector"], + atol=2e-12, + rtol=1e-6, + ) + # Use the upstream injection gate's tolerance for TreeCorr field arithmetic. + old_alpha = np.concatenate([clean["vector"][3 * i : 3 * i + 3] for i in (0, 2, 5)]) + assert not np.allclose( + leaked["vector"][18:] - clean["vector"][18:], + 0.1 * old_alpha, + atol=2e-12, + rtol=1e-6, + ) + for n in (15, 20, 30): + np.testing.assert_allclose( + leaked[f"rows_{n}"][:, 18:] - clean[f"rows_{n}"][:, 18:], + 0.1 * clean[f"alpha_rows_{n}"], + atol=2e-12, + rtol=1e-6, + ) + + +def test_tau_samples_follow_full_field_definition(tmp_path): + """Changing a response or a shear column invalidates samples after tau rebuilds.""" + from types import SimpleNamespace + + output = SimpleNamespace(_output=str(tmp_path)) + fitter = SimpleNamespace(tau_stat_handler=SimpleNamespace(catalogs=output)) + path = tmp_path / "tau_stats_base.fits.json" + definition = rho_tau.tau_field_definition({"e1_col": "e1", "R11": 0.8}) + + def name(): + path.write_text( + json.dumps({"selection": "same-rows", "definition": definition}) + ) + return rho_tau._sample_cache_name(fitter, "base", "lsq", "th") + + original = name() + for key, value in (("e1_col", "other_e1"), ("R11", 0.9)): + previous = definition[key] + definition[key] = value + assert name() != original + definition[key] = previous + assert name() == original diff --git a/src/sp_validation/tests/test_rho_cache_reuse.py b/src/sp_validation/tests/test_rho_cache_reuse.py index 0cf836c8..db0aa3ee 100644 --- a/src/sp_validation/tests/test_rho_cache_reuse.py +++ b/src/sp_validation/tests/test_rho_cache_reuse.py @@ -2,19 +2,19 @@ import numpy as np -from sp_validation.correlation import write_measurement_metadata +from sp_validation.correlation import selection_digest, write_measurement_metadata from sp_validation.rho_tau import rho_cache_valid CONFIG = {"min_sep": 1.0, "max_sep": 100.0, "nbins": 5, "sep_units": "arcmin"} -def _cache(tmp_path, config=None): +def _cache(tmp_path, config=None, selection=None): rho = tmp_path / "rho_stats.fits" cov = tmp_path / "cov_rho_jk.npy" rho.write_bytes(b"") np.save(cov, np.eye(2)) if config is not None: - write_measurement_metadata(rho, config) + write_measurement_metadata(rho, config, selection=selection) return rho, cov @@ -34,3 +34,12 @@ def test_rho_cache_matching_is_reused_unless_forced(tmp_path): assert not rho_cache_valid(rho, cov, CONFIG, force_run=True) cov.unlink() assert not rho_cache_valid(rho, cov, CONFIG) + + +def test_rho_cache_from_another_star_selection_is_remeasured(tmp_path): + """A star veto changes the selection; the cached ρ must not survive it.""" + kept = selection_digest(np.array([True, True, False])) + vetoed = selection_digest(np.array([True, False, False])) + rho, cov = _cache(tmp_path, CONFIG, selection=kept) + assert rho_cache_valid(rho, cov, CONFIG, selection=kept) + assert not rho_cache_valid(rho, cov, CONFIG, selection=vetoed) diff --git a/uv.lock b/uv.lock index 0ef4517b..25860a8d 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 5 +revision = 3 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'linux'", @@ -3844,7 +3844,7 @@ wheels = [ [[package]] name = "shear-psf-leakage" version = "0.2.1" -source = { git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=develop#a1e7e5b97b88544a15e79abd99ec59fb61e42601" } +source = { git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=d0e10d83ecb1978ed56599f0b46326e3d000680e#d0e10d83ecb1978ed56599f0b46326e3d000680e" } dependencies = [ { name = "cs-util" }, { name = "emcee" }, @@ -4060,7 +4060,7 @@ requires-dist = [ { name = "sacc", specifier = ">=2.4,<3" }, { name = "scipy" }, { name = "seaborn" }, - { name = "shear-psf-leakage", git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=develop" }, + { name = "shear-psf-leakage", git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=d0e10d83ecb1978ed56599f0b46326e3d000680e" }, { name = "skyproj" }, { name = "smokescreen", git = "https://github.com/UNIONS-WL/Smokescreen?rev=main" }, { name = "sp-validation", extras = ["test", "docs"], marker = "extra == 'develop'" }, diff --git a/workflow/common.py b/workflow/common.py index 3d98f78b..47a9d8ec 100644 --- a/workflow/common.py +++ b/workflow/common.py @@ -477,6 +477,7 @@ def cv_cov_tau(version, fiducial=None): # src/sp_validation/tests/test_cv_init_params.py, so no default applies silently. CV_INIT_KEYS = ( "rho_tau_method", + "psf_leakage", "cov_estimate_method", "compute_cov_rho", "n_cov", diff --git a/workflow/rules/cosmo_val.smk b/workflow/rules/cosmo_val.smk index b2d3c105..5db23f97 100644 --- a/workflow/rules/cosmo_val.smk +++ b/workflow/rules/cosmo_val.smk @@ -252,13 +252,40 @@ rule cv_plot_tau_stats: "../scripts/cv_plot_diagnostic.py" +# The 'eiv' PSF-leakage fit reads its own joint ρ/τ jackknife measurement; +# 'lsq'/'emcee' fit the published τ statistics with their covariances. +CV_LEAKAGE_FINE = {**CV_FIDUCIAL, "npatch": CV["psf_leakage"]["npatch_fine"]} +CV_LEAKAGE_JACKKNIFE = [ + str(COSMO_VAL / "rho_tau_stats" / f"leakage_eiv_{cv_basename(v, CV_LEAKAGE_FINE)}.npz") + for v in CV_VERSIONS +] + + +rule cv_leakage_jackknife: + """Joint ρ/τ jackknife measurement behind the 'eiv' PSF-leakage fit.""" + input: + catalogues=cv_shear_catalogues() + + [str(Path(catalogue_entry(v)["subdir"]) / catalogue_entry(v)["psf"]["path"]) + for v in CV_VERSIONS], + output: + measurements=CV_LEAKAGE_JACKKNIFE, + params: + cv_init=CV_INIT, + threads: 8 + resources: + mem_mb=128000, + runtime=240, + script: + "../scripts/cv_leakage_jackknife.py" + + rule cv_rho_tau_fits: """Fit the PSF-error model (alpha/beta/eta) and propagate to xi_psf_sys.""" input: rho=[cv_rho_stats(v, CV_FIDUCIAL) for v in CV_VERSIONS], tau=[cv_tau_stats(v, CV_FIDUCIAL) for v in CV_VERSIONS], - covariances=cv_tau_covariances(), - rho_covariances=[str(COSMO_VAL / "rho_tau_stats" / + fit_inputs=CV_LEAKAGE_JACKKNIFE if CV["rho_tau_method"] == "eiv" + else cv_tau_covariances() + [str(COSMO_VAL / "rho_tau_stats" / f"cov_rho_{cv_basename(v, CV_FIDUCIAL)}_jk.npy") for v in CV_VERSIONS], output: products=CV_PSF_PRODUCTS, diff --git a/workflow/scripts/cv_leakage_jackknife.py b/workflow/scripts/cv_leakage_jackknife.py new file mode 100644 index 00000000..3243f7db --- /dev/null +++ b/workflow/scripts/cv_leakage_jackknife.py @@ -0,0 +1,9 @@ +"""Rule cv_leakage_jackknife: the joint ρ/τ jackknife measurement per version.""" + +from cv_runner import _unbuffer_streams, make_cv, verify_outputs + +_unbuffer_streams() +cv = make_cv(snakemake) +for ver in cv.versions: + cv.leakage_jackknife(ver, measure=True) +verify_outputs(snakemake) diff --git a/workflow/scripts/cv_rho_tau_fits.py b/workflow/scripts/cv_rho_tau_fits.py index bcdbff82..abf17ef4 100644 --- a/workflow/scripts/cv_rho_tau_fits.py +++ b/workflow/scripts/cv_rho_tau_fits.py @@ -9,9 +9,21 @@ params = cv.set_params_rho_tau(ver, cv.results[ver]._params, cv.cc[ver]["psf"]) flat_samples = cv.get_samples(ver, params, "all") plus, minus = cv.get_xi_psf_sys_samples(ver, params, "all", "all") + if cv.rho_tau_method == "eiv": + leakage = cv.leakage_jackknife(ver) + theta = leakage["theta"] + fit_summary = cv.leakage_jackknife_path(ver, "all", "leakage_eiv_fit") + extra = dict(leakage_eiv_fit=fit_summary.read_text()) + if cv.psf_leakage["beta_psf_size"]: + # : the beta1 term of the tau0, tau2, tau5 model. + extra["tau_slope_template"] = leakage["slope_xip"][:3] + else: + theta = cv.psf_fitter.rho_stat_handler.rho_stats["theta"] + extra = {} data = dict( - theta=cv.psf_fitter.rho_stat_handler.rho_stats["theta"], + theta=theta, flat_samples=flat_samples, + **extra, theta_edges=np.geomspace( cv.treecorr_config["min_sep"], cv.treecorr_config["max_sep"], diff --git a/workflow/tests/conftest.py b/workflow/tests/conftest.py index af4e4258..6ca9f6eb 100644 --- a/workflow/tests/conftest.py +++ b/workflow/tests/conftest.py @@ -11,8 +11,9 @@ The ``toy`` fixture is a disposable checkout: copies of ``workflow/`` and ``papers/cosmo_val/``, this checkout's ``src/`` symlinked in, a blinded and an -undeclared catalogue, stand-ins for the OneCovariance covariances (their -inputs live on candide), and both output roots in tmp. The host never opens a +undeclared catalogue (each with a star catalogue), stand-ins for the +OneCovariance covariances (their inputs live on candide), and both output +roots in tmp. The host never opens a blind, so there is no blind record. """ @@ -142,6 +143,7 @@ def _cat_config(data): def entry(name, **declaration): catalogue = data / f"{name}.fits" catalogue.touch() + (data / f"{name}_star.fits").touch() return { "subdir": str(data), "pipeline": "SP", @@ -157,6 +159,7 @@ def entry(name, **declaration): "e1_col_corrected": "e1_leak_corrected", "e2_col_corrected": "e2_leak_corrected", }, + "psf": {"path": f"{name}_star.fits"}, **declaration, }