From 626db83ec3516d514e2017ade15e4645750374b4 Mon Sep 17 00:00:00 2001
From: nvill <natvnr@gmail.com>
Date: Fri, 26 Jun 2026 08:55:11 -0500
Subject: [PATCH 1/5] Compute cen_weight=1 in mc_lc_halos.

---
 diffhalos/lightcone_generators/mc_lightcone_halos.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone_halos.py b/diffhalos/lightcone_generators/mc_lightcone_halos.py
index 23bf775..5643054 100644
--- a/diffhalos/lightcone_generators/mc_lightcone_halos.py
+++ b/diffhalos/lightcone_generators/mc_lightcone_halos.py
@@ -229,6 +229,9 @@ def mc_lc_halos(
 
             logt0: float
                 base-10 log of cosmic time at today, in Gyr
+
+            cen_weight: ndarray of ones of shape (n_halos, )
+                one for every object
     """
 
     # generate mc realization of the halo mass function
@@ -262,10 +265,13 @@ def mc_lc_halos(
     # compute MAH values today
     logmp0 = _log_mah_kern(mah_params, 10**logt0, logt0)
 
+    # compute weights = 1 for all objects
+    n_halos = len(z_obs)
+    cen_weight = jnp.ones(n_halos)
+
     # create output namedtuple
-    fields = CenPop._fields[:-1]
-    values = (z_obs, t_obs, logmp_obs, mah_params, logmp0, logt0)
-    cenpop = namedtuple("cenpop", fields)(*values)
+    values = (z_obs, t_obs, logmp_obs, mah_params, logmp0, logt0, cen_weight)
+    cenpop = CenPop(*values)
 
     return cenpop
 

From 67f4091fcc45766fe44d375e6cf1f6dc67afaa69 Mon Sep 17 00:00:00 2001
From: nvill <natvnr@gmail.com>
Date: Fri, 26 Jun 2026 08:55:56 -0500
Subject: [PATCH 2/5] Use global namedtuple SubPop.

---
 .../mc_lightcone_subhalos.py                  | 40 ++++++++++++++-----
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
index 0a9c02d..fce9047 100644
--- a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
+++ b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
@@ -18,6 +18,15 @@
 
 __all__ = ("mc_lc_shmf", "mc_lc_subhalos", "weighted_lc_subhalos")
 
+_SUBPOP_FIELDS = (
+    "sat_weight",
+    "mah_params",
+    "logmu_obs",
+    "logmp_obs",
+    "nsub_per_host",
+)
+SubPop = namedtuple("SubPop", _SUBPOP_FIELDS)
+
 DEFAULT_DIFFMAHNET_SAT_MODEL = "satflow_v2_0_64bit.eqx"
 N_LGMU_PER_HOST = 5
 
@@ -121,14 +130,21 @@ def mc_lc_subhalos(
     -------
     subpop: namedtuple
         subhalo population with fields:
-            nsub_per_host: ndarray of shape (n_host, )
-                number of generated subhalos per host halo
+            sat_weight: ndarray of ones of shape (n_nub, )
+                weights are one for all objects
 
             mah_params_subs: namedtuple of ndarray's with shape (n_subs, n_mah_params)
                 diffmah parameters for each subhalo in the lightcone
 
+            logmp_obs: ndarray of shape (n_subs, )
+                base-10 log of Mpeak/Msun of each subhalo
+
             logmu_obs: ndarray of shape (n_subs, )
                 base-10 log of mu=Msub/Mhost for each subhalo in the lightcone
+
+            nsub_per_host: int
+                number of representative subhalos per host halo
+                equal to the input n_mu_per_host
     """
 
     # two random keys, one for the MC subhalo population and one for diffmahnet
@@ -160,12 +176,16 @@ def mc_lc_subhalos(
     )
 
     # compute the rescaled mu values
-    mc_lg_mu = logmsub_obs - jnp.repeat(cenpop.logmp_obs, n_mu_per_host)
+    logmu_obs = logmsub_obs - jnp.repeat(cenpop.logmp_obs, n_mu_per_host)
+
+    # create array of weights: one for every object
+    n_subs = len(logmu_obs)
+    sat_weight = jnp.ones(n_subs)
+    # TODO: check shape of sat_weight
+    # sat_weight = jnp.repeat(sat_weight, n_mu_per_host) ?
 
-    # add the subhalo data to the halo population namedtuple
-    fields = ("nsub_per_host", "mah_params", "logmu_obs")
-    data = (n_mu_per_host, mah_params_subs, mc_lg_mu)
-    subpop = namedtuple("subpop", fields)(*data)
+    values = (sat_weight, mah_params_subs, logmu_obs, logmsub_obs, n_mu_per_host)
+    subpop = SubPop(*values)
 
     return subpop
 
@@ -281,9 +301,7 @@ def weighted_lc_subhalos(
     # compute the rescaled mu values at t_obs
     logmu_obs = logmsub_obs - jnp.repeat(cenpop.logmp_obs, n_mu_per_host)
 
-    # add subhalo weights to the dictionary
-    fields = ("sat_weight", "mah_params", "logmu_obs", "logmp_obs", "nsub_per_host")
-    data = (sat_weight, mah_params_subs, logmu_obs, logmsub_obs, n_mu_per_host)
-    subpop = namedtuple("subpop", fields)(*data)
+    values = (sat_weight, mah_params_subs, logmu_obs, logmsub_obs, n_mu_per_host)
+    subpop = SubPop(*values)
 
     return subpop

From 5b47bc961a531fd3f8f62fb84bc696aae99a089a Mon Sep 17 00:00:00 2001
From: nvill <natvnr@gmail.com>
Date: Fri, 26 Jun 2026 08:57:54 -0500
Subject: [PATCH 3/5] Use global namedtuple HaloPop.

---
 .../lightcone_generators/mc_lightcone.py      | 233 +++++++++---------
 .../tests/test_mc_lightcone.py                |   4 +-
 2 files changed, 114 insertions(+), 123 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone.py b/diffhalos/lightcone_generators/mc_lightcone.py
index e4ad73e..f108d84 100644
--- a/diffhalos/lightcone_generators/mc_lightcone.py
+++ b/diffhalos/lightcone_generators/mc_lightcone.py
@@ -22,6 +22,111 @@
 
 __all__ = ("mc_lc_mf", "mc_lc", "weighted_lc")
 
+# TODO: I think we could also define halo_weight == gal_weight = cen_weight * sat_weight here, instead of leaving the user to multiply afterwards.
+_HALOPOP_FIELDS = (
+    "z_obs",  # defined for centrals, repeated to be assigned to all objects
+    "t_obs",  # defined for centrals, repeated to be assigned to all objects
+    "logmp_obs",  # join the values of cens and subs
+    "mah_params",  # join the values of cens and subs
+    "logmp0",  # join the values of cens and subs
+    "logt0",  # defined for centrals, repeated to be assigned to all objects
+    "cen_weight",  # defined for centrals, repeated to be assigned to all objects
+    "central",  # host halo idenfier
+    "sat_weight",  # defined for subs, repeated to be assigned to all objects
+    "nsub_per_host",
+    "logmu_obs",  # defined for subs, repeated to be assigned to all objects ?
+    "halo_indx",
+)
+
+HaloPop = namedtuple("HaloPop", _HALOPOP_FIELDS)
+
+
+def _combine_cenpop_subpop(cenpop, subpop):
+    """
+    Auxiliary function to reshape and combine cens and subs quantities, and build an unified halopop namedtuple.
+
+    Returns halopop
+    This is almost the same as simply joining CenPop and SubPop, but some variables defined for each namedtuple separately are combined here into a single one, e.g., MAH params, logmp_obs. So they do not appear twice, as they would if we were simply joining CenPop and SubPop.
+
+    Paste here full description of halopop
+    """
+
+    # Create the halo index array
+    n_host = cenpop.logmp_obs.size
+    n_sub = subpop.logmp_obs.size
+    host_indx = jnp.arange(n_host).astype(int)
+    subhalo_indx = jnp.repeat(host_indx, subpop.nsub_per_host)
+    halo_indx = jnp.concatenate((host_indx, subhalo_indx)).astype(int)
+
+    # Create the central identifier array
+    is_central = jnp.concatenate((jnp.ones(n_host), jnp.zeros(n_sub))).astype(int)
+
+    # -- Combine properties --
+
+    # Reshape z_obs to assign values for all halos
+    z_obs_all = jnp.concatenate(
+        (cenpop.z_obs, jnp.repeat(cenpop.z_obs, subpop.nsub_per_host))
+    )
+
+    # Reshape t_obs to assign values for all halos
+    t_obs_all = jnp.concatenate(
+        (cenpop.t_obs, jnp.repeat(cenpop.t_obs, subpop.nsub_per_host))
+    )
+
+    # Combine logmp_obs from cens and subs
+    logmp_obs_all = jnp.concatenate((cenpop.logmp_obs, subpop.logmp_obs))
+
+    # Combine logmp0 from cens and subs
+    # Compute mah values at z=0 for subs
+    logmp0_subs = _log_mah_kern(subpop.mah_params, 10**cenpop.logt0, cenpop.logt0)
+    logmp0_all = jnp.concatenate((cenpop.logmp0, logmp0_subs))
+    # TODO: logmp0_all = jnp.concatenate((cenpop.logmp0, subpop.logmp0_subs))
+    # This implies adding the logmp0 field to SubPop and running the mah kernel in the subs function instead of doing it here.
+
+    # Reshape cen_weight to assign values for all halos
+    cen_weight_all = np.concatenate(
+        (cenpop.cen_weight, jnp.repeat(cenpop.cen_weight, subpop.nsub_per_host))
+    )
+
+    # Combine halo and subhalo mah_params
+    mah_params_names = cenpop.mah_params._fields
+    mah_params_tot = np.zeros((len(mah_params_names), n_host + n_sub))
+    for i, _param in enumerate(mah_params_names):
+        mah_params_tot[i, :] = np.concatenate(
+            (
+                cenpop.mah_params._asdict()[_param],
+                subpop.mah_params._asdict()[_param],
+            )
+        )
+    mah_params_all = namedtuple("mah_params", cenpop.mah_params._fields)(
+        *mah_params_tot
+    )
+
+    # Reshape logmu_obs to assign values for all halos
+    logmu_obs_all = jnp.concatenate((jnp.zeros(n_host), subpop.logmu_obs))
+
+    # Reshape sat_weight to assign values for all halos
+    sat_weight_all = jnp.concatenate((jnp.ones(n_host), subpop.sat_weight))
+
+    # -- Create the output namedtuple containing host and subhalo information --
+    values = (
+        z_obs_all,
+        t_obs_all,
+        logmp_obs_all,
+        mah_params_all,
+        logmp0_all,
+        cenpop.logt0,
+        cen_weight_all,
+        is_central,
+        sat_weight_all,
+        subpop.nsub_per_host,
+        logmu_obs_all,
+        halo_indx,
+    )
+    halopop = HaloPop(*values)
+
+    return halopop
+
 
 def mc_lc_mf(
     ran_key,
@@ -198,39 +303,14 @@ def mc_lc(
     Returns
     -------
     halopop: namedtuple
-        halo population with fields:
-            z_obs: ndarray of shape (n_host, )
-                lightcone redshift
-
-            logmp_obs: ndarray of shape (n_host, )
-                halo mass at the lightcone redshift, in Msun
-
-            mah_params: namedtuple of ndarray's with shape (n_host+n_sub, n_mah_params)
-                diffmah parameters for each host halo in the lightcone
-
-            logmp0: narray of shape (n_host, )
-                base-10 log of halo mass at z=0, in Msun
 
-            logt0: float
-                base-10 log of cosmic time at today, in Gyr
-
-            nsub_per_host: ndarray of shape (n_host, )
-                number of subhalos generated per host halo
-
-            logmu_obs: ndarray of shape (n_sub, )
-                base-10 log of mu=Msub/Mhost for each generated subhalo
-
-            halo_indx: ndarray of shape (n_host+n_sub, )
-                halo index; for host halos, the index is arange(n_host),
-                while for the subhalos, indeces correspond to the host
-                halo that hosts each generated subhalo
     """
     # two random keys, one for the host and one for the subhalo population
-    host_key, subhalo_key = jran.split(ran_key)
+    cen_key, sub_key = jran.split(ran_key)
 
-    # generate a host halo lightcone
+    # -- Compute cenpop: compute HMF and MAH params. for centrals. --
     cenpop = mclch.mc_lc_halos(
-        host_key,
+        cen_key,
         lgmp_min,
         z_min,
         z_max,
@@ -243,11 +323,10 @@ def mc_lc(
         n_hmf_grid=n_hmf_grid,
         centrals_model_key=centrals_model_key,
     )
-    # fields = ("z_obs", "t_obs", "logmp_obs", "mah_params", "logmp0", "logt0")
 
-    # generate a subhalo lightcone
+    # -- Compute subpop: compute SHMF and MAH params. for subs. --
     subpop = mclcsh.mc_lc_subhalos(
-        subhalo_key,
+        sub_key,
         cenpop,
         lgmsub_min,
         ccshmf_params=ccshmf_params,
@@ -256,34 +335,7 @@ def mc_lc(
         subhalo_model_key=subhalo_model_key,
     )
 
-    # create the index array: [...host_indx..., ...subhalo_indx...]
-    n_host = cenpop.logmp_obs.size
-    host_indx = jnp.arange(n_host).astype(int)
-    n_sub = int(subpop.nsub_per_host.sum())
-    subhalo_indx = jnp.repeat(host_indx, subpop.nsub_per_host)
-    halo_indx = jnp.concatenate((host_indx, subhalo_indx)).astype(int)
-
-    # combine halo and subhalo mah_params
-    mah_params_names = cenpop.mah_params._fields
-    mah_params_tot = np.zeros((len(mah_params_names), n_host + n_sub))
-    for i, _param in enumerate(mah_params_names):
-        mah_params_tot[i, :] = np.concatenate(
-            (
-                cenpop.mah_params._asdict()[_param],
-                subpop.mah_params._asdict()[_param],
-            )
-        )
-    mah_params_ntup = namedtuple("mah_params", cenpop.mah_params._fields)(
-        *mah_params_tot
-    )
-    cenpop = cenpop._replace(mah_params=mah_params_ntup)
-
-    # create the output namedtuple containing host and subhalo information;
-    # this will contain all host halo information, updated to include
-    # the subhalo information and some fields are updated to new shapes
-    halopop = namedtuple(
-        "mc_lc", [*cenpop._fields, "nsub_per_host", "logmu_obs", "halo_indx"]
-    )(*cenpop, subpop.nsub_per_host, subpop.logmu_obs, halo_indx)
+    halopop = _combine_cenpop_subpop(cenpop, subpop)
 
     return halopop
 
@@ -501,67 +553,6 @@ def _weighted_lc_from_grid(
         subhalo_model_key=subhalo_model_key,
     )
 
-    # create the index array
-    n_host = cenpop.logmp_obs.size
-    n_sub = subpop.logmp_obs.size
-    host_indx = jnp.arange(n_host).astype(int)
-    subhalo_indx = jnp.repeat(host_indx, subpop.nsub_per_host)
-    halo_indx = jnp.concatenate((host_indx, subhalo_indx)).astype(int)
-    central = jnp.concatenate((jnp.ones(n_host), jnp.zeros(n_sub))).astype(int)
-
-    z_obs_all = jnp.concatenate(
-        (cenpop.z_obs, jnp.repeat(cenpop.z_obs, subpop.nsub_per_host))
-    )
-    cenpop = cenpop._replace(z_obs=z_obs_all)
-
-    t_obs_all = jnp.concatenate(
-        (cenpop.t_obs, jnp.repeat(cenpop.t_obs, subpop.nsub_per_host))
-    )
-    cenpop = cenpop._replace(t_obs=t_obs_all)
-
-    logmp_obs_all = jnp.concatenate((cenpop.logmp_obs, subpop.logmp_obs))
-    cenpop = cenpop._replace(logmp_obs=logmp_obs_all)
-
-    # compute mah values at z=0 for subs
-    logmp0_subs = _log_mah_kern(subpop.mah_params, 10**cenpop.logt0, cenpop.logt0)
-    logmp0_all = jnp.concatenate((cenpop.logmp0, logmp0_subs))
-    cenpop = cenpop._replace(logmp0=logmp0_all)
-
-    cenpop = cenpop._replace(
-        cen_weight=np.concatenate(
-            (cenpop.cen_weight, jnp.repeat(cenpop.cen_weight, subpop.nsub_per_host))
-        )
-    )
-
-    # combine halo and subhalo mah_params
-    mah_params_names = cenpop.mah_params._fields
-    mah_params_tot = np.zeros((len(mah_params_names), n_host + n_sub))
-    for i, _param in enumerate(mah_params_names):
-        mah_params_tot[i, :] = np.concatenate(
-            (
-                cenpop.mah_params._asdict()[_param],
-                subpop.mah_params._asdict()[_param],
-            )
-        )
-    mah_params_ntup = namedtuple("mah_params", cenpop.mah_params._fields)(
-        *mah_params_tot
-    )
-    cenpop = cenpop._replace(mah_params=mah_params_ntup)
-
-    logmu_obs_all = jnp.concatenate((jnp.zeros(n_host), subpop.logmu_obs))
-
-    sat_weight_all = jnp.concatenate((jnp.ones(n_host), subpop.sat_weight))
-
-    # create the output namedtuple containing host and subhalo information;
-    # this will contain all host halo information, updated to include
-    # the subhalo information and some fields are updated to new shapes
-    halopop = WeightedLightcone(
-        *cenpop, central, sat_weight_all, subpop.nsub_per_host, logmu_obs_all, halo_indx
-    )
+    halopop = _combine_cenpop_subpop(cenpop, subpop)
 
     return halopop
-
-
-SAT_FIELDS = ["central", "sat_weight", "nsub_per_host", "logmu_obs", "halo_indx"]
-_FIELDS = list(mclch.CenPop._fields) + SAT_FIELDS
-WeightedLightcone = namedtuple("WeightedLightcone", _FIELDS)
diff --git a/diffhalos/lightcone_generators/tests/test_mc_lightcone.py b/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
index 38ee97b..fbff8d9 100644
--- a/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
+++ b/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
@@ -74,8 +74,8 @@ def test_mc_lc_behaves_as_expected():
     for _field in halopop._fields:
         assert np.all(np.isfinite(halopop._asdict()[_field]))
 
-    n_host = halopop.logmp_obs.size
-    n_subs = halopop.logmu_obs.size
+    n_host = np.sum(halopop.central == 1)
+    n_subs = np.sum(halopop.central == 0)
     for _param in halopop.mah_params._fields:
         assert halopop.mah_params._asdict()[_param].size == n_host + n_subs
 

From f672f250e49163bda51f47a7268d15b9b9fe7ddf Mon Sep 17 00:00:00 2001
From: nvill <natvnr@gmail.com>
Date: Wed, 1 Jul 2026 09:49:16 -0500
Subject: [PATCH 4/5] Add halo_weight field to HaloPop.

---
 .../lightcone_generators/mc_lightcone.py      | 14 +++++++++-----
 .../tests/test_mc_lightcone.py                | 19 +++++++++++++++----
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone.py b/diffhalos/lightcone_generators/mc_lightcone.py
index f108d84..0a2b8d5 100644
--- a/diffhalos/lightcone_generators/mc_lightcone.py
+++ b/diffhalos/lightcone_generators/mc_lightcone.py
@@ -22,7 +22,6 @@
 
 __all__ = ("mc_lc_mf", "mc_lc", "weighted_lc")
 
-# TODO: I think we could also define halo_weight == gal_weight = cen_weight * sat_weight here, instead of leaving the user to multiply afterwards.
 _HALOPOP_FIELDS = (
     "z_obs",  # defined for centrals, repeated to be assigned to all objects
     "t_obs",  # defined for centrals, repeated to be assigned to all objects
@@ -34,8 +33,9 @@
     "central",  # host halo idenfier
     "sat_weight",  # defined for subs, repeated to be assigned to all objects
     "nsub_per_host",
-    "logmu_obs",  # defined for subs, repeated to be assigned to all objects ?
+    "logmu_obs",  # defined for subs, repeated to be assigned to all objects
     "halo_indx",
+    "halo_weight",  # same as gal_weight = cen_weight * sat_weight
 )
 
 HaloPop = namedtuple("HaloPop", _HALOPOP_FIELDS)
@@ -94,20 +94,23 @@ def _combine_cenpop_subpop(cenpop, subpop):
     for i, _param in enumerate(mah_params_names):
         mah_params_tot[i, :] = np.concatenate(
             (
-                cenpop.mah_params._asdict()[_param],
-                subpop.mah_params._asdict()[_param],
+                getattr(cenpop.mah_params, _param),
+                getattr(subpop.mah_params, _param),
             )
         )
     mah_params_all = namedtuple("mah_params", cenpop.mah_params._fields)(
         *mah_params_tot
     )
 
-    # Reshape logmu_obs to assign values for all halos
+    # Reshape logmu_obs to assign values to all halos
     logmu_obs_all = jnp.concatenate((jnp.zeros(n_host), subpop.logmu_obs))
 
     # Reshape sat_weight to assign values for all halos
     sat_weight_all = jnp.concatenate((jnp.ones(n_host), subpop.sat_weight))
 
+    # Define halo weight
+    halo_weight_all = cen_weight_all * sat_weight_all
+
     # -- Create the output namedtuple containing host and subhalo information --
     values = (
         z_obs_all,
@@ -122,6 +125,7 @@ def _combine_cenpop_subpop(cenpop, subpop):
         subpop.nsub_per_host,
         logmu_obs_all,
         halo_indx,
+        halo_weight_all,
     )
     halopop = HaloPop(*values)
 
diff --git a/diffhalos/lightcone_generators/tests/test_mc_lightcone.py b/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
index fbff8d9..69374e4 100644
--- a/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
+++ b/diffhalos/lightcone_generators/tests/test_mc_lightcone.py
@@ -81,6 +81,12 @@ def test_mc_lc_behaves_as_expected():
 
     assert halopop.halo_indx.size == n_host + n_subs
 
+    # Check halo_weight
+    halo_weight = halopop.halo_weight
+    assert np.allclose(halo_weight, halopop.cen_weight * halopop.sat_weight)
+    # Should be one for all objects in the unweighted (mc) lc
+    assert np.allclose(halo_weight, 1.0)
+
 
 def test_weighted_lc_behaves_as_expected():
     """Check each returned column is finite and has the expected shape"""
@@ -115,6 +121,10 @@ def test_weighted_lc_behaves_as_expected():
     assert halopop.nsub_per_host.shape == ()
     assert np.isfinite(halopop.nsub_per_host)
 
+    # Check halo_weight
+    halo_weight = halopop.halo_weight
+    assert np.allclose(halo_weight, halopop.cen_weight * halopop.sat_weight)
+
 
 def test_weighted_lc_logmp0_is_consistent_with_logmp_obs():
     """Enforce self-consistent behavior for logmp0 and logmp_obs columns
@@ -317,7 +327,7 @@ def test_weighted_lc_logmu_obs():
     assert np.allclose(halopop.logmu_obs[:n_host_halos], 0.0)
 
 
-def test_weighted_lc_gal_weight():
+def test_weighted_lc_halo_weight():
     ran_key = jran.key(0)
 
     n_host_halos = 100
@@ -327,6 +337,7 @@ def test_weighted_lc_gal_weight():
     args = (ran_key, n_host_halos, z_min, z_max, lgmp_min, lgmp_max, sky_area_degsq)
     halopop = mclc.weighted_lc(*args)
 
-    gal_weight = halopop.cen_weight * halopop.sat_weight
-    assert np.allclose(gal_weight[:n_host_halos], halopop.cen_weight[:n_host_halos])
-    assert not np.any(gal_weight[n_host_halos:] == halopop.cen_weight[n_host_halos:])
+    halo_weight = halopop.halo_weight
+
+    assert np.allclose(halo_weight[:n_host_halos], halopop.cen_weight[:n_host_halos])
+    assert not np.any(halo_weight[n_host_halos:] == halopop.cen_weight[n_host_halos:])

From c933daf2bf90e236cf6860607eb761bdc3445829 Mon Sep 17 00:00:00 2001
From: nvill <natvnr@gmail.com>
Date: Wed, 1 Jul 2026 10:16:55 -0500
Subject: [PATCH 5/5] Add logmp0 to SubPop.

---
 .../lightcone_generators/mc_lightcone.py      |  7 +----
 .../mc_lightcone_subhalos.py                  | 28 ++++++++++++++++---
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/diffhalos/lightcone_generators/mc_lightcone.py b/diffhalos/lightcone_generators/mc_lightcone.py
index 0a2b8d5..d7715cc 100644
--- a/diffhalos/lightcone_generators/mc_lightcone.py
+++ b/diffhalos/lightcone_generators/mc_lightcone.py
@@ -10,7 +10,6 @@
 from collections import namedtuple
 
 import numpy as np
-from diffmah.diffmah_kernels import _log_mah_kern
 from jax import numpy as jnp
 from jax import random as jran
 
@@ -77,11 +76,7 @@ def _combine_cenpop_subpop(cenpop, subpop):
     logmp_obs_all = jnp.concatenate((cenpop.logmp_obs, subpop.logmp_obs))
 
     # Combine logmp0 from cens and subs
-    # Compute mah values at z=0 for subs
-    logmp0_subs = _log_mah_kern(subpop.mah_params, 10**cenpop.logt0, cenpop.logt0)
-    logmp0_all = jnp.concatenate((cenpop.logmp0, logmp0_subs))
-    # TODO: logmp0_all = jnp.concatenate((cenpop.logmp0, subpop.logmp0_subs))
-    # This implies adding the logmp0 field to SubPop and running the mah kernel in the subs function instead of doing it here.
+    logmp0_all = jnp.concatenate((cenpop.logmp0, subpop.logmp0))
 
     # Reshape cen_weight to assign values for all halos
     cen_weight_all = np.concatenate(
diff --git a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
index fce9047..4000be6 100644
--- a/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
+++ b/diffhalos/lightcone_generators/mc_lightcone_subhalos.py
@@ -15,6 +15,7 @@
 from ..ccshmf.ccshmf_model import DEFAULT_CCSHMF_PARAMS, subhalo_lightcone_weights
 from ..ccshmf.mc_subs import generate_subhalopop
 from ..mah.utils import apply_mah_rescaling
+from diffmah.diffmah_kernels import _log_mah_kern
 
 __all__ = ("mc_lc_shmf", "mc_lc_subhalos", "weighted_lc_subhalos")
 
@@ -24,6 +25,7 @@
     "logmu_obs",
     "logmp_obs",
     "nsub_per_host",
+    "logmp0",
 )
 SubPop = namedtuple("SubPop", _SUBPOP_FIELDS)
 
@@ -178,13 +180,21 @@ def mc_lc_subhalos(
     # compute the rescaled mu values
     logmu_obs = logmsub_obs - jnp.repeat(cenpop.logmp_obs, n_mu_per_host)
 
+    # Compute mah values at z=0 for subs
+    logmp0_subs = _log_mah_kern(mah_params_subs, 10**cenpop.logt0, cenpop.logt0)
+
     # create array of weights: one for every object
     n_subs = len(logmu_obs)
     sat_weight = jnp.ones(n_subs)
-    # TODO: check shape of sat_weight
-    # sat_weight = jnp.repeat(sat_weight, n_mu_per_host) ?
 
-    values = (sat_weight, mah_params_subs, logmu_obs, logmsub_obs, n_mu_per_host)
+    values = (
+        sat_weight,
+        mah_params_subs,
+        logmu_obs,
+        logmsub_obs,
+        n_mu_per_host,
+        logmp0_subs,
+    )
     subpop = SubPop(*values)
 
     return subpop
@@ -301,7 +311,17 @@ def weighted_lc_subhalos(
     # compute the rescaled mu values at t_obs
     logmu_obs = logmsub_obs - jnp.repeat(cenpop.logmp_obs, n_mu_per_host)
 
-    values = (sat_weight, mah_params_subs, logmu_obs, logmsub_obs, n_mu_per_host)
+    # Compute mah values at z=0 for subs
+    logmp0_subs = _log_mah_kern(mah_params_subs, 10**cenpop.logt0, cenpop.logt0)
+
+    values = (
+        sat_weight,
+        mah_params_subs,
+        logmu_obs,
+        logmsub_obs,
+        n_mu_per_host,
+        logmp0_subs,
+    )
     subpop = SubPop(*values)
 
     return subpop