diff --git a/beast/observationmodel/ast/make_ast_input_list.py b/beast/observationmodel/ast/make_ast_input_list.py index b7f6bea1..ad1d5b74 100755 --- a/beast/observationmodel/ast/make_ast_input_list.py +++ b/beast/observationmodel/ast/make_ast_input_list.py @@ -185,7 +185,7 @@ def pick_models_toothpick_style( else: # ... do not include this model again, since we will reject it # anyway. - include_mask[idxs == rand_idx] = False + include_mask[idxs == rand_idx[r]] = False # Add the approved models chosen_idxs.extend(rand_idx[add_these]) 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..dd056daa 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -12,6 +12,7 @@ flat priors on all the fit parameters. Non-flat priors will be implemented with prior weights. """ + import numpy as np from beast.physicsmodel.grid_weights_stars import compute_distance_grid_weights @@ -24,11 +25,13 @@ PriorMassModel, PriorMetallicityModel, PriorDistanceModel, + PriorDustModel, ) __all__ = [ "compute_age_mass_metallicity_weights", "compute_distance_age_mass_metallicity_weights", + "compute_av_rv_fA_prior_weights", ] @@ -91,7 +94,7 @@ def compute_distance_age_mass_metallicity_weights( dist_prior_weights /= np.sum(dist_prior_weights) dist_weights = dist_grid_weights * dist_prior_weights - # correct for any non-unformity in the number size of the + # correct for any non-uniformity in the number size of the # age-mass grids between metallicity points total_dist_grid_weight /= np.sum(total_dist_grid_weight) total_dist_prior_weight /= np.sum(total_dist_prior_weight) @@ -155,7 +158,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 +174,32 @@ 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) + if isinstance(mass_prior_model, dict): + mass_prior = PriorMassModel(mass_prior_model) + else: + mass_prior = mass_prior_model + + # deal with repeat masses - happens for MegaBEAST + # and have discovred this can happen even for a standard BEAST run + # as sometimes two masses in an isochrone are exactly the same + # new code for MegaBEAST is more correct as then the grid weight + # will be correctly set for any repeated masses + cur_masses = np.unique(_tgrid_single_age["M_ini"]) + n_masses = len(_tgrid_single_age["M_ini"]) + if len(cur_masses) < n_masses: + umass_grid_weights = compute_mass_grid_weights(cur_masses) + umass_prior_weights = mass_prior(cur_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: + cur_masses = _tgrid_single_age["M_ini"] + mass_grid_weights = compute_mass_grid_weights(cur_masses) + 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 @@ -218,3 +246,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 diff --git a/beast/plotting/plot_cmd.py b/beast/plotting/plot_cmd.py index 944c9aa6..e4311b63 100644 --- a/beast/plotting/plot_cmd.py +++ b/beast/plotting/plot_cmd.py @@ -73,7 +73,7 @@ def plot_cmd( # sets up plots to have larger fonts, wider lines, etc. set_params() - plt.plot(col, mag, "k.", ms=0.1) + plt.plot(col, mag, "ko", ms=0.5) plt.gca().invert_yaxis() diff --git a/beast/tests/helpers.py b/beast/tests/helpers.py index b33babce..355a7b74 100644 --- a/beast/tests/helpers.py +++ b/beast/tests/helpers.py @@ -30,7 +30,7 @@ def download_rename(filename, tmpdir=""): return fname -def compare_tables(table_cache, table_new, rtol=1e-7, otag=""): +def compare_tables(table_cache, table_new, rtol=1e-6, otag=""): """ Compare two tables using astropy tables routines.