Change API of primary user-facing functions so that lgt0 and fb are required keyword arguments with no defaults

.github/workflows/linting.yaml

@@ -1,4 +1,4 @@
-name: linting
+name: Flake8 test of source code
 
 on:
   push:
@@ -8,7 +8,7 @@ on:
 
 jobs:
   tests:
-    name: lintest
+    name: flake8 diffstar
     runs-on: "ubuntu-latest"
 
     steps:

.github/workflows/monthly-warning-test.yaml

@@ -1,4 +1,4 @@
-name: Test for Warnings
+name: Monthly test for warnings
 
 on:
   workflow_dispatch: null
@@ -8,7 +8,7 @@ on:
 
 jobs:
   tests:
-    name: tests
+    name: pytest with diffmah/dsps/diffsky@main
     runs-on: "ubuntu-latest"
 
     steps:

.github/workflows/test_releases.yaml

@@ -1,4 +1,4 @@
-name: tests
+name: Test against latest diffstuff releases
 
 on:
   workflow_dispatch: null
@@ -9,7 +9,7 @@ on:
 
 jobs:
   tests:
-    name: tests
+    name: pytest with latest releases on conda-forge
     runs-on: "ubuntu-latest"
 
     steps:

.github/workflows/tests_cron.yaml

@@ -1,4 +1,4 @@
-name: test_main_branch_dependencies
+name: Weekly cron testing
 
 on:
   workflow_dispatch: null
@@ -12,7 +12,7 @@ on:
 
 jobs:
   tests:
-    name: tests
+    name: pytest with diffmah/dsps/diffsky@main
     runs-on: "ubuntu-latest"
 
     steps:

CHANGES.rst

@@ -1,3 +1,8 @@
+1.0.2 (unreleased)
+------------------
+- Monte Carlo SFH generators now have required kwargs lgt0 and fb (https://github.com/ArgonneCPAC/diffstar/pull/108)
+
+
 1.0.1 (2025-11-02)
 ------------------
 - Update scaling relations and recalibrate default parameters (https://github.com/ArgonneCPAC/diffstar/pull/106)

diffstar/diffstarpop/loss_kernels/mstar_ssfr_loss_mgash_anyz.py

@@ -1,18 +1,18 @@
 """ """
 
 from diffsky.diffndhist import tw_ndhist_weighted
-from diffstar.utils import cumulative_mstar_formed
 from jax import jit as jjit
 from jax import numpy as jnp
 from jax import value_and_grad, vmap
 
+from diffstar.utils import cumulative_mstar_formed
+
 from ..kernels.defaults_mgash import (
     DEFAULT_DIFFSTARPOP_U_PARAMS,
     get_bounded_diffstarpop_params,
 )
 from ..mc_diffstarpop_mgash import mc_diffstar_sfh_galpop
 
-
 N_TIMES = 20
 
 _A = (None, 0)
@@ -51,6 +51,8 @@ def _mc_diffstar_sfh_galpop_vmap_kern(
     gyr_since_infall,
     ran_key,
     tobs_target,
+    lgt0,
+    fb,
 ):
     tarr = jnp.logspace(-1, jnp.log10(tobs_target), N_TIMES)
     res = mc_diffstar_sfh_galpop(
@@ -63,11 +65,13 @@ def _mc_diffstar_sfh_galpop_vmap_kern(
         gyr_since_infall,
         ran_key,
         tarr,
+        lgt0=lgt0,
+        fb=fb,
     )
     return res
 
 
-_U = (None, *[0] * 8)
+_U = (None, *[0] * 8, None, None)
 mc_diffstar_sfh_galpop_vmap = jjit(vmap(_mc_diffstar_sfh_galpop_vmap_kern, in_axes=_U))
 
 
@@ -117,6 +121,8 @@ def mstar_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
         logmstar_bins,
         target_mstar_pdf,
     ) = loss_data
@@ -133,6 +139,8 @@ def mstar_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
     )
     diffstar_params_ms, diffstar_params_q, sfh_ms, sfh_q, frac_q, mc_is_q = _res
 
@@ -233,6 +241,8 @@ def mstar_ssfr_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
         ndbins_lo,
         ndbins_hi,
         logmstar_bins,
@@ -255,6 +265,8 @@ def mstar_ssfr_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
     )
     diffstar_params_ms, diffstar_params_q, sfh_ms, sfh_q, frac_q, mc_is_q = _res
 
@@ -347,6 +359,8 @@ def mstar_ssfr_sat_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
         ndbins_lo,
         ndbins_hi,
         logmstar_bins,
@@ -369,6 +383,8 @@ def mstar_ssfr_sat_kern_tobs(u_params, loss_data):
         gyr_since_infall,
         ran_key,
         tobs_target,
+        lgt0,
+        fb,
     )
     diffstar_params_ms, diffstar_params_q, sfh_ms, sfh_q, frac_q, mc_is_q = _res
 

diffstar/diffstarpop/loss_kernels/tests/test_mstar_ssfr_loss_mgash_anyz.py

@@ -91,6 +91,8 @@ def test_h5_data_shapes_and_sanity(loss_data_mstar, loss_data_ssfr, loss_data_ss
         gyr_since_infall_data,
         ran_key_data,
         t_obs_targets,
+        lgt0,
+        fb,
         logmstar_bins_pdf,
         mstar_counts_target,
     ) = loss_data_mstar
@@ -122,6 +124,8 @@ def test_h5_data_shapes_and_sanity(loss_data_mstar, loss_data_ssfr, loss_data_ss
         _,
         _,
         _,
+        _,
+        _,
         ndbins_lo,
         ndbins_hi,
         logmstar_bins_pdf2,
@@ -148,6 +152,8 @@ def test_h5_data_shapes_and_sanity(loss_data_mstar, loss_data_ssfr, loss_data_ss
         _,
         _,
         _,
+        _,
+        _,
         ndbins_lo_s,
         ndbins_hi_s,
         logmstar_bins_pdf_s,

diffstar/diffstarpop/mc_diffstarpop_mgash.py

@@ -6,7 +6,7 @@
 from jax import random as jran
 from jax import vmap
 
-from ..defaults import FB, LGT0, get_bounded_diffstar_params
+from ..defaults import get_bounded_diffstar_params
 from ..sfh_model import calc_sfh_galpop, calc_sfh_singlegal
 from .kernels.diffstarpop_mgash import mc_diffstar_u_params_singlegal_kernel
 
@@ -35,8 +35,9 @@ def mc_diffstar_sfh_singlegal(
     gyr_since_infall,
     ran_key,
     tarr,
-    lgt0=LGT0,
-    fb=FB,
+    *,
+    lgt0,
+    fb,
 ):
     """Monte Carlo realization of a single point in Diffstar parameter space,
     along with the computation of SFH for this point.
@@ -376,8 +377,9 @@ def mc_diffstar_sfh_galpop(
     gyr_since_infall,
     ran_key,
     tarr,
-    lgt0=LGT0,
-    fb=FB,
+    *,
+    lgt0,
+    fb,
 ):
     """Monte Carlo realization of a single point in Diffstar parameter space,
     along with the computation of SFH for this point.

diffstar/diffstarpop/tests/test_gradients.py

@@ -105,7 +105,7 @@ def test_all_diffstarpop_u_param_gradients_are_nonzero():
         default_sfh_q,
         frac_q,
         mc_is_q,
-    ) = mc_diffstar_sfh_galpop(*args)
+    ) = mc_diffstar_sfh_galpop(*args, lgt0=1.14, fb=0.156)
 
     assert default_sfh_q.shape == (n_halos, ntimes)
     assert np.all(np.isfinite(default_sfh_q))
@@ -140,7 +140,7 @@ def test_all_diffstarpop_u_param_gradients_are_nonzero():
         alt_sfh_q,
         alt_frac_q,
         mc_is_q,
-    ) = mc_diffstar_sfh_galpop(*args)
+    ) = mc_diffstar_sfh_galpop(*args, lgt0=1.14, fb=0.156)
     assert alt_sfh_q.shape == (n_halos, ntimes)
     assert np.all(np.isfinite(alt_sfh_q))
 
@@ -167,7 +167,7 @@ def _loss(u_params):
             pred_sfh_q,
             pred_frac_q,
             mc_is_q,
-        ) = mc_diffstar_sfh_galpop(*args)
+        ) = mc_diffstar_sfh_galpop(*args, lgt0=1.14, fb=0.156)
         pred_mean_sfh_total = jnp.mean(
             pred_frac_q[:, None] * pred_sfh_q
             + (1.0 - pred_frac_q[:, None]) * pred_sfh_ms,

diffstar/diffstarpop/tests/test_mc_diffstarpop_mgash.py

@@ -55,7 +55,7 @@ def test_mc_diffstar_sfh_singlegal_evaluates():
         ran_key,
         tarr,
     )
-    _res = mcdsp.mc_diffstar_sfh_singlegal(*args)
+    _res = mcdsp.mc_diffstar_sfh_singlegal(*args, lgt0=1.14, fb=0.156)
     params_ms, params_q, sfh_ms, sfh_q, frac_q, mc_is_q = _res
     assert np.all(frac_q >= 0)
     assert np.all(frac_q <= 1)
@@ -147,6 +147,8 @@ def test_mc_diffstar_sfh_galpop():
         gyr_since_infall,
         ran_key,
         t_table,
+        lgt0=1.14,
+        fb=0.156,
     )
     sfh_q, sfh_ms, frac_q = _res[2:5]
 

diffstar/sfh_model.py

@@ -7,7 +7,6 @@
 from jax import numpy as jnp
 from jax import vmap
 
-from .defaults import FB, LGT0
 from .kernels.history_kernel_builders import _sfh_galpop_kern, _sfh_singlegal_kern
 from .utils import cumulative_mstar_formed
 
@@ -18,7 +17,13 @@
 
 @partial(jjit, static_argnames="return_smh")
 def calc_sfh_singlegal(
-    sfh_params, mah_params, tarr, lgt0=LGT0, fb=FB, return_smh=False
+    sfh_params,
+    mah_params,
+    tarr,
+    *,
+    lgt0,
+    fb,
+    return_smh=False,
 ):
     """Calculate the Diffstar SFH for a single galaxy
 
@@ -72,7 +77,15 @@ def calc_sfh_singlegal(
 
 
 @partial(jjit, static_argnames="return_smh")
-def calc_sfh_galpop(sfh_params, mah_params, tarr, lgt0=LGT0, fb=FB, return_smh=False):
+def calc_sfh_galpop(
+    sfh_params,
+    mah_params,
+    tarr,
+    *,
+    lgt0,
+    fb,
+    return_smh=False,
+):
     """Calculate the Diffstar SFH for a single galaxy
 
     Parameters