diffstar 0.3.3__tar.gz → 0.3.4__tar.gz

diffstar-0.3.4/.github/dependabot.yml

@@ -0,0 +1,10 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"
+    groups:
+      github-actions:
+        patterns:
+          - '*'

{diffstar-0.3.3 → diffstar-0.3.4}/.github/workflows/linting.yaml

@@ -12,9 +12,9 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
 
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: 3.11
           channels: conda-forge,defaults

{diffstar-0.3.3 → diffstar-0.3.4}/.github/workflows/monthly-warning-test.yaml

@@ -12,11 +12,11 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: 3.11
           channels: conda-forge,defaults

{diffstar-0.3.3 → diffstar-0.3.4}/.github/workflows/test_releases.yaml

@@ -16,11 +16,11 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: 3.11
           channels: conda-forge,defaults
@@ -53,7 +53,7 @@ jobs:
           pytest -v diffstar --cov --cov-report=xml
 
       - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v5
 
       - name: test versions
         shell: bash -el {0}

{diffstar-0.3.3 → diffstar-0.3.4}/.github/workflows/tests_cron.yaml

@@ -16,11 +16,11 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           python-version: 3.11
           channels: conda-forge,defaults

{diffstar-0.3.3 → diffstar-0.3.4}/.readthedocs.yml

@@ -6,7 +6,7 @@ sphinx:
 build:
   os: "ubuntu-22.04"
   tools:
-    python: "mambaforge-22.9"
+    python: "miniconda-latest"
 
 conda:
   environment: docs/source/rtd_environment.yaml

{diffstar-0.3.3 → diffstar-0.3.4}/CHANGES.rst

@@ -1,3 +1,8 @@
+0.3.4 (2025-03-03)
+------------------
+- Add convenience function cumulative_mstar_formed_galpop (https://github.com/ArgonneCPAC/diffstar/pull/77)
+
+
 0.3.3 (2025-01-15)
 ------------------
 - Clean out old code so that only tpeak-based diffmah models remain (https://github.com/ArgonneCPAC/diffstar/pull/72)

{diffstar-0.3.3/diffstar.egg-info → diffstar-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: diffstar
-Version: 0.3.3
+Version: 0.3.4
 Summary: Differentiable Star Formation Histories
 Author-email: Alex Alarcon <alexalarcongonzalez@gmail.com>
 License: BSD 3-Clause License

diffstar-0.3.4/diffstar/_version.py

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

{diffstar-0.3.3 → diffstar-0.3.4}/diffstar/diffstarnet/diffstarnet_tdata.py

@@ -34,10 +34,45 @@ _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_KEYS = ["mah_params", "log_mah"] + SFH_KEYS + ["time_arr"]
 TData = namedtuple("TData", TDATA_KEYS)
 
 
+def tdata_generator_dithertarr(
+    ran_key,
+    logm0_sample,
+    n_sfh_table=N_SFH_TABLE,
+    logsm0_min=LGSM0_MIN,
+    n_epochs=float("inf")
+):
+    """
+    Same as tdata_generator, but for each generation, a new t_table_min value
+    is drawn uniformly within the first bin (and same for logm0_sample)
+    """
+    min_m0, max_m0 = logm0_sample.min(), logm0_sample.max()
+    num_m0 = logm0_sample.size
+
+    max_dither_t = (T0 - T_TABLE_MIN) / (n_sfh_table - 1)
+    max_dither_m = (max_m0 - min_m0) / (num_m0 - 1)
+    batchnum = 0
+    while batchnum < n_epochs:
+        ran_key, batch_key, *dither_keys = jran.split(ran_key, 4)
+        dither1, dither2 = jran.uniform(
+            dither_keys[0], (2,), maxval=max_dither_t
+        )
+        tarr = np.linspace(T_TABLE_MIN + dither1, T0 - dither2, n_sfh_table)
+        dither1, dither2 = jran.uniform(
+            dither_keys[1], (2,), maxval=max_dither_m
+        )
+        logm0_new = np.linspace(min_m0 + dither1, max_m0 - dither2, num_m0)
+        tdata = _compute_tdata(
+            batch_key, logm0_new, n_sfh_table, logsm0_min,
+            tarr=tarr
+        )
+        yield tdata
+        batchnum += 1
+
+
 def tdata_generator(
     ran_key,
     logm0_sample,
@@ -78,10 +113,12 @@ def tdata_generator(
 
 
 def _compute_tdata(
-    ran_key, logm0_sample, n_sfh_table=N_SFH_TABLE, logsm0_min=LGSM0_MIN
+    ran_key, logm0_sample, n_sfh_table=N_SFH_TABLE, logsm0_min=LGSM0_MIN,
+    tarr=None,
 ):
     """"""
-    tarr = np.linspace(T_TABLE_MIN, T0, n_sfh_table)
+    if tarr is None:
+        tarr = np.linspace(T_TABLE_MIN, T0, n_sfh_table)
 
     mah_key, early_late_key, sfh_key = jran.split(ran_key, 3)
 
@@ -199,6 +236,7 @@ def _compute_tdata(
         sfh_params_noq_nolag_out,
         sfh_noq_nolag_out,
         smh_noq_nolag_out,
+        tarr
     )
 
 

{diffstar-0.3.3 → diffstar-0.3.4}/diffstar/diffstarnet/tests/test_diffstarnet_tdata.py

@@ -29,6 +29,26 @@ def enforce_good_tdata(tdata, logsm0_min=float("-inf")):
         assert arr.shape == (n_halos, n_times)
 
 
+def test_tdata_generator_dithertarr():
+    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_dithertarr(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)
+
+    # Assert that the time arrays start and end from unique values
+    start_times = [x.time_arr[0] for x in tdata_list]
+    end_times = [x.time_arr[-1] for x in tdata_list]
+    assert len(set(start_times)) == len(start_times)
+    assert len(set(end_times)) == len(end_times)
+
+
 def test_tdata_generator():
     ran_key = jran.key(0)
     n_halos = 5_000

{diffstar-0.3.3 → diffstar-0.3.4}/diffstar/tests/test_utils.py

@@ -1,11 +1,12 @@
 """
 """
+
 import numpy as np
 import pytest
 from jax import random as jran
 
+from .. import utils
 from ..defaults import T_TABLE_MIN
-from ..utils import _get_dt_array, _jax_get_dt_array, cumtrapz, cumulative_mstar_formed
 
 try:
     import dsps
@@ -14,13 +15,21 @@ try:
 except ImportError:
     HAS_DSPS = False
 
+try:
+    from scipy.integrate import trapezoid
+
+    HAS_SCIPY = True
+except ImportError:
+    HAS_SCIPY = False
+
 MSG_HAS_DSPS = "Must have dsps installed to run this test"
+MSG_HAS_SCIPY = "Must have scipy installed to run this test"
 
 
 def test_jax_get_dt_array_linspace():
     tarr = np.linspace(1, 13.8, 50)
-    dtarr_np = _get_dt_array(tarr)
-    dtarr_jnp = _jax_get_dt_array(tarr)
+    dtarr_np = utils._get_dt_array(tarr)
+    dtarr_jnp = utils._jax_get_dt_array(tarr)
     assert np.allclose(dtarr_np, dtarr_jnp, atol=0.01)
 
 
@@ -30,11 +39,12 @@ def test_jax_get_dt_array_random():
     for __ in range(n_tests):
         ran_key, key = jran.split(ran_key, 2)
         tarr = np.sort(jran.uniform(key, minval=0, maxval=14, shape=(50,)))
-        dtarr_np = _get_dt_array(tarr)
-        dtarr_jnp = _jax_get_dt_array(tarr)
+        dtarr_np = utils._get_dt_array(tarr)
+        dtarr_jnp = utils._jax_get_dt_array(tarr)
         assert np.allclose(dtarr_np, dtarr_jnp, atol=0.01)
 
 
+@pytest.mark.skipif(not HAS_SCIPY, reason=MSG_HAS_SCIPY)
 def test_cumtrapz():
     ran_key = jran.PRNGKey(0)
     n_x = 100
@@ -43,16 +53,16 @@ def test_cumtrapz():
         x_key, y_key, ran_key = jran.split(ran_key, 3)
         xarr = np.sort(jran.uniform(x_key, minval=0, maxval=1, shape=(n_x,)))
         yarr = jran.uniform(y_key, minval=0, maxval=1, shape=(n_x,))
-        jax_result = cumtrapz(xarr, yarr)
-        np_result = [np.trapz(yarr[:-i], x=xarr[:-i]) for i in range(1, n_x)][::-1]
+        jax_result = utils.cumtrapz(xarr, yarr)
+        np_result = [trapezoid(yarr[:-i], x=xarr[:-i]) for i in range(1, n_x)][::-1]
         assert np.allclose(jax_result[:-1], np_result, rtol=1e-4)
-        assert np.allclose(jax_result[-1], np.trapz(yarr, x=xarr), rtol=1e-4)
+        assert np.allclose(jax_result[-1], trapezoid(yarr, x=xarr), rtol=1e-4)
 
 
 def test_cumulative_mstar_formed_returns_reasonable_arrays():
     t_table = np.linspace(T_TABLE_MIN, 13.8, 200)
     sfh_table = np.random.uniform(0, 1, t_table.size)
-    smh_table = cumulative_mstar_formed(t_table, sfh_table)
+    smh_table = utils.cumulative_mstar_formed(t_table, sfh_table)
     assert smh_table.shape == t_table.shape
     assert np.all(smh_table > 0)
     assert np.all(np.diff(smh_table) > 0)
@@ -67,6 +77,21 @@ def test_cumulative_mstar_formed_agrees_with_dsps():
     ran_keys = jran.split(ran_key, n_tests)
     for key in ran_keys:
         sfh_table = jran.uniform(key, minval=0, maxval=1, shape=(nt,))
-        smh_table_diffstar = cumulative_mstar_formed(t_table, sfh_table)
+        smh_table_diffstar = utils.cumulative_mstar_formed(t_table, sfh_table)
         smh_table_dsps = dsps.utils.cumulative_mstar_formed(t_table, sfh_table)
         assert np.allclose(smh_table_diffstar, smh_table_dsps, rtol=1e-4)
+
+
+def test_cumulative_mstar_formed_vmap():
+
+    n_t = 200
+    t_table = np.linspace(T_TABLE_MIN, 13.8, n_t)
+    ran_key = jran.PRNGKey(0)
+
+    n_gals = 25
+    sfh_table_galpop = jran.uniform(ran_key, minval=0, maxval=1, shape=(n_gals, n_t))
+    smh_table_galpop = utils.cumulative_mstar_formed_galpop(t_table, sfh_table_galpop)
+
+    for ig in range(n_gals):
+        smh_table_ig = utils.cumulative_mstar_formed(t_table, sfh_table_galpop[ig, :])
+        assert np.allclose(smh_table_ig, smh_table_galpop[ig, :], rtol=1e-5)

{diffstar-0.3.3 → diffstar-0.3.4}/diffstar/utils.py

@@ -5,6 +5,7 @@ import numpy as np
 from jax import jit as jjit
 from jax import lax, nn
 from jax import numpy as jnp
+from jax import vmap
 from jax.lax import scan
 
 from .defaults import SFR_MIN, T_BIRTH_MIN
@@ -273,3 +274,35 @@ def cumulative_mstar_formed(t_table, sfh_table):
     mstar_formed = cumtrapz(padded_t_table, padded_sfh_table)[1:] * YEAR_PER_GYR
 
     return mstar_formed
+
+
+_cuml_mstar_vmap = jjit(vmap(cumulative_mstar_formed, in_axes=(None, 0)))
+
+
+@jjit
+def cumulative_mstar_formed_galpop(t_table, sfh_table):
+    """Compute the cumulative stellar mass formed at each input time
+
+    Parameters
+    ----------
+    t_table : ndarray, shape (n_t, )
+        Age of the Universe in Gyr.
+        Array should be monotonically increasing and
+        t_table[0] >= dsps.constants.T_TABLE_MIN
+
+    sfh_table : ndarray, shape (n_gals, n_t)
+        SFR in Msun/yr at each of the input times
+
+    Returns
+    -------
+    mstar_formed : ndarray, shape (n_gals, n_t)
+        Cumulative stellar mass formed in Msun at each input time
+
+    Notes
+    -----
+    Mstar formed is calculated using trapezoidal integration
+    and assuming that during the interval (dsps.constants.T_BIRTH_MIN, t_table[0]),
+    SFH is constant and equal to dsps.constants.SFR_MIN
+
+    """
+    return _cuml_mstar_vmap(t_table, sfh_table)

{diffstar-0.3.3 → diffstar-0.3.4/diffstar.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: diffstar
-Version: 0.3.3
+Version: 0.3.4
 Summary: Differentiable Star Formation Histories
 Author-email: Alex Alarcon <alexalarcongonzalez@gmail.com>
 License: BSD 3-Clause License

{diffstar-0.3.3 → diffstar-0.3.4}/diffstar.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ pyproject.toml
 requirements.txt
 setup.cfg
 setup.py
+.github/dependabot.yml
 .github/workflows/linting.yaml
 .github/workflows/monthly-warning-test.yaml
 .github/workflows/test_releases.yaml

diffstar-0.3.3/diffstar/_version.py

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