Diffstar without a gas consumption timescale

alt_ms_kernels.py

"""
"""

from diffmah.defaults import MAH_K
from diffmah.individual_halo_assembly import _calc_halo_history_scalar
from jax import jit as jjit
from jax import nn
from jax import numpy as jnp

INDX_K = 9.0  # Main sequence efficiency transition speed
INDX_HI = -1.0


@jjit
def _sigmoid(x, x0, k, ymin, ymax):
    height_diff = ymax - ymin
    return ymin + height_diff * nn.sigmoid(k * (x - x0))


@jjit
def _sfr_eff_plaw(lgm, lgmcrit, lgy_at_mcrit, indx_lo):
    slope = _sigmoid(lgm, lgmcrit, INDX_K, indx_lo, INDX_HI)
    eff = lgy_at_mcrit + slope * (lgm - lgmcrit)
    return 10**eff


@jjit
def _sfr_nolag_scalar_kern(t, mah_params, ms_params, lgt0, fb):
    logmp, logtc, early, late = mah_params
    all_mah_params = lgt0, logmp, logtc, MAH_K, early, late
    lgt = jnp.log10(t)
    res = _calc_halo_history_scalar(lgt, *all_mah_params)
    dmhdt, log_mah = res

    sfr_eff = _sfr_eff_plaw(log_mah, *ms_params)

    dmgdt = fb * dmhdt
    sfr = dmgdt * sfr_eff

    return sfr

diffstar_nolag_demo.ipynb