Include additional diffstar quantities in MC generator

diffsky/mc_diffsky.py

@@ -30,6 +30,7 @@ def mc_diffstar_galpop(
     cosmo_params=DEFAULT_COSMOLOGY,
     diffstarpop_params=DEFAULT_DIFFSTARPOP_PARAMS,
     n_t=N_T,
+    return_internal_quantities=False,
 ):
     """Generate a population of galaxies with diffmah MAH and diffstar SFH
 
@@ -56,6 +57,11 @@ def mc_diffstar_galpop(
         Default is None, in which case volume_com argument must be passed
         and the host halo mass function will be randomly sampled.
 
+    return_internal_quantities : bool, optional
+        If True, returned data will include additional info such as
+        the separate SFHs for the probabilistic main and quenched sequences.
+        Default is False, in which case only a single SFH will be returned.
+
     Returns
     -------
     diffsky_data : dict
@@ -116,6 +122,13 @@ def mc_diffstar_galpop(
     )
     diffstar_data["logssfr_obs"] = logsfh_obs - diffstar_data["logsm_obs"]
 
+    if return_internal_quantities:
+        diffstar_data["sfh_ms"] = sfh_ms
+        diffstar_data["sfh_q"] = sfh_q
+        diffstar_data["frac_q"] = frac_q
+        diffstar_data["sfh_params_ms"] = diffstar_params_ms
+        diffstar_data["sfh_params_q"] = diffstar_params_q
+
     return diffstar_data
 
 
@@ -128,6 +141,7 @@ def mc_diffstar_cenpop(
     cosmo_params=DEFAULT_COSMOLOGY,
     diffstarpop_params=DEFAULT_DIFFSTARPOP_PARAMS,
     n_t=N_T,
+    return_internal_quantities=False,
 ):
     """Generate a population of central galaxies with diffmah MAH and diffstar SFH
 
@@ -154,6 +168,11 @@ def mc_diffstar_cenpop(
         Default is None, in which case volume_com argument must be passed
         and the host halo mass function will be randomly sampled.
 
+    return_internal_quantities : bool, optional
+        If True, returned data will include additional info such as
+        the separate SFHs for the probabilistic main and quenched sequences.
+        Default is False, in which case only a single SFH will be returned.
+
     Returns
     -------
     diffsky_data : dict
@@ -215,4 +234,11 @@ def mc_diffstar_cenpop(
     )
     diffstar_data["logssfr_obs"] = logsfh_obs - diffstar_data["logsm_obs"]
 
+    if return_internal_quantities:
+        diffstar_data["sfh_ms"] = sfh_ms
+        diffstar_data["sfh_q"] = sfh_q
+        diffstar_data["frac_q"] = frac_q
+        diffstar_data["sfh_params_ms"] = diffstar_params_ms
+        diffstar_data["sfh_params_q"] = diffstar_params_q
+
     return diffstar_data

diffsky/tests/test_mc_diffsky.py

@@ -14,7 +14,9 @@ def test_mc_diffstar_galpop():
     lgmp_min = 11.0
     z_obs = 0.01
     args = (ran_key, z_obs, lgmp_min)
-    diffsky_data = mcd.mc_diffstar_galpop(*args, hosts_logmh_at_z=hosts_logmh_at_z)
+    diffsky_data = mcd.mc_diffstar_galpop(
+        *args, hosts_logmh_at_z=hosts_logmh_at_z, return_internal_quantities=True
+    )
 
     for p in diffsky_data["subcat"].mah_params:
         assert np.all(np.isfinite(p))
@@ -27,17 +29,27 @@ def test_mc_diffstar_galpop():
 
     assert diffsky_data["t_obs"] > 13.5
 
+    # Enforce return_internal_quantities supplies additional info
+    for key in ("sfh_params_q", "sfh_ms", "frac_q"):
+        assert key in diffsky_data.keys()
+
 
 def test_mc_diffstar_cenpop():
     ran_key = jran.key(0)
     n_cens = 200
     hosts_logmh_at_z = np.linspace(10, 15, n_cens)
     z_obs = 0.1
     args = (ran_key, z_obs)
-    diffsky_data = mcd.mc_diffstar_cenpop(*args, hosts_logmh_at_z=hosts_logmh_at_z)
+    diffsky_data = mcd.mc_diffstar_cenpop(
+        *args, hosts_logmh_at_z=hosts_logmh_at_z, return_internal_quantities=True
+    )
 
     for p in diffsky_data["subcat"].mah_params:
         assert np.all(np.isfinite(p))
     assert np.all(np.isfinite(diffsky_data["smh"]))
 
     assert diffsky_data["subcat"].logmp0.size == n_cens
+
+    # Enforce return_internal_quantities supplies additional info
+    for key in ("sfh_params_q", "sfh_ms", "frac_q"):
+        assert key in diffsky_data.keys()