diffstar 0.2.4__tar.gz → 0.3.1__tar.gz

{diffstar-0.2.4 → 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.2.4 → 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.2.4 → 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
@@ -31,7 +31,7 @@ jobs:
           use-mamba: true
 
       - name: configure conda and install code
-      # Test against current main branch of diffmah
+      # Test against current main branch of diffmah and dsps
         shell: bash -l {0}
         run: |
           conda config --set always_yes yes
@@ -49,6 +49,7 @@ jobs:
             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

{diffstar-0.2.4 → diffstar-0.3.1}/CHANGES.rst

@@ -1,3 +1,13 @@
+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
+
+
 0.2.4 (2024-01-16)
 ------------------
 - Require diffmah>=0.5.0

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: diffstar
-Version: 0.2.4
+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.2.4 → diffstar-0.3.1}/diffstar/__init__.py

@@ -7,8 +7,11 @@ from .defaults import (
     DEFAULT_DIFFSTAR_U_PARAMS,
     DiffstarParams,
     DiffstarUParams,
+    MSParams,
+    MSUParams,
+    QParams,
+    QUParams,
     get_bounded_diffstar_params,
     get_unbounded_diffstar_params,
 )
-from .kernels import get_ms_sfh_from_mah_kern, get_sfh_from_mah_kern
-from .sfh import sfh_galpop, sfh_singlegal
+from .sfh_model import calc_sfh_galpop, calc_sfh_singlegal

diffstar-0.3.1/diffstar/_version.py

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

{diffstar-0.2.4 → diffstar-0.3.1}/diffstar/defaults.py

@@ -14,6 +14,7 @@ LGT0 = np.log10(TODAY)
 # Constants related to SFH integrals
 SFR_MIN = 1e-14
 T_BIRTH_MIN = 0.001
+T_TABLE_MIN = 0.01
 N_T_LGSM_INTEGRATION = 100
 DEFAULT_N_STEPS = 50
 
@@ -53,6 +54,28 @@ DEFAULT_DIFFSTAR_U_PARAMS = DiffstarUParams(DEFAULT_U_MS_PARAMS, DEFAULT_U_Q_PAR
 
 @jjit
 def get_bounded_diffstar_params(diffstar_u_params):
+    """Calculate diffstar parameters from unbounded counterparts.
+
+    The returned diffstar_params is the input expected by diffstar.calc_sfh_singlegal
+    and diffstar.calc_sfh_galpop.
+
+    Parameters
+    ----------
+    diffstar_u_params : namedtuple, length 2
+        DiffstarUParams = u_ms_params, u_q_params
+            u_ms_params and u_q_params are tuples of floats or ndarrays
+            u_ms_params = u_lgmcrit, u_lgy_at_mcrit, u_indx_lo, u_indx_hi, u_tau_dep
+            u_q_params = u_lg_qt, u_qlglgdt, u_lg_drop, u_lg_rejuv
+
+    Returns
+    -------
+    diffstar_params : namedtuple, length 2
+        DiffstarParams = ms_params, q_params
+            ms_params and q_params are tuples of floats or ndarrays
+            ms_params = lgmcrit, lgy_at_mcrit, indx_lo, indx_hi, tau_dep
+            q_params = lg_qt, qlglgdt, lg_drop, lg_rejuv
+
+    """
     ms_params = MSParams(*_get_bounded_sfr_params(*diffstar_u_params.u_ms_params))
     q_params = QParams(*_get_bounded_q_params(*diffstar_u_params.u_q_params))
     return DiffstarParams(ms_params, q_params)
@@ -60,6 +83,27 @@ def get_bounded_diffstar_params(diffstar_u_params):
 
 @jjit
 def get_unbounded_diffstar_params(diffstar_params):
+    """Calculate unbounded diffstar parameters from standard params.
+
+    This is the inverse function to get_bounded_diffstar_params
+
+    Parameters
+    ----------
+    diffstar_params : namedtuple, length 2
+        DiffstarParams = ms_params, q_params
+            ms_params and q_params are tuples of floats or ndarrays
+            ms_params = lgmcrit, lgy_at_mcrit, indx_lo, indx_hi, tau_dep
+            q_params = lg_qt, qlglgdt, lg_drop, lg_rejuv
+
+    Returns
+    -------
+    diffstar_u_params : namedtuple, length 2
+        DiffstarUParams = u_ms_params, u_q_params
+            u_ms_params and u_q_params are tuples of floats or ndarrays
+            u_ms_params = u_lgmcrit, u_lgy_at_mcrit, u_indx_lo, u_indx_hi, u_tau_dep
+            u_q_params = u_lg_qt, u_qlglgdt, u_lg_drop, u_lg_rejuv
+
+    """
     u_ms_params = MSUParams(*_get_unbounded_sfr_params(*diffstar_params.ms_params))
     u_q_params = QUParams(*_get_unbounded_q_params(*diffstar_params.q_params))
     return DiffstarUParams(u_ms_params, u_q_params)

{diffstar-0.2.4 → 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.1/diffstar/kernels/history_kernel_builders.py

@@ -0,0 +1,268 @@
+"""
+"""
+
+from jax import jit as jjit
+from jax import lax
+from jax import numpy as jnp
+from jax import vmap
+
+from ..defaults import (
+    DEFAULT_DIFFSTAR_PARAMS,
+    DEFAULT_MAH_PARAMS,
+    DEFAULT_N_STEPS,
+    SFR_MIN,
+    T_BIRTH_MIN,
+)
+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",)
+
+N_MAH_PARAMS = len(DEFAULT_MAH_PARAMS)
+N_MS_PARAMS = len(DEFAULT_DIFFSTAR_PARAMS.ms_params)
+N_Q_PARAMS = len(DEFAULT_DIFFSTAR_PARAMS.q_params)
+
+
+def build_sfh_from_mah_kernel(
+    n_steps=DEFAULT_N_STEPS,
+    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.
+
+    Parameters
+    ----------
+    n_steps : int, optional
+        Number of timesteps to use in the tacc integration
+
+    tacc_integration_min : float, optional
+        Earliest time to use in the tacc integrations. Default is 0.01 Gyr.
+
+    tobs_loop : string, optional
+        Argument specifies whether the input time of observation is a scalar or array
+        Default argument is None, for a JAX kernel that assumes scalar input for tobs
+        For a JAX kernel that assumes an array input for tobs,
+        options are either 'vmap' or 'scan', specifying the calculation method
+
+    galpop_loop : string, optional
+        Argument specifies whether the input galaxy/halo parameters assumed by the
+        returned JAX kernel pertain to a single galaxy or a population.
+        Default argument is None, for a single-galaxy JAX 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
+        JAX-jitted function that calculates SFH in accord with the input arguments
+        Function signature is as follows:
+
+        def sfh_from_mah_kern(t, mah_params, ms_params, q_params, lgt0, fb):
+            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
+    def _kern(
+        t_form,
+        logmp,
+        logtc,
+        early_index,
+        late_index,
+        lgmcrit,
+        lgy_at_mcrit,
+        indx_lo,
+        indx_hi,
+        tau_dep,
+        lg_qt,
+        qlglgdt,
+        lg_drop,
+        lg_rejuv,
+        lgt0,
+        fb,
+    ):
+        mah_params = logmp, logtc, early_index, late_index
+        ms_params = lgmcrit, lgy_at_mcrit, indx_lo, indx_hi, tau_dep
+
+        t_min = jnp.max(jnp.array((tacc_integration_min, t_form - tau_dep)))
+        t_table = t_min + uniform_table * (t_form - t_min)
+        args = t_form, mah_params, ms_params, lgt0, fb, t_table
+        ms_sfr = _lax_ms_sfh_scalar_kern(*args)
+        lgt_form = jnp.log10(t_form)
+
+        lg_q_dt = 10**qlglgdt
+        qkern_inputs = lg_qt, lg_q_dt, lg_drop, lg_rejuv
+        qfunc = _quenching_kern(lgt_form, *qkern_inputs)
+        sfr = qfunc * ms_sfr
+        sfr = lax.cond(sfr < SFR_MIN, lambda x: SFR_MIN, lambda x: x, sfr)
+        return sfr
+
+    kern_with_tobs_loop = _get_kern_with_tobs_loop(_kern, tobs_loop)
+    sfh_from_mah_kern = _get_kern_with_galpop_loop(kern_with_tobs_loop, galpop_loop)
+
+    return sfh_from_mah_kern
+
+
+def _get_kern_with_tobs_loop(kern, tobs_loop):
+    if tobs_loop == "vmap":
+        _t = [
+            0,
+            *[None] * N_MAH_PARAMS,
+            *[None] * N_MS_PARAMS,
+            *[None] * N_Q_PARAMS,
+            None,
+            None,
+        ]
+        new_kern = jjit(vmap(kern, in_axes=_t))
+    elif tobs_loop == "scan":
+
+        @jjit
+        def new_kern(
+            tarr,
+            logmp,
+            logtc,
+            early_index,
+            late_index,
+            lgmcrit,
+            lgy_at_mcrit,
+            indx_lo,
+            indx_hi,
+            tau_dep,
+            lg_qt,
+            qlglgdt,
+            lg_drop,
+            lg_rejuv,
+            lgt0,
+            fb,
+        ):
+            @jjit
+            def scan_func_time_array(carryover, el):
+                t_form = el
+                sfr_at_t_form = kern(
+                    t_form,
+                    logmp,
+                    logtc,
+                    early_index,
+                    late_index,
+                    lgmcrit,
+                    lgy_at_mcrit,
+                    indx_lo,
+                    indx_hi,
+                    tau_dep,
+                    lg_qt,
+                    qlglgdt,
+                    lg_drop,
+                    lg_rejuv,
+                    lgt0,
+                    fb,
+                )
+                carryover = sfr_at_t_form
+                accumulated = sfr_at_t_form
+                return carryover, accumulated
+
+            scan_init = 0.0
+            scan_arr = tarr
+            res = lax.scan(scan_func_time_array, scan_init, scan_arr)
+            sfh = res[1]
+            return sfh
+
+    elif tobs_loop is None:
+        new_kern = kern
+    else:
+        msg = "Input `tobs_loop`={0} must be either `vmap` or `scan`"
+        raise ValueError(msg.format(tobs_loop))
+    return new_kern
+
+
+def _get_kern_with_galpop_loop(kern, galpop_loop):
+    if galpop_loop == "vmap":
+        _g = [
+            None,
+            *[0] * N_MAH_PARAMS,
+            *[0] * N_MS_PARAMS,
+            *[0] * N_Q_PARAMS,
+            None,
+            None,
+        ]
+        new_kern = jjit(vmap(kern, in_axes=_g))
+    elif galpop_loop == "scan":
+
+        @jjit
+        def new_kern(
+            t,
+            logmp,
+            logtc,
+            early_index,
+            late_index,
+            lgmcrit,
+            lgy_at_mcrit,
+            indx_lo,
+            indx_hi,
+            tau_dep,
+            lg_qt,
+            qlglgdt,
+            lg_drop,
+            lg_rejuv,
+            lgt0,
+            fb,
+        ):
+            n_gals = logmp.shape[0]
+
+            n_params = N_MAH_PARAMS + N_MS_PARAMS + N_Q_PARAMS
+            galpop_params = jnp.zeros(shape=(n_gals, n_params))
+            galpop_params = galpop_params.at[:, 0].set(logmp)
+            galpop_params = galpop_params.at[:, 1].set(logtc)
+            galpop_params = galpop_params.at[:, 2].set(early_index)
+            galpop_params = galpop_params.at[:, 3].set(late_index)
+
+            galpop_params = galpop_params.at[:, 4].set(lgmcrit)
+            galpop_params = galpop_params.at[:, 5].set(lgy_at_mcrit)
+            galpop_params = galpop_params.at[:, 6].set(indx_lo)
+            galpop_params = galpop_params.at[:, 7].set(indx_hi)
+            galpop_params = galpop_params.at[:, 8].set(tau_dep)
+
+            galpop_params = galpop_params.at[:, 9].set(lg_qt)
+            galpop_params = galpop_params.at[:, 10].set(qlglgdt)
+            galpop_params = galpop_params.at[:, 11].set(lg_drop)
+            galpop_params = galpop_params.at[:, 12].set(lg_rejuv)
+
+            @jjit
+            def scan_func_galpop(carryover, el):
+                params = el
+                mah_params = params[:N_MAH_PARAMS]
+                i, j = N_MAH_PARAMS, N_MAH_PARAMS + N_MS_PARAMS
+                ms_params = params[i:j]
+                i = N_MAH_PARAMS + N_MS_PARAMS
+                q_params = params[i:]
+                sfh_galpop = kern(t, *mah_params, *ms_params, *q_params, lgt0, fb)
+                carryover = sfh_galpop
+                accumulated = sfh_galpop
+                return carryover, accumulated
+
+            scan_init = jnp.zeros_like(t)
+            scan_arr = galpop_params
+            res = lax.scan(scan_func_galpop, scan_init, scan_arr)
+            sfh_galpop = res[1]
+            return sfh_galpop
+
+    elif galpop_loop is None:
+        new_kern = kern
+    else:
+        msg = "Input `galpop_loop`={0} must be either `vmap` or `scan`"
+        raise ValueError(msg.format(galpop_loop))
+    return new_kern

{diffstar-0.2.4 → 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.2.4 → diffstar-0.3.1}/diffstar/kernels/quenching_kernels.py

@@ -78,6 +78,9 @@ def _quenching_kern_u_params(lgt, u_lg_qt, u_qlglgdt, u_lg_drop, u_lg_rejuv):
 def _quenching_kern(lgt, lg_qt, lg_q_dt, q_drop, q_rejuv):
     """Base-10 logarithmic drop and symmetric rise in SFR over a time interval.
 
+    Note the relationship between the input lg_q_dt and q_params[1]=qlglgdt:
+        lg_q_dt = 10**qlglgdt
+
     Parameters
     ----------
     lgt : ndarray