Coverage for src/galsbi/ucat/galaxy_population_models/galaxy_luminosity_function.py: 96%
161 statements
« prev ^ index » next coverage.py v7.6.9, created at 2024-12-13 03:24 +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_max,
139 m_max,
140 parametrization,
141 z_const,
142 alpha,
143 m_star_par,
144 phi_star_par,
145 cosmo,
146 pixarea,
147 ngal_multiplier=1,
148 ):
149 z_density_int = scipy.integrate.quad(
150 func=self._redshift_density,
151 a=0,
152 b=z_max,
153 args=(
154 m_max,
155 parametrization,
156 z_const,
157 alpha,
158 m_star_par,
159 phi_star_par,
160 cosmo,
161 ),
162 )[0]
163 self.n_gal_mean = int(round(z_density_int * pixarea * ngal_multiplier))
165 def __call__(self):
166 if self.n_gal_mean > 0:
167 if self.n_gal_mean < 9e18:
168 n_gal = np.random.poisson(self.n_gal_mean)
169 else:
170 # This is a workaround for the fact that np.random.poisson does not work
171 # for large numbers. We use the fact that the Poisson distribution
172 # converges to a Gaussian for large means and sample from a Gaussian
173 # instead.
174 # It will probably be a crazy galaxy population model anyway...
175 LOGGER.warning("Using Gaussian instead of a Poisson due to large mean.")
176 n = float(self.n_gal_mean)
177 n_gal = int(np.random.normal(n, np.sqrt(n)))
178 else:
179 n_gal = self.n_gal_mean
180 return n_gal
182 def _redshift_density(
183 self, z, m_max, parametrization, z_const, alpha, par_m_star, par_phi_star, cosmo
184 ):
185 m_star = m_star_lum_fct(z, parametrization, z_const, *par_m_star)
186 phi_star = phi_star_lum_fct(z, parametrization, *par_phi_star)
187 e = np.sqrt(cosmo.params.omega_m * (1 + z) ** 3 + cosmo.params.omega_l)
188 d_h = cosmo.params.c / cosmo.params.H0
189 d_m = cosmo.background.dist_trans_a(a=1 / (1 + z))
190 density = (
191 phi_star
192 * d_h
193 * d_m**2
194 / e
195 * upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_star - m_max)))
196 )
197 return density
200class RedshiftAbsMagSampler:
201 """
202 Samples redshifts and absolute magnitudes from the galaxy luminosity function.
203 The sampling is done by first drawing redshifts from the redshift-pdf obtained by
204 integrating out absolute magnitudes. Then, we sample absolute magnitudes from the
205 conditional pdfs obtained by conditioning the luminosity function on the sampled
206 redshifts (the conditional pdf is different for each redshift). See also
207 docs/jupyter_notebooks/sample_redshift_magnitude.ipynb and
208 docs/jupyter_notebooks/test_self_consistency.ipynb.
209 """
211 def __init__(
212 self,
213 z_res,
214 z_max,
215 m_res,
216 m_max,
217 parametrization,
218 z_const,
219 alpha,
220 m_star_par,
221 phi_star_par,
222 cosmo,
223 ):
224 self.z_res = z_res
225 self.z_max = z_max
226 self.m_res = m_res
227 self.m_max = m_max
228 self.parametrization = parametrization
229 self.z_const = z_const
230 self.alpha = alpha
231 self.m_star_par = m_star_par
232 self.phi_star_par = phi_star_par
233 self.cosmo = cosmo
235 # TODO: Do we need to pass all of these parameters? Should be in the class...
236 self._setup_redshift_grid(
237 z_max, z_res, m_max, z_const, alpha, m_star_par, phi_star_par, cosmo
238 )
239 self._setup_mag_grid(
240 z_max, m_max, m_res, parametrization, z_const, alpha, m_star_par
241 )
243 def __call__(self, n_samples):
244 z = np.random.choice(self.z_grid, size=n_samples, replace=True, p=self.nz_grid)
246 m_s = m_star_lum_fct(z, self.parametrization, self.z_const, *self.m_star_par)
247 m_rvs = np.random.uniform(
248 low=0,
249 high=upper_inc_gamma(self.alpha + 1, 10 ** (0.4 * (m_s - self.m_max))),
250 size=n_samples,
251 ) # here, we sample M* - M, where M* is redshift-dependent
252 uig_ind = find_closest_ind(self.uig_grid, m_rvs)
253 m = m_s - self.m_s__m__grid[uig_ind] # now we transform from M* - M to M
254 return z, m
256 def _setup_redshift_grid(
257 self, z_max, z_res, m_max, z_const, alpha, m_star_par, phi_star_par, cosmo
258 ):
259 self.z_grid = np.linspace(
260 z_res, z_max, num=int(round((z_max - z_res) / z_res)) + 1
261 )
263 e = np.sqrt(
264 cosmo.params.omega_m * (1 + self.z_grid) ** 3 + cosmo.params.omega_l
265 )
266 d_h = cosmo.params.c / cosmo.params.H0
267 d_m = cosmo.background.dist_trans_a(a=1 / (1 + self.z_grid))
268 f = d_h * d_m**2 / e
269 m_star = m_star_lum_fct(self.z_grid, self.parametrization, z_const, *m_star_par)
270 phi_star = phi_star_lum_fct(self.z_grid, self.parametrization, *phi_star_par)
272 self.nz_grid = (
273 f * phi_star * upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_star - m_max)))
274 )
275 self.nz_grid /= np.sum(self.nz_grid)
277 def _setup_mag_grid(
278 self, z_max, m_max, m_res, parametrization, z_const, alpha, m_star_par
279 ):
280 m_s_min = scipy.optimize.minimize_scalar(
281 lambda z: m_star_lum_fct(z, parametrization, z_const, *m_star_par),
282 bounds=(0, z_max),
283 method="bounded",
284 ).fun
285 m_s_max = -scipy.optimize.minimize_scalar(
286 lambda z: -m_star_lum_fct(z, parametrization, z_const, *m_star_par),
287 bounds=(0, z_max),
288 method="bounded",
289 ).fun
291 m_min = m_max
292 while upper_inc_gamma(alpha + 1, 10 ** (0.4 * (m_s_min - m_min))) > 0:
293 m_min -= 0.1
294 self.m_s__m__grid = np.linspace(
295 m_s_max - m_min,
296 m_s_min - m_max,
297 num=int(round((m_s_max - m_min - m_s_min + m_max) / m_res)) + 1,
298 )
300 self.uig_grid = upper_inc_gamma(alpha + 1, 10 ** (0.4 * self.m_s__m__grid))
303class LuminosityFunction:
304 """
305 Luminosity function
306 """
308 def __init__(
309 self,
310 name,
311 lum_fct_parametrization,
312 m_star_slope,
313 m_star_intcpt,
314 phi_star_amp,
315 phi_star_exp,
316 z_const,
317 alpha,
318 cosmo,
319 pixarea,
320 galaxy_type,
321 seed_ngal,
322 z_res=0.001,
323 m_res=0.001,
324 z_max=np.inf,
325 m_max=2,
326 z_m_intp=None,
327 ngal_multiplier=1,
328 ):
329 self.parametrization = lum_fct_parametrization
330 self.m_star_slope = m_star_slope
331 self.m_star_intcpt = m_star_intcpt
332 self.phi_star_amp = phi_star_amp
333 self.phi_star_exp = phi_star_exp
334 self.m_star_par = m_star_slope, m_star_intcpt
335 self.phi_star_par = phi_star_amp, phi_star_exp
336 self.z_const = z_const
337 self.alpha = alpha
338 self.z_m_intp = z_m_intp
339 self.m_max = m_max
340 self.cosmo = cosmo
341 self.pixarea = pixarea
342 self.z_res = z_res
343 self.name = name
344 self.galaxy_type = galaxy_type
346 self.z_max = maximum_redshift(
347 z_m_intp=z_m_intp,
348 m_max=m_max,
349 z_max=z_max,
350 parametrization=lum_fct_parametrization,
351 z_const=z_const,
352 alpha=alpha,
353 m_star_par=self.m_star_par,
354 seed_ngal=seed_ngal,
355 )
357 self.n_gal_calc = NumGalCalculator(
358 z_max=self.z_max,
359 m_max=m_max,
360 parametrization=lum_fct_parametrization,
361 z_const=z_const,
362 alpha=alpha,
363 m_star_par=self.m_star_par,
364 phi_star_par=self.phi_star_par,
365 cosmo=cosmo,
366 pixarea=pixarea,
367 ngal_multiplier=ngal_multiplier,
368 )
370 self.z_mabs_sampler = RedshiftAbsMagSampler(
371 z_res=z_res,
372 z_max=self.z_max,
373 parametrization=lum_fct_parametrization,
374 z_const=z_const,
375 m_res=m_res,
376 m_max=m_max,
377 alpha=alpha,
378 m_star_par=self.m_star_par,
379 phi_star_par=self.phi_star_par,
380 cosmo=cosmo,
381 )
383 def sample_z_mabs_and_apply_cut(
384 self,
385 seed_ngal,
386 seed_lumfun,
387 n_gal_max=np.inf,
388 size_chunk=10000,
389 ):
390 """
391 This function gets the abs mag and z using chunking, which uses less memory than
392 the original method. It does not give exactly the same result as before due to
393 different order of random draws in z_mabs_sampler, but it's the same sample.
394 """
395 np.random.seed(seed_ngal)
396 n_gal = self.n_gal_calc()
397 if n_gal == 0:
398 return np.array([]), np.array([])
399 n_chunks = int(np.ceil(float(n_gal) / float(size_chunk)))
400 list_abs_mag = []
401 list_z = []
402 n_gal_final = 0
403 for ic in range(n_chunks):
404 n_gal_sample = n_gal % size_chunk if ic + 1 == n_chunks else size_chunk
405 z_chunk, abs_mag_chunk = self.z_mabs_sampler(n_gal_sample)
406 # Apply cut in z - M - plane
407 if self.z_m_intp is not None:
408 select = abs_mag_chunk < self.z_m_intp(z_chunk)
409 z_chunk = z_chunk[select]
410 abs_mag_chunk = abs_mag_chunk[select]
412 n_gal_final += len(z_chunk)
413 list_z.append(z_chunk)
414 list_abs_mag.append(abs_mag_chunk)
416 if n_gal_final > n_gal_max:
417 break
419 z = np.hstack(list_z)
420 abs_mag = np.hstack(list_abs_mag)
421 return abs_mag, z