diff --git a/beast/physicsmodel/creategrid.py b/beast/physicsmodel/creategrid.py index abbce1c6..3c28a570 100644 --- a/beast/physicsmodel/creategrid.py +++ b/beast/physicsmodel/creategrid.py @@ -14,6 +14,7 @@ * likelihood computations need to be updated to allow computations even if the full grid does not fit in memory """ + import numpy as np import copy @@ -22,9 +23,8 @@ from beast.physicsmodel.stars import stellib from beast.physicsmodel.grid import SpectralGrid, SEDGrid -from beast.physicsmodel.priormodel import PriorDustModel +from beast.physicsmodel.grid_and_prior_weights import compute_av_rv_fA_prior_weights -# from beast.external.eztables import Table from astropy.table import Table from beast.tools.helpers import generator from beast.tools import helpers @@ -321,13 +321,9 @@ def make_extinguished_grid( # Create the sampling mesh # ======================== # basically the dot product from all input 1d vectors - # setup interation over the full dust parameter grid + # setup integration over the full dust parameter grid - # setup the dust prior models - av_prior = PriorDustModel(av_prior_model) - rv_prior = PriorDustModel(rv_prior_model) if with_fA: - fA_prior = PriorDustModel(fA_prior_model) it = np.nditer(np.ix_(avs, rvs, fAs)) niter = np.size(avs) * np.size(rvs) * np.size(fAs) @@ -387,7 +383,7 @@ def make_extinguished_grid( n_filters = len(filter_names) _seds = np.zeros((N, n_filters), dtype=float) if absflux_cov: - n_offdiag = ((n_filters ** 2) - n_filters) / 2 + n_offdiag = ((n_filters**2) - n_filters) / 2 _cov_diag = np.zeros((N, n_filters), dtype=float) _cov_offdiag = np.zeros((N, n_offdiag), dtype=float) @@ -429,34 +425,24 @@ def make_extinguished_grid( cols["Rv"][N0 * count : N0 * (count + 1)] = Rv # compute the dust weights - # moved here in 2023 to support distance based dust priors - dists = g0.grid["distance"].data - if av_prior_model["name"] == "step": - av_weights = av_prior(np.full((len(dists)), Av), y=dists) - else: - av_weights = av_prior(Av) - if rv_prior_model["name"] == "step": - rv_weights = rv_prior(np.full((len(dists)), Rv), y=dists) - else: - rv_weights = rv_prior(Rv) - if fA_prior_model["name"] == "step": - f_A_weights = fA_prior(np.full((len(dists)), f_A), y=dists) - else: - if with_fA: - f_A_weights = fA_prior(f_A) - else: - f_A_weights = 1.0 - - dust_prior_weight = av_weights * rv_weights * f_A_weights + dust_prior_weight = compute_av_rv_fA_prior_weights( + Av, + Rv, + f_A, + g0.grid["distance"].data, + av_prior_model=av_prior_model, + rv_prior_model=rv_prior_model, + fA_prior_model=fA_prior_model, + ) # get new attributes if exist for key in list(temp_results.grid.keys()): if key not in keys: k1 = N0 * count k2 = N0 * (count + 1) - cols.setdefault(key, np.zeros(N, dtype=float))[ - k1:k2 - ] = temp_results.grid[key] + cols.setdefault(key, np.zeros(N, dtype=float))[k1:k2] = ( + temp_results.grid[key] + ) # compute the fractional absflux covariance matrices if absflux_cov: @@ -585,7 +571,7 @@ def add_spectral_properties( def calc_absflux_cov_matrices(specgrid, sedgrid, filter_names): - """ Calculate the absflux covariance matrices for each model + """Calculate the absflux covariance matrices for each model Must be done on the full spectrum of each model to account for the changing combined spectral response due to the model SED and the filter response curve. @@ -611,7 +597,7 @@ def calc_absflux_cov_matrices(specgrid, sedgrid, filter_names): # setup the output quantities n_models = specgrid.seds.shape[0] n_filters = len(filter_names) - n_offdiag = ((n_filters ** 2) - n_filters) / 2 + n_offdiag = ((n_filters**2) - n_filters) / 2 cov_diag = np.zeros((n_models, n_filters), dtype=np.float64) cov_offdiag = np.zeros((n_models, n_offdiag), dtype=np.float64) diff --git a/beast/physicsmodel/grid_and_prior_weights.py b/beast/physicsmodel/grid_and_prior_weights.py index f2713591..825b221c 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -12,6 +12,8 @@ flat priors on all the fit parameters. Non-flat priors will be implemented with prior weights. """ + +from inspect import signature import numpy as np from beast.physicsmodel.grid_weights_stars import compute_distance_grid_weights @@ -24,11 +26,13 @@ PriorMassModel, PriorMetallicityModel, PriorDistanceModel, + PriorDustModel, ) __all__ = [ "compute_age_mass_metallicity_weights", "compute_distance_age_mass_metallicity_weights", + "compute_av_rv_fA_weights", ] @@ -155,7 +159,10 @@ def compute_age_mass_metallicity_weights( # compute the age weights age_grid_weights = compute_age_grid_weights(uniq_ages) - age_prior = PriorAgeModel(age_prior_model) + if isinstance(age_prior_model, dict): + age_prior = PriorAgeModel(age_prior_model) + else: + age_prior = age_prior_model age_prior_weights = age_prior(uniq_ages) for ak, age_val in enumerate(uniq_ages): @@ -168,10 +175,26 @@ def compute_age_mass_metallicity_weights( # 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) + # cur_masses = _tgrid_single_age["M_ini"] + # deal with repeat masses - happens for MegaBEAST + cur_masses = np.unique(_tgrid_single_age["M_ini"]) + umass_grid_weights = compute_mass_grid_weights(cur_masses) + if isinstance(mass_prior_model, dict): + mass_prior = PriorMassModel(mass_prior_model) + else: + mass_prior = mass_prior_model + umass_prior_weights = mass_prior(cur_masses) + n_masses = len(_tgrid_single_age["M_ini"]) + if len(cur_masses) < n_masses: + mass_grid_weights = np.zeros(n_masses, dtype=float) + mass_prior_weights = np.zeros(n_masses, dtype=float) + for k, cmass in enumerate(cur_masses): + gvals = _tgrid_single_age["M_ini"] == cmass + mass_grid_weights[gvals] = umass_grid_weights[k] + mass_prior_weights[gvals] = umass_prior_weights[k] + else: + mass_grid_weights = umass_grid_weights + mass_prior_weights = umass_prior_weights else: # must be a single mass for this age,z combination # set mass weight to zero to remove this point from the grid @@ -218,3 +241,59 @@ def compute_age_mass_metallicity_weights( met_prior_weights[i] * total_z_prior_weight[i] ) _tgrid[zindxs]["weight"] *= met_weights[i] * total_z_weight[i] + + +def compute_av_rv_fA_prior_weights( + Av, + Rv, + f_A, + dists, + av_prior_model={"name": "flat"}, + rv_prior_model={"name": "flat"}, + fA_prior_model={"name": "flat"}, +): + """ + Computes the av, rv, f_A grid and prior weights + on the BEAST model spectra grid + Grid and prior weight columns updated by multiplying by the + existing weights + + Parameters + ---------- + Av : vector + A(V) values + Rv : vector + R(V) values + f_A : vector + f_A values + dists : vector + distance values + av_prior_model : dict + dict including prior model name and parameters + rv_prior_model : dict + dict including prior model name and parameters + fA_prior_model :dict + dict including prior model name and parameters + """ + av_prior = PriorDustModel(av_prior_model) + rv_prior = PriorDustModel(rv_prior_model) + fA_prior = PriorDustModel(fA_prior_model) + if av_prior_model["name"] == "step": + av_weights = av_prior(np.full((len(dists)), Av), y=dists) + else: + av_weights = av_prior(Av) + if rv_prior_model["name"] == "step": + rv_weights = rv_prior(np.full((len(dists)), Rv), y=dists) + else: + rv_weights = rv_prior(Rv) + if fA_prior_model["name"] == "step": + f_A_weights = fA_prior(np.full((len(dists)), f_A), y=dists) + else: + f_A_weights = fA_prior(f_A) + + dust_prior = av_weights * rv_weights * f_A_weights + + # normalize to control for numerical issues + dust_prior /= np.max(dust_prior) + + return dust_prior