Coverage for src/galsbi/ucat_sps/galaxy_population_models/galaxy_stellar_mass

3# Author: Luca Tortorelli, Silvan Fischbacher

4# created: Apr 2025

6from collections import OrderedDict

8import numpy as np

9import scipy.integrate

10import scipy.optimize

11import scipy.special

12from cosmic_toolbox import logger

14LOGGER = logger.get_logger(__file__)

17def find_closest_ind(grid, vals):

18 ind = np.searchsorted(grid, vals)

19 ind[ind == grid.size] -= 1

20 ind[np.fabs(grid[ind] - vals) > np.fabs(grid[ind - 1] - vals)] -= 1

21 return ind

24def single_stellar_mass_function_parametrisation(

25 log10_stellar_masses, log_Mstar, phi_star, alpha

26):

27 """

28 This function parametrises the stellar mass function as a single Schechter

29 function. It computes the number density of galaxies as a function of their stellar

30 mass.

32 :param log10_stellar_masses: (array_like[n_mass_bins,]) log10 of the stellar masses

33 drawn from the SMF.

34 :param log_Mstar: (float) knee of the Schechter function. It corresponds to the

35 stellar mass at which the Schechter transitions from a simple power law with

36 slope alpha at lower masses into an exponential function at higher masses.

37 :param phi_star: (float) number density of galaxies at stellar mass M* (log_Mstar).

38 :param alpha: (float) faint-end slope.

39 :return stellar_mass_function: (array_like[n_mass_bins,]) the number density of

40 galaxies as a function of their stellar mass.

41 """

42 return (

43 np.log(10)

44 * np.exp(-(10 ** (log10_stellar_masses - log_Mstar)))

45 * phi_star

46 * (10 ** (log10_stellar_masses - log_Mstar)) ** (alpha + 1)

47 )

50def double_stellar_mass_function_parametrisation(

51 log10_stellar_masses,

52 log_Mstar,

53 phi_star_low_mass,

54 phi_star_high_mass,

55 alpha_low_mass,

56 alpha_high_mass,

57):

58 """

59 This function parametrises the stellar mass function as a double Schechter

60 function. It computes the number density of galaxies as a function of their stellar

61 mass.

63 :param log10_stellar_masses: (array_like[n_mass_bins,]) log10 of the stellar masses

64 drawn from the SMF.

65 :param log_Mstar: (float) knee of the Schechter function. It corresponds to the

66 stellar mass at which the Schechter transitions from a simple power law with

67 slope alpha at lower masses into an exponential function at higher masses. It

68 has the same value for both the low and high mass SMFs.

69 :param phi_star_low_mass: (float) number density of galaxies at stellar mass M*

70 (log_Mstar) for the low mass SMF.

71 :param phi_star_high_mass: (float) number density of galaxies at stellar mass M*

72 (log_Mstar) for the high mass SMF.

73 :param alpha_low_mass: (float) faint-end slope of the low mass SMF.

74 :param alpha_high_mass: (float) faint-end slope of the high mass SMF.

75 :return stellar_mass_function: (array_like[n_mass_bins,]) the number density of

76 galaxies as a function of their stellar mass.

77 """

78 return (

79 np.log(10)

80 * np.exp(-(10 ** (log10_stellar_masses - log_Mstar)))

81 * (

82 phi_star_low_mass

83 * (10 ** (log10_stellar_masses - log_Mstar)) ** (alpha_low_mass + 1)

84 + phi_star_high_mass

85 * (10 ** (log10_stellar_masses - log_Mstar)) ** (alpha_high_mass + 1)

86 )

87 )

90def stellar_mass_function_m_star(

91 z, parametrization, intercept, slope, quadratic_term=0

92):

93 if parametrization == "linexp":

94 m_s = np.polyval((slope, intercept), z)

95 elif parametrization == "logpower":

96 m_s = intercept + slope * np.log10(1 + z)

97 elif parametrization == "logquadratic_power":

98 m_s = (

99 intercept

100 + slope * np.log10(1 + z)

101 + quadratic_term * (np.log10(1 + z)) ** 2

102 )

103 else:

104 raise ValueError("unknown parametrization value")

105 return m_s

106

107

108def stellar_mass_function_phi_star(z, parametrization, amplitude, exp):

109 if parametrization == "linexp":

110 p_star = amplitude * np.exp(exp * z)

111 elif parametrization == "logpower" or parametrization == "logquadratic_power":

112 p_star = amplitude * (1 + z) ** exp

113 else:

114 raise ValueError("unknown parametrization value")

115 return p_star

116

117

118def upper_inc_gamma(a, x):

119 if a > 0:

120 uig = scipy.special.gamma(a) * scipy.special.gammaincc(a, x)

121

122 elif a == 0:

123 uig = -scipy.special.expi(-x)

124

125 else:

126 uig = 1 / a * (upper_inc_gamma(a + 1, x) - x**a * np.exp(-x))

127

128 return uig

129

130

131def initialize_stellar_mass_functions(

132 par, pixarea, cosmo, z_sm_intp=None, mass_key=None

133):

134 if mass_key is None:

135 raise KeyError

136 kw_smfun = dict(

137 pixarea=pixarea,

138 z_res=par.sm_fct_z_res,

139 sm_res=par.sm_fct_sm_res,

140 z_max=par.sm_fct_z_max,

141 sm_min=par.sm_fct_sm_min,

142 z_sm_intp=z_sm_intp,

143 ngal_multiplier=par.ngal_multiplier,

144 cosmo=cosmo,

145 )

146 if par.ngal_multiplier != 1:

147 LOGGER.warning(f"ngal_multiplier is set to {par.ngal_multiplier}")

148

149 sm_funcs = OrderedDict()

150 for i, g in enumerate(par.galaxy_types):

151 if g in par.galaxy_types:

152 sm_funcs[g] = StellarMassFunction(

153 sm_fct_parametrization=par.sm_fct_parametrization,

154 sm_m_star_0=getattr(par, f"sm_fct_m_star_{g}_0_{mass_key}"),

155 sm_m_star_1=getattr(par, f"sm_fct_m_star_{g}_1_{mass_key}"),

156 sm_m_star_2=getattr(par, f"sm_fct_m_star_{g}_2_{mass_key}"),

157 sm_phi_star_amp=getattr(par, f"sm_fct_phi_star_{g}_amp_{mass_key}"),

158 sm_phi_star_exp=getattr(par, f"sm_fct_phi_star_{g}_exp_{mass_key}"),

159 sm_alpha=getattr(par, f"sm_fct_alpha_{g}_{mass_key}"),

160 name=g,

161 galaxy_type=i,

162 seed_ngal=par.seed_ngal,

163 **kw_smfun,

164 )

165 return sm_funcs

166

167

168def maximum_redshift(

169 z_sm_intp, sm_min, z_max, parametrization, alpha, sm_m_star_par, seed_ngal

170):

171 """Compute the maximum redshift used to sample the stellar-mass function.

172

173 When a redshift-dependent stellar-mass completeness boundary is supplied, sampling

174 stops when fewer than ``1e-5`` of the galaxies above ``sm_min`` are expected to lie

175 above that boundary.

176

177 :param z_sm_intp: (callable) interpolation returning the lower logarithmic stellar-

178 mass boundary as a function of redshift, or ``None`` when no additional cut is

179 required.

180 :param sm_min: (float) minimum logarithmic stellar mass sampled before applying the

181 redshift-dependent cut, in log10 solar-mass units.

182 :param z_max: (float) configured maximum sampling redshift.

183 :param parametrization: (str) parametrization of the redshift evolution of the

184 characteristic stellar mass.

185 :param alpha: (float) low-mass slope of the Schechter function.

186 :param sm_m_star_par: (tuple) parameters describing the characteristic stellar mass.

187 :param seed_ngal: (int) random seed used for the small redshift-grid displacement.

188 :return z_max: (float) maximum redshift used for sampling.

189 """

190 if z_sm_intp is None:

191 return z_max

192

193 def conditional_mass_fraction(z):

194 sm_m_s = stellar_mass_function_m_star(z, parametrization, *sm_m_star_par)

195 sm_cut = max(sm_min, float(z_sm_intp(z)))

196 mass_fraction = (

197 upper_inc_gamma(alpha + 1, 10 ** (sm_cut - sm_m_s))

198 / upper_inc_gamma(alpha + 1, 10 ** (sm_min - sm_m_s))

199 - 1e-5

200 )

201 return mass_fraction

202

203 try:

204 z_max_cutoff = scipy.optimize.brentq(

205 conditional_mass_fraction, 0, z_sm_intp.x[-1]

206 )

207 except ValueError:

208 z_max_cutoff = z_sm_intp.x[-1]

209

210 z_max = min(z_max, z_max_cutoff)

211 np.random.seed(seed_ngal)

212 z_max += np.random.uniform(0, 0.0001)

213

214 return z_max

215

216

217class NumGalCalculator:

218 """

219 Computes galaxy number counts by integrating the galaxy stellar mass function

220 (Schechter function).

221

222 The integral over stellar masses can be done analytically, while the integral

223 over redshifts is computed numerically. See also

224 docs/jupyter_notebooks/sample_redshift_magnitude.ipynb.

225 """

226

227 def __init__(

228 self,

229 z_max,

230 sm_min,

231 parametrization,

232 sm_alpha,

233 sm_m_star_par,

234 sm_phi_star_par,

235 cosmo,

236 pixarea,

237 ngal_multiplier=1,

238 ):

239 z_density_int = scipy.integrate.quad(

240 func=self._redshift_density,

241 a=0,

242 b=z_max,

243 args=(

244 sm_min,

245 parametrization,

246 sm_alpha,

247 sm_m_star_par,

248 sm_phi_star_par,

249 cosmo,

250 ),

251 )[0]

252 self.n_gal_mean = float(z_density_int * pixarea * ngal_multiplier)

253 if not np.isfinite(self.n_gal_mean) or self.n_gal_mean < 0:

254 raise ValueError(

255 "The expected number of galaxies must be finite and nonnegative, "

256 f"but received {self.n_gal_mean}."

257 )

258

259 def __call__(self):

260 return np.random.poisson(self.n_gal_mean) if self.n_gal_mean > 0 else 0

261

262 def _redshift_density(

263 self,

264 z,

265 sm_min,

266 parametrization,

267 sm_alpha,

268 par_sm_m_star,

269 par_sm_phi_star,

270 cosmo,

271 ):

272 sm_m_star = stellar_mass_function_m_star(z, parametrization, *par_sm_m_star)

273 sm_phi_star = stellar_mass_function_phi_star(

274 z, parametrization, *par_sm_phi_star

275 )

276 e = np.sqrt(cosmo.params.omega_m * (1 + z) ** 3 + cosmo.params.omega_l)

277 d_h = cosmo.params.c / cosmo.params.H0

278 d_m = cosmo.background.dist_trans_a(a=1 / (1 + z))

279 density = (

280 sm_phi_star

281 * d_h

282 * d_m**2

283 / e

284 * upper_inc_gamma(sm_alpha + 1, 10 ** (sm_min - sm_m_star))

285 )

286

287 return density

288

289

290class RedshiftSMSampler:

291 """

292 Samples redshifts and stellar masses from the galaxy stellar mass function.

293

294 The sampling is done by first drawing redshifts from the redshift-pdf obtained by

295 integrating out stellar masses. Then, we sample stellar masses from the

296 conditional pdfs obtained by conditioning the stellar mass function on the sampled

297 redshifts (the conditional pdf is different for each redshift). See also

298 docs/jupyter_notebooks/sample_redshift_magnitude.ipynb and

299 docs/jupyter_notebooks/test_self_consistency.ipynb.

300 """

301

302 def __init__(

303 self,

304 z_res,

305 z_max,

306 sm_res,

307 sm_min,

308 parametrization,

309 sm_alpha,

310 sm_m_star_par,

311 sm_phi_star_par,

312 cosmo,

313 ):

314 self.z_res = z_res

315 self.z_max = z_max

316 self.sm_res = sm_res

317 self.sm_min = sm_min

318 self.parametrization = parametrization

319 self.sm_alpha = sm_alpha

320 self.sm_m_star_par = sm_m_star_par

321 self.sm_phi_star_par = sm_phi_star_par

322 self.cosmo = cosmo

323

324 self._setup_redshift_grid(

325 z_max, z_res, sm_min, sm_alpha, sm_m_star_par, sm_phi_star_par, cosmo

326 )

327 self._setup_mag_grid(

328 z_max, sm_min, sm_res, parametrization, sm_alpha, sm_m_star_par

329 )

330

331 def __call__(self, n_samples):

332 z = np.random.choice(self.z_grid, size=n_samples, replace=True, p=self.nz_grid)

333

334 sm_m_s = stellar_mass_function_m_star(

335 z, self.parametrization, *self.sm_m_star_par

336 )

337 sm_m_rvs = np.random.uniform(

338 low=0,

339 high=upper_inc_gamma(self.sm_alpha + 1, 10 ** (self.sm_min - sm_m_s)),

340 size=n_samples,

341 ) # here, we sample M - M*, where M* is redshift-dependent

342 uig_ind = find_closest_ind(self.uig_grid, sm_m_rvs)

343 sm = sm_m_s + self.m_s__m__grid[uig_ind] # now we transform from M - M* to M

344 return z, sm

345

346 def _setup_redshift_grid(

347 self, z_max, z_res, sm_min, sm_alpha, sm_m_star_par, sm_phi_star_par, cosmo

348 ):

349 self.z_grid = np.linspace(

350 z_res, z_max, num=int(round((z_max - z_res) / z_res)) + 1

351 )

352 e = np.sqrt(

353 cosmo.params.omega_m * (1 + self.z_grid) ** 3 + cosmo.params.omega_l

354 )

355 d_h = cosmo.params.c / cosmo.params.H0

356 d_m = cosmo.background.dist_trans_a(a=1 / (1 + self.z_grid))

357 f = d_h * d_m**2 / e

358

359 sm_m_star = stellar_mass_function_m_star(

360 self.z_grid, self.parametrization, *sm_m_star_par

361 )

362 sm_phi_star = stellar_mass_function_phi_star(

363 self.z_grid, self.parametrization, *sm_phi_star_par

364 )

365

366 self.nz_grid = (

367 f * sm_phi_star * upper_inc_gamma(sm_alpha + 1, 10 ** (sm_min - sm_m_star))

368 )

369 self.nz_grid /= np.sum(self.nz_grid)

370

371 def _setup_mag_grid(

372 self, z_max, sm_min, sm_res, parametrization, sm_alpha, sm_m_star_par

373 ):

374 sm_m_s_max = -scipy.optimize.minimize_scalar(

375 lambda z: -stellar_mass_function_m_star(z, parametrization, *sm_m_star_par),

376 bounds=(0, z_max),

377 method="bounded",

378 ).fun

379 sm_m_s_min = scipy.optimize.minimize_scalar(

380 lambda z: stellar_mass_function_m_star(z, parametrization, *sm_m_star_par),

381 bounds=(0, z_max),

382 method="bounded",

383 ).fun

384

385 sm_max = sm_min

386 while upper_inc_gamma(sm_alpha + 1, 10 ** (sm_max - sm_m_s_min)) > 0:

387 sm_max += 0.1

388

389 self.m_s__m__grid = np.linspace(

390 sm_max - sm_m_s_min,

391 sm_min - sm_m_s_max,

392 num=int(round((sm_max - sm_m_s_min - sm_min + sm_m_s_max) / sm_res)) + 1,

393 )

394 self.uig_grid = upper_inc_gamma(sm_alpha + 1, 10**self.m_s__m__grid)

395

396

397class StellarMassFunction:

398 """

399 Stellar Mass function.

400 """

401

402 def __init__(

403 self,

404 name,

405 sm_fct_parametrization,

406 sm_m_star_0,

407 sm_m_star_1,

408 sm_m_star_2,

409 sm_phi_star_amp,

410 sm_phi_star_exp,

411 sm_alpha,

412 cosmo,

413 pixarea,

414 galaxy_type,

415 seed_ngal,

416 z_res=0.001,

417 sm_res=0.01,

418 z_max=3,

419 sm_min=8,

420 z_sm_intp=None,

421 ngal_multiplier=1,

422 ):

423 self.parametrization = sm_fct_parametrization

424 self.sm_m_star_0 = sm_m_star_0

425 self.sm_m_star_1 = sm_m_star_1

426 self.sm_m_star_2 = sm_m_star_2

427 self.sm_phi_star_amp = sm_phi_star_amp

428 self.sm_phi_star_exp = sm_phi_star_exp

429 self.sm_m_star_par = sm_m_star_0, sm_m_star_1, sm_m_star_2

430 self.sm_phi_star_par = sm_phi_star_amp, sm_phi_star_exp

431 self.sm_alpha = sm_alpha

432 self.z_sm_intp = z_sm_intp

433 self.sm_min = sm_min

434 self.cosmo = cosmo

435 self.pixarea = pixarea

436 self.z_res = z_res

437 self.name = name

438 self.galaxy_type = galaxy_type

439

440 self.z_max = maximum_redshift(

441 z_sm_intp=z_sm_intp,

442 sm_min=self.sm_min,

443 z_max=z_max,

444 parametrization=self.parametrization,

445 alpha=self.sm_alpha,

446 sm_m_star_par=self.sm_m_star_par,

447 seed_ngal=seed_ngal,

448 )

449

450 self.n_gal_calc = NumGalCalculator(

451 z_max=self.z_max,

452 sm_min=sm_min,

453 parametrization=sm_fct_parametrization,

454 sm_alpha=sm_alpha,

455 sm_m_star_par=self.sm_m_star_par,

456 sm_phi_star_par=self.sm_phi_star_par,

457 cosmo=cosmo,

458 pixarea=pixarea,

459 ngal_multiplier=ngal_multiplier,

460 )

461

462 self.z_sm_sampler = RedshiftSMSampler(

463 z_res=z_res,

464 z_max=self.z_max,

465 parametrization=sm_fct_parametrization,

466 sm_res=sm_res,

467 sm_min=sm_min,

468 sm_alpha=sm_alpha,

469 sm_m_star_par=self.sm_m_star_par,

470 sm_phi_star_par=self.sm_phi_star_par,

471 cosmo=cosmo,

472 )

473

474 def sample_z_sm_and_apply_cut(self, seed_ngal, n_gal_max=np.inf, size_chunk=10000):

475 """

476 This function gets the abs mag and z using chunking, which uses less memory

477 than the original method.

478

479 It does not give exactly the same result as before due to different order of

480 random draws in z_mabs_sampler, but it's the same sample.

481 """

482 np.random.seed(seed_ngal)

483 n_gal = self.n_gal_calc()

484 if n_gal == 0:

485 return np.array([]), np.array([])

486 n_chunks = int(np.ceil(float(n_gal) / float(size_chunk)))

487 list_sm = []

488 list_z = []

489 n_gal_final = 0

490 for ic in range(n_chunks):

491 n_gal_sample = min(size_chunk, n_gal - ic * size_chunk)

492 z_chunk, sm_chunk = self.z_sm_sampler(n_gal_sample)

493

494 # Apply cut in z - M - plane

495 if self.z_sm_intp is not None:

496 select = sm_chunk >= self.z_sm_intp(z_chunk)

497 z_chunk = z_chunk[select]

498 sm_chunk = sm_chunk[select]

499

500 n_gal_final += len(z_chunk)

501 list_z.append(z_chunk)

502 list_sm.append(sm_chunk)

503

504 if n_gal_final > n_gal_max:

505 break

506

507 z = np.hstack(list_z)

508 sm = np.hstack(list_sm)

509 return sm, z

Coverage for src/galsbi/ucat_sps/galaxy_population_models/galaxy_stellar_mass_function.py: 99%

165 statements