Merge pull request #795 from karllark/better_distance_dust_priors

Adding in step and absexp priors of dust and dist

CHANGES.rst

@@ -1,6 +1,7 @@
 2.1 (unreleased)
 ================
-- No changes yet
+- added 2D dust prior step model that also depends on distance
+- added absexponential distance prior
 
 2.0 (2022-11-11)
 ================

beast/physicsmodel/creategrid.py

@@ -395,7 +395,6 @@ def make_extinguished_grid(
 
             if with_fA:
                 Av, Rv, f_A = pt
-                dust_prior_weight = av_prior(Av) * rv_prior(Rv) * fA_prior(f_A)
                 Rv_MW = extLaw.get_Rv_A(Rv, f_A)
                 r = g0.applyExtinctionLaw(extLaw, Av=Av, Rv=Rv, f_A=f_A, inplace=False)
                 # add extra "spectral bands" if requested
@@ -415,7 +414,6 @@ def make_extinguished_grid(
 
             else:
                 Av, Rv = pt
-                dust_prior_weight = av_prior(Av) * rv_prior(Rv)
                 r = g0.applyExtinctionLaw(extLaw, Av=Av, Rv=Rv, inplace=False)
 
                 if add_spectral_properties_kwargs is not None:
@@ -430,6 +428,27 @@ def make_extinguished_grid(
                 cols["Av"][N0 * count : N0 * (count + 1)] = Av
                 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
+
             # get new attributes if exist
             for key in list(temp_results.grid.keys()):
                 if key not in keys:

beast/physicsmodel/grid_and_prior_weights.py

@@ -82,7 +82,6 @@ def compute_distance_age_mass_metallicity_weights(
         total_dist_prior_weight[dz] = np.sum(_tgrid[dindxs]["prior_weight"])
         total_dist_weight[dz] = np.sum(_tgrid[dindxs]["weight"])
 
-    # ensure that the distance prior is uniform
     if n_dist > 1:
         # get the distance weights
         dist_grid_weights = compute_distance_grid_weights(uniq_dists)

beast/physicsmodel/priormodel.py

@@ -1,6 +1,7 @@
 import numpy as np
 from scipy.interpolate import interp1d
 from scipy.integrate import quad
+import astropy.units as u
 
 from beast.physicsmodel.grid_weights_stars import compute_bin_boundaries
 import beast.physicsmodel.priormodel_functions as pmfuncs
@@ -36,14 +37,16 @@ def __init__(self, model, allowed_models=None):
         # save the model
         self.model = model
 
-    def __call__(self, x):
+    def __call__(self, x, y=None):
         """
         Weights based on input model choice
 
         Parameters
         ----------
         x : float
             values for model evaluation
+        y : float
+            secondary values for 2D priors
         """
         if self.model["name"] == "flat":
             if "amp" in self.model.keys():
@@ -73,9 +76,7 @@ def __call__(self, x):
             else:
                 # interpolate model to grid ages
                 return np.interp(
-                    x,
-                    np.array(self.model["x"]),
-                    np.array(self.model["values"]),
+                    x, np.array(self.model["x"]), np.array(self.model["values"]),
                 )
         elif self.model["name"] == "lognormal":
             for ckey in ["mean", "sigma"]:
@@ -95,11 +96,35 @@ def __call__(self, x):
                 N1=self.model["N1_to_N2"],
                 N2=1.0,
             )
-        elif self.model["name"] == "exponential":
-            for ckey in ["tau"]:
+        elif self.model["name"] == "absexponential":
+            for ckey in ["dist0", "tau", "amp"]:
+                if ckey not in self.model.keys():
+                    raise ValueError(f"{ckey} not in prior model keys")
+            return pmfuncs._absexponential(
+                x,
+                dist0=self.model["dist0"].to(u.pc).value,
+                tau=self.model["tau"].to(u.pc).value,
+                amp=self.model["amp"],
+            )
+        elif self.model["name"] == "step":
+            for ckey in ["dist0", "amp1", "damp2", "lgsigma1", "lgsigma2"]:
                 if ckey not in self.model.keys():
                     raise ValueError(f"{ckey} not in prior model keys")
-            return pmfuncs._exponential(x, tau=self.model["tau"])
+            if y is None:
+                raise ValueError("y values not passed required for 2D priors")
+            if len(x) != len(y):
+                raise ValueError(
+                    "x and y values not the same length, required for 2D priors"
+                )
+            return pmfuncs._step(
+                x,
+                y,
+                dist0=self.model["dist0"].to(u.pc).value,
+                amp1=self.model["amp1"],
+                damp2=self.model["damp2"],
+                lgsigma1=self.model["lgsigma1"],
+                lgsigma2=self.model["lgsigma2"],
+            )
         else:
             modname = self.model["name"]
             raise NotImplementedError(f"{modname} is not an allowed model")
@@ -120,7 +145,7 @@ def __init__(self, model):
           Possible choices are flat, lognormal, two_lognormal, and exponential
         """
         super().__init__(
-            model, allowed_models=["flat", "lognormal", "two_lognormal", "exponential"]
+            model, allowed_models=["flat", "lognormal", "two_lognormal", "step"]
         )
 
 
@@ -198,7 +223,7 @@ def __init__(self, model):
         model : dict
           Possible choices are flat
         """
-        super().__init__(model, allowed_models=["flat"])
+        super().__init__(model, allowed_models=["flat", "absexponential"])
 
 
 class PriorMassModel(PriorModel):

beast/physicsmodel/priormodel_functions.py

@@ -4,6 +4,8 @@
     "_lognorm",
     "_two_lognorm",
     "_exponential",
+    "_absexponential",
+    "_step",
     "_imf_salpeter",
     "_imf_kroupa",
     "_imf_flat",
@@ -109,6 +111,64 @@ def _exponential(x, tau=2.0, amp=1.0):
     return amp * np.exp(-1.0 * x / tau)
 
 
+def _absexponential(x, dist0, tau=2.0, amp=1.0):
+    """
+    Absolute value of exponential distribution
+    Used for stellar density distriction versus distance
+
+    Parameters
+    ----------
+    x : vector
+       x values
+    dist0 : float
+       distance at peak amplitude
+    tau : float
+       Decay Rate parameter in exp: e^-(abs(d - dist0)/tau)
+    amp : float
+       Amplitude for dist0
+
+    Returns
+    -------
+    absolute value of exponential computed on the x grid
+    """
+    return amp * np.exp(-1.0 * np.absolute(x - dist0) / tau)
+
+
+def _step(x, y, dist0, amp1=0.0, damp2=1.0, lgsigma1=0.1, lgsigma2=0.1):
+    """
+    Step function
+
+    Parameters
+    ----------
+    x : vector
+       x values
+    dist0 : float
+       distance of step
+    amp1 : float
+       Amplitude before dist0
+    damp2 : float
+       Delta amplitude after dist0 (e.g. afterwards amp = amp1 + damp2)
+    lgsigma1 : float
+       log-normal sigma for amp1
+    lgsigma2 : float
+       log-normal sigma for amp2
+
+    Returns
+    -------
+    Step function evaluted on the x grid
+    """
+    probs = x * 0.0
+    gvals = y < dist0
+    probs[gvals] = _lognorm(x[gvals], amp1, sigma=lgsigma1)
+    # amp1 includes as all are affected
+    #  not quite correct, likely the 2nd log-normal should be convolved by the 1st
+    #  but how to do this analytically?
+    gvals = y >= dist0
+    probs[gvals] += _lognorm(x[gvals], amp1 + damp2, sigma=lgsigma2)
+
+    return probs
+
+
 def _imf_flat(x):
     """
     Compute a flat IMF (useful for simulations, not for normal BEAST runs)
@@ -174,10 +234,10 @@ def _imf_kroupa(in_x, alpha0=0.3, alpha1=1.3, alpha2=2.3, alpha3=2.3):
     m2 = 0.5
     m3 = 1.0
 
-    ialpha0 = -1. * alpha0
-    ialpha1 = -1. * alpha1
-    ialpha2 = -1. * alpha2
-    ialpha3 = -1. * alpha3
+    ialpha0 = -1.0 * alpha0
+    ialpha1 = -1.0 * alpha1
+    ialpha2 = -1.0 * alpha2
+    ialpha3 = -1.0 * alpha3
 
     imf = np.full((len(x)), 0.0)
 

beast/physicsmodel/tests/test_dust_prior_weights.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+import astropy.units as u
 
 from beast.physicsmodel.priormodel import PriorDustModel
 
@@ -24,14 +25,12 @@ def test_av_prior_weights():
             "sigma2": 0.2,
             "N1_to_N2": 1.0 / 5.0,
         },
-        {"name": "exponential", "tau": 1.0},
     ]
     # fmt: off
     expected_vals = [
         [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
         [0.0, 0.120985, 0.095149, 0.066168, 0.046282, 0.033133],
         [0.000000e+00, 4.969014e-01, 2.563834e-03, 3.719773e-04, 1.341047e-04, 5.743044e-05],
-        [1.000000e+00, 1.353353e-01, 1.831564e-02, 2.478752e-03, 3.354626e-04, 4.539993e-05],
     ]
     # fmt: on
 
@@ -48,9 +47,38 @@ def test_av_prior_weights():
     # test error output when unsupported model is given
     with pytest.raises(NotImplementedError) as exc_av:
         PriorDustModel({"name": "exp"})
-    assert (
-        str(exc_av.value)
-        == "exp is not an allowed model"
+    assert str(exc_av.value) == "exp is not an allowed model"
+
+
+def test_av_prior_weights_2d():
+    """
+    Test the A(V) dust prior weights for 2D priors
+    """
+
+    dist_vals = np.array([50.0, 59.0, 61.0, 70.0]) * 1e3
+    av_vals = np.array([0.1, 0.9, 1.1, 2.0])
+
+    # Test A(V) priors
+    model = {
+        "name": "step",
+        "dist0": 60 * u.kpc,
+        "amp1": 0.1,
+        "damp2": 1.0,
+        "lgsigma1": 0.05,
+        "lgsigma2": 0.05,
+    }
+
+    # fmt: off
+    expected_vals = [7.968878e+01, 0.000000e+00, 7.244435e+00, 6.544045e-31]
+    # fmt: on
+
+    av_prior = PriorDustModel(model)
+    mname = model["name"]
+    np.testing.assert_allclose(
+        av_prior(av_vals, y=dist_vals),
+        expected_vals,
+        atol=1e-6,
+        err_msg=f"A problem occurred while setting the 2D A(V) priors to {mname}.",
     )
 
 
@@ -94,6 +122,38 @@ def test_rv_prior_weights():
         )
 
 
+def test_rv_prior_weights_2d():
+    """
+    Test the R(V) dust prior weights for 2D priors
+    """
+
+    dist_vals = np.array([50.0, 59.0, 61.0, 70.0]) * 1e3
+    rv_vals = np.array([3.1, 3.8, 4.6, 5.5])
+
+    # Test A(V) priors
+    model = {
+        "name": "step",
+        "dist0": 60 * u.kpc,
+        "amp1": 3.1,
+        "damp2": 1.5,
+        "lgsigma1": 0.05,
+        "lgsigma2": 0.05,
+    }
+
+    # fmt: off
+    expected_vals = [2.570606e+00, 6.449322e-04, 1.732365e+00, 2.918564e-03]
+    # fmt: on
+
+    av_prior = PriorDustModel(model)
+    mname = model["name"]
+    np.testing.assert_allclose(
+        av_prior(rv_vals, y=dist_vals),
+        expected_vals,
+        atol=1e-6,
+        err_msg=f"A problem occurred while setting the 2D R(V) priors to {mname}.",
+    )
+
+
 def test_fA_prior_weights():
     """
     Test the f_A dust prior weights
@@ -132,3 +192,35 @@ def test_fA_prior_weights():
             atol=1e-6,
             err_msg=f"A problem occurred while setting the fA priors to {mname}.",
         )
+
+
+def test_fA_prior_weights_2d():
+    """
+    Test the fA dust prior weights for 2D priors
+    """
+
+    dist_vals = np.array([50.0, 59.0, 61.0, 70.0]) * 1e3
+    fA_vals = np.array([0.1, 0.5, 0.9, 0.3])
+
+    # Test A(V) priors
+    model = {
+        "name": "step",
+        "dist0": 60 * u.kpc,
+        "amp1": 0.1,
+        "damp2": 0.8,
+        "lgsigma1": 0.05,
+        "lgsigma2": 0.1,
+    }
+
+    # fmt: off
+    expected_vals = [7.968878e+001, 8.158912e-224, 4.410584e+000, 2.728553e-026]
+    # fmt: on
+
+    av_prior = PriorDustModel(model)
+    mname = model["name"]
+    np.testing.assert_allclose(
+        av_prior(fA_vals, y=dist_vals),
+        expected_vals,
+        atol=1e-6,
+        err_msg=f"A problem occurred while setting the 2D f_A priors to {mname}.",
+    )