diffstar 0.3.0__tar.gz → 0.3.2__tar.gz

{diffstar-0.3.0 → diffstar-0.3.2}/.github/workflows/linting.yml

@@ -16,7 +16,7 @@ jobs:
 
       - uses: conda-incubator/setup-miniconda@v2
         with:
-          python-version: 3.9
+          python-version: 3.11
           channels: conda-forge,defaults
           channel-priority: strict
           show-channel-urls: true

diffstar-0.3.2/.github/workflows/monthly-warning-test.yml

@@ -0,0 +1,55 @@
+name: Test for Warnings
+
+on:
+  workflow_dispatch: null
+  schedule:
+    # Runs "First of every month at 3:15am Central"
+    - cron: '15 8 1 * *'
+
+jobs:
+  tests:
+    name: tests
+    runs-on: "ubuntu-latest"
+
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - uses: conda-incubator/setup-miniconda@v2
+        with:
+          python-version: 3.11
+          channels: conda-forge,defaults
+          channel-priority: strict
+          show-channel-urls: true
+          miniforge-version: latest
+          miniforge-variant: Mambaforge
+          use-mamba: true
+
+      - name: configure conda and install code
+      # Test against current main branch of diffmah and dsps
+        shell: bash -l {0}
+        run: |
+          conda config --set always_yes yes
+          mamba install --quiet \
+            --file=requirements.txt
+          python -m pip install --no-deps -e .
+          mamba install -y -q \
+            flake8 \
+            pytest \
+            pytest-xdist \
+            pytest-cov \
+            pip \
+            setuptools \
+            "setuptools_scm>=7,<8" \
+            python-build
+          pip uninstall diffmah --yes
+          pip install --no-deps git+https://github.com/ArgonneCPAC/diffmah.git
+          pip install --no-deps git+https://github.com/ArgonneCPAC/dsps.git
+          python -m pip install --no-build-isolation --no-deps -e .
+
+      - name: test that no warnings are raised
+        shell: bash -l {0}
+        run: |
+          export PYTHONWARNINGS=error
+          pytest -v diffstar --cov --cov-report=xml

{diffstar-0.3.0 → diffstar-0.3.2}/.github/workflows/test_releases.yml

@@ -22,7 +22,7 @@ jobs:
 
       - uses: conda-incubator/setup-miniconda@v2
         with:
-          python-version: 3.9
+          python-version: 3.11
           channels: conda-forge,defaults
           channel-priority: strict
           show-channel-urls: true

{diffstar-0.3.0 → diffstar-0.3.2}/.github/workflows/tests_cron.yml

@@ -22,7 +22,7 @@ jobs:
 
       - uses: conda-incubator/setup-miniconda@v2
         with:
-          python-version: 3.9
+          python-version: 3.11
           channels: conda-forge,defaults
           channel-priority: strict
           show-channel-urls: true

{diffstar-0.3.0 → diffstar-0.3.2}/CHANGES.rst

@@ -1,3 +1,13 @@
+0.3.2 (2024-10-25)
+------------------
+- Adapt sfh kernels to be compatible with diffmah 0.6.1
+
+
+0.3.1 (2024-6-19)
+------------------
+- Performance improvements for calc_sfh_galpop and calc_sfh_singlegal
+
+
 0.3.0 (2024-01-17)
 ------------------
 - Implement new API for primary user-facing functions calc_sfh_galpop and calc_sfh_singlegal

{diffstar-0.3.0/diffstar.egg-info → diffstar-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diffstar
-Version: 0.3.0
+Version: 0.3.2
 Summary: Differentiable Star Formation Histories
 Author-email: Alex Alarcon <alexalarcongonzalez@gmail.com>
 License: BSD 3-Clause License
@@ -35,10 +35,10 @@ License: BSD 3-Clause License
         
 Project-URL: home, https://github.com/ArgonneCPAC/diffstar
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE.rst
-Requires-Dist: diffmah>=0.5.0
+Requires-Dist: diffmah>=0.6.1
 Requires-Dist: numpy
 Requires-Dist: jax
 Requires-Dist: h5py

diffstar-0.3.2/diffstar/_version.py

@@ -0,0 +1 @@
+__version__ = '0.3.2'

{diffstar-0.3.0 → diffstar-0.3.2}/diffstar/data_loaders/load_smah_data.py

@@ -1,5 +1,6 @@
 """
 """
+
 import os
 import warnings
 
@@ -9,21 +10,31 @@ import numpy as np
 from ..defaults import SFR_MIN
 from ..utils import _get_dt_array
 
+try:
+    from umachine_pyio.load_mock import load_mock_from_binaries
+
+    HAS_UM_LOADER = True
+except ImportError:
+    HAS_UM_LOADER = False
+
 TASSO = "/Users/aphearin/work/DATA/diffmah_data"
 BEBOP = "/lcrc/project/halotools/diffmah_data"
 LAPTOP = "/Users/alarcon/Documents/diffmah_data"
-
+BEBOP_SMDPL = os.path.join(
+    "/lcrc/project/galsampler/SMDPL/",
+    "dr1_no_merging_upidh/sfh_binary_catalogs/a_1.000000/",
+)
 H_BPL = 0.678
 H_TNG = 0.6774
 H_MDPL = H_BPL
 
 
-def load_fit_mah(filename, data_drn=BEBOP):
+def load_fit_mah(basename, data_drn=BEBOP):
     """Load the best fit diffmah parameter data
 
     Parameters
     ----------
-    filename : string
+    basename : string
         Name of the h5 file where the diffmah best fit parameters are stored
 
     data_drn : string
@@ -38,30 +49,107 @@ def load_fit_mah(filename, data_drn=BEBOP):
     logmp:  ndarray of shape (n_gal, )
         Base-10 logarithm of the present day peak halo mass
 
+    logmp:  ndarray of shape (n_gal, )
+        Base-10 logarithm of the present day peak halo mass
     """
-    fitting_data = dict()
 
-    fn = os.path.join(data_drn, filename)
+    fn = os.path.join(data_drn, basename)
     with h5py.File(fn, "r") as hdf:
-        for key in hdf.keys():
-            if key == "halo_id":
-                fitting_data[key] = hdf[key][...]
-            else:
-                fitting_data["fit_" + key] = hdf[key][...]
-
-    mah_fit_params = np.array(
-        [
-            fitting_data["fit_mah_logtc"],
-            fitting_data["fit_mah_k"],
-            fitting_data["fit_early_index"],
-            fitting_data["fit_late_index"],
-        ]
-    ).T
-    logmp = fitting_data["fit_logmp_fit"]
+        mah_fit_params = np.array(
+            [
+                hdf["logm0"][:],
+                hdf["logtc"][:],
+                hdf["early_index"][:],
+                hdf["late_index"][:],
+            ]
+        ).T
+        logmp = hdf["logm0"][:]
 
     return mah_fit_params, logmp
 
 
+def load_fit_mah_tpeak(basename, data_drn=BEBOP):
+    """Load the best fit diffmah parameter data
+
+    Parameters
+    ----------
+    basename : string
+        Name of the h5 file where the diffmah best fit parameters are stored
+
+    data_drn : string
+        Filepath where the Diffstar best-fit parameters are stored
+
+    Returns
+    -------
+    mah_fit_params:  ndarray of shape (n_gal, 4)
+        Best fit parameters for each halo:
+            (logtc, k, early_index, late_index)
+
+    logmp:  ndarray of shape (n_gal, )
+        Base-10 logarithm of the present day peak halo mass
+
+    logmp:  ndarray of shape (n_gal, )
+        Base-10 logarithm of the present day peak halo mass
+    """
+
+    fn = os.path.join(data_drn, basename)
+    with h5py.File(fn, "r") as hdf:
+        mah_fit_params = np.array(
+            [
+                hdf["logm0"][:],
+                hdf["logtc"][:],
+                hdf["early_index"][:],
+                hdf["late_index"][:],
+            ]
+        ).T
+        t_peak = hdf["t_peak"][:]
+        logmp = hdf["logm0"][:]
+
+    return mah_fit_params, logmp, t_peak
+
+
+def load_fit_sfh(basename, data_drn=BEBOP):
+    """Load the best fit diffmah parameter data
+
+    Parameters
+    ----------
+    basename : string
+        Name of the h5 file where the diffmah best fit parameters are stored
+
+    data_drn : string
+        Filepath where the Diffstar best-fit parameters are stored
+
+    Returns
+    -------
+    sfh_fit_params:  ndarray of shape (n_gal, 4)
+        Best fit parameters for each halo:
+            (logtc, k, early_index, late_index)
+
+    """
+
+    fn = os.path.join(data_drn, basename)
+    with h5py.File(fn, "r") as hdf:
+        ms_fit_params = np.array(
+            [
+                hdf["lgmcrit"][:],
+                hdf["lgy_at_mcrit"][:],
+                hdf["indx_lo"][:],
+                hdf["indx_hi"][:],
+                hdf["tau_dep"][:],
+            ]
+        ).T
+        q_fit_params = np.array(
+            [
+                hdf["lg_qt"][:],
+                hdf["qlglgdt"][:],
+                hdf["lg_drop"][:],
+                hdf["lg_rejuv"][:],
+            ]
+        ).T
+
+    return ms_fit_params, q_fit_params
+
+
 def load_bolshoi_data(gal_type, data_drn=BEBOP):
     """Load the stellar mass histories from UniverseMachine simulation
     applied to the Bolshoi-Planck (BPL) simulation.
@@ -397,3 +485,67 @@ def load_mdpl_small_data(gal_type, data_drn=BEBOP):
         log_smahs = np.where(sm_cumsum == 0, 0, np.log10(sm_cumsum))
 
     return halo_ids, log_smahs, sfrh, mdpl_t, dt
+
+
+def load_SMDPL_data(subvols, data_drn=BEBOP_SMDPL):
+    """Load the stellar mass histories from UniverseMachine simulation
+    applied to the Bolshoi-Planck (BPL) simulation.
+
+    The loaded stellar mass data has units of Msun assuming the h = H_BPL
+    from the cosmology of the underlying simulation.
+
+    The output stellar mass data has units of Msun/h, or units of
+    Mstar[h=H_BPL] using the h value of the simulation.
+
+    H_BPL is defined at the top of the module.
+
+    Parameters
+    ----------
+    gal_type : string
+        Name of the galaxy type of the file being loaded. Options are
+            'cens': central galaxies
+            'sats': satellite galaxies
+            'orphans': orphan galaxies
+    data_drn : string
+        Filepath where the Diffstar best-fit parameters are stored.
+
+    Returns
+    -------
+    halo_ids:  ndarray of shape (n_gal, )
+        IDs of the halos in the file.
+    log_smahs: ndarray of shape (n_gal, n_times)
+        Cumulative stellar mass history in units of Msun assuming h=1.
+    sfrh: ndarray of shape (n_gal, n_times)
+        Star formation rate history in units of Msun/yr assuming h=1.
+    bpl_t : ndarray of shape (n_times, )
+        Cosmic time of each simulated snapshot in Gyr
+    dt : ndarray of shape (n_times, )
+        Cosmic time steps between each simulated snapshot in Gyr
+    """
+    if not HAS_UM_LOADER:
+        raise ImportError("Must have umachine_pyio installed to load this dataset")
+
+    galprops = ["halo_id", "sfr_history_main_prog", "mpeak_history_main_prog"]
+    mock = load_mock_from_binaries(subvols, root_dirname=data_drn, galprops=galprops)
+
+    SMDPL_t = np.loadtxt(os.path.join(data_drn, "smdpl_cosmic_time.txt"))
+
+    halo_ids = mock["halo_id"]
+    dt = _get_dt_array(SMDPL_t)
+    sfrh = mock["sfr_history_main_prog"]
+    sm_cumsum = np.cumsum(sfrh * dt, axis=1) * 1e9
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        log_smahs = np.where(sm_cumsum == 0, 0, np.log10(sm_cumsum))
+
+    lgmh_min = 7.0
+    mh_min = 10**lgmh_min
+    msk = mock["mpeak_history_main_prog"] < mh_min
+    clipped_mahs = np.where(msk, 1.0, mock["mpeak_history_main_prog"])
+    log_mahs = np.log10(clipped_mahs)
+    log_mahs = np.maximum.accumulate(log_mahs, axis=1)
+
+    logmp = log_mahs[:, -1]
+
+    return halo_ids, log_smahs, sfrh, SMDPL_t, dt, log_mahs, logmp

diffstar-0.3.2/diffstar/data_loaders/tests/test_load_smah_data.py

@@ -0,0 +1,28 @@
+"""
+"""
+
+import os
+
+import numpy as np
+
+from ..load_smah_data import load_fit_mah_tpeak
+
+_THIS_DRNAME = os.path.dirname(os.path.abspath(__file__))
+
+
+def test_load_fit_mah_tpeak():
+    basename = "subvol_000_diffmah_fits.h5"
+    data_drn = os.path.join(_THIS_DRNAME, "testing_data")
+    _res = load_fit_mah_tpeak(basename, data_drn=data_drn)
+    for x in _res:
+        assert np.all(np.isfinite(x))
+    mah_fit_params, logmp, t_peak = _res
+    n_halos, n_params = mah_fit_params.shape
+
+    assert logmp.shape == (n_halos,)
+    assert np.all(logmp > 10)
+    assert np.all(logmp < 16)
+
+    assert t_peak.shape == (n_halos,)
+    assert np.all(t_peak > 0)
+    assert np.all(t_peak < 14)

diffstar-0.3.2/diffstar/diffstarnet/diffstarnet_tdata.py

@@ -0,0 +1,226 @@
+"""
+"""
+
+from collections import namedtuple
+
+import numpy as np
+from diffmah.diffmah_kernels import DEFAULT_MAH_PARAMS
+from diffmah.diffmahpop_kernels.bimod_censat_params import DEFAULT_DIFFMAHPOP_PARAMS
+from diffmah.diffmahpop_kernels.mc_bimod_cens import _mc_diffmah_singlecen_vmap_kern
+from jax import numpy as jnp
+from jax import random as jran
+
+from ..defaults import (
+    DEFAULT_DIFFSTAR_U_PARAMS,
+    DEFAULT_Q_PARAMS_UNQUENCHED,
+    DEFAULT_U_MS_PARAMS,
+    DEFAULT_U_Q_PARAMS,
+    FB,
+    LGT0,
+    T_TABLE_MIN,
+    get_bounded_diffstar_params,
+)
+from ..kernels.main_sequence_kernels_tpeak import MS_PARAM_BOUNDS_PDICT
+from ..sfh_model_tpeak import calc_sfh_galpop
+
+T0 = 10**LGT0
+N_SFH_TABLE = 200
+LGMH_MIN = 10.5
+LGSM0_MIN = 5.0
+
+TAU_INST = MS_PARAM_BOUNDS_PDICT["tau_dep"][0] + 1e-4
+
+_TDATA_SFH_ROOTKEYS = ["sfh_params", "sfh", "smh"]
+_TDATA_NOQ_KEYS = [key + "_noq" for key in _TDATA_SFH_ROOTKEYS]
+_TDATA_NOQ_NOLAG_KEYS = [key + "_nolag" for key in _TDATA_NOQ_KEYS]
+SFH_KEYS = _TDATA_SFH_ROOTKEYS + _TDATA_NOQ_KEYS + _TDATA_NOQ_NOLAG_KEYS
+TDATA_KEYS = ["mah_params", "log_mah"] + SFH_KEYS
+TData = namedtuple("TData", TDATA_KEYS)
+
+
+def tdata_generator(
+    ran_key,
+    logm0_sample,
+    n_sfh_table=N_SFH_TABLE,
+    logsm0_min=LGSM0_MIN,
+    n_epochs=float("inf"),
+):
+    """Training data generator for diffstarnet
+
+    Parameters
+    ----------
+    ran_key : jax.random.key
+
+    logm0_sample : array, shape (n_halos, )
+        Array of values of diffmah parameter logm0
+
+    logsm0_min : float, optional
+        Minimum z=0 stellar mass in the training data
+        Default is set by LGSM0_MIN at top of module
+
+    n_epochs : int, optional
+        maximum number of batches to yield. Default is infinite
+
+    Yields
+    ------
+    tdata : namedtuple
+        mah_params, t_peak, log_mah, sfh_params, sfh, smh
+        sfh_params_noq, sfh_noq, smh_noq
+        sfh_params_noq_nolag, sfh_noq_nolag, smh_noq_nolag
+
+    """
+    batchnum = 0
+    while batchnum < n_epochs:
+        ran_key, batch_key = jran.split(ran_key, 2)
+        tdata = _compute_tdata(batch_key, logm0_sample, n_sfh_table, logsm0_min)
+        yield tdata
+        batchnum += 1
+
+
+def _compute_tdata(
+    ran_key, logm0_sample, n_sfh_table=N_SFH_TABLE, logsm0_min=LGSM0_MIN
+):
+    """"""
+    tarr = np.linspace(T_TABLE_MIN, T0, n_sfh_table)
+
+    mah_key, early_late_key, sfh_key = jran.split(ran_key, 3)
+
+    _reslist = mc_diffmah_halo_sample(mah_key, tarr, logm0_sample)
+    mah_params_early, dmhdt_early, log_mah_early = _reslist[:3]
+    mah_params_late, dmhdt_late, log_mah_late = _reslist[3:6]
+    frac_early = _reslist[6]
+
+    n_halos = mah_params_early.logm0.size
+    uran_mah = jran.uniform(early_late_key, minval=0, maxval=1, shape=(n_halos,))
+    msk_mah = frac_early < uran_mah
+    mah_params = DEFAULT_MAH_PARAMS._make(
+        [
+            jnp.where(
+                msk_mah, getattr(mah_params_early, key), getattr(mah_params_late, key)
+            )
+            for key in mah_params_late._fields
+        ]
+    )
+
+    log_mah = jnp.where(msk_mah.reshape((-1, 1)), log_mah_early, log_mah_late)
+
+    ZZ = np.zeros(n_halos)
+
+    uran = jran.uniform(sfh_key, minval=-100, maxval=100, shape=(8, n_halos))
+
+    u_ms_late_index = np.zeros(n_halos) + DEFAULT_U_MS_PARAMS.u_indx_hi
+    u_ms_params = [uran[0], uran[1], uran[2], u_ms_late_index, uran[3]]
+    u_ms_params = DEFAULT_U_MS_PARAMS._make(u_ms_params)
+    u_q_params = DEFAULT_U_Q_PARAMS._make([uran[i, :] for i in range(4, 8)])
+
+    sfh_u_params = DEFAULT_DIFFSTAR_U_PARAMS._make((u_ms_params, u_q_params))
+    sfh_params = get_bounded_diffstar_params(sfh_u_params)
+
+    q_params_noq = u_q_params._make([ZZ + x for x in DEFAULT_Q_PARAMS_UNQUENCHED])
+    sfh_params_noq = sfh_params._replace(q_params=q_params_noq)
+
+    ms_params_nolag = sfh_params.ms_params._replace(tau_dep=TAU_INST + ZZ)
+    sfh_params_noq_nolag = sfh_params_noq._replace(ms_params=ms_params_nolag)
+
+    sfh, smh = calc_sfh_galpop(
+        sfh_params,
+        mah_params,
+        tarr,
+        lgt0=LGT0,
+        fb=FB,
+        return_smh=True,
+    )
+    logsm0 = np.log10(smh)[:, -1]
+
+    sfh_noq, smh_noq = calc_sfh_galpop(
+        sfh_params_noq,
+        mah_params,
+        tarr,
+        lgt0=LGT0,
+        fb=FB,
+        return_smh=True,
+    )
+    logsm0_noq = np.log10(smh_noq)[:, -1]
+
+    sfh_noq_nolag, smh_noq_nolag = calc_sfh_galpop(
+        sfh_params_noq_nolag,
+        mah_params,
+        tarr,
+        lgt0=LGT0,
+        fb=FB,
+        return_smh=True,
+    )
+    logsm0_noq_nolag = np.log10(smh_noq_nolag)[:, -1]
+
+    # Implement stellar mass cut
+    msk = (
+        (logsm0 > logsm0_min)
+        & (logsm0_noq > logsm0_min)
+        & (logsm0_noq_nolag > logsm0_min)
+    )
+
+    mah_params_out = mah_params._make([x[msk] for x in mah_params])
+    log_mah_out = log_mah[msk]
+    sfh_out = sfh[msk]
+    smh_out = smh[msk]
+    sfh_noq_out = sfh_noq[msk]
+    smh_noq_out = smh_noq[msk]
+    sfh_noq_nolag_out = sfh_noq_nolag[msk]
+    smh_noq_nolag_out = smh_noq_nolag[msk]
+
+    ms_params = sfh_params.ms_params._make([x[msk] for x in sfh_params.ms_params])
+    q_params = sfh_params.q_params._make([x[msk] for x in sfh_params.q_params])
+    sfh_params_out = sfh_params._make((ms_params, q_params))
+
+    ms_params_noq_out = sfh_params_noq.ms_params._make(
+        [x[msk] for x in sfh_params_noq.ms_params]
+    )
+    q_params_noq_out = sfh_params_noq.q_params._make(
+        [x[msk] for x in sfh_params_noq.q_params]
+    )
+    sfh_params_noq_out = sfh_params_noq._make((ms_params_noq_out, q_params_noq_out))
+
+    ms_params_noq_nolag_out = sfh_params_noq_nolag.ms_params._make(
+        [x[msk] for x in sfh_params_noq_nolag.ms_params]
+    )
+    sfh_params_noq_nolag_out = sfh_params_noq_out._replace(
+        ms_params=ms_params_noq_nolag_out
+    )
+
+    return TData(
+        mah_params_out,
+        log_mah_out,
+        sfh_params_out,
+        sfh_out,
+        smh_out,
+        sfh_params_noq_out,
+        sfh_noq_out,
+        smh_noq_out,
+        sfh_params_noq_nolag_out,
+        sfh_noq_nolag_out,
+        smh_noq_nolag_out,
+    )
+
+
+def mc_diffmah_halo_sample(ran_key, tarr, logm0_sample):
+    n_halos = logm0_sample.size
+    ZZ = np.zeros(n_halos)
+    t_0 = tarr[-1]
+    t_obs = t_0 + ZZ
+    lgt0 = np.log10(t_0)
+    ran_keys = jran.split(ran_key, n_halos)
+    _reslist = _mc_diffmah_singlecen_vmap_kern(
+        DEFAULT_DIFFMAHPOP_PARAMS, tarr, logm0_sample, t_obs, ran_keys, lgt0
+    )
+    mah_params_early, dmhdt_early, log_mah_early = _reslist[:3]
+    mah_params_late, dmhdt_late, log_mah_late = _reslist[3:6]
+    frac_early = _reslist[6]
+    return (
+        mah_params_early,
+        dmhdt_early,
+        log_mah_early,
+        mah_params_late,
+        dmhdt_late,
+        log_mah_late,
+        frac_early,
+    )

diffstar-0.3.2/diffstar/diffstarnet/tests/test_diffstarnet_tdata.py

@@ -0,0 +1,80 @@
+"""
+"""
+
+import numpy as np
+from jax import random as jran
+
+from .. import diffstarnet_tdata as dtg
+
+
+def enforce_good_tdata(tdata, logsm0_min=float("-inf")):
+    for x in tdata:
+        try:
+            assert np.all(np.isfinite(x))
+        except ValueError:  # x is a namedtuple
+            for y in x:
+                assert np.all(np.isfinite(y))
+
+    logsm0 = np.log10(tdata.smh[:, -1])
+    assert np.all(logsm0 >= logsm0_min)
+
+    n_halos, n_times = tdata.sfh.shape
+    n_halos2 = tdata.mah_params.logm0.size
+    assert n_halos2 == n_halos
+
+    history_keys = [x for x in dtg.SFH_KEYS if "params" not in x]
+
+    for key in history_keys:
+        arr = getattr(tdata, key)
+        assert arr.shape == (n_halos, n_times)
+
+
+def test_tdata_generator():
+    ran_key = jran.key(0)
+    n_halos = 5_000
+    logm0_sample = np.linspace(10, 15, n_halos)
+
+    # generate 5 epochs of data
+    n_epochs = 5
+    gen = dtg.tdata_generator(ran_key, logm0_sample, n_epochs=n_epochs)
+    tdata_list = list(gen)
+    assert len(tdata_list) == n_epochs
+    for tdata in tdata_list:
+        enforce_good_tdata(tdata, logsm0_min=dtg.LGSM0_MIN)
+
+    # demo typical usage
+    LOGSM0_MIN = 6.0
+    gen = dtg.tdata_generator(ran_key, logm0_sample, n_epochs=2, logsm0_min=LOGSM0_MIN)
+    tdata0 = next(gen)
+    tdata1 = next(gen)
+    try:
+        next(gen)
+    except StopIteration:
+        pass  # expected because we tried to iterate for longer than n_epochs
+
+    # tdata should not contain Mstar exceeding logsm0_min
+    assert np.all(tdata0.smh[:, -1] >= 10**LOGSM0_MIN)
+    assert np.all(tdata1.smh[:, -1] >= 10**LOGSM0_MIN)
+
+    # tdata generator should yield different tdata with each iteration
+    assert not np.allclose(tdata0.log_mah[0, :], tdata1.log_mah[0, :])
+    assert not np.allclose(tdata0.sfh[0, :], tdata1.sfh[0, :])
+    assert not np.allclose(tdata0.sfh_noq[0, :], tdata1.sfh_noq[0, :])
+    assert not np.allclose(tdata0.sfh_noq_nolag[0, :], tdata1.sfh_noq_nolag[0, :])
+
+
+def test_mc_diffmah_halo_sample():
+    ran_key = jran.key(0)
+    n_halos_init = 2_000
+    logm0_sample = np.linspace(10, 15, n_halos_init)
+
+    gen = dtg.tdata_generator(ran_key, logm0_sample)
+    tdata = next(gen)
+
+    n_halos = tdata.mah_params.logm0.size
+    assert n_halos_init >= n_halos
+    assert tdata.mah_params.t_peak.size == n_halos
+
+    diff = tdata.mah_params.logm0 - tdata.log_mah[:, -1]
+    assert np.abs(diff).mean() < 0.1
+    assert np.std(diff) < 0.3