# Copyright (C) 2025 ETH Zurich
# Institute for Particle Physics and Astrophysics
# Author: Silvan Fischbacher
# created: Tue Jun 24 2025
import pickle
import re
import numpy as np
from cosmic_toolbox.logger import get_logger
from scipy.special import erf
from galsbi.sham_ot import utils
LOGGER = get_logger(__file__)
[docs]
class HaloHistogram:
"""
Class to handle histograms of halos for optimal use in the SHAM-OT framework.
The class saves the following histograms:
- 2D histogram of redshift vs halo mass for all halos,
shape (n_z_bins, n_mass_bins)
- 2D histogram of redshift vs halo mass for central halos,
shape (n_z_bins, n_mass_bins)
- 3D histogram of redshift vs halo mass vs time since last merger for subhalos,
shape (n_z_bins, n_mass_bins, n_time_bins)
The indexing follows: hist[z_bin_index, mass_bin_index, time_bin_index],
i.e. histogram[0] gives the 0th redshift bin as a histogram over log masses.
"""
def __init__(
self,
z_min=None,
z_max=None,
min_mass=None,
max_mass=6e15,
min_time_since_merger=0,
max_time_since_merger=6,
n_bins_mass=500,
n_bins_z=500,
n_bins_time=100,
):
"""
Initializes the HaloHistogram class with specified mass and redshift ranges
and the number of bins for each.
Defaults for redshift and minimum mass are taken directly from the simulation
if not specified.
:param z_min: Minimum redshift of the histogram
:param z_max: Maximum redshift of the histogram
:param min_mass: Minimum mass of the histogram
:param max_mass: Maximum mass of the histogram (default: 6e15)
:param min_time_since_merger: Minimum time since last merger (default: 0)
:param max_time_since_merger: Maximum time since last merger (default: 10)
:param n_bins_mass: Number of bins in the mass dimension (default: 500)
:param n_bins_z: Number of bins in the redshift dimension (default: 500)
:param n_bins_time: Number of bins in the time since last merger dimension
(default: 100)
"""
self.z_min = z_min
self.z_max = z_max
self.min_mass = min_mass
self.max_mass = max_mass
self.min_time_since_merger = min_time_since_merger
self.max_time_since_merger = max_time_since_merger
self.n_bins_mass = n_bins_mass
self.n_bins_z = n_bins_z
self.n_bins_time = n_bins_time
# Flag to enforce initialization of bins and histograms the first time
# a halo catalog is added
self.no_catalogs_added = True
@property
def min_log_mass(self):
"""
Returns the minimum log mass of the histogram.
If min_mass is not set, it returns None.
"""
if self.min_mass is not None:
return np.log10(self.min_mass)
return None
@property
def max_log_mass(self):
"""
Returns the maximum log mass of the histogram.
If max_mass is not set, it returns None.
"""
return np.log10(self.max_mass)
@property
def log_mass_bins(self):
"""
Returns the logarithmic mass bins for the histogram.
If min_log_mass or max_log_mass is not set, it returns None.
"""
if hasattr(self, "mass_bins"):
return np.log10(self.mass_bins)
return None
[docs]
def add_halo_catalog(
self,
halo_catalog,
mass_field="m",
redshift_field="z_true",
is_central_field="is_halo",
time_since_merger_field="delta_t",
pinocchio_output_filename=None,
):
"""
Adds a halo catalog to the histograms. The catalog should be a recarray or
dictionary-like object containing the necessary fields:
- Halo mass
- Redshift
- Is central halo
- Time since last merger
:param halo_catalog: A recarray or dictionary-like object containing halo data.
:param mass_field: Field name for halo mass
(default: 'm')
:param redshift_field: Field name for redshift
(default: 'z_true')
:param is_central_field: Field name for central halo indicator
(default: 'is_halo')
:param time_since_merger_field: Field name for time since last merger
(default: 'delta_t')
"""
if self.no_catalogs_added:
if self.z_min is None or self.z_max is None or self.min_mass is None:
m_part, min_num_part, z_min, z_max = read_pinocchio_config_details(
pinocchio_output_filename
)
if self.z_min is None:
self.z_min = z_min
if self.z_max is None:
self.z_max = z_max
if self.min_mass is None:
self.min_mass = 0.999 * m_part * min_num_part
log_mass_bins = np.linspace(
self.min_log_mass, self.max_log_mass, self.n_bins_mass + 1
)
self.mass_bins = np.power(10, log_mass_bins)
self.z_bins = np.linspace(self.z_min, self.z_max, self.n_bins_z + 1)
self.time_bins = np.linspace(
self.min_time_since_merger,
self.max_time_since_merger,
self.n_bins_time + 1,
)
self.histogram_all = np.zeros((self.n_bins_z, self.n_bins_mass), dtype=int)
self.histogram_central = np.zeros(
(self.n_bins_z, self.n_bins_mass), dtype=int
)
self.histogram_subhalo = np.zeros(
(self.n_bins_z, self.n_bins_mass, self.n_bins_time), dtype=int
)
self.no_catalogs_added = False
is_central = halo_catalog[is_central_field].astype(bool)
self.histogram_all += np.histogram2d(
halo_catalog[redshift_field],
np.log10(halo_catalog[mass_field]),
bins=(self.z_bins, self.log_mass_bins),
)[0].astype(int)
self.histogram_central += np.histogram2d(
halo_catalog[redshift_field][is_central],
np.log10(halo_catalog[mass_field][is_central]),
bins=(self.z_bins, self.log_mass_bins),
)[0].astype(int)
self.histogram_subhalo += np.histogramdd(
(
halo_catalog[redshift_field][~is_central],
np.log10(halo_catalog[mass_field][~is_central]),
halo_catalog[time_since_merger_field][~is_central],
),
bins=(self.z_bins, self.log_mass_bins, self.time_bins),
)[0].astype(int)
if np.any(halo_catalog[mass_field] > self.max_mass):
too_high = halo_catalog[mass_field] > self.max_mass
LOGGER.warning(
f"{sum(too_high)} halos have mass > {self.max_mass}. "
"They are added to the heaviest mass bin"
)
# Add overflow halos to the correct redshift bins in the last mass bin
overflow_hist_all = np.histogram(
halo_catalog[redshift_field][too_high], bins=self.z_bins
)[0].astype(int)
overflow_hist_central = np.histogram(
halo_catalog[redshift_field][too_high & is_central], bins=self.z_bins
)[0].astype(int)
overflow_hist_subhalo = np.histogramdd(
(
halo_catalog[redshift_field][too_high & ~is_central],
halo_catalog[time_since_merger_field][too_high & ~is_central],
),
bins=(self.z_bins, self.time_bins),
)[0].astype(int)
self.histogram_all[:, -1] += overflow_hist_all
self.histogram_central[:, -1] += overflow_hist_central
self.histogram_subhalo[:, -1, :] += overflow_hist_subhalo
if np.any(halo_catalog[time_since_merger_field] > self.max_time_since_merger):
too_late = (
halo_catalog[time_since_merger_field] > self.max_time_since_merger
)
LOGGER.warning(
f"{sum(too_late)} subhalos have time since last merger > "
"{self.max_time_since_merger}. They are added to the latest time bin."
)
subhalos_too_late = too_late & ~is_central
if np.any(subhalos_too_late):
overflow_hist_subhalo = np.histogram2d(
halo_catalog[redshift_field][subhalos_too_late],
np.log10(halo_catalog[mass_field][subhalos_too_late]),
bins=(self.z_bins, self.log_mass_bins),
)[0].astype(int)
self.histogram_subhalo[:, :, -1] += overflow_hist_subhalo
[docs]
def add_halo_histogram(self, halo_histogram):
"""
Adds another HaloHistogram instance to the current one.
:param halo_histogram: Another HaloHistogram instance to add.
"""
if not isinstance(halo_histogram, HaloHistogram):
raise TypeError("halo_histogram must be an instance of HaloHistogram")
if (
not np.array_equal(self.mass_bins, halo_histogram.mass_bins)
or not np.array_equal(self.z_bins, halo_histogram.z_bins)
or not np.array_equal(self.time_bins, halo_histogram.time_bins)
):
raise ValueError("Mass, redshift, or time bins do not match.")
self.histogram_all += halo_histogram.histogram_all
self.histogram_central += halo_histogram.histogram_central
self.histogram_subhalo += halo_histogram.histogram_subhalo
[docs]
def get_histogram(self, galaxy_type="blue", model="Berner+24", **kwargs):
"""
Returns the histogram for red or blue galaxies for the specified SHAM model.
Currently, the following models are supported:
- "Berner+24": free parameters m_limit and quenching time, see arXiv:2310.15223
- "Peng+10": quenching fractions based on mass, see arXiv:1003.4747
:param galaxy_type: Type of galaxy ("blue" or "red")
:param model: Model name (default: "Berner+24")
:param m_limit: Mass limit for central halos in Berner+24 (default: 9e12)
:param quenching_time: Quenching time for subhalos in Berner+24 (default: 2)
:param mlim_sigma: Scatter in mass limit for extended Berner+24 (default: 0.3)
:param initial_quenched_fraction: Initial quenched fraction for subhalos in
extended Berner+24 (default: 0.1)
:param quenching_time_sigma: Scatter in quenching time for extended Berner+24
(default: 0.5)
:param M_mu: Characteristic mass for Peng+10 (default: 11.935)
:param sigma: Mass scatter for Peng+10 (default: 0.51)
:param b: Satellite quenching efficiency for Peng+10 (default: 0.41)
:return: Histogram for the specified galaxy type
"""
if model == "Berner+24":
kwargs.setdefault("m_limit", 8e12)
kwargs.setdefault("quenching_time", 2.0)
return self._get_histogram_berner(
galaxy_type, kwargs["m_limit"], kwargs["quenching_time"]
)
elif model == "Berner+24_ext":
kwargs.setdefault("m_limit", 8e12)
kwargs.setdefault("quenching_time", 2.0)
kwargs.setdefault("mlim_sigma", 0.3)
kwargs.setdefault("initial_quenched_fraction", 0.1)
kwargs.setdefault("quenching_time_sigma", 0.5)
return self._get_histogram_berner_ext(
galaxy_type,
kwargs["m_limit"],
kwargs["quenching_time"],
mlim_sigma=kwargs["mlim_sigma"],
initial_quenched_fraction=kwargs["initial_quenched_fraction"],
quenching_time_sigma=kwargs["quenching_time_sigma"],
)
elif model == "Peng+10":
kwargs.setdefault("M_mu", 11.935)
kwargs.setdefault("sigma", 0.51)
kwargs.setdefault("b", 0.41)
return self._get_histogram_peng(
galaxy_type, kwargs["M_mu"], kwargs["sigma"], kwargs["b"]
)
else:
raise ValueError(
f"Model '{model}' is not supported. Currently only the following "
"models are supported: 'Berner+24', 'Berner+24_ext, 'Peng+10'."
)
def _get_histogram_berner(self, galaxy_type, m_limit, quenching_time):
"""
Get histogram using Berner+24 model with hard cuts.
"""
histogram = np.zeros_like(self.histogram_all)
# Create fractional masks for mass and time bins
mass_bin_centers = self.get_mass_bins()[1]
time_bin_centers = self.get_time_bins()[1]
# For mass bins: fraction of bin that is <= m_limit
mass_blue_mask = np.zeros(len(mass_bin_centers))
for i, (bin_start, bin_end) in enumerate(
zip(self.mass_bins[:-1], self.mass_bins[1:], strict=True)
):
if bin_end <= m_limit:
mass_blue_mask[i] = 1.0 # Entire bin is blue
elif bin_start > m_limit:
mass_blue_mask[i] = 0.0 # Entire bin is red
else:
# Limit is within bin, calculate fraction
mass_blue_mask[i] = (m_limit - bin_start) / (bin_end - bin_start)
# For time bins: fraction of bin that is <= quenching_time
time_blue_mask = np.zeros(len(time_bin_centers))
for i, (bin_start, bin_end) in enumerate(
zip(self.time_bins[:-1], self.time_bins[1:], strict=True)
):
if bin_end <= quenching_time:
time_blue_mask[i] = 1.0 # Entire bin is blue
elif bin_start > quenching_time:
time_blue_mask[i] = 0.0 # Entire bin is red
else:
# Limit is within bin, calculate fraction
time_blue_mask[i] = (quenching_time - bin_start) / (bin_end - bin_start)
if galaxy_type == "blue":
# Blue galaxies:
# central halos with mass <= m_limit
# subhalos with delta_t <= quenching_time
histogram += np.round(
self.histogram_central * mass_blue_mask[None, :]
).astype(int)
histogram += (
np.round(self.histogram_subhalo * time_blue_mask[None, None, :])
.sum(axis=2)
.astype(int)
)
elif galaxy_type == "red":
# Red galaxies:
# central halos with mass > m_limit
# subhalos with delta_t > quenching_time
histogram += np.round(
self.histogram_central * (1 - mass_blue_mask)[None, :]
).astype(int)
histogram += (
np.round(self.histogram_subhalo * (1 - time_blue_mask)[None, None, :])
.sum(axis=2)
.astype(int)
)
else:
raise ValueError(
f"Galaxy type '{galaxy_type}' not supported. Use 'blue' or 'red'."
)
return histogram
def _get_histogram_berner_ext(
self,
galaxy_type,
m_limit,
quenching_time,
mlim_sigma,
initial_quenched_fraction,
quenching_time_sigma,
):
"""
Get histogram using extended Berner+24 model with mass-dependent quenching
fractions.
"""
histogram = np.zeros_like(self.histogram_all, dtype=float)
mass_centers = self.get_mass_bins()[1]
time_centers = self.get_time_bins()[1]
red_centrals = utils.adapted_err_function(
np.log10(mass_centers), np.log10(m_limit), mlim_sigma, min_val=0, max_val=1
)
blue_centrals = 1 - red_centrals
red_subhalos = utils.adapted_err_function(
time_centers,
quenching_time,
quenching_time_sigma,
min_val=initial_quenched_fraction,
max_val=1,
)
blue_subhalos = 1 - red_subhalos
histogram_central = self.histogram_central.astype(float)
histogram_subhalo = self.histogram_subhalo.astype(float)
if galaxy_type == "blue":
# Blue galaxies = (1 - quenched fraction) * histogram
histogram += histogram_central * blue_centrals[None, :]
# subhalo histogram is 3D, the selection is over the 3rd dimension
histogram += (histogram_subhalo * blue_subhalos[None, None, :]).sum(axis=2)
elif galaxy_type == "red":
# Red galaxies = quenched fraction * histogram
histogram += histogram_central * red_centrals[None, :]
# subhalo histogram is 3D, the selection is over the 3rd dimension
histogram += (histogram_subhalo * red_subhalos[None, None, :]).sum(axis=2)
else:
raise ValueError(
f"Galaxy type '{galaxy_type}' not supported. Use 'blue' or 'red'."
)
return histogram
def _get_histogram_peng(self, galaxy_type, M_mu, sigma, b):
"""
Get histogram using Peng+10 model with mass-dependent quenching fractions.
Attention: M_mu is in log10 mass units.
"""
histogram = np.zeros_like(self.histogram_all, dtype=float)
mass_centers = self.get_mass_bins()[1]
def _r0(M_mu, sigma, masses):
return np.log10(masses / (10**M_mu)) / (np.sqrt(2) * sigma)
def _quenched_fraction_central(M_mu, sigma, masses):
return (1 / 2) * (1 + erf(_r0(M_mu, sigma, masses)))
def _quenched_fraction_subhalo(M_mu, sigma, b, masses):
return ((1 - b) / 2) * (1 + erf(_r0(M_mu, sigma, masses))) + b
# Calculate quenching fractions for each mass bin
quenched_frac_central = _quenched_fraction_central(M_mu, sigma, mass_centers)
quenched_frac_subhalo = 1 - _quenched_fraction_subhalo(
M_mu, sigma, b, mass_centers
)
histogram_central = self.histogram_central.astype(float)
histogram_subhalo = self.histogram_subhalo.astype(float).sum(axis=2)
if galaxy_type == "blue":
# Blue galaxies = (1 - quenched fraction) * histogram
histogram += histogram_central * (1 - quenched_frac_central)[None, :]
histogram += histogram_subhalo * (1 - quenched_frac_subhalo)[None, :]
elif galaxy_type == "red":
# Red galaxies = quenched fraction * histogram
histogram += histogram_central * quenched_frac_central[None, :]
histogram += histogram_subhalo * quenched_frac_subhalo[None, :]
else:
raise ValueError(
f"Galaxy type '{galaxy_type}' not supported. Use 'blue' or 'red'."
)
return histogram
[docs]
def save(self, filename):
"""
Saves the class instance to a file using pickle.
:param filename: Name of the file to save the instance to.
"""
with open(filename, "wb") as f:
pickle.dump(self, f)
LOGGER.info(f"Saved HaloHistogram to {filename}")
[docs]
def get_total_counts(self):
"""
Returns the total counts of all histograms.
:return: Tuple of total counts for all, central, and subhalo histograms.
"""
total_all = np.sum(self.histogram_all)
total_central = np.sum(self.histogram_central)
total_subhalo = np.sum(self.histogram_subhalo)
return total_all, total_central, total_subhalo
[docs]
def get_z_bins(self):
"""
Returns the redshift bin edges and centers.
:return: Array of redshift bin edges and centers.
"""
return self.z_bins, (self.z_bins[:-1] + self.z_bins[1:]) / 2
[docs]
def get_mass_bins(self):
"""
Returns the mass bin edges and log centers.
:return: Array of mass bin edges and log centers.
"""
centers_log = (self.log_mass_bins[:-1] + self.log_mass_bins[1:]) / 2
return self.mass_bins, np.power(10, centers_log)
[docs]
def get_time_bins(self):
"""
Returns the time since last merger bin edges and centers.
:return: Array of time since last merger bin edges and centers.
"""
return self.time_bins, (self.time_bins[:-1] + self.time_bins[1:]) / 2
[docs]
def load_halo_histogram(filename):
"""
Loads a HaloHistogram instance from a file.
:param filename: Name of the file to load the instance from.
:return: Loaded HaloHistogram instance.
"""
with open(filename, "rb") as f:
halo_hist = pickle.load(f)
LOGGER.info(f"Loaded HaloHistogram from {filename}")
return halo_hist
[docs]
def read_pinocchio_config_details(pinocchio_output_filename):
"""
Reads the Pinocchio output file to extract details about the simulation,
such as particle mass, minimum number of particles, and redshift range.
:param pinocchio_output_filename: Path to the Pinocchio output file.
:return: Tuple containing particle mass, minimum number of particles,
minimum redshift, and maximum redshift.
"""
with open(pinocchio_output_filename) as f:
text = f.read()
m_part = float(
re.search("Particle Mass \\(Msun/h\\)\\s+(.+)\n", text).groups()[0],
)
min_num_part = int(
re.search("MinHaloMass \\(particles\\)\\s+(\\d+)\n", text).groups()[0],
)
z_max, z_min = (
float(i)
for i in re.search(
"The Past Light Cone will be reconstruct from z=(.+) to z=(.+)\n", text
).groups()
)
return (
m_part,
min_num_part,
z_min,
z_max,
)
[docs]
def rebin_halo_histogram(
halo_hist, new_n_bins_mass=None, new_n_bins_z=None, new_n_bins_time=None
):
"""
Rebins a HaloHistogram to a new number of bins by combining neighboring bins.
The new number of bins must be a divisor of the current number of bins.
:param halo_hist: HaloHistogram instance to rebin
:param new_n_bins_mass: New number of mass bins (must divide current n_bins_mass)
:param new_n_bins_z: New number of redshift bins (must divide current n_bins_z)
:param new_n_bins_time: New number of time bins (must divide current n_bins_time)
:return: New HaloHistogram instance with rebinned histograms
"""
if halo_hist.no_catalogs_added:
raise ValueError("Cannot rebin histogram before any catalogs have been added")
# Use current values if new values are not specified
if new_n_bins_mass is None:
new_n_bins_mass = halo_hist.n_bins_mass
if new_n_bins_z is None:
new_n_bins_z = halo_hist.n_bins_z
if new_n_bins_time is None:
new_n_bins_time = halo_hist.n_bins_time
# Check if new bin numbers are valid divisors
if halo_hist.n_bins_mass % new_n_bins_mass != 0:
raise ValueError(
f"new_n_bins_mass ({new_n_bins_mass}) must be a divisor of "
f"current n_bins_mass ({halo_hist.n_bins_mass})"
)
if halo_hist.n_bins_z % new_n_bins_z != 0:
raise ValueError(
f"new_n_bins_z ({new_n_bins_z}) must be a divisor of "
f"current n_bins_z ({halo_hist.n_bins_z})"
)
if halo_hist.n_bins_time % new_n_bins_time != 0:
raise ValueError(
f"new_n_bins_time ({new_n_bins_time}) must be a divisor of "
f"current n_bins_time ({halo_hist.n_bins_time})"
)
# Calculate rebin factors
mass_factor = halo_hist.n_bins_mass // new_n_bins_mass
z_factor = halo_hist.n_bins_z // new_n_bins_z
time_factor = halo_hist.n_bins_time // new_n_bins_time
# Create new histogram instance
new_hist = HaloHistogram(
z_min=halo_hist.z_min,
z_max=halo_hist.z_max,
min_mass=halo_hist.min_mass,
max_mass=halo_hist.max_mass,
min_time_since_merger=halo_hist.min_time_since_merger,
max_time_since_merger=halo_hist.max_time_since_merger,
n_bins_mass=new_n_bins_mass,
n_bins_z=new_n_bins_z,
n_bins_time=new_n_bins_time,
)
# Manually initialize the bins (same as what would happen in add_halo_catalog)
log_mass_bins = np.linspace(
new_hist.min_log_mass, new_hist.max_log_mass, new_n_bins_mass + 1
)
new_hist.mass_bins = np.power(10, log_mass_bins)
new_hist.z_bins = np.linspace(new_hist.z_min, new_hist.z_max, new_n_bins_z + 1)
new_hist.time_bins = np.linspace(
new_hist.min_time_since_merger,
new_hist.max_time_since_merger,
new_n_bins_time + 1,
)
new_hist.no_catalogs_added = False
# Rebin the histograms
new_hist.histogram_all = _rebin_2d(halo_hist.histogram_all, z_factor, mass_factor)
new_hist.histogram_central = _rebin_2d(
halo_hist.histogram_central, z_factor, mass_factor
)
new_hist.histogram_subhalo = _rebin_3d(
halo_hist.histogram_subhalo, z_factor, mass_factor, time_factor
)
return new_hist
def _rebin_2d(histogram, factor_axis0, factor_axis1):
"""
Helper function to rebin a 2D histogram by combining neighboring bins.
"""
old_shape = histogram.shape
new_shape = (old_shape[0] // factor_axis0, old_shape[1] // factor_axis1)
# Reshape and sum over the dimensions to be combined
reshaped = histogram.reshape(new_shape[0], factor_axis0, new_shape[1], factor_axis1)
return reshaped.sum(axis=(1, 3))
def _rebin_3d(histogram, factor_axis0, factor_axis1, factor_axis2):
"""
Helper function to rebin a 3D histogram by combining neighboring bins.
"""
old_shape = histogram.shape
new_shape = (
old_shape[0] // factor_axis0,
old_shape[1] // factor_axis1,
old_shape[2] // factor_axis2,
)
# Reshape and sum over the dimensions to be combined
reshaped = histogram.reshape(
new_shape[0],
factor_axis0,
new_shape[1],
factor_axis1,
new_shape[2],
factor_axis2,
)
return reshaped.sum(axis=(1, 3, 5))