Coverage for src/galsbi/ucat/galaxy_population_models/galaxy_position.py: 98%
200 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-09 18:05 +0000
1# Copyright (C) 2018 ETH Zurich, Institute for Particle Physics and Astrophysics
3"""
4Created 2021
5author: Tomasz Kacprzak
6"""
8import hashlib
10import h5py
11import healpy as hp
12import numpy as np
13from cosmic_toolbox import arraytools as at
14from cosmic_toolbox import logger
15from ufig import coordinate_util
17from galsbi.sham_ot import histograms, sham_ot
18from galsbi.ucat.galaxy_population_models.galaxy_luminosity_function import (
19 LuminosityFunction,
20)
22SEED_OFFSET_LUMFUN = 123491
24##########################################################
25#
26# Sampling position
27#
28##########################################################
30LOGGER = logger.get_logger(__file__)
32# Cache OT solutions within process to avoid repeated expensive solves
33_OT_SOLUTION_CACHE = {}
35# Consistent with PINOCCHIO/SHAM pipeline
36GALAXY_TYPE_FLAGS = {
37 "blue": 1,
38 "red": 0,
39}
42def _hash_array(array):
43 """Hash numpy array content without keeping extra copies in memory."""
44 arr = np.ascontiguousarray(array)
45 h = hashlib.sha1()
46 h.update(np.asarray(arr.shape, dtype=np.int64).tobytes())
47 h.update(str(arr.dtype).encode())
48 h.update(arr.tobytes())
49 return h.hexdigest()
52def _safe_key_val(value):
53 """Return a hashable, comparable representation without forcing float casts."""
54 if value is None:
55 return None
56 if isinstance(value, str | bytes | bool):
57 return value
58 try:
59 return float(value)
60 except (TypeError, ValueError):
61 return repr(value)
64def _build_ot_cache_key(
65 par,
66 lumfunc,
67 galaxy_type,
68 zmin,
69 zmax,
70 logmass_bin_centers,
71 histogram_slice,
72 ot_epsilon,
73):
74 """Construct cache key covering all OT inputs affecting a redshift bin."""
75 return (
76 "bin",
77 float(zmin),
78 float(zmax),
79 galaxy_type,
80 "sham",
81 par.sham_model,
82 _safe_key_val(par.mlim_sham),
83 _safe_key_val(par.mlim_sigma_sham),
84 _safe_key_val(par.quenching_time_sham),
85 _safe_key_val(par.quenching_time_sigma_sham),
86 _safe_key_val(par.initial_quenched_fraction_sham),
87 _safe_key_val(par.m_mu_peng),
88 _safe_key_val(par.sigma_peng),
89 _safe_key_val(par.b_peng),
90 "ot",
91 _safe_key_val(ot_epsilon),
92 _safe_key_val(par.ot_mu_0),
93 _safe_key_val(par.ot_alpha),
94 _safe_key_val(par.ot_regularizer),
95 "lum",
96 lumfunc.name,
97 lumfunc.parametrization,
98 _safe_key_val(lumfunc.m_star_slope),
99 _safe_key_val(lumfunc.m_star_intcpt),
100 _safe_key_val(lumfunc.phi_star_amp),
101 _safe_key_val(lumfunc.phi_star_exp),
102 _safe_key_val(lumfunc.z_const),
103 _safe_key_val(lumfunc.alpha),
104 _safe_key_val(lumfunc.z_res),
105 _safe_key_val(lumfunc.m_res),
106 _safe_key_val(lumfunc.z_min),
107 _safe_key_val(lumfunc.z_max),
108 _safe_key_val(lumfunc.m_min),
109 _safe_key_val(lumfunc.m_max),
110 _safe_key_val(lumfunc.ngal_multiplier),
111 _safe_key_val(lumfunc.pixarea),
112 "data",
113 _hash_array(logmass_bin_centers),
114 _hash_array(histogram_slice),
115 "pixarea_halo_cat",
116 _safe_key_val(par.pixarea_of_halo_catalog),
117 )
120def sample_position_uniform(numobj, w, pixel_index, nside):
121 """
122 Sample a RING-ordered HEALPix pixel uniformly in solid angle.
124 Longitude bounds are unwrapped around the pixel centre when they cross
125 the signed-longitude boundary. Sampling is uniform in longitude and
126 cos(theta), followed by rejection outside the requested pixel. The
127 caller controls reproducibility through NumPy's existing random state.
129 :param numobj: Number of uniform samples
130 :param w: wcs-object containing all the relevant wcs-information
131 :param pixel_index: Index of the Healpix pixels sampled
132 :param nside: NSIDE of the Healpix map
133 :return: Uniformly drawn x-coordinate (in pixels)
134 :return: Uniformly drawn y-coordinate (in pixels)
135 """
137 corners = hp.boundaries(nside, pixel_index, 1)
138 theta, phi = hp.vec2ang(np.transpose(corners))
140 phi[phi > np.pi] -= 2 * np.pi # To deal with periodic boundaries
141 if np.ptp(phi) > np.pi:
142 # A fixed longitude cut can split a small pixel across -pi and +pi.
143 # Use its centre to place every corner on the same local interval.
144 _, centre_phi = hp.pix2ang(nside, pixel_index)
145 phi = centre_phi + (phi - centre_phi + np.pi) % (2 * np.pi) - np.pi
147 cos_theta_sample = np.random.uniform(
148 low=np.min(np.cos(theta)), high=np.max(np.cos(theta)), size=numobj
149 )
150 phi_sample = np.random.uniform(low=np.min(phi), high=np.max(phi), size=numobj)
152 n_reps = 0
154 while True:
155 mask = hp.ang2pix(
156 nside, np.arccos(cos_theta_sample), phi_sample % (2 * np.pi)
157 ) != pixel_index
159 if np.sum(mask) == 0:
160 break
162 cos_theta_sample[mask] = np.random.uniform(
163 low=np.min(np.cos(theta)), high=np.max(np.cos(theta)), size=np.sum(mask)
164 )
166 phi_sample[mask] = np.random.uniform(
167 low=np.min(phi), high=np.max(phi), size=np.sum(mask)
168 )
170 n_reps += 1
172 if n_reps > 10_000:
173 raise RuntimeError("more than 10_000 many iterations")
174 if w is None:
175 # If no wcs is provided, ra and dec are returned
176 return coordinate_util.thetaphi2radec(np.arccos(cos_theta_sample), phi_sample)
177 else:
178 # If a wcs is provided, x and y are returned
179 return coordinate_util.thetaphi2xy(w, np.arccos(cos_theta_sample), phi_sample)
182def sample_with_sham_ot(par, halo_catalog, lumfunc, pixels, w=None):
183 """
184 Samples absolute magnitude, redshift and position of galaxies using the SHAM-OT
185 framework. If the halo catalog does not contain enough halos to populate the faint
186 end of the luminosity function or the high redshift part, the rest of the luminosity
187 function is sampled with a uniform distribution in position.
189 :param par: ucat parameters from config file
190 :param halo_catalog: Halo catalog with halo mass, redshift, ra, dec and is_blue
191 :param lumfunc: Luminosity function instance
192 :param pixels: Healpix pixels to sample from
193 :param w: wcs-object containing all the relevant wcs-information
194 :return: Absolute magnitude, redshift, position of galaxies, logmasses,
195 halo_or_subhalo, time_since_merger
196 """
197 # transform galaxy type from flag to string
198 galaxy_type = "blue" if lumfunc.galaxy_type == GALAXY_TYPE_FLAGS["blue"] else "red"
199 if galaxy_type == "blue":
200 cat_select = halo_catalog["is_blue"].astype(bool)
201 else:
202 cat_select = ~halo_catalog["is_blue"].astype(bool)
204 halo_catalog = halo_catalog[cat_select]
206 ot_epsilon = get_ot_epsilon(par, galaxy_type)
207 use_ot_cache = par.ot_cache_enable
208 # Pre-sort halo catalog by redshift for efficient binary search
209 z_sort_indices = np.argsort(halo_catalog["z_obs"])
210 halo_catalog = halo_catalog[z_sort_indices]
212 halo_histogram = histograms.load_halo_histogram(par.halo_histograms_filename)
213 LOGGER.info(f"Sampling {galaxy_type} galaxies with SHAM-OT model {par.sham_model}")
214 LOGGER.debug(f"OT epsilon: {ot_epsilon}")
215 LOGGER.debug(f"mlim SHAM: {par.mlim_sham}")
216 LOGGER.debug(f"quenching time: {par.quenching_time_sham}")
217 LOGGER.debug(f"mlim sigma: {par.mlim_sigma_sham}")
218 LOGGER.debug(f"initial quenched fraction: {par.initial_quenched_fraction_sham}")
219 LOGGER.debug(f"quenching time sigma: {par.quenching_time_sigma_sham}")
220 LOGGER.debug(f"M_mu (Peng): {par.m_mu_peng}")
221 LOGGER.debug(f"sigma (Peng): {par.sigma_peng}")
222 LOGGER.debug(f"b (Peng): {par.b_peng}")
223 histogram = halo_histogram.get_histogram(
224 galaxy_type=galaxy_type,
225 model=par.sham_model,
226 m_limit=par.mlim_sham,
227 quenching_time=par.quenching_time_sham,
228 mlim_sigma=par.mlim_sigma_sham,
229 initial_quenched_fraction=par.initial_quenched_fraction_sham,
230 quenching_time_sigma=par.quenching_time_sigma_sham,
231 M_mu=par.m_mu_peng,
232 sigma=par.sigma_peng,
233 b=par.b_peng,
234 )
235 redshift_bin_edges, redshift_bin_centers = halo_histogram.get_z_bins()
236 _, mass_bin_centers = halo_histogram.get_mass_bins()
237 logmass_bin_centers = np.log10(mass_bin_centers)
239 absmags = np.array([])
240 redshift = np.array([])
241 ra = np.array([])
242 dec = np.array([])
243 logmasses = np.array([])
244 halo_or_subhalo = np.array([])
245 time_since_merger = np.array([])
247 faintest_mags = []
248 if par.save_galaxy_halo_relation:
249 # save for each redshift bin the absolute magnitude to logmass relation
250 # we have the ot_solution which interpolates from absmags to logmass
251 logmasses_to_save = np.linspace(
252 par.galaxy_halo_relation_logmass_min,
253 par.galaxy_halo_relation_logmass_max,
254 num=par.galaxy_halo_relation_nbins,
255 )
256 gh_relation = []
257 redshifts_to_save = []
258 # loop trough redshift bins
259 faintest_mag_initial_guess = -20
260 for i, z_center in enumerate(redshift_bin_centers):
261 LOGGER.debug(f"Sampling galaxies at redshift {z_center:.2f}.")
262 if sum(histogram[i]) == 0:
263 LOGGER.debug(f"No halos in the histogram in redshift bin {z_center:.2f}.")
264 faintest_mags.append(np.inf)
265 continue # skip redshift bin if no halos are present
267 zmin = redshift_bin_edges[i]
268 zmax = redshift_bin_edges[i + 1]
269 # Use binary search to find start and end indices for this redshift bin
270 start_idx = np.searchsorted(halo_catalog["z_obs"], zmin, side="left")
271 end_idx = np.searchsorted(halo_catalog["z_obs"], zmax, side="left")
272 if start_idx == end_idx:
273 LOGGER.debug(f"No halos in halo catalog in redshift bin {z_center:.2f}.")
274 faintest_mags.append(np.inf)
275 continue # skip redshift bin if no halos are present
277 cache_key = None
278 cached = None
279 if use_ot_cache:
280 cache_key = _build_ot_cache_key(
281 par,
282 lumfunc,
283 galaxy_type,
284 zmin,
285 zmax,
286 logmass_bin_centers,
287 histogram[i],
288 ot_epsilon,
289 )
290 cached = _OT_SOLUTION_CACHE.get(cache_key)
292 if cached is not None:
293 ot_solution, faintest_mag = cached
294 LOGGER.debug(
295 "Reusing cached OT solution for "
296 f"z=[{zmin:.3f}, {zmax:.3f}) ({galaxy_type})."
297 )
298 else:
299 ot_solution, faintest_mag = sham_ot.find_ot_solution(
300 logmass_bin_centers,
301 histogram[i],
302 lumfunc,
303 zmin,
304 zmax,
305 par.pixarea_of_halo_catalog,
306 faintest_mag_initial_guess=faintest_mag_initial_guess,
307 epsilon=ot_epsilon,
308 mu_0=par.ot_mu_0,
309 alpha=par.ot_alpha,
310 regularizer=par.ot_regularizer,
311 )
312 if cache_key is not None:
313 _OT_SOLUTION_CACHE[cache_key] = (ot_solution, faintest_mag)
315 LOGGER.debug(f"Faintest magnitude: {faintest_mag:.2f}")
316 # Update the initial guess for the next iteration
317 faintest_mag_initial_guess = faintest_mag
318 if par.save_galaxy_halo_relation:
319 gh_relation.append(ot_solution(logmasses_to_save))
320 redshifts_to_save.append(z_center)
322 am, z, r, d, ish, logm, dt = sham_ot.sample_catalog_from_simulation(
323 halo_catalog[start_idx:end_idx],
324 ot_solution,
325 )
326 absmags = np.concatenate((absmags, am))
327 redshift = np.concatenate((redshift, z))
328 ra = np.concatenate((ra, r))
329 dec = np.concatenate((dec, d))
330 logmasses = np.concatenate((logmasses, logm))
331 halo_or_subhalo = np.concatenate((halo_or_subhalo, ish))
332 time_since_merger = np.concatenate((time_since_merger, dt))
334 faintest_mags.append(faintest_mag)
336 faintest_mags = np.array(faintest_mags)
337 if par.save_galaxy_halo_relation:
338 # save the galaxy-halo relation for each redshift bin
339 gh_relation = np.array(gh_relation)
340 filename = par.galaxy_halo_relation_filename_format.format(galaxy_type)
342 # Save as HDF5 file
343 with h5py.File(filename, "w") as f:
344 f.create_dataset("galaxy_halo_relation", data=gh_relation)
345 f.create_dataset("logmasses", data=logmasses_to_save)
346 f.create_dataset("redshifts", data=redshifts_to_save)
348 LOGGER.info(
349 f"Saved galaxy-halo relation for {galaxy_type} galaxies to {filename}"
350 )
352 if np.any(faintest_mags < par.lum_fct_m_max):
353 LOGGER.warning(
354 "Not enough halos to populate the faint end of LF for all redshift bins."
355 )
356 am, z, r, d = sample_rest_of_lf(
357 lumfunc, 0, redshift_bin_edges[-1], np.min(faintest_mags), pixels, par
358 )
359 for i, _ in enumerate(redshift_bin_centers):
360 zmin = redshift_bin_edges[i]
361 zmax = redshift_bin_edges[i + 1]
362 faintest_mag = faintest_mags[i]
363 # select objects in redshift bin that are fainter than the faintest
364 # magnitude from the previous loop
365 select = (z >= zmin) & (z < zmax) & (am > faintest_mag)
366 absmags = np.concatenate((absmags, am[select]))
367 redshift = np.concatenate((redshift, z[select]))
368 ra = np.concatenate((ra, r[select]))
369 dec = np.concatenate((dec, d[select]))
370 logmasses = np.concatenate((logmasses, np.nan * np.ones_like(am[select])))
371 halo_or_subhalo = np.concatenate(
372 (halo_or_subhalo, np.nan * np.ones_like(am[select]))
373 )
374 time_since_merger = np.concatenate(
375 (time_since_merger, np.nan * np.ones_like(am[select]))
376 )
378 if redshift_bin_edges[-1] < lumfunc.z_max:
379 LOGGER.info(
380 f"Sampling the high redshift part of the LF (z > {redshift_bin_edges[-1]})"
381 )
382 am, z, r, d = sample_rest_of_lf(
383 lumfunc, redshift_bin_edges[-1], None, None, pixels, par
384 )
385 logm = np.nan * np.ones_like(am)
386 absmags = np.concatenate((absmags, am))
387 redshift = np.concatenate((redshift, z))
388 ra = np.concatenate((ra, r))
389 dec = np.concatenate((dec, d))
390 logmasses = np.concatenate((logmasses, logm))
391 halo_or_subhalo = np.concatenate((halo_or_subhalo, np.nan * np.ones_like(am)))
392 time_since_merger = np.concatenate(
393 (time_since_merger, np.nan * np.ones_like(am))
394 )
396 if w is None:
397 return absmags, redshift, ra, dec, logmasses, halo_or_subhalo, time_since_merger
398 else:
399 x, y = coordinate_util.radec2xy(w, ra, dec)
400 return absmags, redshift, x, y, logmasses, halo_or_subhalo, time_since_merger
403def get_ot_epsilon(par, galaxy_type):
404 """
405 Get the OT epsilon parameter based on the configuration parameters.
406 If ot_epsilon is not None, it is used for both galaxy types.
407 If ot_epsilon is None, the ot_epsilon_red and ot_epsilon_blue parameters are used
409 :param par: ucat parameters from config file
410 :param galaxy_type: Type of galaxy ("red" or "blue")
411 :return: OT epsilon parameter
412 """
413 if par.ot_epsilon is not None:
414 return par.ot_epsilon
415 else:
416 if galaxy_type == "red":
417 return par.ot_epsilon_red
418 else:
419 return par.ot_epsilon_blue
422def sample_rest_of_lf(lumfunc, zmin, zmax, m_min, pixels, par):
423 """
424 Samples the rest of the luminosity function in case there are not enough halos.
425 Positions are sampled uniformly.
427 :param lumfunc: Luminosity function instance
428 :param zmin: Minimum redshift
429 :param zmax: Maximum redshift
430 :param m_min: Minimum absolute magnitude (brightest magnitude still to be sampled)
431 :param pixels: Healpix pixels to sample from
432 :param par: ucat parameters from config file
433 :return: Absolute magnitude, redshift, ra and dec of galaxies
434 """
435 lf = LuminosityFunction(
436 name=lumfunc.name,
437 lum_fct_parametrization=lumfunc.parametrization,
438 m_star_slope=lumfunc.m_star_slope,
439 m_star_intcpt=lumfunc.m_star_intcpt,
440 phi_star_amp=lumfunc.phi_star_amp,
441 phi_star_exp=lumfunc.phi_star_exp,
442 z_const=lumfunc.z_const,
443 alpha=lumfunc.alpha,
444 cosmo=lumfunc.cosmo,
445 pixarea=lumfunc.pixarea,
446 galaxy_type=lumfunc.galaxy_type,
447 seed_ngal=par.seed_ngal,
448 z_res=lumfunc.z_res,
449 m_res=lumfunc.m_res,
450 z_min=zmin if zmin is not None else lumfunc.z_min,
451 z_max=zmax if zmax is not None else lumfunc.z_max,
452 m_min=m_min if m_min is not None else lumfunc.m_min,
453 m_max=lumfunc.m_max,
454 z_m_intp=lumfunc.z_m_intp,
455 ngal_multiplier=lumfunc.ngal_multiplier,
456 )
458 abs_mag = np.array([])
459 redshift = np.array([])
460 ra = np.array([])
461 dec = np.array([])
462 for i, pix in enumerate(pixels):
463 seed_ngal = par.seed + pix + par.gal_num_seed_offset + SEED_OFFSET_LUMFUN * i
464 seed_lumfunc = (
465 par.seed + pix + par.gal_lum_fct_seed_offset + SEED_OFFSET_LUMFUN * i
466 )
467 am, z = lf.sample_z_mabs_and_apply_cut(
468 seed_ngal=seed_ngal, seed_lumfun=seed_lumfunc
469 )
471 np.random.seed(par.seed + pix + par.gal_dist_seed_offset)
472 r, d = sample_position_uniform(len(am), None, pix, par.nside_sampling)
473 abs_mag = np.concatenate((abs_mag, am))
474 redshift = np.concatenate((redshift, z))
475 ra = np.concatenate((ra, r))
476 dec = np.concatenate((dec, d))
477 return abs_mag, redshift, ra, dec
480def get_halo_catalog(par):
481 """
482 Loads the halo catalog from the specified file and adds a column indicating
483 whether the galaxy is blue or not based on the SHAM-OT model parameters.
485 :param par: ucat parameters from config file
486 :return: halo catalog with an additional column 'is_blue'
487 """
488 halo_catalog = at.load_hdf_cols(par.halo_catalog_filename)
489 is_blue = sham_ot.select_type(
490 halo_catalog,
491 is_blue=True,
492 model=par.sham_model,
493 m_limit=par.mlim_sham,
494 mlim_sigma=par.mlim_sigma_sham,
495 initial_quenched_fraction=par.initial_quenched_fraction_sham,
496 quenching_time=par.quenching_time_sham,
497 quenching_time_sigma=par.quenching_time_sigma_sham,
498 M_mu=par.m_mu_peng,
499 sigma=par.sigma_peng,
500 b=par.b_peng,
501 )
502 if "is_blue" in halo_catalog.dtype.names:
503 LOGGER.warning(
504 "Column 'is_blue' already exists in the halo catalog. "
505 "Updating the existing column."
506 )
507 halo_catalog = at.delete_cols(halo_catalog, ["is_blue"])
509 halo_catalog = at.add_cols(
510 halo_catalog,
511 ["is_blue"],
512 data=is_blue,
513 )
514 return halo_catalog