From 4728dbd84f1b2a6211e543107eb6bd21f0c2418c Mon Sep 17 00:00:00 2001
From: Andrew Hearin <ahearin@anl.gov>
Date: Wed, 5 Mar 2025 09:37:05 -0600
Subject: [PATCH] Fix bug in calculation of time-since-infall argument

---
 diffsky/mc_diffsky.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/diffsky/mc_diffsky.py b/diffsky/mc_diffsky.py
index a9a6d1b..89bd445 100644
--- a/diffsky/mc_diffsky.py
+++ b/diffsky/mc_diffsky.py
@@ -77,13 +77,14 @@ def mc_diffstar_galpop(
     )
 
     logmu_infall = subcat.logmp_ult_inf - subcat.logmhost_ult_inf
+    t_obs = flat_wcdm._age_at_z_kern(z_obs, *cosmo_params)
     args = (
         diffstarpop_params,
         subcat.mah_params,
         subcat.logmp0,
         logmu_infall,
         subcat.logmhost_ult_inf,
-        subcat.t_ult_inf,
+        t_obs - subcat.t_ult_inf,
         sfh_key,
         t_table,
     )
@@ -93,8 +94,6 @@ def mc_diffstar_galpop(
     sfh_table = jnp.where(mc_is_q.reshape((-1, 1)), sfh_q, sfh_ms)
     smh_table = cumulative_mstar_formed_galpop(t_table, sfh_table)
 
-    t_obs = flat_wcdm._age_at_z_kern(z_obs, *cosmo_params)
-
     diffstar_data = dict()
     diffstar_data["subcat"] = subcat
     diffstar_data["t_table"] = t_table
@@ -172,13 +171,14 @@ def mc_diffstar_cenpop(
     )
 
     logmu_infall = subcat.logmp_ult_inf - subcat.logmhost_ult_inf
+    t_obs = flat_wcdm._age_at_z_kern(z_obs, *cosmo_params)
     args = (
         diffstarpop_params,
         subcat.mah_params,
         subcat.logmp0,
         logmu_infall,
         subcat.logmhost_ult_inf,
-        subcat.t_ult_inf,
+        t_obs - subcat.t_ult_inf,
         sfh_key,
         t_table,
     )
@@ -188,8 +188,6 @@ def mc_diffstar_cenpop(
     sfh_table = jnp.where(mc_is_q.reshape((-1, 1)), sfh_q, sfh_ms)
     smh_table = cumulative_mstar_formed_galpop(t_table, sfh_table)
 
-    t_obs = flat_wcdm._age_at_z_kern(z_obs, *cosmo_params)
-
     diffstar_data = dict()
     diffstar_data["subcat"] = subcat
     diffstar_data["t_table"] = t_table