Coverage for src/galsbi/sham_ot/sham_ot.py: 99%
193 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-09 18:05 +0000
1# Copyright (C) 2025 ETH Zurich
2# Institute for Particle Physics and Astrophysics
3# Author: Silvan Fischbacher
4# created: Thu Mar 06 2025
6import numpy as np
7import ot
8from cosmic_toolbox import logger
9from scipy.integrate import cumulative_trapezoid
10from scipy.interpolate import interp1d
11from scipy.optimize import fsolve
12from scipy.special import erf
14from galsbi.sham_ot import utils
16LOGGER = logger.get_logger(__name__)
18BRIGHTEST_GALAXY = -27
21def find_ot_solution(
22 logmasses,
23 halo_histogram,
24 luminosity_function,
25 z0,
26 z1,
27 pixarea,
28 magnitude_bins=100,
29 m_max=-5,
30 faintest_mag_initial_guess=-20,
31 epsilon=None,
32 mu_0=9.94,
33 alpha=-1.7,
34 regularizer="sinkhorn",
35):
36 """
37 Finds the OT solution for the given mass distribution and luminosity function.
38 Returns an interpolator predicting the absolute magnitude given the log halo mass.
40 :param logmasses: log-masses of the halos at which the histogram is evaluated
41 :param halo_histogram: histogram of the halos at the given redshift
42 :param luminosity_function: luminosity function instance
43 :param z0: lower redshift limit
44 :param z1: upper redshift limit
45 :param pixarea: sky area in square radians of the simulation
46 :param magnitude_bins: number of bins in the absolute magnitude histogram
47 :param m_max: maximum absolute magnitude to consider (default: -5)
48 :param faintest_mag_initial_guess: initial guess for faintest magnitude
49 :param epsilon: regularization parameter for OT (None for unregularized)
50 :param mu_0: halo mass scale parameter for cost matrix modification (Fischbacher+25)
51 :param alpha: exponent for cost matrix scaling (Fischbacher+25)
52 :param regularizer: regularizer type for OT
53 :return: OT solution interpolator, faintest magnitude
54 """
55 n_halos = np.sum(halo_histogram)
56 faintest_absmag = compute_faintest_absmag(
57 luminosity_function,
58 z0,
59 z1,
60 n_halos,
61 pixarea,
62 initial_guess=faintest_mag_initial_guess,
63 )
64 if faintest_absmag > m_max:
65 LOGGER.debug(
66 f"Faintest absolute magnitude {faintest_absmag} is larger than the maximum "
67 f"absolute magnitude {m_max}. Using m_max instead."
68 )
69 faintest_absmag = m_max
70 n_halos_adjusted = compute_n_halos_given_faintest_absmag(
71 luminosity_function, z0, z1, faintest_absmag, pixarea
72 )
73 # adjust halo histogram such that the sum is corresponding to adjusted halo n
74 # this means we keep the high mass halos and remove the low mass bins
75 if n_halos_adjusted < 1:
76 # avoid empty histogram
77 n_halos_to_remove = n_halos - 1
78 else:
79 n_halos_to_remove = np.round(n_halos - n_halos_adjusted)
80 while n_halos_to_remove > 0:
81 # remove the lowest mass bin
82 if halo_histogram[0] > n_halos_to_remove:
83 halo_histogram[0] -= n_halos_to_remove
84 n_halos_to_remove = 0
85 else:
86 n_halos_to_remove -= np.sum(halo_histogram[0])
87 halo_histogram = halo_histogram[1:]
88 logmasses = logmasses[1:]
90 absmags, absmags_histogram = generate_galaxy_histogram(
91 luminosity_function,
92 faintest_absmag,
93 brightest_absmag=BRIGHTEST_GALAXY,
94 nbins=magnitude_bins,
95 redshift=(z0 + z1) / 2,
96 )
97 empty_mass_bins = halo_histogram == 0
99 _, ot = solve_ot(
100 halo_histogram[~empty_mass_bins],
101 absmags_histogram,
102 logmasses=logmasses[~empty_mass_bins],
103 absmags=absmags,
104 epsilon=epsilon,
105 mu_0=mu_0,
106 alpha=alpha,
107 regularizer=regularizer,
108 )
110 # Choose between regularized (stochastic) and unregularized (deterministic) OT
111 if epsilon is None:
112 # Unregularized OT: use deterministic interpolator
113 ot_interp = get_interpolated_ot_solution(
114 ot, logmasses[~empty_mass_bins], absmags
115 )
116 else:
117 # Regularized OT: use stochastic sampler
118 ot_interp = get_regularized_ot_solution(
119 ot, logmasses[~empty_mass_bins], absmags
120 )
122 return ot_interp, faintest_absmag
125def compute_faintest_absmag(
126 luminosity_function, z0, z1, n_halos, pixarea, initial_guess=-20, grid_size=100
127):
128 """
129 Compute the faintest absolute magnitude for a given number of halos in a given
130 redshift range. This is done by computing the integral of the luminosity function
131 over the redshift range and from the maximum absolute magnitude to the faintest
132 absolute magnitude.
134 :param luminosity_function: luminosity function instance
135 :param z0: lower redshift limit
136 :param z1: upper redshift limit
137 :param n_halos: number of halos in the redshift range
138 :param pixarea: sky area in square radians of the simulation
139 :param initial_guess: initial guess for the faintest absolute magnitude
140 :param grid_size: number of points in the redshift grid for integration
141 :return: faintest absolute magnitude
142 """
144 def _func(m_max):
145 # Create a grid of redshift points
146 z_points = np.linspace(z0, z1, grid_size)
148 z_density_values = z_density(luminosity_function, z_points, m_max)
150 integral = cumulative_trapezoid(z_density_values, z_points, initial=0)
151 integral = integral[-1] # Take the final value for total integral
152 return integral * pixarea - n_halos
154 faintest_absmag = fsolve(_func, x0=initial_guess)[0]
156 return faintest_absmag
159def compute_n_halos_given_faintest_absmag(
160 luminosity_function, z0, z1, faintest_absmag, pixarea, grid_size=100
161):
162 """
163 Compute the number of halos given the faintest absolute magnitude in a given
164 redshift range. This is done by computing the integral of the luminosity function
165 over the redshift range and from the maximum absolute magnitude to the faintest
166 absolute magnitude.
168 :param luminosity_function: luminosity function instance
169 :param z0: lower redshift limit
170 :param z1: upper redshift limit
171 :param faintest_absmag: faintest absolute magnitude
172 :param pixarea: sky area in square radians of the simulation
173 :param grid_size: number of points in the redshift grid for integration
174 :return: number of halos in the redshift range
175 """
176 # Create a grid of redshift points (same as in compute_faintest_absmag)
177 z_points = np.linspace(z0, z1, grid_size)
179 z_density_values = z_density(luminosity_function, z_points, faintest_absmag)
181 integral = cumulative_trapezoid(z_density_values, z_points, initial=0)
182 integral = integral[-1] # Take the final value for total integral
184 n_halos = integral * pixarea
186 return n_halos
189def z_density(lumfunction, z, m_max):
190 m_star_par = lumfunction.m_star_slope, lumfunction.m_star_intcpt
191 phi_star_par = (
192 lumfunction.phi_star_amp,
193 lumfunction.phi_star_exp,
194 )
196 return lumfunction.n_gal_calc._redshift_density(
197 z=z,
198 m_min=None,
199 m_max=m_max,
200 parametrization=lumfunction.parametrization,
201 z_const=lumfunction.z_const,
202 alpha=lumfunction.alpha,
203 par_m_star=m_star_par,
204 par_phi_star=phi_star_par,
205 cosmo=lumfunction.cosmo,
206 )
209def generate_galaxy_histogram(
210 luminosity_function,
211 faintest_absmag,
212 brightest_absmag=BRIGHTEST_GALAXY,
213 nbins=100,
214 redshift=0,
215):
216 """
217 Generate a histogram of galaxies in absolute magnitude space based on the given
218 luminosity function and redshift.
220 :param luminosity_function: luminosity function instance
221 :param faintest_absmag: faintest absolute magnitude
222 :param brightest_absmag: brightest absolute magnitude
223 :param nbins: number of bins in the histogram
224 :param redshift: redshift of the galaxies
225 :return: absolute magnitudes, histogram
226 """
227 absmags = np.linspace(faintest_absmag, brightest_absmag, nbins)
228 absmags_histogram = luminosity_function.analytical_number_density(absmags, redshift)
230 return absmags, absmags_histogram
233def solve_ot(
234 masses_histogram,
235 absmags_histogram,
236 logmasses=None,
237 absmags=None,
238 epsilon=None,
239 mu_0=9.94,
240 alpha=-1.7,
241 regularizer="sinkhorn",
242 norm="sqeuclidean",
243):
244 """
245 Solve the Optimal Transport problem between two histograms.
246 This function assumes that the histograms are already consistently normalized.
247 If unregularized OT is used (epsilon=None), it is not necessary to define the
248 logmasses and absmags. This makes only a difference if the regularized OT is used.
250 :param masses_histogram: Histogram of the halo masses
251 :param absmags_histogram: Histogram of the galaxy absolute magnitudes
252 :param logmasses: log-masses of the halos
253 :param absmags: absolute magnitudes of the galaxies
254 :param epsilon: Regularization parameter
255 :param mu_0: halo mass scale parameter for cost matrix modification (Fischbacher+25)
256 :param alpha: exponent for cost matrix scaling (Fischbacher+25)
257 :param regularizer: Regularizer to use
258 :param norm: Norm to use
259 :return: cost matrix, ot solution
260 """
262 masses_histogram = masses_histogram / np.sum(masses_histogram)
263 absmags_histogram = absmags_histogram / np.sum(absmags_histogram)
265 x = np.arange(len(masses_histogram)) if logmasses is None else logmasses
266 y = np.arange(len(absmags_histogram)) if absmags is None else -absmags
268 M = ot.dist(np.expand_dims(x, 1), np.expand_dims(y, 1), metric=norm).astype(
269 np.float64
270 )
272 # Apply Fischbacher+25 cost matrix modification when regularized OT is used
273 if epsilon is not None:
274 M *= (x - mu_0)[:, np.newaxis] ** alpha
276 M /= M.max()
278 if epsilon is None:
279 G = ot.emd(masses_histogram, absmags_histogram, M)
280 else:
281 if regularizer == "sinkhorn":
282 G = ot.sinkhorn(masses_histogram, absmags_histogram, M, epsilon)
283 elif regularizer == "sinkhorn_stabilized":
284 G = ot.bregman.sinkhorn_stabilized(
285 masses_histogram, absmags_histogram, M, epsilon
286 )
287 elif regularizer == "quadratic":
288 G = ot.smooth.smooth_ot_dual(
289 masses_histogram, absmags_histogram, M, epsilon, reg_type="kl"
290 )
291 else:
292 raise ValueError(f"Regularizer {regularizer} not recognized.")
294 return M, G
297def get_regularized_ot_solution(ot_solution, x1, x2):
298 """
299 Returns a callable for regularized OT solution that samples from probability
300 distributions. For each query logmass, it interpolates between nearby probability
301 distributions and samples. Vectorized implementation for efficient processing of
302 large arrays.
304 :param ot_solution: OT solution matrix (n x m)
305 :param x1: x values of the source distribution (logmasses) (n)
306 :param x2: x values of the target distribution (absolute magnitudes) (m)
307 :return: callable that takes logmass and returns sampled absolute magnitude
308 """
310 def sample_magnitude(logmass):
311 """Sample absolute magnitude for given logmass from interpolated probability
312 distribution."""
313 # Handle scalar input
314 scalar_input = np.isscalar(logmass)
315 logmass = np.atleast_1d(logmass)
316 n_queries = len(logmass)
318 # Vectorized interpolation index finding
319 indices_right = np.searchsorted(x1, logmass)
320 indices_left = np.maximum(indices_right - 1, 0)
321 indices_right = np.minimum(indices_right, len(x1) - 1)
323 # Handle boundary cases
324 below_grid = logmass <= x1[0]
325 above_grid = logmass >= x1[-1]
326 in_grid = ~(below_grid | above_grid)
328 # Initialize probability distributions array
329 prob_dists = np.zeros((n_queries, len(x2)))
331 # Below grid: use first distribution
332 if np.any(below_grid):
333 prob_dists[below_grid] = ot_solution[0, :]
335 # Above grid: use last distribution
336 if np.any(above_grid):
337 prob_dists[above_grid] = ot_solution[-1, :]
339 # In grid: interpolate between distributions
340 if np.any(in_grid):
341 # Get indices for in-grid points
342 left_idx = indices_left[in_grid]
343 right_idx = indices_right[in_grid]
345 # Compute interpolation weights vectorized
346 x_left = x1[left_idx]
347 x_right = x1[right_idx]
348 x_query = logmass[in_grid]
350 # Avoid division by zero for identical points
351 dx = x_right - x_left
352 safe_dx = np.where(
353 dx == 0, 1, dx
354 ) # Replace 0 with 1 to avoid division by zero
355 weight_right = np.where(dx == 0, 0, (x_query - x_left) / safe_dx)
356 weight_left = 1.0 - weight_right
358 # Vectorized interpolation of probability distributions
359 left_dists = ot_solution[left_idx, :]
360 right_dists = ot_solution[right_idx, :]
361 prob_dists[in_grid] = (
362 weight_left[:, np.newaxis] * left_dists
363 + weight_right[:, np.newaxis] * right_dists
364 )
366 # Normalize probability distributions
367 row_sums = prob_dists.sum(axis=1)
368 valid_dists = row_sums > 0
369 prob_dists[valid_dists] = (
370 prob_dists[valid_dists] / row_sums[valid_dists, np.newaxis]
371 )
373 # Handle invalid distributions (fallback to uniform or first magnitude)
374 invalid_dists = ~valid_dists
375 if np.any(invalid_dists):
376 prob_dists[invalid_dists, 0] = 1.0 # All probability on first magnitude
378 # Vectorized sampling using cumulative probabilities
379 cumprobs = np.cumsum(prob_dists, axis=1)
380 random_vals = np.random.random(n_queries)
382 # Interpolate within the cumulative probability bins for continuous sampling
383 # Find which bin each random value falls into
384 sampled_indices = (cumprobs >= random_vals[:, np.newaxis]).argmax(axis=1)
386 # Ensure indices are within bounds
387 sampled_indices = np.minimum(sampled_indices, len(x2) - 1)
389 # Get cumulative probabilities for interpolation
390 # For each query, get cumprob at sampled_index and previous index
391 cumprob_right = cumprobs[np.arange(n_queries), sampled_indices]
392 sampled_indices_left = np.maximum(sampled_indices - 1, 0)
393 cumprob_left = cumprobs[np.arange(n_queries), sampled_indices_left]
395 # Linear interpolation within the cumulative probability bin
396 # Find fractional position of random_val between cumprob_left and cumprob_right
397 cumprob_diff = cumprob_right - cumprob_left
398 # Avoid division by zero
399 safe_diff = np.where(cumprob_diff == 0, 1.0, cumprob_diff)
400 fractions = np.where(
401 cumprob_diff == 0, 0.5, (random_vals - cumprob_left) / safe_diff
402 )
404 # Interpolate between magnitude bin centers
405 mag_left = x2[sampled_indices_left]
406 mag_right = x2[sampled_indices]
408 # Get interpolated magnitudes
409 results = mag_left + fractions * (mag_right - mag_left)
411 return results[0] if scalar_input else results
413 return sample_magnitude
416def get_interpolated_ot_solution(ot_solution, x1, x2):
417 """
418 Returns a deterministic interpolator for unregularized OT solution.
419 To ensure a smooth interpolator, all constant segments in the OT solution are
420 removed.
422 :param ot_solution: OT solution matrix (n x m)
423 :param x1: x values of the source distribution (n)
424 :param x2: x values of the target distribution (m)
425 :return: interpolator
426 """
427 # Get target indices
428 target_idx = np.argmax(ot_solution, axis=1)
429 target_vals = x2[target_idx]
431 # Find unique steps (keep first point of each constant segment)
432 changes = np.where(np.diff(target_idx) != 0)[0]
433 keep_idx = np.concatenate(([0], changes + 1, [len(x1) - 1]))
434 keep_idx = np.unique(keep_idx)
436 # Create reduced arrays for interpolation
437 x1_reduced = x1[keep_idx]
438 target_reduced = target_vals[keep_idx]
440 # remove non monotonic segments
441 monotonic_mask = np.diff(target_reduced) < 0
442 monotonic_mask = np.concatenate(([True], monotonic_mask))
443 x1_reduced = x1_reduced[monotonic_mask]
444 target_reduced = target_reduced[monotonic_mask]
445 if sum(~monotonic_mask) > 0.25 * len(x1_reduced):
446 LOGGER.warning(
447 "More than 25% of the OT solution is non-monotonic. "
448 "The galaxy-halo connection might not be well defined."
449 )
451 return interp1d(
452 x1_reduced,
453 target_reduced,
454 fill_value=(target_reduced[0], target_reduced[-1]),
455 bounds_error=False,
456 )
459def sample_catalog_from_simulation(
460 halo_catalog,
461 ot_solution,
462):
463 """
464 Sample a galaxy catalog from a halo catalog using the OT solution.
466 :param halo_catalog: Halo catalog (masked for redshift)
467 :param ot_solution: OT solution interpolator
468 :param zmin: lower redshift limit
469 :param zmax: upper redshift limit
470 :return: abs. magnitudes, redshifts, ra, dec, is_halo, logmasses, time_since_merger
471 """
473 logmasses = np.log10(halo_catalog["m"])
474 absmags = ot_solution(logmasses)
475 return (
476 absmags,
477 halo_catalog["z_obs"],
478 halo_catalog["ra"],
479 halo_catalog["dec"],
480 halo_catalog["is_halo"],
481 logmasses,
482 halo_catalog["delta_t"], # time since merger
483 )
486def select_type(halo_catalog, is_blue, model="Berner+24", **kwargs):
487 """
488 Select galaxies based on the SHAM model and type (blue/red).
489 Depending on the model, different additional arguments are needed.
491 - For the Berner+24 model: m_limit (mass limit in Msun/h) and quenching_time
492 (time after which a galaxy is considered red in Gyr).
493 - For the Berner+24_ext model: m_limit, quenching_time,
494 mlim_sigma (transition width for the mass limit),
495 initial_quenched_fraction (baseline quenched fraction at low mass for subhalos),
496 quenching_time_sigma (transition width for the quenched fraction).
498 :param halo_catalog: Halo catalog
499 :param is_blue: flag to select blue galaxies
500 :param model: SHAM model to use for selecting galaxies
501 :param Mlimit: mass limit for the SHAM model (in Msun/h)
502 :param quenching_time: time after which a galaxy is considered red (in Gyr)
503 :param M_mu: mass limit for the Peng model (in log Msun/h)
504 :param sigma: transition width from star-forming to quenched (only for Peng model
505 :param b: baseline quenched fraction at low mass for subhalos (only for Peng model)
506 :return: boolean mask for selected galaxies
507 """
508 if model == "Berner+24":
509 kwargs.setdefault("m_limit", 8e12)
510 kwargs.setdefault("quenching_time", 2.0)
511 return select_type_berner(
512 halo_catalog,
513 is_blue,
514 m_limit=kwargs["m_limit"],
515 quenching_time=kwargs["quenching_time"],
516 )
517 elif model == "Berner+24_ext":
518 kwargs.setdefault("m_limit", 8e12)
519 kwargs.setdefault("quenching_time", 2.0)
520 kwargs.setdefault("mlim_sigma", 0.3)
521 kwargs.setdefault("initial_quenched_fraction", 0.1)
522 kwargs.setdefault("quenching_time_sigma", 0.5)
523 return select_type_berner_ext(
524 halo_catalog,
525 is_blue,
526 m_limit=kwargs["m_limit"],
527 quenching_time=kwargs["quenching_time"],
528 mlim_sigma=kwargs["mlim_sigma"],
529 initial_quenched_fraction=kwargs["initial_quenched_fraction"],
530 quenching_time_sigma=kwargs["quenching_time_sigma"],
531 )
532 elif model == "Peng+10":
533 kwargs.setdefault("M_mu", 11.935)
534 kwargs.setdefault("sigma", 0.51)
535 kwargs.setdefault("b", 0.41)
536 return select_type_peng(
537 halo_catalog,
538 is_blue,
539 M_mu=kwargs["M_mu"],
540 sigma=kwargs["sigma"],
541 b=kwargs["b"],
542 )
543 else:
544 raise ValueError(f"SHAM model {model} not recognized.")
547def select_type_berner(halo_catalog, is_blue, m_limit=8e12, quenching_time=2):
548 """
549 Select galaxies based on the Berner SHAM model. See arXiv:2310.15223 for details.
551 :param halo_catalog: Halo catalog
552 :param is_blue: flag to select blue galaxies (True) or red galaxies (False)
553 :param Mlimit: mass limit for the SHAM model (in Msun/h)
554 :param quenching_time: time after which a galaxy is considered red (in Gyr)
555 :return: boolean mask for selected galaxies
556 """
558 if is_blue:
559 central_blue = (halo_catalog["is_halo"] == 1) & (halo_catalog["m"] <= m_limit)
560 satellite_blue = (halo_catalog["is_halo"] == 0) & (
561 halo_catalog["delta_t"] <= quenching_time
562 )
563 type_mask = central_blue | satellite_blue
564 else:
565 central_red = (halo_catalog["is_halo"] == 1) & (halo_catalog["m"] > m_limit)
566 satellite_red = (halo_catalog["is_halo"] == 0) & (
567 halo_catalog["delta_t"] > quenching_time
568 )
569 type_mask = central_red | satellite_red
570 return type_mask
573def select_type_berner_ext(
574 halo_catalog,
575 is_blue,
576 m_limit=8e12,
577 quenching_time=2,
578 mlim_sigma=0.2,
579 initial_quenched_fraction=0.0,
580 quenching_time_sigma=0.2,
581):
582 """
583 Select galaxies based on the extended Berner SHAM model.
585 :param halo_catalog: Halo catalog
586 :param is_blue: flag to select blue galaxies (True) or red galaxies (False)
587 :param m_limit: mass limit for the SHAM model (in Msun/h)
588 :param quenching_time: time after which a galaxy is considered red (in Gyr)
589 :param mlim_sigma: transition width for the mass limit
590 :param initial_quenched_fraction: baseline quenched fraction at low mass
591 for subhalos
592 :param quenching_time_sigma: transition width for the quenched fraction
593 :return: boolean mask for selected galaxies
594 """
596 random_draw = np.random.uniform(size=len(halo_catalog))
597 centrals = halo_catalog["is_halo"] == 1
599 blue_centrals = centrals & (
600 random_draw
601 > utils.adapted_err_function(
602 np.log10(halo_catalog["m"]),
603 np.log10(m_limit),
604 mlim_sigma,
605 min_val=0,
606 max_val=1,
607 )
608 )
609 blue_satellites = ~centrals & (
610 random_draw
611 > utils.adapted_err_function(
612 halo_catalog["delta_t"],
613 quenching_time,
614 quenching_time_sigma,
615 min_val=initial_quenched_fraction,
616 max_val=1,
617 )
618 )
619 blues = blue_centrals | blue_satellites
620 type_mask = blues if is_blue else ~blues
621 return type_mask
624def select_type_peng(halo_catalog, is_blue, M_mu=11.935, sigma=0.51, b=0.41):
625 """ "
626 Select galaxies based on the Peng SHAM model. See arXiv:1003.4747 for details
628 :param halo_catalog: Halo catalog
629 :param is_blue: flag to select blue galaxies (True) or red galaxies (False)
630 :param M_mu: mass limit for the SHAM model (in log Msun/h)
631 :param sigma: transition width from star-forming to quenched
632 :param b: baseline quenched fraction at low mass for subhalos
633 :return: boolean mask for selected galaxies
634 """
636 centrals = halo_catalog["is_halo"] == 1
638 def _r0(M_mu, sigma):
639 return np.log10(halo_catalog["m"] / np.power(10, M_mu)) / (np.sqrt(2) * sigma)
641 def _quenched_fraction(b):
642 return (1 - b) / 2 * (1 + erf(_r0(M_mu, sigma))) + b
644 random_draw = np.random.uniform(size=len(halo_catalog))
645 blue_centrals = centrals & (random_draw > _quenched_fraction(0))
646 blue_subhalos = ~centrals & (random_draw < _quenched_fraction(b))
647 blues = blue_centrals | blue_subhalos
648 type_mask = blues if is_blue else ~blues
649 return type_mask