Merge pull request #821 from karllark/fix_imf_prior

modifications to support megabeast simulations

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])

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)
 

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

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()
 

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.