Skip to content

refactor(cosmo_val): tighten the mixin modules (behavior-preserving cleanup) #208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cailmdaley merged 1 commit into from
Jun 20, 2026
Merged

refactor(cosmo_val): tighten the mixin modules (behavior-preserving cleanup) #208

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 6 additions & 8 deletions src/sp_validation/cosmo_val/catalog_characterization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,11 +121,9 @@ def compute_survey_stats(
}

if overwrite_config:
if "cov_th" not in self.cc[ver]:
self.cc[ver]["cov_th"] = {}
self.cc[ver]["cov_th"]["A"] = float(area_deg2)
self.cc[ver]["cov_th"]["n_e"] = float(n_eff)
self.cc[ver]["cov_th"]["sigma_e"] = float(sigma_e)
self.cc[ver].setdefault("cov_th", {}).update(
A=float(area_deg2), n_e=float(n_eff), sigma_e=float(sigma_e)
)
self._write_catalog_config()

return results
Expand Down Expand Up @@ -268,7 +266,7 @@ def plot_ellipticity(self, nbins=200):
else:
self.print_start("Computing ellipticity histograms:")

fig, axs = plt.subplots(1, 2, figsize=(22, 7))
_fig, axs = plt.subplots(1, 2, figsize=(22, 7))
bins = np.linspace(-1.1, 1.1, nbins + 1)
for ver in self.versions:
self.print_magenta(ver)
Expand Down Expand Up @@ -317,7 +315,7 @@ def plot_weights(self, nbins=200):
else:
self.print_start("Computing weight histograms:")

fig, ax = plt.subplots(1, 1, figsize=(10, 7))
plt.figure(figsize=(10, 7))
for ver in self.versions:
self.print_magenta(ver)
with self.results[ver].temporarily_read_data():
Expand Down Expand Up @@ -345,7 +343,7 @@ def plot_weights(self, nbins=200):
def plot_separation(self, nbins=200):
self.print_start("Separation histograms")
if "SP_matched_MP_v1.0" in self.versions:
fig, axs = plt.subplots(1, 1, figsize=(10, 7))
_fig, axs = plt.subplots(1, 1, figsize=(10, 7))
with self.results["SP_matched_MP_v1.0"].temporarily_read_data():
sep = self.results["SP_matched_MP_v1.0"].dat_shear["Separation"]
axs.hist(
Expand Down
75 changes: 38 additions & 37 deletions src/sp_validation/cosmo_val/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# %%
import copy
import os
import re
from pathlib import Path

import colorama
Expand Down Expand Up @@ -249,19 +250,17 @@ def __init__(
"cell_method must be 'map' or 'catalog'"
)
assert self.noise_bias_method in ["analytic", "randoms"], (
"noise_bias_method must be 'analytical' or 'randoms'"
"noise_bias_method must be 'analytic' or 'randoms'"
)
assert self.fiducial_input_inka in ["coupled", "decoupled"], (
"fiducial_input_inka must be 'coupled' or 'decoupled'"
)

# For theory calculations:
# Create cosmology object using new functionality
if cosmo_params is not None:
self.cosmo = get_cosmo(**cosmo_params)
else:
# Use Planck 2018 defaults
self.cosmo = get_cosmo()
# Cosmology for theory predictions: caller-supplied params, else the
# get_cosmo() Planck 2018 defaults.
self.cosmo = (
get_cosmo(**cosmo_params) if cosmo_params is not None else get_cosmo()
)

self.treecorr_config = {
"ra_units": "degrees",
Expand All @@ -276,15 +275,15 @@ def __init__(

self.catalog_config_path = Path(catalog_config)
with self.catalog_config_path.open("r") as file:
self.cc = cc = yaml.load(file.read(), Loader=yaml.FullLoader)
self.cc = cc = yaml.load(file, Loader=yaml.FullLoader)

def resolve_paths_for_version(ver):
"""Resolve relative paths for a version using its subdir."""
subdir = Path(cc[ver]["subdir"])
for key in cc[ver]:
if "path" in cc[ver][key]:
path = Path(cc[ver][key]["path"])
cc[ver][key]["path"] = (
for section in cc[ver].values():
if "path" in section:
path = Path(section["path"])
section["path"] = (
str(path) if path.is_absolute() else str(subdir / path)
)

Expand Down Expand Up @@ -381,8 +380,6 @@ def get_redshift(self, version):
specified blind (A, B, or C) by replacing the blind suffix in the
configured path.
"""
import re

redshift_path = self.cc[version]["shear"]["redshift_path"]

# Override blind if specified
Expand All @@ -398,10 +395,14 @@ def _write_catalog_config(self):
def color_reset(self):
print(colorama.Fore.BLACK, end="")

def print_blue(self, msg, end="\n"):
print(colorama.Fore.BLUE + msg, end=end)
def _print_color(self, color, msg, end="\n"):
"""Print ``msg`` in ``color``, then restore the default foreground."""
print(color + msg, end=end)
self.color_reset()

def print_blue(self, msg, end="\n"):
self._print_color(colorama.Fore.BLUE, msg, end=end)

def print_start(self, msg, end="\n"):
print()
self.print_blue(msg, end=end)
Expand All @@ -410,30 +411,29 @@ def print_done(self, msg):
self.print_blue(msg)

def print_magenta(self, msg):
print(colorama.Fore.MAGENTA + msg)
self.color_reset()
self._print_color(colorama.Fore.MAGENTA, msg)

def print_green(self, msg):
print(colorama.Fore.GREEN + msg)
self.color_reset()
self._print_color(colorama.Fore.GREEN, msg)

def print_cyan(self, msg):
print(colorama.Fore.CYAN + msg)
self.color_reset()
self._print_color(colorama.Fore.CYAN, msg)

def init_results(self, objectwise=False):
# Branch is loop-invariant: pick the leakage class and its parameter
# builder once, then apply per version.
make_leakage, set_params = (
(run_object.LeakageObject, self.set_params_leakage_object)
if objectwise
else (run_scale.LeakageScale, self.set_params_leakage_scale)
)

results = {}
for ver in self.versions:
# Set parameters depending on the type of leakage
if objectwise:
results[ver] = run_object.LeakageObject()
results[ver]._params.update(self.set_params_leakage_object(ver))
else:
results[ver] = run_scale.LeakageScale()
results[ver]._params.update(self.set_params_leakage_scale(ver))

results[ver].check_params()
results[ver].prepare_output()
leakage = results[ver] = make_leakage()
leakage._params.update(set_params(ver))
leakage.check_params()
leakage.prepare_output()

return results

Expand Down Expand Up @@ -501,10 +501,11 @@ def summarize_bmodes(self, fiducial_scale_cut=(12, 83), versions=None):
)
except (KeyError, RuntimeError):
pass
if "eb_samples" in res:
cov_methods.add("semi-analytic")
else:
cov_methods.add(f"jackknife ({gg.npatch1} patches)")
cov_methods.add(
"semi-analytic"
if "eb_samples" in res
else f"jackknife ({gg.npatch1} patches)"
)

# COSEBIs PTE from stored results
if ver in self._cosebis_results:
Expand Down
30 changes: 12 additions & 18 deletions src/sp_validation/cosmo_val/cosebis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,7 +138,7 @@ def calculate_cosebis(
gg=gg, nmodes=nmodes, scale_cuts=None, cov_path=cov_path
)
# Extract single results dict from scale_cuts dictionary
results = list(results.values())[0]
results = next(iter(results.values()))

return results

Expand Down Expand Up @@ -237,19 +237,17 @@ def plot_cosebis(

# Generate plots using specialized plotting functions
# Extract single result for plotting if multiple scale cuts were evaluated
if isinstance(results, dict) and all(
isinstance(k, tuple) for k in results.keys()
):
multiple_scale_cuts = isinstance(results, dict) and all(
isinstance(k, tuple) for k in results
)
if multiple_scale_cuts:
# Multiple scale cuts: use fiducial_scale_cut if provided, otherwise use
# full range
if fiducial_scale_cut is not None:
plot_results = results[
find_conservative_scale_cut_key(results, fiducial_scale_cut)
]
else:
# Use full range result (largest scale cut)
max_range_key = max(results.keys(), key=lambda x: x[1] - x[0])
plot_results = results[max_range_key]
# full range (largest scale cut)
plot_results = results[
find_conservative_scale_cut_key(results, fiducial_scale_cut)
if fiducial_scale_cut is not None
else max(results, key=lambda x: x[1] - x[0])
]
else:
# Single result
plot_results = results
Expand All @@ -266,11 +264,7 @@ def plot_cosebis(
)

# Generate scale cut heatmap if we have multiple scale cuts
if (
isinstance(results, dict)
and all(isinstance(k, tuple) for k in results.keys())
and len(results) > 1
):
if multiple_scale_cuts and len(results) > 1:
# Create temporary gg object with correct binning for mapping
treecorr_config_temp = {
**self.treecorr_config,
Expand Down
74 changes: 22 additions & 52 deletions src/sp_validation/cosmo_val/pseudo_cl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,6 @@ def calculate_pseudo_cl_eb_cov(self):
self.print_cyan(f"Extracting the fiducial power spectrum for {ver}")

lmax = 2 * self.nside
z, dndz = self.get_redshift(ver)
ell = np.arange(1, lmax + 1)
pw = hp.pixwin(nside, lmax=lmax)
if pw.shape[0] != len(ell) + 1:
Expand All @@ -158,7 +157,6 @@ def calculate_pseudo_cl_eb_cov(self):
# Load redshift distribution and calculate theory C_ell
path_redshift_distr = self.cc[ver]["shear"]["redshift_path"]
z, dndz = np.loadtxt(path_redshift_distr, unpack=True)
ell = np.arange(1, lmax + 1)
fiducial_cl = (
get_theo_c_ell(
ell=ell,
Expand All @@ -181,7 +179,7 @@ def calculate_pseudo_cl_eb_cov(self):
# Load data and create shear and noise maps
cat_gal = fits.getdata(self.cc[ver]["shear"]["path"])

n_gal, unique_pix, idx, idx_rep = self.get_n_gal_map(
n_gal, unique_pix, _idx, idx_rep = self.get_n_gal_map(
params, nside, cat_gal
)

Expand All @@ -202,14 +200,7 @@ def calculate_pseudo_cl_eb_cov(self):
idx_rep,
)

cl_noise = np.mean(cl_noise, axis=0)
noise_bias_cl = cl_noise
noise_bias_cl = b.unbin_cell(noise_bias_cl) # Unbin
lowest_ell = b.get_ell_list(0)[0]
for i in range(4):
noise_bias_cl[i, :lowest_ell] = noise_bias_cl[
i, lowest_ell
] # Fill the data vector below lmin
noise_bias_cl = np.mean(cl_noise, axis=0)

elif self.noise_bias_method == "analytic":
self.print_cyan("Getting analytic noise bias.")
Expand All @@ -230,18 +221,17 @@ def calculate_pseudo_cl_eb_cov(self):
noise_bias_cl[3, :] = noise_bias

noise_bias_cl = wsp.decouple_cell(noise_bias_cl) # Decouple
noise_bias_cl = b.unbin_cell(noise_bias_cl) # Unbin
lowest_ell = b.get_ell_list(0)[0]
for i in range(4):
noise_bias_cl[i, :lowest_ell] = noise_bias_cl[
i, lowest_ell
] # Fill the data vector below lmin

else:
raise ValueError(
f"Noise bias method {self.noise_bias_method} not recognized. It should be 'randoms' or 'analytic'."
)

# Unbin, then fill the data vector below lmin with the lowest-ell value
noise_bias_cl = b.unbin_cell(noise_bias_cl)
lowest_ell = b.get_ell_list(0)[0]
noise_bias_cl[:, :lowest_ell] = noise_bias_cl[:, [lowest_ell]]

self.print_cyan("Adding noise bias to the fiducial Cls.")

fiducial_cl = (
Expand Down Expand Up @@ -288,43 +278,21 @@ def calculate_pseudo_cl_eb_cov(self):
wb=wsp,
).reshape([n_ell_actual, 4, n_ell_actual, 4])

covar_EE_EE = covar_22_22[:, 0, :, 0]
covar_EE_EB = covar_22_22[:, 0, :, 1]
covar_EE_BE = covar_22_22[:, 0, :, 2]
covar_EE_BB = covar_22_22[:, 0, :, 3]
covar_EB_EE = covar_22_22[:, 1, :, 0]
covar_EB_EB = covar_22_22[:, 1, :, 1]
covar_EB_BE = covar_22_22[:, 1, :, 2]
covar_EB_BB = covar_22_22[:, 1, :, 3]
covar_BE_EE = covar_22_22[:, 2, :, 0]
covar_BE_EB = covar_22_22[:, 2, :, 1]
covar_BE_BE = covar_22_22[:, 2, :, 2]
covar_BE_BB = covar_22_22[:, 2, :, 3]
covar_BB_EE = covar_22_22[:, 3, :, 0]
covar_BB_EB = covar_22_22[:, 3, :, 1]
covar_BB_BE = covar_22_22[:, 3, :, 2]
covar_BB_BB = covar_22_22[:, 3, :, 3]

self.print_cyan("Saving Pseudo-Cl covariance")

# covar_22_22 is indexed [ell, pol_a, ell, pol_b]; store each of the
# 16 EE/EB/BE/BB cross-blocks as a named HDU (row-major pol order).
# Append rather than construct from a list so astropy promotes the
# first HDU to a PrimaryHDU on write.
pols = ["EE", "EB", "BE", "BB"]
hdu = fits.HDUList()

hdu.append(fits.ImageHDU(covar_EE_EE, name="COVAR_EE_EE"))
hdu.append(fits.ImageHDU(covar_EE_EB, name="COVAR_EE_EB"))
hdu.append(fits.ImageHDU(covar_EE_BE, name="COVAR_EE_BE"))
hdu.append(fits.ImageHDU(covar_EE_BB, name="COVAR_EE_BB"))
hdu.append(fits.ImageHDU(covar_EB_EE, name="COVAR_EB_EE"))
hdu.append(fits.ImageHDU(covar_EB_EB, name="COVAR_EB_EB"))
hdu.append(fits.ImageHDU(covar_EB_BE, name="COVAR_EB_BE"))
hdu.append(fits.ImageHDU(covar_EB_BB, name="COVAR_EB_BB"))
hdu.append(fits.ImageHDU(covar_BE_EE, name="COVAR_BE_EE"))
hdu.append(fits.ImageHDU(covar_BE_EB, name="COVAR_BE_EB"))
hdu.append(fits.ImageHDU(covar_BE_BE, name="COVAR_BE_BE"))
hdu.append(fits.ImageHDU(covar_BE_BB, name="COVAR_BE_BB"))
hdu.append(fits.ImageHDU(covar_BB_EE, name="COVAR_BB_EE"))
hdu.append(fits.ImageHDU(covar_BB_EB, name="COVAR_BB_EB"))
hdu.append(fits.ImageHDU(covar_BB_BE, name="COVAR_BB_BE"))
hdu.append(fits.ImageHDU(covar_BB_BB, name="COVAR_BB_BB"))
for i, pa in enumerate(pols):
for j, pb in enumerate(pols):
hdu.append(
fits.ImageHDU(
covar_22_22[:, i, :, j], name=f"COVAR_{pa}_{pb}"
)
)

hdu.writeto(out_path, overwrite=True)

Expand Down Expand Up @@ -565,7 +533,9 @@ def calculate_pseudo_cl_map(self, ver, nside, out_path):

w = cat_gal[params["w_col"]]
self.print_cyan("Creating maps and computing Cl's...")
n_gal_map, unique_pix, idx, idx_rep = self.get_n_gal_map(params, nside, cat_gal)
n_gal_map, unique_pix, _idx, idx_rep = self.get_n_gal_map(
params, nside, cat_gal
)
mask = n_gal_map != 0

shear_map_e1 = np.zeros(hp.nside2npix(nside))
Expand Down
Loading
Loading