From 0d5c65dec2b35ab6fa5979a8282a4e1c7674b8df Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 09:11:05 -0400 Subject: [PATCH 01/16] leakage_fit: joint errors-in-variables rho/tau fit, jackknife rows, sky mask Pure-numerics module for the PSF-leakage fit: the rho->tau design, the k-template joint EIV profile fit (residual covariance from joint delete-one rows including the rho-tau cross term, Hartlap-scaled chi2, profile-Hessian errors), Ward grouping of fine patches, O(pairs) TreeCorr match-jackknife rows with optional grouping, the circle sky mask, and xi_sys = a^T M b. Reproduces the leakadq reference fits on its bound9clean rows to 1e-12 in chi2. Co-Authored-By: Claude Opus 5.5 --- src/sp_validation/leakage_fit.py | 441 ++++++++++++++++++++ src/sp_validation/tests/test_leakage_fit.py | 245 +++++++++++ 2 files changed, 686 insertions(+) create mode 100644 src/sp_validation/leakage_fit.py create mode 100644 src/sp_validation/tests/test_leakage_fit.py diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py new file mode 100644 index 00000000..43a7878e --- /dev/null +++ b/src/sp_validation/leakage_fit.py @@ -0,0 +1,441 @@ +"""Numerical tools for joint errors-in-variables PSF-leakage fits. + +This module has no catalogue I/O. It builds the production rho-to-tau design, +profiles noisy templates from a joint delete-one covariance, and provides +TreeCorr jackknife, patch-grouping, and spherical-mask 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 or centres.shape[0] < 3: + raise ValueError("centres must be a two-dimensional array of patch centres") + 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. The + returned tuple is ``(rows, weights, max_checked_error)``; the weight vector + is diagnostic and matches TreeCorr's patch-weight convention. + """ + blocks = [] + weights = None + 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}" + ) + + treecorr_factor = len(corr.getWeight()) / len(corr.weight) + block_weights = denominator.sum(axis=1) * treecorr_factor + weights = block_weights if weights is None else weights + block_weights + blocks.append(rows) + + if not blocks: + raise ValueError("at least one TreeCorr correlation is required") + return np.concatenate(blocks, axis=1), weights, max_error + + +def keep_sky(ra, dec, circles, chunk_size=1_000_000): + """Return a boolean mask outside all spherical circles. + + ``ra`` and ``dec`` are in degrees. Each circle is a mapping with + ``ra_deg``, ``dec_deg`` and ``radius_deg`` fields. Angular separation is + computed on the sphere, including across the RA 0/360-degree boundary. + """ + ra = np.asarray(ra, dtype=float) + dec = np.asarray(dec, dtype=float) + if ra.shape != dec.shape or ra.ndim != 1: + raise ValueError("ra and dec must be one-dimensional arrays of equal length") + circles = tuple(circles) + keep = np.ones(ra.size, dtype=bool) + for circle in circles: + radius = float(circle["radius_deg"]) + if not 0 < radius < 180: + raise ValueError("circle radius must be between 0 and 180 degrees") + centre_ra = np.deg2rad(float(circle["ra_deg"])) + centre_dec_deg = float(circle["dec_deg"]) + centre_dec = np.deg2rad(centre_dec_deg) + cos_radius = np.cos(np.deg2rad(radius)) + for start in range(0, ra.size, chunk_size): + stop = min(start + chunk_size, ra.size) + candidate = np.flatnonzero( + np.abs(dec[start:stop] - centre_dec_deg) <= radius + ) + if not candidate.size: + continue + point_dec = np.deg2rad(dec[start:stop][candidate]) + point_ra = np.deg2rad(ra[start:stop][candidate]) + cosine = ( + np.sin(point_dec) * np.sin(centre_dec) + + np.cos(point_dec) + * np.cos(centre_dec) + * np.cos(point_ra - centre_ra) + ) + keep[start + candidate[cosine >= cos_radius]] = False + return keep + + +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. + + ``x_rows``/``y_rows`` are delete-one jackknife rows of ``x`` and ``y``. With + jackknife covariances C_xx, C_xy, C_yy built from them, the residual + ``r = y - x p`` has covariance V(p) = C_yy - (B C_xy + C_yx B^T) + + B C_xx B^T, B the linear map p -> x p, and the objective is + r^T V(p)^-1 r (no log-determinant term). + """ + 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") + + 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 residual_covariance(parameters): + return cov_yy + np.einsum("i,iab->ab", parameters, linear) + np.einsum( + "i,j,ijab->ab", parameters, parameters, cov_xx + ) + + def objective(scaled_parameters): + parameters = scaled_parameters * scales + residual = y - x @ parameters + covariance = residual_covariance(parameters) + factor = cho_factor(covariance, lower=True, check_finite=False) + weighted_residual = cho_solve(factor, residual, check_finite=False) + value = float(residual @ weighted_residual) + covariance_derivative = linear + np.einsum( + "j,ijab->iab", parameters, cov_xx + ) + np.einsum("j,jiab->iab", parameters, cov_xx) + gradient = ( + -2 * x.T @ weighted_residual + - np.einsum( + "a,iab,b->i", + weighted_residual, + covariance_derivative, + weighted_residual, + ) + ) * scales + return value, gradient + + def generalized_least_squares(covariance): + factor = cho_factor(covariance, lower=True, check_finite=False) + weighted_x = cho_solve(factor, x, check_finite=False) + weighted_y = cho_solve(factor, y, check_finite=False) + return np.linalg.solve(x.T @ weighted_x, x.T @ weighted_y) + + fixed_start = generalized_least_squares(cov_yy) + initial = fixed_start if start is None else np.asarray(start, dtype=float) + if initial.shape != (n_param,): + raise ValueError("start must contain one value per fit parameter") + 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: + starts.append(initial) + rng = np.random.default_rng(1028) + starts.extend(fixed_start + rng.normal(size=n_param) * scales for _ in range(3)) + + trials = [] + fitted = [] + for initial_parameters in starts: + optimization = minimize( + objective, + initial_parameters / scales, + jac=True, + method="BFGS", + options={"gtol": 1e-7, "maxiter": 500}, + ) + fitted.append(optimization) + trials.append( + { + "parameters": (optimization.x * scales).tolist(), + "chi2_raw": float(optimization.fun), + "success": bool(optimization.success), + "gradient_norm": float(np.linalg.norm(optimization.jac)), + } + ) + best = min(fitted, key=lambda fit: fit.fun) + parameters = best.x * scales + covariance = residual_covariance(parameters) + residual = y - x @ parameters + precision_residual = np.linalg.solve(covariance, residual) + chi2_raw = float(residual @ precision_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 + eigenvalues = np.linalg.eigvalsh(hessian) + parameter_covariance = ( + 2 + * np.linalg.inv(hessian) + * scales[:, None] + * scales[None, :] + / hartlap + ) + dof = n_data - n_param + result = { + "parameters": parameters.tolist(), + "sigma": np.sqrt(np.maximum(np.diag(parameter_covariance), 0)).tolist(), + "parameter_covariance": parameter_covariance.tolist(), + "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": np.sqrt(np.diag(covariance)).tolist(), + "residual_sigma_normalized": ( + residual / np.sqrt(np.diag(covariance)) + ).tolist(), + "profile_hessian_eigenvalues": eigenvalues.tolist(), + "profile_hessian_positive": bool(eigenvalues.min() > 0), + "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), + } + return result, covariance + + +def fit_joint_eiv( + vector, + jk_rows, + *, + nbin=None, + template_vector=None, + template_rows=None, + scales=None, + start=None, + cross_term=True, +): + """Fit rho0..rho5 to tau0/tau2/tau5 with the joint EIV covariance. + + ``vector`` is ordered ``[rho0..rho5, tau0, tau2, tau5]`` and ``jk_rows`` + contains the same delete-one measurements, one row per patch. Optional + template columns add measured predictors such as ``u * delta_e``; their + point values and rows have shapes ``(ndata, ntemplate)`` and + ``(npatch, ndata, ntemplate)``. Returns ``(result, residual_covariance)``. + + 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. + """ + 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 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/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py new file mode 100644 index 00000000..ae0bcffc --- /dev/null +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -0,0 +1,245 @@ +"""Synthetic and TreeCorr checks for the joint PSF-leakage estimator.""" + +import numpy as np +import pytest +import treecorr + +from sp_validation.leakage_fit import ( + fit_joint_eiv, + keep_sky, + rho_tau_design, + systematic_correlation_matrix, + treecorr_jackknife_rows, + ward_group_map, + xi_psf_sys, +) + + +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, weights, _ = treecorr_jackknife_rows([fine_corr]) + native, native_weights = 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) + np.testing.assert_allclose(weights, native_weights, rtol=1e-12) + 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, grouped_weights, _ = treecorr_jackknife_rows( + [fine_corr], groups=groups, check_rows=False + ) + native_grouped, native_group_weights = 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) + np.testing.assert_allclose(grouped_weights, native_group_weights, rtol=1e-12) + + +def test_spherical_circle_mask_handles_inside_outside_and_ra_wrap(): + """Circle excision follows great-circle distance across RA=0 degrees.""" + circles = [ + {"name": "wrap", "ra_deg": 359.8, "dec_deg": 0.0, "radius_deg": 0.5} + ] + keep = keep_sky( + ra=np.array([359.9, 0.1, 359.0, 180.0]), + dec=np.array([0.0, 0.0, 0.0, 0.0]), + circles=circles, + ) + np.testing.assert_array_equal(keep, [False, False, True, True]) + + +def test_sky_mask_rejects_mismatched_coordinates(): + with pytest.raises(ValueError, match="equal length"): + keep_sky([1.0], [1.0, 2.0], []) + + +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)) From f012de470d0ba33f53b508e16af020f5912a7067 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 09:23:42 -0400 Subject: [PATCH 02/16] rho_tau: joint rho/tau jackknife measurement and EIV fit measure_leakage_jackknife builds stars and galaxies as get_rho_tau does, with npatch_fine shared k-means centres, and in one pass of match-weighted GG correlations measures rho0..rho5, tau0/2/5 and optionally the u*de PSF-size templates; delete-one rows at the fine count and at Ward unions for each patch count come from the per-patch-pair results. Products are cached as npz + json keyed by the measurement settings. fit_leakage and get_samples_eiv run the joint EIV fit and draw Gaussian samples. Co-Authored-By: Claude Opus 5.5 --- src/sp_validation/leakage_fit.py | 45 ++-- src/sp_validation/rho_tau.py | 255 +++++++++++++++++++- src/sp_validation/tests/test_leakage_fit.py | 9 +- 3 files changed, 279 insertions(+), 30 deletions(-) diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py index 43a7878e..38263ec8 100644 --- a/src/sp_validation/leakage_fit.py +++ b/src/sp_validation/leakage_fit.py @@ -44,7 +44,9 @@ def ward_group_map(centres, ngroups): raise ValueError("centres must be a two-dimensional array of patch centres") 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 + 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}" @@ -68,7 +70,9 @@ def treecorr_jackknife_rows(correlations, groups=None, check_rows=True): 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") + raise ValueError( + "each correlation must have matching patches and npatch > 1" + ) if group_map is None: n_patch = fine_n else: @@ -117,9 +121,7 @@ def treecorr_jackknife_rows(correlations, groups=None, check_rows=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 - ): + 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}" ) @@ -164,12 +166,9 @@ def keep_sky(ra, dec, circles, chunk_size=1_000_000): continue point_dec = np.deg2rad(dec[start:stop][candidate]) point_ra = np.deg2rad(ra[start:stop][candidate]) - cosine = ( - np.sin(point_dec) * np.sin(centre_dec) - + np.cos(point_dec) - * np.cos(centre_dec) - * np.cos(point_ra - centre_ra) - ) + cosine = np.sin(point_dec) * np.sin(centre_dec) + np.cos( + point_dec + ) * np.cos(centre_dec) * np.cos(point_ra - centre_ra) keep[start + candidate[cosine >= cos_radius]] = False return keep @@ -205,8 +204,10 @@ def _profile_fit(x, y, x_rows, y_rows, scales, cross_term, start): cov_xx = jk_factor * np.einsum("nai,nbj->ijab", dx, dx) def residual_covariance(parameters): - return cov_yy + np.einsum("i,iab->ab", parameters, linear) + np.einsum( - "i,j,ijab->ab", parameters, parameters, cov_xx + return ( + cov_yy + + np.einsum("i,iab->ab", parameters, linear) + + np.einsum("i,j,ijab->ab", parameters, parameters, cov_xx) ) def objective(scaled_parameters): @@ -216,9 +217,11 @@ def objective(scaled_parameters): factor = cho_factor(covariance, lower=True, check_finite=False) weighted_residual = cho_solve(factor, residual, check_finite=False) value = float(residual @ weighted_residual) - covariance_derivative = linear + np.einsum( - "j,ijab->iab", parameters, cov_xx - ) + np.einsum("j,jiab->iab", parameters, cov_xx) + covariance_derivative = ( + linear + + np.einsum("j,ijab->iab", parameters, cov_xx) + + np.einsum("j,jiab->iab", parameters, cov_xx) + ) gradient = ( -2 * x.T @ weighted_residual - np.einsum( @@ -294,11 +297,7 @@ def generalized_least_squares(covariance): hessian = (hessian + hessian.T) / 2 eigenvalues = np.linalg.eigvalsh(hessian) parameter_covariance = ( - 2 - * np.linalg.inv(hessian) - * scales[:, None] - * scales[None, :] - / hartlap + 2 * np.linalg.inv(hessian) * scales[:, None] * scales[None, :] / hartlap ) dof = n_data - n_param result = { @@ -313,9 +312,7 @@ def generalized_least_squares(covariance): "residual": residual.tolist(), "prediction": (x @ parameters).tolist(), "residual_sigma": np.sqrt(np.diag(covariance)).tolist(), - "residual_sigma_normalized": ( - residual / np.sqrt(np.diag(covariance)) - ).tolist(), + "residual_sigma_normalized": (residual / np.sqrt(np.diag(covariance))).tolist(), "profile_hessian_eigenvalues": eigenvalues.tolist(), "profile_hessian_positive": bool(eigenvalues.min() > 0), "gradient_norm": float(np.linalg.norm(best.jac)), diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 454c0891..142ba032 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -1,16 +1,18 @@ import gc +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, @@ -398,6 +400,257 @@ def get_rho_tau( 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, sky_mask +): + """The settings that identify a joint leakage measurement. + + ``sky_mask`` is the list of circles applied to stars and galaxies, or None. + """ + config = { + key: value + for key, value in treecorr_config.items() + if key not in ("num_threads", "var_method") + } + return { + "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), + "sky_mask": sky_mask, + } + + +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`) holds the settings: TreeCorr + configuration, ``npatch_fine``, ``patch_counts``, ``beta_psf_size`` and + the sky-mask circles, which ``mask_star``/``mask_gal`` must already apply. + Stars and galaxies are built into TreeCorr catalogues exactly as + `get_rho_tau` builds them (``shear_psf_leakage`` field definitions, + galaxy ellipticities centred on their 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 ``path`` (``.npz``) and its ``.json`` metadata, and 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 + ) + 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"] + 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, + 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:]) + + 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:], + ) + result, _ = leakage_fit.fit_joint_eiv( + products["vector"], rows[:, :n_joint], cross_term=cross_term, **template + ) + return result + + +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 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. + """ + fits = { + int(n): fit_leakage(products, n, beta_psf_size=beta_psf_size) + for n in patch_counts + } + for n, result in fits.items(): + print( + f"EIV leakage fit, {n} patches: chi2 = {result['chi2']:.2f}/" + f"{result['dof']} (PTE {result['pte']:.3g}), parameters " + + ", ".join( + f"{mu:.4g} +/- {sigma:.2g}" + for mu, sigma in zip(result["parameters"], result["sigma"]) + ) + ) + fiducial = fits[int(fiducial_npatch)] + 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, diff --git a/src/sp_validation/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py index ae0bcffc..93e311e1 100644 --- a/src/sp_validation/tests/test_leakage_fit.py +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -14,7 +14,6 @@ xi_psf_sys, ) - TRUE_PARAMETERS = np.array([0.013, 1.28, 0.96, 1.9]) @@ -187,16 +186,16 @@ def coarse(catalogue): 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( + 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) np.testing.assert_allclose(grouped_weights, native_group_weights, rtol=1e-12) def test_spherical_circle_mask_handles_inside_outside_and_ra_wrap(): """Circle excision follows great-circle distance across RA=0 degrees.""" - circles = [ - {"name": "wrap", "ra_deg": 359.8, "dec_deg": 0.0, "radius_deg": 0.5} - ] + circles = [{"name": "wrap", "ra_deg": 359.8, "dec_deg": 0.0, "radius_deg": 0.5}] keep = keep_sky( ra=np.array([359.9, 0.1, 359.0, 180.0]), dec=np.array([0.0, 0.0, 0.0, 0.0]), From 4033e1c18e8bd9fe41200e6e4b46a1e8a401944b Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 09:23:50 -0400 Subject: [PATCH 03/16] cosmo_val: eiv PSF-leakage fit by default, star-contamination sky mask rho_tau_method 'eiv' (now the default) fits the joint rho/tau jackknife measurement and reports alpha, beta, eta with covariance, chi2/dof and PTE at each psf_leakage patch count, drawing downstream samples from the fiducial fit. psf_leakage settings: npatch_fine 800, patch_counts [200, 400], fiducial_npatch 400, beta_psf_size false, and the sky mask config/masks/star_contamination.json (M33, M101, M13, M92), applied to stars and galaxies in every rho/tau measurement. With beta_psf_size, xi_sys includes the u*de terms. Co-Authored-By: Claude Opus 5.5 --- config/masks/star_contamination.json | 33 +++ src/sp_validation/cosmo_val/core.py | 29 ++- .../cosmo_val/psf_systematics.py | 221 +++++++++++++++--- .../test_rho_tau_ignores_hsm_flags.py | 1 + .../tests/test_leakage_measurement.py | 213 +++++++++++++++++ 5 files changed, 462 insertions(+), 35 deletions(-) create mode 100644 config/masks/star_contamination.json create mode 100644 src/sp_validation/tests/test_leakage_measurement.py diff --git a/config/masks/star_contamination.json b/config/masks/star_contamination.json new file mode 100644 index 00000000..18820652 --- /dev/null +++ b/config/masks/star_contamination.json @@ -0,0 +1,33 @@ +{ + "description": "Sky circles excised from PSF stars and galaxies in the rho/tau PSF-leakage measurement. A row is removed when its angular separation from a circle centre (ICRS degrees) is less than radius_deg. M33 contaminates the PSF stars (CosmoStat/sp_validation#452); M101, M13 and M92 are conservative crowding or extended-source excisions.", + "circles": [ + { + "name": "M33", + "ra_deg": 23.462, + "dec_deg": 30.66, + "radius_deg": 1.0, + "reason": "Resolved galaxy contaminating the PSF star sample (CosmoStat/sp_validation#452)" + }, + { + "name": "M101", + "ra_deg": 210.8023, + "dec_deg": 54.3489, + "radius_deg": 0.6, + "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" + }, + { + "name": "M13", + "ra_deg": 250.4235, + "dec_deg": 36.4613, + "radius_deg": 0.5, + "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" + }, + { + "name": "M92", + "ra_deg": 259.2808, + "dec_deg": 43.1365, + "radius_deg": 0.5, + "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" + } + ] +} diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index 5e5576ca..34ed08b8 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,20 @@ 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)), and ``sky_mask`` + (a JSON file of sky circles excised from stars and galaxies in every + rho/tau measurement, relative to the repository root, or None). 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 +296,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 +331,13 @@ 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 {})} + if self.psf_leakage["fiducial_npatch"] not in self.psf_leakage["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/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index ef0e004b..5d4146ff 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -5,6 +5,7 @@ object-wise) for catalogue validation. """ +import json import os from pathlib import Path @@ -18,14 +19,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 ..rho_tau import ( get_rho_tau_w_cov, get_samples, + get_samples_eiv, + leakage_cache_key, + measure_leakage_jackknife, + read_leakage_jackknife, ) from .sacc_writers import rho_tau_to_sacc +REPO_ROOT = Path(__file__).resolve().parents[3] + +# 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, + "sky_mask": "config/masks/star_contamination.json", +} +RHO_TAU_METHODS = ("eiv", "lsq", "emcee", "none") + # TODO: Reorganise the order of functions so it is more readable class PSFSystematicsMixin: @@ -97,16 +115,11 @@ def calculate_rho_tau_stats(self, tomography=True): else: tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] - # Get the selection for stars - mask_star = self._get_star_mask(ver) - for tomo_bin_id in tomo_bin_ids: self.print_cyan(f"Computing for the tomographic bin: {tomo_bin_id}") base_rho = self.basename(ver) base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) - - # Get the selection for galaxies - mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) + mask_star, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) # Compute rho and tau statistics rho_stat_handler, tau_stat_handler = get_rho_tau_w_cov( @@ -370,6 +383,117 @@ def _get_star_mask(self, ver): star_mask = np.ones(len(cat_star), dtype=bool) return star_mask + def _sky_mask_circles(self): + """Circles of the configured ``psf_leakage`` sky mask, or None.""" + path = self.psf_leakage["sky_mask"] + if path is None: + return None + path = Path(path) if Path(path).is_absolute() else REPO_ROOT / path + return json.loads(path.read_text())["circles"] + + def _get_rho_tau_masks(self, ver, tomo_bin_id): + """Star and galaxy selections of every rho/tau measurement. + + Stars pass both HSM flags (`_get_star_mask`), galaxies the tomographic + selection (`_get_galaxy_mask`), and both lie outside the configured + sky-mask circles. + """ + mask_star = self._get_star_mask(ver) + mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) + circles = self._sky_mask_circles() + if circles is not None: + for mask, block, ra_default, dec_default in ( + (mask_star, "psf", "RA", "DEC"), + (mask_gal, "shear", "RA", "Dec"), + ): + entry = self.cc[ver][block] + cat = open_entry(entry) + mask &= leakage_fit.keep_sky( + np.asarray(cat[entry.get("ra_col", ra_default)], dtype=float), + np.asarray(cat[entry.get("dec_col", dec_default)], dtype=float), + circles, + ) + return mask_star, mask_gal + + def leakage_jackknife(self, ver, tomo_bin_id="all"): + """Joint rho/tau jackknife measurement of one version and bin. + + Read from ``rho_tau_stats/leakage_eiv_.npz`` when cached + with the current settings, else measured there by + `rho_tau.measure_leakage_jackknife` with the ``psf_leakage`` settings. + """ + 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 + key = leakage_cache_key( + self.treecorr_config, + settings["npatch_fine"], + settings["patch_counts"], + settings["beta_psf_size"], + self._sky_mask_circles(), + ) + products = None if self.force_run else read_leakage_jackknife(path, key) + if products is None: + mask_star, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) + 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. + """ + 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, + ) + summary = { + str(n): { + key: fit[key] + for key in ( + "parameters", + "sigma", + "parameter_covariance", + "chi2", + "dof", + "pte", + "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} @@ -451,8 +575,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): @@ -462,15 +589,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 @@ -479,16 +605,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(): @@ -506,14 +644,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") diff --git a/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py b/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py index 4562c81a..be5a4b07 100644 --- a/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py +++ b/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py @@ -89,6 +89,7 @@ def __init__(self, cc): self.cov_estimate_method = "sim" self.compute_cov_rho = False self.npatch = 4 + self.psf_leakage = {"sky_mask": None} def basename(self, version, tomo_bin_a="all", **kwargs): return f"{version}_{tomo_bin_a}" 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..ca4ae32e --- /dev/null +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -0,0 +1,213 @@ +"""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" +CIRCLE = {"name": "hole", "ra_deg": 11.0, "dec_deg": 1.0, "radius_deg": 0.3} + + +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", "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() + mask = tmp_path / "mask.json" + mask.write_text(json.dumps({"circles": [CIRCLE]})) + 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_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, + "sky_mask": str(mask), + } + 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 test_sky_mask_removes_stars_and_galaxies_inside_the_circles(tmp_path): + val = _Validation(tmp_path) + mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") + stars = fits.getdata(val.cc[VERSION]["psf"]["path"]) + gals = fits.getdata(val.cc[VERSION]["shear"]["path"]) + + def inside(ra, dec): + return np.hypot((ra - 11.0) * np.cos(np.deg2rad(1.0)), dec - 1.0) < 0.299 + + assert not np.any(mask_star & inside(stars["RA"], stars["DEC"])) + assert not np.any(mask_gal & inside(gals["RA"], gals["Dec"])) + assert not np.any(mask_star & (stars["HSM_FLAG_PSF"] != 0)) + outside_stars = ~inside(stars["RA"], stars["DEC"]) & (stars["HSM_FLAG_PSF"] == 0) + assert mask_star.sum() >= 0.99 * outside_stars.sum() + + +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) + mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") + (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 + 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 + ) + fresh = _Validation.__new__(_Validation) + fresh.__dict__.update( + {k: v for k, v in val.__dict__.items() if not k.startswith("_leakage")} + ) + np.testing.assert_array_equal( + fresh.leakage_jackknife(VERSION)["vector"], products["vector"] + ) From aeb68e3001c55a960b3edf88dda23b792e3ce6ba Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 09:23:52 -0400 Subject: [PATCH 04/16] workflow: cv_leakage_jackknife rule; eiv fits read its measurement The cosmo_val config selects rho_tau_method eiv with the psf_leakage settings, forwarded to every cosmo_val rule and to rho_tau_stats for the sky mask. cv_rho_tau_fits depends on the joint measurement under eiv and stores the per-patch-count fit summary in the PSF product. Co-Authored-By: Claude Opus 5.5 --- papers/cosmo_val/config/config.yaml | 16 +++++++++++- workflow/common.py | 1 + workflow/rules/cosmo_val.smk | 31 ++++++++++++++++++++++-- workflow/rules/twopoint.smk | 2 ++ workflow/scripts/cv_leakage_jackknife.py | 9 +++++++ workflow/scripts/cv_rho_tau_fits.py | 10 +++++++- workflow/scripts/run_rho_tau.py | 1 + 7 files changed, 66 insertions(+), 4 deletions(-) create mode 100644 workflow/scripts/cv_leakage_jackknife.py diff --git a/papers/cosmo_val/config/config.yaml b/papers/cosmo_val/config/config.yaml index 99f553c9..e226b45f 100644 --- a/papers/cosmo_val/config/config.yaml +++ b/papers/cosmo_val/config/config.yaml @@ -51,7 +51,21 @@ 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 + # sky circles excised from stars and galaxies in every rho/tau measurement + sky_mask: config/masks/star_contamination.json # Fiducial scale cut (arcmin) for B-mode PTE extraction. Paper II: [12, 83]. fiducial_scale_cut: [12, 83] 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 fb05f7ef..2ab0f594 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=96000, + 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/rules/twopoint.smk b/workflow/rules/twopoint.smk index ae9bc17d..7f5bc705 100644 --- a/workflow/rules/twopoint.smk +++ b/workflow/rules/twopoint.smk @@ -98,6 +98,8 @@ rule rho_tau_stats: cat_config=CAT_CONFIG, output_dir=str(COSMO_VAL), b_target=config.get("cosmo_val", {}).get("b_target", 0.01), + # Its sky mask applies to the ρ/τ stars and galaxies. + psf_leakage=config.get("cosmo_val", {}).get("psf_leakage"), blind=lambda w: blind_of(w.version), resources: mem_mb=64000, diff --git a/workflow/scripts/cv_leakage_jackknife.py b/workflow/scripts/cv_leakage_jackknife.py new file mode 100644 index 00000000..dd7cb39e --- /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) +verify_outputs(snakemake) diff --git a/workflow/scripts/cv_rho_tau_fits.py b/workflow/scripts/cv_rho_tau_fits.py index bcdbff82..166d7692 100644 --- a/workflow/scripts/cv_rho_tau_fits.py +++ b/workflow/scripts/cv_rho_tau_fits.py @@ -9,9 +9,17 @@ 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": + theta = cv.leakage_jackknife(ver)["theta"] + fit_summary = cv.leakage_jackknife_path(ver, "all", "leakage_eiv_fit") + extra = dict(leakage_eiv_fit=fit_summary.read_text()) + 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/scripts/run_rho_tau.py b/workflow/scripts/run_rho_tau.py index 1a495b7e..9077b6e2 100644 --- a/workflow/scripts/run_rho_tau.py +++ b/workflow/scripts/run_rho_tau.py @@ -20,6 +20,7 @@ catalog_config=params["cat_config"], output_dir=params["output_dir"], cov_estimate_method=params.get("cov_estimate_method", "th"), + psf_leakage=params.get("psf_leakage"), ) cv.calculate_rho_tau_stats(tomography=False) verify_outputs(snakemake) From cbeaf7de49229968ab83af663ba8a1339be2ea6f Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 10:32:28 -0400 Subject: [PATCH 05/16] workflow: cv_leakage_jackknife memory for the full-area measurement The bound-9 measurement (170M galaxies, 800 patches, slope templates) peaked at 83 GiB. Co-Authored-By: Claude Opus 5.5 --- workflow/rules/cosmo_val.smk | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workflow/rules/cosmo_val.smk b/workflow/rules/cosmo_val.smk index 2ab0f594..f3ca1449 100644 --- a/workflow/rules/cosmo_val.smk +++ b/workflow/rules/cosmo_val.smk @@ -273,7 +273,7 @@ rule cv_leakage_jackknife: cv_init=CV_INIT, threads: 8 resources: - mem_mb=96000, + mem_mb=128000, runtime=240, script: "../scripts/cv_leakage_jackknife.py" From 9cbb7eae2d2a3806373566a9ea11914db9d39775 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 11:47:23 -0400 Subject: [PATCH 06/16] rho_tau: key cached rho/tau products on the star and galaxy selection A cached rho_stats/tau_stats file made before the sky mask (or with any other selection) no longer satisfies measurement_matches; the metadata records a digest of the row masks it was measured on. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- src/sp_validation/correlation.py | 37 +++++++++++++---- src/sp_validation/rho_tau.py | 45 +++++++++++++++++---- src/sp_validation/tests/test_correlation.py | 16 ++++++++ 3 files changed, 82 insertions(+), 16 deletions(-) diff --git a/src/sp_validation/correlation.py b/src/sp_validation/correlation.py index 131a81bc..4f5f5306 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. before or after a sky mask, 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,8 +33,8 @@ 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): + """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 @@ -23,6 +42,8 @@ def measurement_matches(path, config): 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 +54,12 @@ 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): + """Record both TreeCorr configurations and the row selection beside a product.""" + record = {"means": "unpatched", "configs": _measurement_configs(config)} + 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/rho_tau.py b/src/sp_validation/rho_tau.py index 142ba032..2b342511 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -16,6 +16,7 @@ from sp_validation.correlation import ( measure_with_patches, measurement_matches, + selection_digest, write_measurement_metadata, ) @@ -31,7 +32,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. @@ -78,7 +79,9 @@ 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 + ) return means @@ -347,7 +350,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 @@ -358,7 +365,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: @@ -375,7 +387,10 @@ 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) + and not force_run + ): print( f"Skipping tau statistics computation, file {tau_path} already exists." ) @@ -394,7 +409,13 @@ 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 @@ -731,12 +752,18 @@ 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 + ) and not force_run ): print(f"Skipping covariance computation, file {tau_cov_path} already exists.") @@ -843,6 +870,7 @@ def chunks(i): rho_filename, save_cov=True, means=rho_means, + selection=rho_selection, ) tau_means = _compute_stats_with_patches( tau_stat_handler, @@ -851,6 +879,7 @@ def chunks(i): tau_filename, save_cov=True, means=tau_means, + selection=tau_selection, ) # Update the keys in the dictionaries 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( From 99095c4fd9e4cab958bb1aaeefabfab4989f867e Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:14:48 -0400 Subject: [PATCH 07/16] leakage_fit: flag unconverged and non-positive-definite fits; V(p) testable The best trial is the lowest objective among the BFGS runs that converged (bfgs_success False when none did); sigma and the parameter covariance are None when the profile Hessian is not positive definite (Cholesky test). residual_covariance_model exposes V(p) and dV/dp, checked against L C L^T of an explicit joint jackknife covariance with a template column, and the noisy fit against a direct minimisation of the dense chi2. The design and constant-beta xi_sys are checked against PSFErrorFit. Unused returns (jackknife weights, residual covariance) are dropped. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- src/sp_validation/leakage_fit.py | 190 ++++++++++++-------- src/sp_validation/rho_tau.py | 7 +- src/sp_validation/tests/test_leakage_fit.py | 148 +++++++++++++-- 3 files changed, 254 insertions(+), 91 deletions(-) diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py index 38263ec8..77f07fa2 100644 --- a/src/sp_validation/leakage_fit.py +++ b/src/sp_validation/leakage_fit.py @@ -40,8 +40,12 @@ def ward_group_map(centres, ngroups): patch index to a contiguous coarse group index. """ centres = np.asarray(centres, dtype=float) - if centres.ndim != 2 or centres.shape[0] < 3: + 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 = ( @@ -58,12 +62,12 @@ 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. The - returned tuple is ``(rows, weights, max_checked_error)``; the weight vector - is diagnostic and matches TreeCorr's patch-weight convention. + 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 = [] - weights = None max_error = 0.0 group_map = None if groups is None else np.asarray(groups, dtype=int) @@ -126,14 +130,11 @@ def treecorr_jackknife_rows(correlations, groups=None, check_rows=True): f"fast jackknife row {patch} differs from TreeCorr by {error}" ) - treecorr_factor = len(corr.getWeight()) / len(corr.weight) - block_weights = denominator.sum(axis=1) * treecorr_factor - weights = block_weights if weights is None else weights + block_weights blocks.append(rows) if not blocks: raise ValueError("at least one TreeCorr correlation is required") - return np.concatenate(blocks, axis=1), weights, max_error + return np.concatenate(blocks, axis=1), max_error def keep_sky(ra, dec, circles, chunk_size=1_000_000): @@ -173,26 +174,18 @@ def keep_sky(ra, dec, circles, chunk_size=1_000_000): return keep -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. +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``/``y_rows`` are delete-one jackknife rows of ``x`` and ``y``. With - jackknife covariances C_xx, C_xy, C_yy built from them, the residual - ``r = y - x p`` has covariance V(p) = C_yy - (B C_xy + C_yx B^T) - + B C_xx B^T, B the linear map p -> x p, and the objective is - r^T V(p)^-1 r (no log-determinant term). + ``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, 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") - + 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) @@ -203,46 +196,70 @@ def _profile_fit(x, y, x_rows, y_rows, scales, cross_term, start): linear = -cov_xy - cov_xy.transpose(0, 2, 1) cov_xx = jk_factor * np.einsum("nai,nbj->ijab", dx, dx) - def residual_covariance(parameters): + 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 - covariance = residual_covariance(parameters) - factor = cho_factor(covariance, lower=True, check_finite=False) + 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) - covariance_derivative = ( - linear - + np.einsum("j,ijab->iab", parameters, cov_xx) - + np.einsum("j,jiab->iab", parameters, cov_xx) - ) gradient = ( -2 * x.T @ weighted_residual - np.einsum( "a,iab,b->i", weighted_residual, - covariance_derivative, + covariance_derivative(parameters), weighted_residual, ) ) * scales return value, gradient - def generalized_least_squares(covariance): - factor = cho_factor(covariance, lower=True, check_finite=False) - weighted_x = cho_solve(factor, x, check_finite=False) - weighted_y = cho_solve(factor, y, check_finite=False) - return np.linalg.solve(x.T @ weighted_x, x.T @ weighted_y) - - fixed_start = generalized_least_squares(cov_yy) - initial = fixed_start if start is None else np.asarray(start, dtype=float) - if initial.shape != (n_param,): - raise ValueError("start must contain one value per fit parameter") + # 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( @@ -252,35 +269,37 @@ def generalized_least_squares(covariance): ] ) if start is not None: - starts.append(initial) + 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)) - trials = [] - fitted = [] - for initial_parameters in starts: - optimization = minimize( + fitted = [ + minimize( objective, - initial_parameters / scales, + initial / scales, jac=True, method="BFGS", options={"gtol": 1e-7, "maxiter": 500}, ) - fitted.append(optimization) - trials.append( - { - "parameters": (optimization.x * scales).tolist(), - "chi2_raw": float(optimization.fun), - "success": bool(optimization.success), - "gradient_norm": float(np.linalg.norm(optimization.jac)), - } - ) - best = min(fitted, key=lambda fit: fit.fun) + 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 - precision_residual = np.linalg.solve(covariance, residual) - chi2_raw = float(residual @ precision_residual) + chi2_raw = float(residual @ np.linalg.solve(covariance, residual)) chi2_value = hartlap * chi2_raw step = 1e-4 @@ -295,15 +314,26 @@ def generalized_least_squares(covariance): ] ) hessian = (hessian + hessian.T) / 2 - eigenvalues = np.linalg.eigvalsh(hessian) - parameter_covariance = ( - 2 * np.linalg.inv(hessian) * scales[:, None] * scales[None, :] / hartlap - ) + 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 - result = { + residual_sigma = np.sqrt(np.diag(covariance)) + return { "parameters": parameters.tolist(), - "sigma": np.sqrt(np.maximum(np.diag(parameter_covariance), 0)).tolist(), - "parameter_covariance": parameter_covariance.tolist(), + "sigma": sigma, + "parameter_covariance": parameter_covariance, "chi2": chi2_value, "chi2_raw": chi2_raw, "dof": dof, @@ -311,16 +341,16 @@ def generalized_least_squares(covariance): "hartlap": hartlap, "residual": residual.tolist(), "prediction": (x @ parameters).tolist(), - "residual_sigma": np.sqrt(np.diag(covariance)).tolist(), - "residual_sigma_normalized": (residual / np.sqrt(np.diag(covariance))).tolist(), - "profile_hessian_eigenvalues": eigenvalues.tolist(), - "profile_hessian_positive": bool(eigenvalues.min() > 0), + "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), } - return result, covariance def fit_joint_eiv( @@ -340,7 +370,7 @@ def fit_joint_eiv( contains the same delete-one measurements, one row per patch. Optional template columns add measured predictors such as ``u * delta_e``; their point values and rows have shapes ``(ndata, ntemplate)`` and - ``(npatch, ndata, ntemplate)``. Returns ``(result, residual_covariance)``. + ``(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 @@ -349,6 +379,12 @@ def fit_joint_eiv( 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) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 2b342511..d264bc62 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -578,11 +578,11 @@ def measure_leakage_jackknife( 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]) + 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( + grouped, _ = leakage_fit.treecorr_jackknife_rows( [corr], groups=groups[n], check_rows=False ) rows[n].append(grouped) @@ -636,10 +636,9 @@ def fit_leakage(products, npatch, beta_psf_size=False, cross_term=True): template_vector=products["template_vector"], template_rows=rows[:, n_joint:], ) - result, _ = leakage_fit.fit_joint_eiv( + return leakage_fit.fit_joint_eiv( products["vector"], rows[:, :n_joint], cross_term=cross_term, **template ) - return result def get_samples_eiv( diff --git a/src/sp_validation/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py index 93e311e1..5b4c41e5 100644 --- a/src/sp_validation/tests/test_leakage_fit.py +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -3,10 +3,13 @@ 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, keep_sky, + residual_covariance_model, rho_tau_design, systematic_correlation_matrix, treecorr_jackknife_rows, @@ -70,7 +73,7 @@ def test_joint_eiv_is_exact_without_noise(slope): vector, rows, template, template_rows, expected = _synthetic_joint( seed=204, slope=slope, noise=False ) - result, _ = fit_joint_eiv( + 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) @@ -88,7 +91,7 @@ def test_joint_eiv_recovers_parameters_within_errors(slope): vector, rows, template, template_rows, expected = _synthetic_joint( seed=seed, slope=slope ) - result, _ = fit_joint_eiv( + result = fit_joint_eiv( vector, rows, template_vector=template, template_rows=template_rows ) pulls.append((np.array(result["parameters"]) - expected) / result["sigma"]) @@ -103,8 +106,8 @@ def test_joint_eiv_recovers_parameters_within_errors(slope): 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) + 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 @@ -152,15 +155,14 @@ def coarse(catalogue): for left, right in ((first, second), (first, first)): fine_corr = treecorr.GGCorrelation(config) fine_corr.process(left, right) - rows, weights, _ = treecorr_jackknife_rows([fine_corr]) - native, native_weights = treecorr.build_multi_cov_design_matrix( + 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) - np.testing.assert_allclose(weights, native_weights, rtol=1e-12) centered = rows - rows.mean(axis=0) covariance = (len(rows) - 1) / len(rows) * centered.T @ centered np.testing.assert_allclose( @@ -177,10 +179,10 @@ def coarse(catalogue): coarse_corr = treecorr.GGCorrelation(config) coarse_corr.process(coarse(left), coarse(right)) - grouped, grouped_weights, _ = treecorr_jackknife_rows( + grouped, _ = treecorr_jackknife_rows( [fine_corr], groups=groups, check_rows=False ) - native_grouped, native_group_weights = treecorr.build_multi_cov_design_matrix( + native_grouped, _ = treecorr.build_multi_cov_design_matrix( [coarse_corr], "jackknife", func=lambda correlations: correlations[0].xip, @@ -190,7 +192,6 @@ def coarse(catalogue): 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) - np.testing.assert_allclose(grouped_weights, native_group_weights, rtol=1e-12) def test_spherical_circle_mask_handles_inside_outside_and_ra_wrap(): @@ -242,3 +243,130 @@ def test_ward_groups_are_contiguous_unions(): 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 + + +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, + ) From 913cc1597d4d333ee87a14900163f6645aab0ac5 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:17:16 -0400 Subject: [PATCH 08/16] cosmo_val: measure the joint leakage jackknife only when asked leakage_jackknife measures only with measure=True, which the rule cv_leakage_jackknife passes; the fits, get_samples, xi_psf_sys and interactive use read the cache and raise, naming the rule, when it is missing or stale. The cache key records the star and star+galaxy selection digests and leaves out cross_patch_weight, which the measurement overrides. The fit summary JSON carries the Hartlap factor, BFGS success, gradient norm and the profile Hessian's positivity and eigenvalues; get_samples_eiv raises when the fiducial Hessian is not positive definite. Tomographic bins draw samples from seeds derived from the bin id. The u de template must have the de catalogue's rows (and weights). The docstring says where galaxy ellipticities are centred. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- .../cosmo_val/psf_systematics.py | 36 +++++++-- src/sp_validation/rho_tau.py | 76 ++++++++++++++----- src/sp_validation/tests/test_leakage_fit.py | 3 + .../tests/test_leakage_measurement.py | 56 ++++++++++++-- workflow/scripts/cv_leakage_jackknife.py | 2 +- 5 files changed, 141 insertions(+), 32 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 5d4146ff..0e045142 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -7,6 +7,7 @@ import json import os +import zlib from pathlib import Path import matplotlib.pyplot as plt @@ -415,27 +416,41 @@ def _get_rho_tau_masks(self, ver, tomo_bin_id): ) return mask_star, mask_gal - def leakage_jackknife(self, ver, tomo_bin_id="all"): + 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`` when cached - with the current settings, else measured there by - `rho_tau.measure_leakage_jackknife` with the ``psf_leakage`` settings. + 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, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) key = leakage_cache_key( self.treecorr_config, settings["npatch_fine"], settings["patch_counts"], settings["beta_psf_size"], self._sky_mask_circles(), + mask_star=mask_star, + mask_gal=mask_gal, ) - products = None if self.force_run else read_leakage_jackknife(path, key) + 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: - mask_star, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) products = measure_leakage_jackknife( self.cc, ver, @@ -460,7 +475,8 @@ 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. + 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: @@ -471,6 +487,7 @@ def _leakage_eiv_fit(self, ver, tomo_bin_id, nsamples): settings["fiducial_npatch"], beta_psf_size=settings["beta_psf_size"], nsamples=nsamples, + seed=zlib.crc32(str(tomo_bin_id).encode()), ) summary = { str(n): { @@ -482,6 +499,11 @@ def _leakage_eiv_fit(self, ver, tomo_bin_id, nsamples): "chi2", "dof", "pte", + "hartlap", + "bfgs_success", + "gradient_norm", + "profile_hessian_positive", + "profile_hessian_eigenvalues", "residual", "residual_sigma", ) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index d264bc62..18d93cd2 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -443,16 +443,26 @@ def get_rho_tau( def leakage_cache_key( - treecorr_config, npatch_fine, patch_counts, beta_psf_size, sky_mask + treecorr_config, + npatch_fine, + patch_counts, + beta_psf_size, + sky_mask, + mask_star=None, + mask_gal=None, ): - """The settings that identify a joint leakage measurement. + """The settings and row selections that identify a joint leakage measurement. - ``sky_mask`` is the list of circles applied to stars and galaxies, or None. + ``sky_mask`` is the list of circles applied to stars and galaxies, or + None; ``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") + if key not in ("num_threads", "var_method", "cross_patch_weight") } return { "treecorr_config": config, @@ -460,6 +470,10 @@ def leakage_cache_key( "patch_counts": sorted(int(n) for n in patch_counts), "beta_psf_size": bool(beta_psf_size), "sky_mask": sky_mask, + "selection": { + "stars": selection_digest(mask_star), + "stars_galaxies": selection_digest(mask_star, mask_gal), + }, } @@ -486,12 +500,14 @@ def measure_leakage_jackknife( ): """Measure the rho/tau vector and its joint delete-one jackknife rows. - ``key`` (from `leakage_cache_key`) holds the settings: TreeCorr - configuration, ``npatch_fine``, ``patch_counts``, ``beta_psf_size`` and - the sky-mask circles, which ``mask_star``/``mask_gal`` must already apply. - Stars and galaxies are built into TreeCorr catalogues exactly as - `get_rho_tau` builds them (``shear_psf_leakage`` field definitions, - galaxy ellipticities centred on their weighted mean), sharing one set of + ``key`` (from `leakage_cache_key` with these masks) holds the settings: + TreeCorr configuration, ``npatch_fine``, ``patch_counts``, + ``beta_psf_size`` and the sky-mask circles, which ``mask_star``/ + ``mask_gal`` must already apply. 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 weighted mean of the full, unmasked shear catalogue before + ``mask_gal`` selects rows), 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 @@ -550,6 +566,11 @@ def measure_leakage_jackknife( 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, @@ -557,6 +578,7 @@ def measure_leakage_jackknife( dec_units="rad", g1=u * de.g1, g2=u * de.g2, + w=de.w, patch_centers=centres, ) del size, u @@ -646,24 +668,42 @@ def get_samples_eiv( ): """Joint EIV fits at each patch count and Gaussian samples at the fiducial. - Returns ``(samples, fits, quantiles)``: ``nsamples`` draws 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. + 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 " - + ", ".join( - f"{mu:.4g} +/- {sigma:.2g}" - for mu, sigma in zip(result["parameters"], result["sigma"]) - ) + 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 ) diff --git a/src/sp_validation/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py index 5b4c41e5..44c7d63f 100644 --- a/src/sp_validation/tests/test_leakage_fit.py +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -16,6 +16,7 @@ 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]) @@ -337,6 +338,8 @@ def test_unconstrained_parameter_is_flagged_not_given_an_error(): 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}, [200], 200) def _production_fitter(tmp_path, rho): diff --git a/src/sp_validation/tests/test_leakage_measurement.py b/src/sp_validation/tests/test_leakage_measurement.py index ca4ae32e..31c03bda 100644 --- a/src/sp_validation/tests/test_leakage_measurement.py +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -115,6 +115,16 @@ def _output_path(self, *parts): return os.path.join(self.output, *parts) +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_sky_mask_removes_stars_and_galaxies_inside_the_circles(tmp_path): val = _Validation(tmp_path) mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") @@ -134,7 +144,7 @@ def inside(ra, dec): 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) + products = val.leakage_jackknife(VERSION, measure=True) mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") (tmp_path / "production").mkdir() rho_handler, tau_handler = rho_tau.get_rho_tau( @@ -167,6 +177,7 @@ def test_joint_vector_equals_production_rho_tau(tmp_path): 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) @@ -204,10 +215,43 @@ def fail(*args, **kwargs): monkeypatch.setattr( "sp_validation.cosmo_val.psf_systematics.measure_leakage_jackknife", fail ) - fresh = _Validation.__new__(_Validation) - fresh.__dict__.update( - {k: v for k, v in val.__dict__.items() if not k.startswith("_leakage")} - ) np.testing.assert_array_equal( - fresh.leakage_jackknife(VERSION)["vector"], products["vector"] + _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", ["sky_mask", "beta_psf_size", "selection"]) +def test_cache_is_stale_after_a_change(tmp_path, change): + """Sky mask, PSF-size slope and star selection each invalidate the cache.""" + val = _Validation(tmp_path) + val.leakage_jackknife(VERSION, measure=True) + if change == "sky_mask": + mask = tmp_path / "wider.json" + mask.write_text(json.dumps({"circles": [{**CIRCLE, "radius_deg": 0.4}]})) + fresh = _fresh(val, sky_mask=str(mask)) + elif 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): + # A strip far from the masked circle, 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 + diff --git a/workflow/scripts/cv_leakage_jackknife.py b/workflow/scripts/cv_leakage_jackknife.py index dd7cb39e..3243f7db 100644 --- a/workflow/scripts/cv_leakage_jackknife.py +++ b/workflow/scripts/cv_leakage_jackknife.py @@ -5,5 +5,5 @@ _unbuffer_streams() cv = make_cv(snakemake) for ver in cv.versions: - cv.leakage_jackknife(ver) + cv.leakage_jackknife(ver, measure=True) verify_outputs(snakemake) From 9eb7839bed43baf981d7bdb050f30232310d35f9 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:17:17 -0400 Subject: [PATCH 09/16] cosmo_val: validate psf_leakage patch counts Every patch count must lie strictly between 1 and npatch_fine, and the fiducial count must be one of them. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- src/sp_validation/cosmo_val/core.py | 9 ++++++++- .../tests/test_leakage_measurement.py | 14 ++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index 34ed08b8..c3d694cd 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -336,7 +336,14 @@ def __init__( if unknown: raise ValueError(f"Unknown psf_leakage settings: {sorted(unknown)}") self.psf_leakage = {**PSF_LEAKAGE_DEFAULTS, **(psf_leakage or {})} - if self.psf_leakage["fiducial_npatch"] not in self.psf_leakage["patch_counts"]: + 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 diff --git a/src/sp_validation/tests/test_leakage_measurement.py b/src/sp_validation/tests/test_leakage_measurement.py index 31c03bda..8a46e9cd 100644 --- a/src/sp_validation/tests/test_leakage_measurement.py +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -255,3 +255,17 @@ def fewer_stars(self, ver): 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) From 3d6b8e264aefbc18170d1d273aa9dd0736429d96 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:23:18 -0400 Subject: [PATCH 10/16] rho_tau: theory covariance and fit samples follow the row selection get_theory_cov records the star/galaxy selection_digest beside cov_tau_*_th.npy and reuses it only for the same selection. The lsq and emcee sample caches carry the selection recorded beside the tau statistics in their names, so remeasured statistics on other rows are refitted. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- src/sp_validation/rho_tau.py | 34 +++++++++++++++--- .../test_lsq_sample_cache_ignores_cov_type.py | 33 +++++++++++++++++ .../tests/test_leakage_measurement.py | 36 +++++++++++++++++++ 3 files changed, 99 insertions(+), 4 deletions(-) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 18d93cd2..52ada6d1 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -728,7 +728,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]) @@ -736,8 +737,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 @@ -764,6 +771,7 @@ 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() @@ -953,6 +961,24 @@ 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 + selection = json.loads(metadata.read_text()).get("selection") + return name if selection is None else f"{name}_selection={selection[:12]}" + + def get_samples( psf_fitter, base_rho, @@ -1041,7 +1067,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.") @@ -1093,7 +1119,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/test_leakage_measurement.py b/src/sp_validation/tests/test_leakage_measurement.py index 8a46e9cd..2fdcf6e1 100644 --- a/src/sp_validation/tests/test_leakage_measurement.py +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -269,3 +269,39 @@ def test_psf_leakage_settings_are_validated(settings, message): 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 = val._get_rho_tau_masks(VERSION, "all") + 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 From 106b6cbb593864b04ea7c832b5c6d8c410c012e0 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:25:03 -0400 Subject: [PATCH 11/16] PSF-leakage plots follow the parameter count of the fit With beta_psf_size the samples are (alpha, beta0, eta, beta1): the getdist contours (cosmo_val and the cv_plot_rho_tau_fits renderer) and the dashboard's coefficient draws take names, labels and markers from the sample width, and the dashboard tau model adds beta1 times the template, which cv_rho_tau_fits now stores as tau_slope_template (it raises without it). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- .../cosmo_val/diagnostic_plots.py | 27 ++++++++++++--- .../cosmo_val/psf_systematics.py | 8 +++-- src/sp_validation/dashboard/psf.py | 33 +++++++++++++++---- src/sp_validation/tests/test_dashboard_psf.py | 30 +++++++++++++++++ workflow/scripts/cv_rho_tau_fits.py | 6 +++- 5 files changed, 90 insertions(+), 14 deletions(-) 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 0e045142..4d54b5fc 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -30,6 +30,7 @@ measure_leakage_jackknife, read_leakage_jackknife, ) +from .diagnostic_plots import leakage_contour_axes from .sacc_writers import rho_tau_to_sacc REPO_ROOT = Path(__file__).resolve().parents[3] @@ -1239,15 +1240,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/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/workflow/scripts/cv_rho_tau_fits.py b/workflow/scripts/cv_rho_tau_fits.py index 166d7692..abf17ef4 100644 --- a/workflow/scripts/cv_rho_tau_fits.py +++ b/workflow/scripts/cv_rho_tau_fits.py @@ -10,9 +10,13 @@ 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": - theta = cv.leakage_jackknife(ver)["theta"] + 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 = {} From ef70a7293e21247bb7a341cae23cf09c1d0b036d Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:25:56 -0400 Subject: [PATCH 12/16] workflow: sky-mask file is an input of the rho/tau rules rho_tau_stats and cv_leakage_jackknife list the psf_leakage sky-mask JSON as input, so editing its circles reruns them, and rho_tau_stats receives only the sky-mask setting rather than the whole psf_leakage dict. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- .../tests/test_cv_init_params.py | 14 +++++++++++ workflow/common.py | 23 +++++++++++++++++++ workflow/rules/cosmo_val.smk | 1 + workflow/rules/twopoint.smk | 6 +++-- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/src/sp_validation/tests/test_cv_init_params.py b/src/sp_validation/tests/test_cv_init_params.py index 6af38de3..5d17215a 100644 --- a/src/sp_validation/tests/test_cv_init_params.py +++ b/src/sp_validation/tests/test_cv_init_params.py @@ -75,3 +75,17 @@ def test_cv_basename_is_the_class_basename(): assert _load_common().cv_basename( "SP_v1.4.6.3", fiducial ) == CosmologyValidation.basename(cv, "SP_v1.4.6.3") + + +def test_workflow_sky_mask_default_is_the_class_default(): + """Rules list the sky-mask file as input; unset, it is the class default.""" + from sp_validation.cosmo_val.psf_systematics import PSF_LEAKAGE_DEFAULTS + + common = _load_common() + assert common.PSF_LEAKAGE_SKY_MASK == PSF_LEAKAGE_DEFAULTS["sky_mask"] + assert common.psf_leakage_sky_mask({}) == PSF_LEAKAGE_DEFAULTS["sky_mask"] + (path,) = common.psf_leakage_sky_mask_input({}) + assert Path(path) == REPO / PSF_LEAKAGE_DEFAULTS["sky_mask"] + assert Path(path).exists() + off = {"cosmo_val": {"psf_leakage": {"sky_mask": None}}} + assert common.psf_leakage_sky_mask_input(off) == [] diff --git a/workflow/common.py b/workflow/common.py index 47a9d8ec..33e70223 100644 --- a/workflow/common.py +++ b/workflow/common.py @@ -472,6 +472,29 @@ def cv_cov_tau(version, fiducial=None): return str(COSMO_VAL / "rho_tau_stats" / f"cov_tau_{base}_th.npy") +# CosmologyValidation's default ``psf_leakage`` sky mask (PSF_LEAKAGE_DEFAULTS +# in sp_validation.cosmo_val.psf_systematics; test_cv_init_params checks they agree). +PSF_LEAKAGE_SKY_MASK = "config/masks/star_contamination.json" + + +def psf_leakage_sky_mask(config): + """The ρ/τ sky-mask setting: ``cosmo_val.psf_leakage.sky_mask`` or the default. + + A JSON file of circles relative to the repository root, or None for none. + """ + leakage = (config.get("cosmo_val") or {}).get("psf_leakage") or {} + return leakage.get("sky_mask", PSF_LEAKAGE_SKY_MASK) + + +def psf_leakage_sky_mask_input(config): + """The sky-mask file as rule input, so editing its circles reruns the rule.""" + path = psf_leakage_sky_mask(config) + if path is None: + return [] + path = Path(path) + return [str(path if path.is_absolute() else REPO_ROOT / path)] + + # CosmologyValidation constructor kwargs read from config["cosmo_val"]. Every # keyword with a default is either here or explicitly exempted in # src/sp_validation/tests/test_cv_init_params.py, so no default applies silently. diff --git a/workflow/rules/cosmo_val.smk b/workflow/rules/cosmo_val.smk index f3ca1449..a4b89873 100644 --- a/workflow/rules/cosmo_val.smk +++ b/workflow/rules/cosmo_val.smk @@ -267,6 +267,7 @@ rule cv_leakage_jackknife: catalogues=cv_shear_catalogues() + [str(Path(catalogue_entry(v)["subdir"]) / catalogue_entry(v)["psf"]["path"]) for v in CV_VERSIONS], + sky_mask=psf_leakage_sky_mask_input(config), output: measurements=CV_LEAKAGE_JACKKNIFE, params: diff --git a/workflow/rules/twopoint.smk b/workflow/rules/twopoint.smk index 7f5bc705..ee2c59c2 100644 --- a/workflow/rules/twopoint.smk +++ b/workflow/rules/twopoint.smk @@ -80,6 +80,8 @@ rule xi_onecov: rule rho_tau_stats: + input: + sky_mask=psf_leakage_sky_mask_input(config), output: rho_stats=str(COSMO_VAL / "rho_tau_stats/rho_stats_{version}_tomo_bin_all_minsep={min_sep}_maxsep={max_sep}_nbins={nbins}_npatch={npatch}.fits"), tau_stats=str(COSMO_VAL / "rho_tau_stats/tau_stats_{version}_tomo_bin_all_minsep={min_sep}_maxsep={max_sep}_nbins={nbins}_npatch={npatch}.fits"), @@ -98,8 +100,8 @@ rule rho_tau_stats: cat_config=CAT_CONFIG, output_dir=str(COSMO_VAL), b_target=config.get("cosmo_val", {}).get("b_target", 0.01), - # Its sky mask applies to the ρ/τ stars and galaxies. - psf_leakage=config.get("cosmo_val", {}).get("psf_leakage"), + # The psf_leakage sky mask applies to the ρ/τ stars and galaxies. + psf_leakage={"sky_mask": psf_leakage_sky_mask(config)}, blind=lambda w: blind_of(w.version), resources: mem_mb=64000, From af40b4859e5b58beef426e20cbfa5d4b2e3ee9a6 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Fri, 30 Oct 2026 12:41:16 -0400 Subject: [PATCH 13/16] workflow tests: the toy checkout carries star catalogues and the sky mask cv_leakage_jackknife lists each version's PSF catalogue and the rho/tau rules the sky-mask file as inputs; the toy catalogue config had no psf block and the toy checkout no config/masks, so every DAG test failed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9EfMQEn7kKMQzWV5o2 --- workflow/tests/conftest.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/workflow/tests/conftest.py b/workflow/tests/conftest.py index af4e4258..139ca41b 100644 --- a/workflow/tests/conftest.py +++ b/workflow/tests/conftest.py @@ -9,9 +9,10 @@ parse with the standard library and Snakemake alone -- the condition a host Snakemake is in. The ``candide`` tests need candide itself; CI deselects them. -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 +The ``toy`` fixture is a disposable checkout: copies of ``workflow/``, +``papers/cosmo_val/`` and ``config/masks/``, this checkout's ``src/`` +symlinked in, a blinded and an 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, } @@ -175,6 +178,7 @@ def toy(tmp_path_factory): shutil.copytree( REPO / "papers" / "cosmo_val", root / "papers" / "cosmo_val", ignore=skip ) + shutil.copytree(REPO / "config" / "masks", root / "config" / "masks") (root / "src").symlink_to(REPO / "src") (root / "data").mkdir() From 0661334beec27505b30e313a4a8afeecfdded103 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Sat, 31 Oct 2026 00:05:24 +0100 Subject: [PATCH 14/16] rho/tau: drop the star-contamination sky mask Masking of contaminated PSF stars belongs upstream in ShapePipe, so the rho/tau measurements here use only the star (HSM flags) and galaxy (tomographic bin) selections. Removes config/masks/star_contamination.json, the psf_leakage sky_mask setting, leakage_fit.keep_sky, the workflow sky-mask inputs and params, and the toy checkout's mask copy. The selection-keyed caching of rho/tau products stays. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JK7o9H5tLNFzcwRnYwMFEA --- config/masks/star_contamination.json | 33 ------------- papers/cosmo_val/config/config.yaml | 2 - src/sp_validation/correlation.py | 2 +- src/sp_validation/cosmo_val/core.py | 4 +- .../cosmo_val/psf_systematics.py | 46 ++++--------------- src/sp_validation/leakage_fit.py | 39 +--------------- src/sp_validation/rho_tau.py | 18 +++----- .../test_rho_tau_ignores_hsm_flags.py | 1 - .../tests/test_cv_init_params.py | 14 ------ src/sp_validation/tests/test_leakage_fit.py | 17 ------- .../tests/test_leakage_measurement.py | 41 +++++------------ workflow/common.py | 23 ---------- workflow/rules/cosmo_val.smk | 1 - workflow/rules/twopoint.smk | 4 -- workflow/scripts/run_rho_tau.py | 1 - workflow/tests/conftest.py | 11 ++--- 16 files changed, 34 insertions(+), 223 deletions(-) delete mode 100644 config/masks/star_contamination.json diff --git a/config/masks/star_contamination.json b/config/masks/star_contamination.json deleted file mode 100644 index 18820652..00000000 --- a/config/masks/star_contamination.json +++ /dev/null @@ -1,33 +0,0 @@ -{ - "description": "Sky circles excised from PSF stars and galaxies in the rho/tau PSF-leakage measurement. A row is removed when its angular separation from a circle centre (ICRS degrees) is less than radius_deg. M33 contaminates the PSF stars (CosmoStat/sp_validation#452); M101, M13 and M92 are conservative crowding or extended-source excisions.", - "circles": [ - { - "name": "M33", - "ra_deg": 23.462, - "dec_deg": 30.66, - "radius_deg": 1.0, - "reason": "Resolved galaxy contaminating the PSF star sample (CosmoStat/sp_validation#452)" - }, - { - "name": "M101", - "ra_deg": 210.8023, - "dec_deg": 54.3489, - "radius_deg": 0.6, - "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" - }, - { - "name": "M13", - "ra_deg": 250.4235, - "dec_deg": 36.4613, - "radius_deg": 0.5, - "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" - }, - { - "name": "M92", - "ra_deg": 259.2808, - "dec_deg": 43.1365, - "radius_deg": 0.5, - "reason": "Known extended galaxy or dense globular cluster with positive size-residual, ellipticity-residual or density outlier" - } - ] -} diff --git a/papers/cosmo_val/config/config.yaml b/papers/cosmo_val/config/config.yaml index e226b45f..084509ce 100644 --- a/papers/cosmo_val/config/config.yaml +++ b/papers/cosmo_val/config/config.yaml @@ -64,8 +64,6 @@ cosmo_val: fiducial_npatch: 400 # beta(T_PSF) = beta0 + beta1 (T_PSF / - 1) beta_psf_size: false - # sky circles excised from stars and galaxies in every rho/tau measurement - sky_mask: config/masks/star_contamination.json # Fiducial scale cut (arcmin) for B-mode PTE extraction. Paper II: [12, 83]. fiducial_scale_cut: [12, 83] diff --git a/src/sp_validation/correlation.py b/src/sp_validation/correlation.py index 4f5f5306..97811a2b 100644 --- a/src/sp_validation/correlation.py +++ b/src/sp_validation/correlation.py @@ -12,7 +12,7 @@ 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. before or after a sky mask, is never reused. + e.g. another star flag or tomographic bin, is never reused. """ digest = hashlib.sha1() for mask in masks: diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index c3d694cd..7c9fff12 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -139,9 +139,7 @@ class CosmologyValidation( 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)), and ``sky_mask`` - (a JSON file of sky circles excised from stars and galaxies in every - rho/tau measurement, relative to the repository root, or None). + 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. diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 4d54b5fc..f5de041e 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -33,8 +33,6 @@ from .diagnostic_plots import leakage_contour_axes from .sacc_writers import rho_tau_to_sacc -REPO_ROOT = Path(__file__).resolve().parents[3] - # 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 = { @@ -42,7 +40,6 @@ "patch_counts": [200, 400], "fiducial_npatch": 400, "beta_psf_size": False, - "sky_mask": "config/masks/star_contamination.json", } RHO_TAU_METHODS = ("eiv", "lsq", "emcee", "none") @@ -117,11 +114,16 @@ def calculate_rho_tau_stats(self, tomography=True): else: tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] + # Get the selection for stars + mask_star = self._get_star_mask(ver) + for tomo_bin_id in tomo_bin_ids: self.print_cyan(f"Computing for the tomographic bin: {tomo_bin_id}") base_rho = self.basename(ver) base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) - mask_star, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) + + # Get the selection for galaxies + mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) # Compute rho and tau statistics rho_stat_handler, tau_stat_handler = get_rho_tau_w_cov( @@ -385,38 +387,6 @@ def _get_star_mask(self, ver): star_mask = np.ones(len(cat_star), dtype=bool) return star_mask - def _sky_mask_circles(self): - """Circles of the configured ``psf_leakage`` sky mask, or None.""" - path = self.psf_leakage["sky_mask"] - if path is None: - return None - path = Path(path) if Path(path).is_absolute() else REPO_ROOT / path - return json.loads(path.read_text())["circles"] - - def _get_rho_tau_masks(self, ver, tomo_bin_id): - """Star and galaxy selections of every rho/tau measurement. - - Stars pass both HSM flags (`_get_star_mask`), galaxies the tomographic - selection (`_get_galaxy_mask`), and both lie outside the configured - sky-mask circles. - """ - mask_star = self._get_star_mask(ver) - mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) - circles = self._sky_mask_circles() - if circles is not None: - for mask, block, ra_default, dec_default in ( - (mask_star, "psf", "RA", "DEC"), - (mask_gal, "shear", "RA", "Dec"), - ): - entry = self.cc[ver][block] - cat = open_entry(entry) - mask &= leakage_fit.keep_sky( - np.asarray(cat[entry.get("ra_col", ra_default)], dtype=float), - np.asarray(cat[entry.get("dec_col", dec_default)], dtype=float), - circles, - ) - return mask_star, mask_gal - def leakage_jackknife(self, ver, tomo_bin_id="all", measure=False): """Joint rho/tau jackknife measurement of one version and bin. @@ -431,13 +401,13 @@ def leakage_jackknife(self, ver, tomo_bin_id="all", measure=False): if (ver, tomo_bin_id) not in cache: path = self.leakage_jackknife_path(ver, tomo_bin_id) settings = self.psf_leakage - mask_star, mask_gal = self._get_rho_tau_masks(ver, tomo_bin_id) + 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"], - self._sky_mask_circles(), mask_star=mask_star, mask_gal=mask_gal, ) diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py index 77f07fa2..aff62367 100644 --- a/src/sp_validation/leakage_fit.py +++ b/src/sp_validation/leakage_fit.py @@ -2,7 +2,7 @@ This module has no catalogue I/O. It builds the production rho-to-tau design, profiles noisy templates from a joint delete-one covariance, and provides -TreeCorr jackknife, patch-grouping, and spherical-mask helpers. +TreeCorr jackknife and patch-grouping helpers. """ import numpy as np @@ -137,43 +137,6 @@ def treecorr_jackknife_rows(correlations, groups=None, check_rows=True): return np.concatenate(blocks, axis=1), max_error -def keep_sky(ra, dec, circles, chunk_size=1_000_000): - """Return a boolean mask outside all spherical circles. - - ``ra`` and ``dec`` are in degrees. Each circle is a mapping with - ``ra_deg``, ``dec_deg`` and ``radius_deg`` fields. Angular separation is - computed on the sphere, including across the RA 0/360-degree boundary. - """ - ra = np.asarray(ra, dtype=float) - dec = np.asarray(dec, dtype=float) - if ra.shape != dec.shape or ra.ndim != 1: - raise ValueError("ra and dec must be one-dimensional arrays of equal length") - circles = tuple(circles) - keep = np.ones(ra.size, dtype=bool) - for circle in circles: - radius = float(circle["radius_deg"]) - if not 0 < radius < 180: - raise ValueError("circle radius must be between 0 and 180 degrees") - centre_ra = np.deg2rad(float(circle["ra_deg"])) - centre_dec_deg = float(circle["dec_deg"]) - centre_dec = np.deg2rad(centre_dec_deg) - cos_radius = np.cos(np.deg2rad(radius)) - for start in range(0, ra.size, chunk_size): - stop = min(start + chunk_size, ra.size) - candidate = np.flatnonzero( - np.abs(dec[start:stop] - centre_dec_deg) <= radius - ) - if not candidate.size: - continue - point_dec = np.deg2rad(dec[start:stop][candidate]) - point_ra = np.deg2rad(ra[start:stop][candidate]) - cosine = np.sin(point_dec) * np.sin(centre_dec) + np.cos( - point_dec - ) * np.cos(centre_dec) * np.cos(point_ra - centre_ra) - keep[start + candidate[cosine >= cos_radius]] = False - return keep - - 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``. diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 52ada6d1..a589a512 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -447,14 +447,12 @@ def leakage_cache_key( npatch_fine, patch_counts, beta_psf_size, - sky_mask, mask_star=None, mask_gal=None, ): """The settings and row selections that identify a joint leakage measurement. - ``sky_mask`` is the list of circles applied to stars and galaxies, or - None; ``mask_star``/``mask_gal`` are the star and galaxy selections the + ``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. @@ -469,7 +467,6 @@ def leakage_cache_key( "npatch_fine": int(npatch_fine), "patch_counts": sorted(int(n) for n in patch_counts), "beta_psf_size": bool(beta_psf_size), - "sky_mask": sky_mask, "selection": { "stars": selection_digest(mask_star), "stars_galaxies": selection_digest(mask_star, mask_gal), @@ -502,13 +499,12 @@ def measure_leakage_jackknife( ``key`` (from `leakage_cache_key` with these masks) holds the settings: TreeCorr configuration, ``npatch_fine``, ``patch_counts``, - ``beta_psf_size`` and the sky-mask circles, which ``mask_star``/ - ``mask_gal`` must already apply. 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 weighted mean of the full, unmasked shear catalogue before - ``mask_gal`` selects rows), sharing one set of - ``npatch_fine`` k-means patch centres found on the selected stars. One + ``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 weighted mean of the full, + unmasked shear catalogue before ``mask_gal`` selects rows), 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) diff --git a/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py b/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py index be5a4b07..4562c81a 100644 --- a/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py +++ b/src/sp_validation/tests/regression/test_rho_tau_ignores_hsm_flags.py @@ -89,7 +89,6 @@ def __init__(self, cc): self.cov_estimate_method = "sim" self.compute_cov_rho = False self.npatch = 4 - self.psf_leakage = {"sky_mask": None} def basename(self, version, tomo_bin_a="all", **kwargs): return f"{version}_{tomo_bin_a}" diff --git a/src/sp_validation/tests/test_cv_init_params.py b/src/sp_validation/tests/test_cv_init_params.py index 5d17215a..6af38de3 100644 --- a/src/sp_validation/tests/test_cv_init_params.py +++ b/src/sp_validation/tests/test_cv_init_params.py @@ -75,17 +75,3 @@ def test_cv_basename_is_the_class_basename(): assert _load_common().cv_basename( "SP_v1.4.6.3", fiducial ) == CosmologyValidation.basename(cv, "SP_v1.4.6.3") - - -def test_workflow_sky_mask_default_is_the_class_default(): - """Rules list the sky-mask file as input; unset, it is the class default.""" - from sp_validation.cosmo_val.psf_systematics import PSF_LEAKAGE_DEFAULTS - - common = _load_common() - assert common.PSF_LEAKAGE_SKY_MASK == PSF_LEAKAGE_DEFAULTS["sky_mask"] - assert common.psf_leakage_sky_mask({}) == PSF_LEAKAGE_DEFAULTS["sky_mask"] - (path,) = common.psf_leakage_sky_mask_input({}) - assert Path(path) == REPO / PSF_LEAKAGE_DEFAULTS["sky_mask"] - assert Path(path).exists() - off = {"cosmo_val": {"psf_leakage": {"sky_mask": None}}} - assert common.psf_leakage_sky_mask_input(off) == [] diff --git a/src/sp_validation/tests/test_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py index 44c7d63f..e25ebfc1 100644 --- a/src/sp_validation/tests/test_leakage_fit.py +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -8,7 +8,6 @@ from sp_validation.leakage_fit import ( fit_joint_eiv, - keep_sky, residual_covariance_model, rho_tau_design, systematic_correlation_matrix, @@ -195,22 +194,6 @@ def coarse(catalogue): np.testing.assert_allclose(grouped, native_grouped, rtol=1e-11, atol=1e-14) -def test_spherical_circle_mask_handles_inside_outside_and_ra_wrap(): - """Circle excision follows great-circle distance across RA=0 degrees.""" - circles = [{"name": "wrap", "ra_deg": 359.8, "dec_deg": 0.0, "radius_deg": 0.5}] - keep = keep_sky( - ra=np.array([359.9, 0.1, 359.0, 180.0]), - dec=np.array([0.0, 0.0, 0.0, 0.0]), - circles=circles, - ) - np.testing.assert_array_equal(keep, [False, False, True, True]) - - -def test_sky_mask_rejects_mismatched_coordinates(): - with pytest.raises(ValueError, match="equal length"): - keep_sky([1.0], [1.0, 2.0], []) - - 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) diff --git a/src/sp_validation/tests/test_leakage_measurement.py b/src/sp_validation/tests/test_leakage_measurement.py index 2fdcf6e1..c6e3add1 100644 --- a/src/sp_validation/tests/test_leakage_measurement.py +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -11,7 +11,6 @@ from sp_validation.cosmo_val.psf_systematics import PSFSystematicsMixin VERSION = "SP_synthetic" -CIRCLE = {"name": "hole", "ra_deg": 11.0, "dec_deg": 1.0, "radius_deg": 0.3} def _write_catalogues(tmp_path, n_star=4000, n_gal=6000): @@ -56,8 +55,6 @@ class _Validation(PSFSystematicsMixin): def __init__(self, tmp_path, beta_psf_size=False): psf, shear = _write_catalogues(tmp_path) (tmp_path / "rho_tau_stats").mkdir() - mask = tmp_path / "mask.json" - mask.write_text(json.dumps({"circles": [CIRCLE]})) self.output = tmp_path self.cc = { "paths": {"output": str(tmp_path)}, @@ -103,7 +100,6 @@ def __init__(self, tmp_path, beta_psf_size=False): "patch_counts": [15, 20], "fiducial_npatch": 20, "beta_psf_size": beta_psf_size, - "sky_mask": str(mask), } self.rho_tau_method = "eiv" self.force_run = False @@ -115,6 +111,11 @@ 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) @@ -125,27 +126,11 @@ def _fresh(val, **psf_leakage): return fresh -def test_sky_mask_removes_stars_and_galaxies_inside_the_circles(tmp_path): - val = _Validation(tmp_path) - mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") - stars = fits.getdata(val.cc[VERSION]["psf"]["path"]) - gals = fits.getdata(val.cc[VERSION]["shear"]["path"]) - - def inside(ra, dec): - return np.hypot((ra - 11.0) * np.cos(np.deg2rad(1.0)), dec - 1.0) < 0.299 - - assert not np.any(mask_star & inside(stars["RA"], stars["DEC"])) - assert not np.any(mask_gal & inside(gals["RA"], gals["Dec"])) - assert not np.any(mask_star & (stars["HSM_FLAG_PSF"] != 0)) - outside_stars = ~inside(stars["RA"], stars["DEC"]) & (stars["HSM_FLAG_PSF"] == 0) - assert mask_star.sum() >= 0.99 * outside_stars.sum() - - 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 = val._get_rho_tau_masks(VERSION, "all") + mask_star, mask_gal = _masks(val) (tmp_path / "production").mkdir() rho_handler, tau_handler = rho_tau.get_rho_tau( val.cc, @@ -230,23 +215,19 @@ def test_reading_without_a_measurement_names_the_rule(tmp_path): assert not val.leakage_jackknife_path(VERSION).exists() -@pytest.mark.parametrize("change", ["sky_mask", "beta_psf_size", "selection"]) +@pytest.mark.parametrize("change", ["beta_psf_size", "selection"]) def test_cache_is_stale_after_a_change(tmp_path, change): - """Sky mask, PSF-size slope and star selection each invalidate the cache.""" + """PSF-size slope and star selection each invalidate the cache.""" val = _Validation(tmp_path) val.leakage_jackknife(VERSION, measure=True) - if change == "sky_mask": - mask = tmp_path / "wider.json" - mask.write_text(json.dumps({"circles": [{**CIRCLE, "radius_deg": 0.4}]})) - fresh = _fresh(val, sky_mask=str(mask)) - elif change == "beta_psf_size": + 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): - # A strip far from the masked circle, so the selection changes. + # 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) @@ -284,7 +265,7 @@ def build_cov(self, **kwargs): return np.eye(2) * len(built) monkeypatch.setattr(rho_tau, "CovTauTh", FakeCovTauTh) - mask_star, mask_gal = val._get_rho_tau_masks(VERSION, "all") + mask_star, mask_gal = _masks(val) outdir = tmp_path / "rho_tau_stats" def covariance(star, gal): diff --git a/workflow/common.py b/workflow/common.py index 33e70223..47a9d8ec 100644 --- a/workflow/common.py +++ b/workflow/common.py @@ -472,29 +472,6 @@ def cv_cov_tau(version, fiducial=None): return str(COSMO_VAL / "rho_tau_stats" / f"cov_tau_{base}_th.npy") -# CosmologyValidation's default ``psf_leakage`` sky mask (PSF_LEAKAGE_DEFAULTS -# in sp_validation.cosmo_val.psf_systematics; test_cv_init_params checks they agree). -PSF_LEAKAGE_SKY_MASK = "config/masks/star_contamination.json" - - -def psf_leakage_sky_mask(config): - """The ρ/τ sky-mask setting: ``cosmo_val.psf_leakage.sky_mask`` or the default. - - A JSON file of circles relative to the repository root, or None for none. - """ - leakage = (config.get("cosmo_val") or {}).get("psf_leakage") or {} - return leakage.get("sky_mask", PSF_LEAKAGE_SKY_MASK) - - -def psf_leakage_sky_mask_input(config): - """The sky-mask file as rule input, so editing its circles reruns the rule.""" - path = psf_leakage_sky_mask(config) - if path is None: - return [] - path = Path(path) - return [str(path if path.is_absolute() else REPO_ROOT / path)] - - # CosmologyValidation constructor kwargs read from config["cosmo_val"]. Every # keyword with a default is either here or explicitly exempted in # src/sp_validation/tests/test_cv_init_params.py, so no default applies silently. diff --git a/workflow/rules/cosmo_val.smk b/workflow/rules/cosmo_val.smk index a4b89873..f3ca1449 100644 --- a/workflow/rules/cosmo_val.smk +++ b/workflow/rules/cosmo_val.smk @@ -267,7 +267,6 @@ rule cv_leakage_jackknife: catalogues=cv_shear_catalogues() + [str(Path(catalogue_entry(v)["subdir"]) / catalogue_entry(v)["psf"]["path"]) for v in CV_VERSIONS], - sky_mask=psf_leakage_sky_mask_input(config), output: measurements=CV_LEAKAGE_JACKKNIFE, params: diff --git a/workflow/rules/twopoint.smk b/workflow/rules/twopoint.smk index ee2c59c2..ae9bc17d 100644 --- a/workflow/rules/twopoint.smk +++ b/workflow/rules/twopoint.smk @@ -80,8 +80,6 @@ rule xi_onecov: rule rho_tau_stats: - input: - sky_mask=psf_leakage_sky_mask_input(config), output: rho_stats=str(COSMO_VAL / "rho_tau_stats/rho_stats_{version}_tomo_bin_all_minsep={min_sep}_maxsep={max_sep}_nbins={nbins}_npatch={npatch}.fits"), tau_stats=str(COSMO_VAL / "rho_tau_stats/tau_stats_{version}_tomo_bin_all_minsep={min_sep}_maxsep={max_sep}_nbins={nbins}_npatch={npatch}.fits"), @@ -100,8 +98,6 @@ rule rho_tau_stats: cat_config=CAT_CONFIG, output_dir=str(COSMO_VAL), b_target=config.get("cosmo_val", {}).get("b_target", 0.01), - # The psf_leakage sky mask applies to the ρ/τ stars and galaxies. - psf_leakage={"sky_mask": psf_leakage_sky_mask(config)}, blind=lambda w: blind_of(w.version), resources: mem_mb=64000, diff --git a/workflow/scripts/run_rho_tau.py b/workflow/scripts/run_rho_tau.py index 9077b6e2..1a495b7e 100644 --- a/workflow/scripts/run_rho_tau.py +++ b/workflow/scripts/run_rho_tau.py @@ -20,7 +20,6 @@ catalog_config=params["cat_config"], output_dir=params["output_dir"], cov_estimate_method=params.get("cov_estimate_method", "th"), - psf_leakage=params.get("psf_leakage"), ) cv.calculate_rho_tau_stats(tomography=False) verify_outputs(snakemake) diff --git a/workflow/tests/conftest.py b/workflow/tests/conftest.py index 139ca41b..6ca9f6eb 100644 --- a/workflow/tests/conftest.py +++ b/workflow/tests/conftest.py @@ -9,11 +9,11 @@ parse with the standard library and Snakemake alone -- the condition a host Snakemake is in. The ``candide`` tests need candide itself; CI deselects them. -The ``toy`` fixture is a disposable checkout: copies of ``workflow/``, -``papers/cosmo_val/`` and ``config/masks/``, this checkout's ``src/`` -symlinked in, a blinded and an 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 +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 (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. """ @@ -178,7 +178,6 @@ def toy(tmp_path_factory): shutil.copytree( REPO / "papers" / "cosmo_val", root / "papers" / "cosmo_val", ignore=skip ) - shutil.copytree(REPO / "config" / "masks", root / "config" / "masks") (root / "src").symlink_to(REPO / "src") (root / "data").mkdir() From 8f26957917faaf346ce4f27e1fb2ae5584779ab2 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Sun, 1 Nov 2026 06:44:10 +0100 Subject: [PATCH 15/16] Use centred galaxy PSFs for the leakage alpha predictor --- pyproject.toml | 7 +- src/sp_validation/correlation.py | 8 +- .../cosmo_val/psf_systematics.py | 2 + src/sp_validation/leakage_fit.py | 20 ++- src/sp_validation/rho_tau.py | 106 ++++++++++++-- ...test_tomo_tau_centering_before_bin_mask.py | 10 -- src/sp_validation/tests/test_leakage_fit.py | 60 +++++++- .../tests/test_leakage_measurement.py | 130 ++++++++++++++++++ uv.lock | 4 +- 9 files changed, 315 insertions(+), 32 deletions(-) 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 97811a2b..80de8370 100644 --- a/src/sp_validation/correlation.py +++ b/src/sp_validation/correlation.py @@ -33,12 +33,14 @@ def _measurement_configs(config): return {"means": means, "patched": patched} -def measurement_matches(path, config, selection=None): +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 @@ -54,9 +56,11 @@ def measurement_matches(path, config, selection=None): return recorded_configs == expected_configs -def write_measurement_metadata(path, config, selection=None): +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") diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 871db96c..ce263194 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -23,6 +23,7 @@ from .. import blinding, leakage_fit, sacc_io from ..io import open_entry from ..rho_tau import ( + get_params_rho_tau, get_rho_tau_w_cov, get_samples, get_samples_eiv, @@ -427,6 +428,7 @@ def leakage_jackknife(self, ver, tomo_bin_id="all", measure=False): 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 diff --git a/src/sp_validation/leakage_fit.py b/src/sp_validation/leakage_fit.py index aff62367..4a945633 100644 --- a/src/sp_validation/leakage_fit.py +++ b/src/sp_validation/leakage_fit.py @@ -1,6 +1,6 @@ """Numerical tools for joint errors-in-variables PSF-leakage fits. -This module has no catalogue I/O. It builds the production rho-to-tau design, +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. """ @@ -321,16 +321,20 @@ def fit_joint_eiv( jk_rows, *, nbin=None, + alpha_vector=None, + alpha_rows=None, template_vector=None, template_rows=None, scales=None, start=None, cross_term=True, ): - """Fit rho0..rho5 to tau0/tau2/tau5 with the joint EIV covariance. + """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)``. @@ -369,6 +373,18 @@ def fit_joint_eiv( [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) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index c3862825..f0cdc180 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -1,4 +1,5 @@ import gc +import hashlib import json import os import time @@ -25,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.""" @@ -80,7 +111,12 @@ def measure(cats, measurement_config): getattr(handler, f"save_{kind}_stats")(filename) getattr(handler, f"load_{kind}_stats")(filename) write_measurement_metadata( - Path(handler.catalogs._output) / filename, config, selection=selection + Path(handler.catalogs._output) / filename, + config, + selection=selection, + definition=tau_field_definition(handler.catalogs._params) + if kind == "tau" + else None, ) return means @@ -127,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}") @@ -206,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") @@ -388,7 +434,12 @@ def get_rho_tau( ) if ( - measurement_matches(tau_path, treecorr_config, selection=tau_selection) + measurement_matches( + tau_path, + treecorr_config, + selection=tau_selection, + definition=tau_field_definition(params), + ) and not force_run ): print( @@ -449,6 +500,8 @@ def leakage_cache_key( beta_psf_size, mask_star=None, mask_gal=None, + *, + params, ): """The settings and row selections that identify a joint leakage measurement. @@ -463,6 +516,12 @@ def leakage_cache_key( 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), @@ -502,8 +561,7 @@ def measure_leakage_jackknife( ``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 weighted mean of the full, - unmasked shear catalogue before ``mask_gal`` selects rows), sharing one + 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 @@ -512,8 +570,10 @@ def measure_leakage_jackknife( per-patch-pair results at ``npatch_fine`` and, for every count in ``patch_counts``, at Ward-linkage unions of the fine patches. - Writes ``path`` (``.npz``) and its ``.json`` metadata, and returns the - products as a dict: ``theta``, ``vector`` (rho0..rho5, tau0, tau2, tau5 + 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``, @@ -547,6 +607,9 @@ def measure_leakage_jackknife( 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() @@ -621,6 +684,16 @@ def measure_leakage_jackknife( products["slope_xip"] = np.array(xip[n_joint:]) products["slope_xim"] = np.array(xim[n_joint:]) + 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 = { @@ -655,7 +728,12 @@ def fit_leakage(products, npatch, beta_psf_size=False, cross_term=True): template_rows=rows[:, n_joint:], ) return leakage_fit.fit_joint_eiv( - products["vector"], rows[:, :n_joint], cross_term=cross_term, **template + products["vector"], + rows[:, :n_joint], + alpha_vector=products["alpha_vector"], + alpha_rows=products[f"alpha_rows_{npatch}"], + cross_term=cross_term, + **template, ) @@ -820,7 +898,10 @@ def get_jackknife_cov( Path(outdir) / rho_filename, treecorr_config, selection=rho_selection ) and measurement_matches( - Path(outdir) / tau_filename, treecorr_config, selection=tau_selection + Path(outdir) / tau_filename, + treecorr_config, + selection=tau_selection, + definition=tau_field_definition(get_params_rho_tau(config[version])), ) and not force_run ): @@ -998,7 +1079,12 @@ def _sample_cache_name(psf_fitter, base_tau, sampler, cov_type): ) if not metadata.exists(): return name - selection = json.loads(metadata.read_text()).get("selection") + 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]}" 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_leakage_fit.py b/src/sp_validation/tests/test_leakage_fit.py index e25ebfc1..bfb6b3fb 100644 --- a/src/sp_validation/tests/test_leakage_fit.py +++ b/src/sp_validation/tests/test_leakage_fit.py @@ -322,7 +322,18 @@ def test_unconstrained_parameter_is_flagged_not_given_an_error(): 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}, [200], 200) + 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): @@ -356,3 +367,50 @@ def test_design_and_xi_sys_match_shear_psf_leakage(tmp_path): 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 index c6e3add1..bac80e40 100644 --- a/src/sp_validation/tests/test_leakage_measurement.py +++ b/src/sp_validation/tests/test_leakage_measurement.py @@ -41,6 +41,8 @@ def _write_catalogues(tmp_path, n_star=4000, n_gal=6000): [ 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)), @@ -78,6 +80,8 @@ def __init__(self, tmp_path, beta_psf_size=False): "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", }, @@ -286,3 +290,129 @@ def covariance(star, gal): 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/uv.lock b/uv.lock index 804f7447..2e6f0b4c 100644 --- a/uv.lock +++ b/uv.lock @@ -3821,7 +3821,7 @@ wheels = [ [[package]] name = "shear-psf-leakage" version = "0.2.1" -source = { git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=develop#b380ea778277b5cff6c2d5d10c2943890f593ef2" } +source = { git = "https://github.com/CosmoStat/shear_psf_leakage.git?rev=d0e10d83ecb1978ed56599f0b46326e3d000680e#d0e10d83ecb1978ed56599f0b46326e3d000680e" } dependencies = [ { name = "cs-util" }, { name = "emcee" }, @@ -4037,7 +4037,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'" }, From 7587907dab0918a710a55651f2b214e171a93da6 Mon Sep 17 00:00:00 2001 From: Cail Daley Date: Sun, 1 Nov 2026 07:14:22 +0100 Subject: [PATCH 16/16] Release the shear field tree before alpha measurement --- src/sp_validation/rho_tau.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index f0cdc180..ae7535b5 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -684,6 +684,8 @@ def measure_leakage_jackknife( 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: