Source code for megabeast.helpers

import numpy as np
from numpy.random import default_rng

from beast.tools.read_beast_data import read_lnp_data
from beast.physicsmodel.grid_weights_stars import compute_bin_boundaries

__all__ = [
    "get_likelihoods",
    "precompute_mass_multipliers",
    "get_predicted_num_stars",
    "get_predicted_num_stars_simulate",
]


[docs] def get_likelihoods(ppdf_file, beast_model_data): """ Read in the saved BEAST sparse posterior PDFs and divide by the BEAST priors to get the spare liklilhoods. Parameters ---------- ppdf_file : string filename of the saved BEAST posterior PDFs Returns ------- lnp_data : dictonary contains arrays of the likelihood values (vals) and indices to the BEAST model grid (indxs) """ # BEAST saves posterior PDFs labeled as log(pPDF) lnpdata = read_lnp_data(ppdf_file) # divide by the BEAST prior weights to return to recover the likelihoods n_lnps, n_stars = lnpdata["indxs"].shape for i in range(n_stars): indxs = lnpdata["indxs"][:, i] gmask = np.isfinite(lnpdata["vals"][:, i]) lnpdata["vals"][gmask, i] = ( np.exp(lnpdata["vals"][gmask, i]) / beast_model_data["prior_weight"][indxs[gmask]] ) return lnpdata
[docs] def precompute_mass_multipliers(bphysparams, physmodmass): """ Calculate the value to mulitply the SFR to get the total mass in stars at all ages and masses given the physics model on the BEAST model grid. Parameters ---------- bphysparams : astropy.table table giving the physical parameters, weights, and completeness on the BEAST physical grid physmodmass : beast.physicmodel.priormodel physics model of for initial mass Returns ------- sfhmassinfo : dict "massmult" gives the value to muliply the SFH at each age and metallicity "ages" gives the ages and "Zs" gives the metallicities """ mass_range = [min(bphysparams["M_ini"]), max(bphysparams["M_ini"])] # compute the total mass and average mass of a star given the mass_prior_model nmass = 100 masspts = np.logspace(np.log10(mass_range[0]), np.log10(mass_range[1]), nmass) massprior = physmodmass(masspts) totmass = np.trapz(massprior, masspts) avemass = np.trapz(masspts * massprior, masspts) / totmass # loop over all the ages and compute the mass to simulate # ***need to add metallicity as well*** grid_ages = np.unique(bphysparams["logA"]) bin_boundaries = compute_bin_boundaries(grid_ages) bin_widths = np.diff(10 ** (bin_boundaries)) grid_mets = np.unique(bphysparams["Z"]) massmults = np.full((len(grid_ages), len(grid_mets)), 0.0) for i, cage in enumerate(grid_ages): for j, cmet in enumerate(grid_mets): gmods = (bphysparams["logA"] == cage) & (bphysparams["Z"] == cmet) cur_mass_range = [ min(bphysparams["M_ini"][gmods]), max(bphysparams["M_ini"][gmods]), ] gmass = (masspts >= cur_mass_range[0]) & (masspts <= cur_mass_range[1]) curmasspts = masspts[gmass] curmassprior = massprior[gmass] totcurmass = np.trapz(curmassprior, curmasspts) # compute the mass remaining at each age -> this is the mass to simulate massmults[i, j] = bin_widths[i] * totcurmass / totmass return { "massmult": massmults, "ages": grid_ages, "Zs": grid_mets, "avemass": avemass, }
[docs] def get_predicted_num_stars( massmult_info, bphysparams, bphysmod, physmodage ): """ Calculate the expected number of stars based on the physics model as given on the BEAST model grid including completeness. Parameters ---------- massmult_info : dict "massmult" gives the value to muliply the SFH at each age and metallicity "ages" gives the ages and "Zs" gives the metallicities bphysparams : astropy.table table giving the physical parameters, weights, and completeness on the BEAST physical grid bphysmod : array probability of the physics model on the BEAST physics grid physmodage : beast.physicmodel.priormodel physics model of for age Returns ------- n_stars : float number of stars expected """ gridweights = bphysmod * bphysparams["grid_weight"] gridweights = gridweights / np.sum(gridweights) ageprior = physmodage(massmult_info["ages"]) # assumes a flat metallicity prior -> will revise when metallicity metprior = np.full(len(massmult_info["Zs"]), 1.0) metprior /= np.sum(metprior) # loop over the age and mass points and compute the number of stars expected # at each point including the completeness n_totstars = 0.0 for i, cage in enumerate(massmult_info["ages"]): for j, cmet in enumerate(massmult_info["Zs"]): gmods = (bphysparams["logA"] == cage) & (bphysparams["Z"] == cmet) # compute the mass remaining at each age -> this is the mass to simulate simmass = ageprior[i] * metprior[j] * massmult_info["massmult"][i, j] # number of stars (can be fractional) nstars_curage = simmass / massmult_info["avemass"] # normalize the total probability to have number of stars curweights = gridweights[gmods] if np.sum(curweights) > 0: curweights *= nstars_curage / np.sum(curweights) # account for the completeness curweights *= bphysparams["completeness"][gmods] nstars_curage_wcomp = np.sum(curweights) n_totstars += nstars_curage_wcomp return n_totstars
[docs] def get_predicted_num_stars_simulate( massmult_info, bphysparams, bphysmod, physmodage ): """ Calculate the expected number of stars based on the physics model as given on the BEAST model grid including completeness. Parameters ---------- massmult_info : dict "massmult" gives the value to muliply the SFH at each age and metallicity "ages" gives the ages and "Zs" gives the metallicities bphysparams : astropy.table table giving the physical parameters, weights, and completeness on the BEAST physical grid bphysmod : array probability of the physics model on the BEAST physics grid physmodage : beast.physicmodel.priormodel physics model of for age Returns ------- n_stars : float number of stars expected """ # initialize the random number generator rangen = default_rng() # setup model_indx = np.arange(len(bphysparams["M_ini"])) nsim = 0 gridweights = bphysmod * bphysparams["grid_weight"] gridweights = gridweights / np.sum(gridweights) ageprior = physmodage(massmult_info["ages"]) # assumes a flat metallicity prior -> will revise when metallicity metprior = np.full(len(massmult_info["Zs"]), 1.0) metprior /= np.sum(metprior) # loop over the age and mass points and compute the number of stars expected # at each point including the completeness totsim_indx = np.array([], dtype=int) for i, cage in enumerate(massmult_info["ages"]): for j, cmet in enumerate(massmult_info["Zs"]): gmods = (bphysparams["logA"] == cage) & (bphysparams["Z"] == cmet) # compute the mass remaining at each age -> this is the mass to simulate simmass = ageprior[i] * metprior[j] * massmult_info["massmult"][i, j] nsim_curage = int(simmass / massmult_info["avemass"]) # simluate the stars at the current age and metallicity curweights = gridweights[gmods] if np.sum(curweights) > 0: curweights /= np.sum(curweights) cursim_indx = rangen.choice( model_indx[gmods], size=nsim_curage, p=curweights ) totsim_indx = np.concatenate((totsim_indx, cursim_indx)) nsim += nsim_curage # totsimcurmass = np.sum(sedgrid["M_ini"][cursim_indx]) # print(cage, totcurmass / totmass, simmass, totsimcurmass, nsim_curage) # totsimmass = np.sum(bphysparams["M_ini"][totsim_indx]) # print(f"number total simulated stars = {nsim}; mass = {totsimmass}") compl_choice = rangen.random(nsim) compl_indx = bphysparams["completeness"][totsim_indx] >= compl_choice sim_indx = totsim_indx[compl_indx] # totcompsimmass = np.sum(bphysparams["M_ini"][sim_indx]) # print( # f"number of simulated stars w/ completeness = {len(sim_indx)}; mass = {totcompsimmass}" # ) return len(sim_indx)