diff --git a/diffsky/mc_diffsky.py b/diffsky/mc_diffsky.py index 3778f7f..b6d65d6 100644 --- a/diffsky/mc_diffsky.py +++ b/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 diff --git a/diffsky/tests/test_mc_diffsky.py b/diffsky/tests/test_mc_diffsky.py index 1b400bd..8310fd0 100644 --- a/diffsky/tests/test_mc_diffsky.py +++ b/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,6 +29,10 @@ 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) @@ -34,10 +40,16 @@ def test_mc_diffstar_cenpop(): 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()