From 8b1cf11cef0575a3003d191993bf0eed9893b47e Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Tue, 13 Aug 2024 15:40:31 -0400 Subject: [PATCH 1/6] modifications to support megabeast simulations --- beast/physicsmodel/creategrid.py | 50 ++++------- beast/physicsmodel/grid_and_prior_weights.py | 89 ++++++++++++++++++-- 2 files changed, 102 insertions(+), 37 deletions(-) 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 From 78f20bd99781901e15152c635a5a646da8854dfc Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Tue, 13 Aug 2024 15:50:11 -0400 Subject: [PATCH 2/6] fixing errors --- beast/physicsmodel/grid_and_prior_weights.py | 5 ++--- beast/plotting/plot_cmd.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/beast/physicsmodel/grid_and_prior_weights.py b/beast/physicsmodel/grid_and_prior_weights.py index 825b221c..e7d7bffd 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -13,7 +13,6 @@ with prior weights. """ -from inspect import signature import numpy as np from beast.physicsmodel.grid_weights_stars import compute_distance_grid_weights @@ -32,7 +31,7 @@ __all__ = [ "compute_age_mass_metallicity_weights", "compute_distance_age_mass_metallicity_weights", - "compute_av_rv_fA_weights", + "compute_av_rv_fA_prior_weights", ] @@ -294,6 +293,6 @@ def compute_av_rv_fA_prior_weights( dust_prior = av_weights * rv_weights * f_A_weights # normalize to control for numerical issues - dust_prior /= np.max(dust_prior) + # 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..e59f2d75 100644 --- a/beast/plotting/plot_cmd.py +++ b/beast/plotting/plot_cmd.py @@ -71,9 +71,9 @@ def plot_cmd( fig = plt.figure(figsize=(9, 9)) # sets up plots to have larger fonts, wider lines, etc. - set_params() + #set_params() - plt.plot(col, mag, "k.", ms=0.1) + plt.plot(col, mag, "ko", ms=0.5) plt.gca().invert_yaxis() From 98524a5e0c823335a5d3ff5a0702d3135adc3594 Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Wed, 14 Aug 2024 11:34:19 -0400 Subject: [PATCH 3/6] fixing errors --- beast/plotting/plot_cmd.py | 2 +- beast/tests/helpers.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/beast/plotting/plot_cmd.py b/beast/plotting/plot_cmd.py index e59f2d75..e4311b63 100644 --- a/beast/plotting/plot_cmd.py +++ b/beast/plotting/plot_cmd.py @@ -71,7 +71,7 @@ def plot_cmd( fig = plt.figure(figsize=(9, 9)) # sets up plots to have larger fonts, wider lines, etc. - #set_params() + set_params() plt.plot(col, mag, "ko", ms=0.5) 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. From 5c09bb944b166856828316fde27b9ca93f7fa4ef Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Wed, 14 Aug 2024 11:52:47 -0400 Subject: [PATCH 4/6] fixing mass grid weights --- beast/physicsmodel/grid_and_prior_weights.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/beast/physicsmodel/grid_and_prior_weights.py b/beast/physicsmodel/grid_and_prior_weights.py index e7d7bffd..5105d7cb 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -174,17 +174,17 @@ def compute_age_mass_metallicity_weights( # compute the mass weights if len(aindxs) > 1: - # 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) + + # deal with repeat masses - happens for MegaBEAST + 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): @@ -192,8 +192,9 @@ def compute_age_mass_metallicity_weights( 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 + 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 From de8c8aefa7e7828042bbff60cc1b147732411024 Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Wed, 14 Aug 2024 13:27:15 -0400 Subject: [PATCH 5/6] minor update --- beast/physicsmodel/grid_and_prior_weights.py | 1 + 1 file changed, 1 insertion(+) diff --git a/beast/physicsmodel/grid_and_prior_weights.py b/beast/physicsmodel/grid_and_prior_weights.py index 5105d7cb..4978efe0 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -195,6 +195,7 @@ def compute_age_mass_metallicity_weights( 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 From 31df91929b07c4791e5a750ad866f4705107e4c0 Mon Sep 17 00:00:00 2001 From: Karl Gordon Date: Fri, 16 Aug 2024 15:55:07 -0400 Subject: [PATCH 6/6] updates needed for regression test file creation --- beast/observationmodel/ast/make_ast_input_list.py | 2 +- beast/physicsmodel/grid_and_prior_weights.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) 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/grid_and_prior_weights.py b/beast/physicsmodel/grid_and_prior_weights.py index 4978efe0..dd056daa 100644 --- a/beast/physicsmodel/grid_and_prior_weights.py +++ b/beast/physicsmodel/grid_and_prior_weights.py @@ -94,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) @@ -180,6 +180,10 @@ def compute_age_mass_metallicity_weights( 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: