def compute_age_prior_weights(logages, age_prior_model):
"""
Computes the age prior for the specified model
Parameters
----------
logages : numpy vector
log(ages)
age_prior_model: dict
dict including prior model name and parameters
Returns
-------
age_weights : numpy vector
weights needed according to the prior model
"""
if age_prior_model["name"] == "flat" or age_prior_model[
"name"] == "flat_linear":
age_weights = np.full(len(logages), 1.0)
elif age_prior_model["name"] == "flat_log":
# flat in log space means use the native log(age) grid spacing
# thus the priors weights are the inverse of the grid weights
# assumes the logace spacing is uniform
age_weights = 1.0 / compute_age_grid_weights(logages)
elif age_prior_model["name"] == "bins_histo":
# interpolate according to bins, assuming SFR constant from i to i+1
# and allow for bin edges input
if len(age_prior_model["values"]) == len(
age_prior_model["logages"]) - 1:
age_prior_model["values"].append(0.0)
ageND = interp1d(age_prior_model["logages"],
age_prior_model["values"],
kind="zero")
age_weights = ageND(logages)
elif age_prior_model["name"] == "bins_interp":
# interpolate model to grid ages
age_weights = np.interp(
logages,
np.array(age_prior_model["logages"]),
np.array(age_prior_model["values"]),
)
elif age_prior_model["name"] == "exp":
# assumes SFR(t) \propto e**(-t/tau) \propto e**(age/tau)
# where age \propto -t (for age=t0-t) and tau in Gyr
vals = (10**logages) / (age_prior_model["tau"] * 1e9)
age_weights = np.exp(vals)
else:
raise NotImplementedError(
"input age prior ''{}'' function not supported".format(
age_prior_model["name"]))
# normalize to avoid numerical issues (too small or too large)
age_weights /= np.average(age_weights)
return age_weights
def test_age_grid_weights():
"""
Test age grid weights
"""
ages = np.array([6, 7, 8, 9, 10])
weights = compute_age_grid_weights(ages)
expected_weights = [
4.500045e-04,
4.500045e-03,
4.500045e-02,
4.500045e-01,
4.500045e00,
]
np.testing.assert_allclose(
weights, expected_weights, err_msg=("Stellar grid age weights error")
)
def compute_age_mass_metallicity_weights(
_tgrid,
indxs,
age_prior_model={"name": "flat"},
mass_prior_model={"name": "kroupa"},
met_prior_model={"name": "flat"},
**kwargs
):
"""
Computes the age-mass-metallicity grid and prior weights
on the BEAST model spectra grid
Grid and prior weight columns updated by multiplying by the
age-mass-metallicity weight.
Parameters
----------
_tgrid : SpectralGrid
BEAST models spectral grid
age_prior_model : dict
dict including prior model name and parameters
mass_prior_model : dict
dict including prior model name and parameters
met_prior_model : dict
dict including prior model name and parameters
"""
# get the unique metallicities
uniq_Zs = np.unique(_tgrid[indxs]["Z"])
# setup the vector to hold the z weight vector
total_z_grid_weight = np.zeros(len(uniq_Zs))
total_z_prior_weight = np.zeros(len(uniq_Zs))
total_z_weight = np.zeros(len(uniq_Zs))
for az, z_val in enumerate(uniq_Zs):
print("computing the age-mass-metallicity grid weight for Z = ", z_val)
# get the grid for a single metallicity
(zindxs,) = np.where(_tgrid[indxs]["Z"] == z_val)
# get the unique ages for this metallicity
zindxs = indxs[zindxs]
uniq_ages = np.unique(_tgrid[zindxs]["logA"])
# compute the age weights
age_grid_weights = compute_age_grid_weights(uniq_ages)
age_prior = PriorAgeModel(age_prior_model)
age_prior_weights = age_prior(uniq_ages)
for ak, age_val in enumerate(uniq_ages):
# get the grid for a single age
(aindxs,) = np.where(
(_tgrid[indxs]["logA"] == age_val) & (_tgrid[indxs]["Z"] == z_val)
)
aindxs = indxs[aindxs]
_tgrid_single_age = _tgrid[aindxs]
# compute the mass weights
if len(aindxs) > 1:
cur_masses = _tgrid_single_age["M_ini"]
mass_grid_weights = compute_mass_grid_weights(cur_masses)
mass_prior = PriorMassModel(mass_prior_model)
mass_prior_weights = mass_prior(cur_masses)
else:
# must be a single mass for this age,z combination
# set mass weight to zero to remove this point from the grid
mass_grid_weights = np.zeros(1)
mass_prior_weights = np.zeros(1)
# apply both the mass and age weights
for i, k in enumerate(aindxs):
comb_grid_weights = mass_grid_weights[i] * age_grid_weights[ak]
comb_prior_weights = mass_prior_weights[i] * age_prior_weights[ak]
_tgrid[k]["grid_weight"] *= comb_grid_weights
_tgrid[k]["prior_weight"] *= comb_prior_weights
_tgrid[k]["weight"] *= comb_grid_weights * comb_prior_weights
# compute the current total weight at each metallicity
total_z_grid_weight[az] = np.sum(_tgrid[zindxs]["grid_weight"])
total_z_prior_weight[az] = np.sum(_tgrid[zindxs]["prior_weight"])
total_z_weight[az] = np.sum(_tgrid[zindxs]["weight"])
# ensure that the metallicity prior is uniform
if len(uniq_Zs) > 1:
# get the metallicity weights
met_grid_weights = compute_metallicity_grid_weights(uniq_Zs)
met_grid_weights /= np.sum(met_grid_weights)
met_prior = PriorMetallicityModel(met_prior_model)
met_prior_weights = met_prior(uniq_Zs)
met_prior_weights /= np.sum(met_prior_weights)
met_weights = met_grid_weights * met_prior_weights
# correct for any non-unformity in the number size of the
# age-mass grids between metallicity points
total_z_grid_weight /= np.sum(total_z_grid_weight)
total_z_prior_weight /= np.sum(total_z_prior_weight)
total_z_weight /= np.sum(total_z_weight)
for i, z_val in enumerate(uniq_Zs):
# get the grid for this metallicity
(zindxs,) = np.where(_tgrid[indxs]["Z"] == z_val)
zindxs = indxs[zindxs]
_tgrid[zindxs]["grid_weight"] *= (
met_grid_weights[i] * total_z_grid_weight[i]
)
_tgrid[zindxs]["prior_weight"] *= (
met_prior_weights[i] * total_z_prior_weight[i]
)
_tgrid[zindxs]["weight"] *= met_weights[i] * total_z_weight[i]
def compute_age_prior_weights(logages, age_prior_model):
"""
Computes the age prior for the specified model
Parameters
----------
logages : numpy vector
log(ages)
age_prior_model: dict
dict including prior model name and parameters
Returns
-------
age_weights : numpy vector
weights needed according to the prior model
"""
if age_prior_model["name"] == "flat" or age_prior_model[
"name"] == "flat_linear":
if "sfr" in age_prior_model.keys():
sfr = age_prior_model["sfr"]
else:
sfr = 1.0
age_weights = np.full(len(logages), sfr)
elif age_prior_model["name"] == "flat_log":
# flat in log space means use the native log(age) grid spacing
# thus the priors weights are the inverse of the grid weights
# assumes the logace spacing is uniform
age_weights = 1.0 / compute_age_grid_weights(logages)
elif age_prior_model["name"] == "bins_histo":
# check if all ages within interpolation range
if np.all([
np.max(logages) <= x <= np.min(logages)
for x in age_prior_model["values"]
]):
raise ValueError(
"Age prior weight error: Requested ages outside of model range"
)
# interpolate according to bins, assuming SFR constant from i to i+1
# and allow for bin edges input
if len(age_prior_model["values"]) == len(
age_prior_model["logages"]) - 1:
age_prior_model["values"].append(0.0)
ageND = interp1d(age_prior_model["logages"],
age_prior_model["values"],
kind="zero")
age_weights = ageND(logages)
elif age_prior_model["name"] == "bins_interp":
# interpolate model to grid ages
age_weights = np.interp(
logages,
np.array(age_prior_model["logages"]),
np.array(age_prior_model["values"]),
)
elif age_prior_model["name"] == "exp":
# assumes SFR(t) \propto e**(-t/tau) \propto e**(age/tau)
# where age \propto -t (for age=t0-t) and tau in Gyr
vals = (10**logages) / (age_prior_model["tau"] * 1e9)
age_weights = np.exp(-1.0 * vals)
else:
raise NotImplementedError(
"input age prior ''{}'' function not supported".format(
age_prior_model["name"]))
# normalize to avoid numerical issues (too small or too large)
# do not normalize as the absolute level now matters for simulations and megaBEAST
# age_weights /= np.average(age_weights)
return age_weights