Coverage for src/galsbi/ucat/galaxy_population_models/galaxy_luminosity_function.py: 100%
180 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 on Sept 2021
5author: Tomasz Kacprzak
6using code from: Joerg Herbel
7"""
9from collections import OrderedDict
11import numpy as np
12import scipy.integrate
13import scipy.optimize
14import scipy.special
15from cosmic_toolbox import logger
17LOGGER = logger.get_logger(__file__)
20def find_closest_ind(grid, vals):
21 ind = np.searchsorted(grid, vals)
22 ind[ind == grid.size] -= 1
23 ind[np.fabs(grid[ind] - vals) > np.fabs(grid[ind - 1] - vals)] -= 1
24 return ind
27def initialize_luminosity_functions(par, pixarea, cosmo, z_m_intp=None):
28 kw_lumfun = dict(
29 pixarea=pixarea,
30 z_res=par.lum_fct_z_res,
31 m_res=par.lum_fct_m_res,
32 z_max=par.lum_fct_z_max,
33 m_max=par.lum_fct_m_max,
34 z_m_intp=z_m_intp,
35 ngal_multiplier=par.ngal_multiplier,
36 cosmo=cosmo,
37 )
38 if par.ngal_multiplier != 1:
39 LOGGER.warning(f"ngal_multiplier is set to {par.ngal_multiplier}")
41 lum_funcs = OrderedDict()
42 for i, g in enumerate(par.galaxy_types):
43 lum_funcs[g] = LuminosityFunction(
44 lum_fct_parametrization=par.lum_fct_parametrization,
45 m_star_slope=getattr(par, f"lum_fct_m_star_{g}_slope"),
46 m_star_intcpt=getattr(par, f"lum_fct_m_star_{g}_intcpt"),
47 phi_star_amp=getattr(par, f"lum_fct_phi_star_{g}_amp"),
48 phi_star_exp=getattr(par, f"lum_fct_phi_star_{g}_exp"),
49 z_const=getattr(par, f"lum_fct_z_const_{g}"),
50 alpha=getattr(par, f"lum_fct_alpha_{g}"),
51 name=g,
52 galaxy_type=i,
53 seed_ngal=par.seed_ngal,
54 **kw_lumfun,
55 )
56 return lum_funcs
59def maximum_redshift(
60 z_m_intp, m_max, z_max, parametrization, z_const, alpha, m_star_par, seed_ngal
61):
62 """
63 Computes the maximum redshift up to which we sample objects from the luminosity
64 function. The cutoff is based on the criterion that the CDF for absolute magnitudes
65 is larger than 1e-5, i.e. that there is a reasonable probability of actually
66 obtaining objects at this redshift and absolute magnitude which still pass the cut
67 on par.gals_mag_max.
68 """
69 if z_m_intp is None:
70 return z_max
72 def cond_mag_cdf_lim(z):
73 m_s = m_star_lum_fct(z, parametrization, z_const, *m_star_par)
74 cdf_lim = (
75 upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_s - z_m_intp(z))))
76 / upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_s - m_max)))
77 - 1e-5
78 )
79 return cdf_lim
81 try:
82 z_max_cutoff = scipy.optimize.brentq(cond_mag_cdf_lim, 0, z_m_intp.x[-1])
83 except ValueError:
84 z_max_cutoff = z_m_intp.x[-1]
86 z_max = min(z_max, z_max_cutoff)
87 np.random.seed(seed_ngal)
88 z_max += np.random.uniform(0, 0.0001)
90 return z_max
93def m_star_lum_fct(z, parametrization, z_const, slope, intercept):
94 if parametrization == "linexp":
95 # Eq. 3.3. from http://arxiv.org/abs/1705.05386
96 m_s = np.polyval((slope, intercept), z)
97 elif parametrization == "logpower":
98 # Eq. 21 + 22 from https://arxiv.org/abs/1106.2039
99 m_s = intercept + slope * np.log10(1 + z)
100 elif parametrization == "truncated_logexp":
101 m_s = intercept + slope * np.log10(1 + np.where(z < z_const, z, z_const))
102 return m_s
105def phi_star_lum_fct(z, parametrization, amplitude, exp):
106 if (parametrization == "linexp") | (parametrization == "truncated_logexp"):
107 # Eq. 3.4. from http://arxiv.org/abs/1705.05386
108 p_star = amplitude * np.exp(exp * z)
109 elif parametrization == "logpower":
110 # Eq. 25 from https://arxiv.org/abs/1106.2039
111 p_star = amplitude * (1 + z) ** exp
112 return p_star
115def upper_inc_gamma(a, x):
116 if a > 0:
117 uig = scipy.special.gamma(a) * scipy.special.gammaincc(a, x)
119 elif a == 0:
120 uig = -scipy.special.expi(-x)
122 else:
123 uig = 1 / a * (upper_inc_gamma(a + 1, x) - x**a * np.exp(-x))
125 return uig
128class NumGalCalculator:
129 """
130 Computes galaxy number counts by integrating the galaxy luminosity function.
131 The integral over absolute magnitudes can be done analytically, while the integral
132 over redshifts is computed numerically. See also
133 docs/jupyter_notebooks/sample_redshift_magnitude.ipynb.
134 """
136 def __init__(
137 self,
138 z_min,
139 z_max,
140 m_min,
141 m_max,
142 parametrization,
143 z_const,
144 alpha,
145 m_star_par,
146 phi_star_par,
147 cosmo,
148 pixarea,
149 ngal_multiplier=1,
150 ):
151 """
152 Initializes the galaxy number count calculator.
154 :param z_min: Minimum redshift
155 :param z_max: Maximum redshift
156 :param m_min: Minimum absolute magnitude (brightest)
157 :param m_max: Maximum absolute magnitude (faintest)
158 :param parametrization: Luminosity function parametrization
159 :param z_const: Redshift constant for truncated logexp parametrization
160 :param alpha: Luminosity function slope
161 :param m_star_par: Parameters for the M* - z relation (slope, intercept)
162 :param phi_star_par: Parameters for the phi* - z relation (amplitude, exponent)
163 :param cosmo: PyCosmo cosmology object
164 :param pixarea: Sky area in square radians
165 :param ngal_multiplier: multiplier for the number of galaxies
166 """
168 z_density_int = scipy.integrate.quad(
169 func=self._redshift_density,
170 a=z_min,
171 b=z_max,
172 args=(
173 m_min,
174 m_max,
175 parametrization,
176 z_const,
177 alpha,
178 m_star_par,
179 phi_star_par,
180 cosmo,
181 ),
182 )[0]
183 self.n_gal_mean = z_density_int * pixarea * ngal_multiplier
185 def __call__(self):
186 """
187 Randomly samples the number of galaxies from the Poisson distribution with the
188 mean number of galaxies.
189 """
190 if self.n_gal_mean > 0:
191 if self.n_gal_mean < 9e18:
192 n_gal = np.random.poisson(self.n_gal_mean)
193 else:
194 # This is a workaround for the fact that np.random.poisson does not work
195 # for large numbers. We use the fact that the Poisson distribution
196 # converges to a Gaussian for large means and sample from a Gaussian
197 # instead.
198 # It will probably be a crazy galaxy population model anyway...
199 LOGGER.warning("Using Gaussian instead of a Poisson due to large mean.")
200 n = float(self.n_gal_mean)
201 n_gal = int(np.random.normal(n, np.sqrt(n)))
202 else:
203 n_gal = self.n_gal_mean
204 return n_gal
206 def _redshift_density(
207 self,
208 z,
209 m_min,
210 m_max,
211 parametrization,
212 z_const,
213 alpha,
214 par_m_star,
215 par_phi_star,
216 cosmo,
217 ):
218 m_star = m_star_lum_fct(z, parametrization, z_const, *par_m_star)
219 phi_star = phi_star_lum_fct(z, parametrization, *par_phi_star)
220 e = np.sqrt(cosmo.params.omega_m * (1 + z) ** 3 + cosmo.params.omega_l)
221 d_h = cosmo.params.c / cosmo.params.H0
222 d_m = cosmo.background.dist_trans_a(a=1 / (1 + z))
223 if m_min is not None:
224 m_min_subtract = upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_star - m_min)))
225 else:
226 m_min_subtract = 0
227 density = (
228 phi_star
229 * d_h
230 * d_m**2
231 / e
232 * (
233 upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_star - m_max)))
234 - m_min_subtract
235 )
236 )
237 return density
240class RedshiftAbsMagSampler:
241 """
242 Samples redshifts and absolute magnitudes from the galaxy luminosity function.
243 The sampling is done by first drawing redshifts from the redshift-pdf obtained by
244 integrating out absolute magnitudes. Then, we sample absolute magnitudes from the
245 conditional pdfs obtained by conditioning the luminosity function on the sampled
246 redshifts (the conditional pdf is different for each redshift). See also
247 docs/jupyter_notebooks/sample_redshift_magnitude.ipynb and
248 docs/jupyter_notebooks/test_self_consistency.ipynb.
249 """
251 def __init__(
252 self,
253 z_res,
254 z_min,
255 z_max,
256 m_res,
257 m_min,
258 m_max,
259 parametrization,
260 z_const,
261 alpha,
262 m_star_par,
263 phi_star_par,
264 cosmo,
265 ):
266 """
267 Initializes the redshift and absolute magnitude grid for sampling.
269 :param z_res: Redshift resolution
270 :param z_min: Minimum redshift
271 :param z_max: Maximum redshift
272 :param m_res: Absolute magnitude resolution
273 :param m_min: Minimum absolute magnitude (brightest)
274 :param m_max: Maximum absolute magnitude (faintest)
275 :param parametrization: Luminosity function parametrization
276 :param z_const: Redshift constant for truncated logexp parametrization
277 :param alpha: Luminosity function slope
278 :param m_star_par: Parameters for the M* - z relation (slope, intercept)
279 :param phi_star_par: Parameters for the phi* - z relation (amplitude, exponent)
280 :param cosmo: PyCosmo cosmology object
281 """
282 self.z_res = z_res
283 self.z_min = z_min
284 self.z_max = z_max
285 self.m_res = m_res
286 self.m_min = m_min
287 self.m_max = m_max
288 self.parametrization = parametrization
289 self.z_const = z_const
290 self.alpha = alpha
291 self.m_star_par = m_star_par
292 self.phi_star_par = phi_star_par
293 self.cosmo = cosmo
295 # TODO: Do we need to pass all of these parameters? Should be in the class...
296 self._setup_redshift_grid()
297 self._setup_mag_grid()
299 def __call__(self, n_samples):
300 """
301 Returns n_samples redshifts and absolute magnitudes sampled from the galaxy
302 luminosity function.
304 :param n_samples: Number of samples
305 :return: Redshifts and absolute magnitudes
306 """
307 z = np.random.choice(self.z_grid, size=n_samples, replace=True, p=self.nz_grid)
309 m_s = m_star_lum_fct(z, self.parametrization, self.z_const, *self.m_star_par)
310 if self.m_min is None:
311 low = 0
312 else:
313 low = upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * (m_s - self.m_min)))
314 m_rvs = np.random.uniform(
315 low=low,
316 high=upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * (m_s - self.m_max))),
317 size=n_samples,
318 ) # here, we sample M* - M, where M* is redshift-dependent
319 uig_ind = find_closest_ind(self.uig_grid, m_rvs)
320 m = m_s - self.m_s__m__grid[uig_ind] # now we transform from M* - M to M
321 return z, m
323 def _setup_redshift_grid(self):
324 self.z_grid = np.linspace(
325 self.z_min,
326 self.z_max,
327 num=int(round((self.z_max - self.z_min) / self.z_res)) + 1,
328 )
330 e = np.sqrt(
331 self.cosmo.params.omega_m * (1 + self.z_grid) ** 3
332 + self.cosmo.params.omega_l
333 )
334 d_h = self.cosmo.params.c / self.cosmo.params.H0
335 d_m = self.cosmo.background.dist_trans_a(a=1 / (1 + self.z_grid))
336 f = d_h * d_m**2 / e
337 m_star = m_star_lum_fct(
338 self.z_grid, self.parametrization, self.z_const, *self.m_star_par
339 )
340 phi_star = phi_star_lum_fct(
341 self.z_grid, self.parametrization, *self.phi_star_par
342 )
344 if self.m_min is not None:
345 m_min_subtract = upper_inc_gamma(
346 self.alpha + 1, 10 ** (0.4 * (m_star - self.m_min))
347 )
348 else:
349 m_min_subtract = 0
350 self.nz_grid = (
351 f
352 * phi_star
353 * (
354 upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * (m_star - self.m_max)))
355 - m_min_subtract
356 )
357 )
358 self.nz_grid /= np.sum(self.nz_grid)
360 def _setup_mag_grid(self):
361 m_s_min = scipy.optimize.minimize_scalar(
362 lambda z: m_star_lum_fct(
363 z, self.parametrization, self.z_const, *self.m_star_par
364 ),
365 bounds=(self.z_min, self.z_max),
366 method="bounded",
367 ).fun
368 m_s_max = -scipy.optimize.minimize_scalar(
369 lambda z: -m_star_lum_fct(
370 z, self.parametrization, self.z_const, *self.m_star_par
371 ),
372 bounds=(self.z_min, self.z_max),
373 method="bounded",
374 ).fun
376 m_min = self.m_max
377 while upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * (m_s_min - m_min))) > 0:
378 m_min -= 0.1
379 self.m_s__m__grid = np.linspace(
380 m_s_max - m_min,
381 m_s_min - self.m_max,
382 num=int(round((m_s_max - m_min - m_s_min + self.m_max) / self.m_res)) + 1,
383 )
385 self.uig_grid = upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * self.m_s__m__grid))
388class LuminosityFunction:
389 """
390 Luminosity function
391 """
393 def __init__(
394 self,
395 name,
396 lum_fct_parametrization,
397 m_star_slope,
398 m_star_intcpt,
399 phi_star_amp,
400 phi_star_exp,
401 z_const,
402 alpha,
403 cosmo,
404 pixarea,
405 galaxy_type,
406 seed_ngal,
407 z_res=0.001,
408 m_res=0.001,
409 z_min=0,
410 z_max=6,
411 m_min=None,
412 m_max=2,
413 z_m_intp=None,
414 ngal_multiplier=1,
415 ):
416 """
417 Initializes the luminosity function.
419 :param lum_fct_parametrization: Luminosity function parametrization
420 :param m_star_slope: redshift evolution parameter for M* - z relation
421 :param m_star_intcpt: amplitude of M* at z=0
422 :param phi_star_amp: amplitude of phi* at z=0
423 :param phi_star_exp: redshift evolution parameter for phi* - z relation
424 :param z_const: redshift constant for truncated logexp parametrization
425 :param alpha: luminosity function slope
426 :param cosmo: PyCosmo cosmology object
427 :param pixarea: Sky area in square radians
428 :param galaxy_type: Galaxy type
429 :param seed_ngal: Seed for random number generator
430 :param z_res: Redshift resolution
431 :param m_res: Absolute magnitude resolution
432 :param z_min: Minimum redshift
433 :param z_max: Maximum redshift
434 :param m_min: Minimum absolute magnitude (brightest)
435 :param m_max: Maximum absolute magnitude (faintest)
436 :param z_m_intp: Redshift - absolute magnitude interpolation function
437 :param ngal_multiplier: Multiplier for the number of galaxies
438 """
439 self.parametrization = lum_fct_parametrization
440 self.m_star_slope = m_star_slope
441 self.m_star_intcpt = m_star_intcpt
442 self.phi_star_amp = phi_star_amp
443 self.phi_star_exp = phi_star_exp
444 self.m_star_par = m_star_slope, m_star_intcpt
445 self.phi_star_par = phi_star_amp, phi_star_exp
446 self.z_const = z_const
447 self.alpha = alpha
448 self.z_m_intp = z_m_intp
449 self.m_min = m_min
450 self.m_max = m_max
451 self.cosmo = cosmo
452 self.pixarea = pixarea
453 self.z_res = z_res
454 self.m_res = m_res
455 self.z_min = z_min
456 self.name = name
457 self.galaxy_type = galaxy_type
458 self.ngal_multiplier = ngal_multiplier
460 self.z_max = maximum_redshift(
461 z_m_intp=z_m_intp,
462 m_max=m_max,
463 z_max=z_max,
464 parametrization=lum_fct_parametrization,
465 z_const=z_const,
466 alpha=alpha,
467 m_star_par=self.m_star_par,
468 seed_ngal=seed_ngal,
469 )
471 self.n_gal_calc = NumGalCalculator(
472 z_min=self.z_min,
473 z_max=self.z_max,
474 m_min=m_min,
475 m_max=m_max,
476 parametrization=lum_fct_parametrization,
477 z_const=z_const,
478 alpha=alpha,
479 m_star_par=self.m_star_par,
480 phi_star_par=self.phi_star_par,
481 cosmo=cosmo,
482 pixarea=pixarea,
483 ngal_multiplier=ngal_multiplier,
484 )
486 self.z_mabs_sampler = RedshiftAbsMagSampler(
487 z_res=z_res,
488 z_min=self.z_min,
489 z_max=self.z_max,
490 parametrization=lum_fct_parametrization,
491 z_const=z_const,
492 m_res=m_res,
493 m_min=m_min,
494 m_max=m_max,
495 alpha=alpha,
496 m_star_par=self.m_star_par,
497 phi_star_par=self.phi_star_par,
498 cosmo=cosmo,
499 )
501 def sample_z_mabs_and_apply_cut(
502 self,
503 seed_ngal,
504 seed_lumfun,
505 n_gal_max=np.inf,
506 size_chunk=10000,
507 ):
508 """
509 This function gets the abs mag and z using chunking, which uses less memory than
510 the original method. It does not give exactly the same result as before due to
511 different order of random draws in z_mabs_sampler, but it's the same sample.
512 """
513 np.random.seed(seed_ngal)
514 n_gal = self.n_gal_calc()
515 if n_gal == 0:
516 return np.array([]), np.array([])
517 n_chunks = int(np.ceil(float(n_gal) / float(size_chunk)))
518 list_abs_mag = []
519 list_z = []
520 n_gal_final = 0
521 for ic in range(n_chunks):
522 n_gal_sample = n_gal % size_chunk if ic + 1 == n_chunks else size_chunk
523 z_chunk, abs_mag_chunk = self.z_mabs_sampler(n_gal_sample)
524 # Apply cut in z - M - plane
525 if self.z_m_intp is not None:
526 select = abs_mag_chunk < self.z_m_intp(z_chunk)
527 z_chunk = z_chunk[select]
528 abs_mag_chunk = abs_mag_chunk[select]
530 n_gal_final += len(z_chunk)
531 list_z.append(z_chunk)
532 list_abs_mag.append(abs_mag_chunk)
534 if n_gal_final > n_gal_max:
535 break
537 z = np.hstack(list_z)
538 abs_mag = np.hstack(list_abs_mag)
539 return abs_mag, z
541 def analytical_number_density(self, m, z):
542 """
543 Returns the number density of galaxies at redshift z and absolute magnitude m.
544 Classic Schechter function with units of Mpc^-3 mag^-1.
546 :param m: Absolute magnitude
547 :param z: Redshift
548 :return: Number density of galaxies at redshift z and absolute magnitude m
549 """
550 m_star = m_star_lum_fct(z, self.parametrization, self.z_const, *self.m_star_par)
551 phi_star = phi_star_lum_fct(z, self.parametrization, *self.phi_star_par)
552 return (
553 0.4
554 * np.log(10)
555 * (
556 phi_star
557 * (10 ** (0.4 * (m_star - m))) ** (self.alpha + 1)
558 * np.exp(-(10 ** (0.4 * (m_star - m))))
559 )
560 )