diffstar 0.3.0__tar.gz → 0.3.1__tar.gz

{diffstar-0.3.0 → diffstar-0.3.1}/.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.1/.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.1}/.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.1}/.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.1}/CHANGES.rst

@@ -1,3 +1,8 @@
+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.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diffstar
-Version: 0.3.0
+Version: 0.3.1
 Summary: Differentiable Star Formation Histories
 Author-email: Alex Alarcon <alexalarcongonzalez@gmail.com>
 License: BSD 3-Clause License
@@ -35,7 +35,7 @@ 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

diffstar-0.3.1/diffstar/_version.py

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

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar/kernels/gas_consumption.py

@@ -1,5 +1,6 @@
 """Calculate mass of gas available for star formation from the freshly accreted gas.
 """
+
 from jax import jit as jjit
 from jax import lax
 from jax import numpy as jnp
@@ -28,6 +29,11 @@ def _gas_conversion_kern(t_form, t_acc, dt, tau_dep, tau_dep_max):
     return tri_kern
 
 
+_vmap_gas_conversion_kern = jjit(
+    vmap(_gas_conversion_kern, in_axes=(None, 0, None, None, None))
+)
+
+
 _a, _b = (0, None, 0, None, None), (None, 0, None, None, None)
 _depletion_kernel = jjit(vmap(vmap(_gas_conversion_kern, in_axes=_b), in_axes=_a))
 

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar/kernels/history_kernel_builders.py

@@ -1,5 +1,6 @@
 """
 """
+
 from jax import jit as jjit
 from jax import lax
 from jax import numpy as jnp
@@ -12,7 +13,10 @@ from ..defaults import (
     SFR_MIN,
     T_BIRTH_MIN,
 )
-from .main_sequence_kernels import _lax_ms_sfh_scalar_kern
+from .main_sequence_kernels import (
+    _lax_ms_sfh_scalar_kern_scan,
+    _lax_ms_sfh_scalar_kern_sum,
+)
 from .quenching_kernels import _quenching_kern
 
 __all__ = ("build_sfh_from_mah_kernel",)
@@ -27,6 +31,7 @@ def build_sfh_from_mah_kernel(
     tacc_integration_min=T_BIRTH_MIN,
     tobs_loop=None,
     galpop_loop=None,
+    tform_loop="sum",
 ):
     """Build a JAX-jitted kernel to calculate SFHs of a galaxy population.
 
@@ -51,6 +56,10 @@ def build_sfh_from_mah_kernel(
         For a JAX kernel that assumes galaxy population,
         options are either 'vmap' or 'scan', specifying the calculation method
 
+    tform_loop : string
+        Use 'sum' for faster vmap-based calculation and 'scan' for slower alternative.
+        Default is 'sum'
+
     Returns
     -------
     sfh_from_mah_kern : function
@@ -61,6 +70,11 @@ def build_sfh_from_mah_kernel(
             return sfh
 
     """
+    if tform_loop == "sum":
+        _lax_ms_sfh_scalar_kern = _lax_ms_sfh_scalar_kern_sum
+    elif tform_loop == "scan":
+        _lax_ms_sfh_scalar_kern = _lax_ms_sfh_scalar_kern_scan
+
     uniform_table = jnp.linspace(0, 1, n_steps)
 
     @jjit

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar/kernels/main_sequence_kernels.py

@@ -1,10 +1,12 @@
 """
 """
+
 from collections import OrderedDict, namedtuple
 
 import numpy as np
 from diffmah.defaults import MAH_K
 from diffmah.individual_halo_assembly import (
+    _calc_halo_history,
     _calc_halo_history_scalar,
     _rolling_plaw_vs_logt,
 )
@@ -14,7 +16,7 @@ from jax import numpy as jnp
 from jax import vmap
 
 from ..utils import _inverse_sigmoid, _jax_get_dt_array, _sigmoid
-from .gas_consumption import _gas_conversion_kern
+from .gas_consumption import _gas_conversion_kern, _vmap_gas_conversion_kern
 
 DEFAULT_MS_PDICT = OrderedDict(
     lgmcrit=12.0,
@@ -53,7 +55,7 @@ MS_BOUNDING_SIGMOID_PDICT = calculate_sigmoid_bounds(MS_PARAM_BOUNDS_PDICT)
 
 
 @jjit
-def _lax_ms_sfh_scalar_kern(t_form, mah_params, ms_params, lgt0, fb, t_table):
+def _lax_ms_sfh_scalar_kern_scan(t_form, mah_params, ms_params, lgt0, fb, t_table):
     logmp, logtc, early, late = mah_params
     all_mah_params = lgt0, logmp, logtc, MAH_K, early, late
     lgt_form = jnp.log10(t_form)
@@ -93,6 +95,38 @@ def _lax_ms_sfh_scalar_kern(t_form, mah_params, ms_params, lgt0, fb, t_table):
     return sfr
 
 
+@jjit
+def _lax_ms_sfh_scalar_kern_sum(t_form, mah_params, ms_params, lgt0, fb, t_table):
+    logmp, logtc, early, late = mah_params
+    all_mah_params = lgt0, logmp, logtc, MAH_K, early, late
+    lgt_form = jnp.log10(t_form)
+    log_mah_at_tform = _rolling_plaw_vs_logt(lgt_form, *all_mah_params)
+
+    sfr_eff_params = ms_params[:4]
+    sfr_eff = _sfr_eff_plaw(log_mah_at_tform, *sfr_eff_params)
+
+    tau_dep = ms_params[4]
+    tau_dep_max = MS_BOUNDING_SIGMOID_PDICT["tau_dep"][3]
+
+    # compute inst. gas accretion
+    lgtacc = jnp.log10(t_table)
+    res = _calc_halo_history(lgtacc, *all_mah_params)
+    dmhdt_at_tacc, log_mah_at_tacc = res
+    dmgdt_inst = fb * dmhdt_at_tacc
+
+    # compute the consumption kernel
+    dt = t_table[1] - t_table[0]
+    kern = _vmap_gas_conversion_kern(t_form, t_table, dt, tau_dep, tau_dep_max)
+
+    # convolve
+    dmgas_dt = jnp.sum(dmgdt_inst * kern * dt)
+    sfr = dmgas_dt * sfr_eff
+    return sfr
+
+
+_lax_ms_sfh_scalar_kern = _lax_ms_sfh_scalar_kern_sum
+
+
 @jjit
 def _sfr_eff_plaw(lgm, lgmcrit, lgy_at_mcrit, indx_lo, indx_hi):
     """Instantaneous baryon conversion efficiency of main sequence galaxies

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar/sfh_model.py

@@ -1,5 +1,6 @@
 """
 """
+
 from collections import namedtuple
 from functools import partial
 
@@ -14,6 +15,7 @@ _sfh_singlegal_kern = build_sfh_from_mah_kernel(
     n_steps=DEFAULT_N_STEPS,
     tacc_integration_min=T_BIRTH_MIN,
     tobs_loop="scan",
+    tform_loop="sum",
 )
 
 _sfh_galpop_kern = build_sfh_from_mah_kernel(
@@ -21,6 +23,7 @@ _sfh_galpop_kern = build_sfh_from_mah_kernel(
     tacc_integration_min=T_BIRTH_MIN,
     tobs_loop="scan",
     galpop_loop="vmap",
+    tform_loop="sum",
 )
 
 _cumulative_mstar_formed_vmap = jjit(vmap(cumulative_mstar_formed, in_axes=(None, 0)))

diffstar-0.3.1/diffstar/tests/test_main_sequence_kernels.py

@@ -0,0 +1,28 @@
+import numpy as np
+import jax.numpy as jnp
+
+from diffstar.kernels.main_sequence_kernels import (
+    _lax_ms_sfh_scalar_kern_scan,
+    _lax_ms_sfh_scalar_kern_sum,
+)
+from diffstar.kernels.main_sequence_kernels import DEFAULT_MS_PARAMS
+from diffmah.defaults import DEFAULT_MAH_PARAMS
+from diffstar.defaults import T_TABLE_MIN, TODAY
+from diffstar.defaults import FB
+
+
+def test_main_sequence_kernels_lax_ms_sfh_scalar_kern_scan_vs_sum():
+    lgt0 = jnp.log10(TODAY)
+    t_form = 12.0
+    t_table = jnp.linspace(T_TABLE_MIN, t_form, 20)
+
+    np.testing.assert_allclose(
+        _lax_ms_sfh_scalar_kern_scan(
+            t_form, DEFAULT_MAH_PARAMS, DEFAULT_MS_PARAMS, lgt0, FB, t_table
+        ),
+        _lax_ms_sfh_scalar_kern_sum(
+            t_form, DEFAULT_MAH_PARAMS, DEFAULT_MS_PARAMS, lgt0, FB, t_table
+        ),
+        rtol=1e-6,
+        atol=1e-6,
+    )

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar/utils.py

@@ -1,8 +1,9 @@
 """
 """
+
 import numpy as np
 from jax import jit as jjit
-from jax import lax
+from jax import lax, nn
 from jax import numpy as jnp
 from jax.lax import scan
 
@@ -91,7 +92,7 @@ def _sigmoid(x, x0, k, ymin, ymax):
 
     """
     height_diff = ymax - ymin
-    return ymin + height_diff * lax.logistic(k * (x - x0))
+    return ymin + height_diff * nn.sigmoid(k * (x - x0))
 
 
 @jjit

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diffstar
-Version: 0.3.0
+Version: 0.3.1
 Summary: Differentiable Star Formation Histories
 Author-email: Alex Alarcon <alexalarcongonzalez@gmail.com>
 License: BSD 3-Clause License
@@ -35,7 +35,7 @@ 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

{diffstar-0.3.0 → diffstar-0.3.1}/diffstar.egg-info/SOURCES.txt

@@ -11,6 +11,7 @@ requirements.txt
 setup.cfg
 setup.py
 .github/workflows/linting.yml
+.github/workflows/monthly-warning-test.yml
 .github/workflows/test_releases.yml
 .github/workflows/tests_cron.yml
 diffstar/__init__.py
@@ -63,6 +64,7 @@ diffstar/tests/test_defaults.py
 diffstar/tests/test_gas.py
 diffstar/tests/test_lax_main_sequence.py
 diffstar/tests/test_lax_sfh.py
+diffstar/tests/test_main_sequence_kernels.py
 diffstar/tests/test_quenching.py
 diffstar/tests/test_sfh.py
 diffstar/tests/test_sfh_model.py

{diffstar-0.3.0 → diffstar-0.3.1}/pyproject.toml

@@ -26,7 +26,7 @@ authors = [
 ]
 description = "Differentiable Star Formation Histories"
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 license = {file = "LICENSE.rst"}
 classifiers = [
     "Programming Language :: Python :: 3",

diffstar-0.3.0/diffstar/_version.py

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