modifications to support megabeast simulations

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,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