diffsky 0.2.3__tar.gz → 0.2.4__tar.gz

{diffsky-0.2.3 → diffsky-0.2.4}/.github/workflows/linting.yaml

@@ -12,7 +12,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - uses: conda-incubator/setup-miniconda@v3
         with:

{diffsky-0.2.3 → diffsky-0.2.4}/.github/workflows/monthly-warning-test.yaml

@@ -12,7 +12,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           fetch-depth: 0
 

{diffsky-0.2.3 → diffsky-0.2.4}/.github/workflows/test_latest_releases.yaml

@@ -12,7 +12,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           fetch-depth: 0
 

{diffsky-0.2.3 → diffsky-0.2.4}/.github/workflows/test_main_branch.yaml

@@ -12,7 +12,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           fetch-depth: 0
 

{diffsky-0.2.3 → diffsky-0.2.4}/.github/workflows/tests_cron.yaml

@@ -16,7 +16,7 @@ jobs:
     runs-on: "ubuntu-latest"
 
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
         with:
           fetch-depth: 0
 

{diffsky-0.2.3 → diffsky-0.2.4}/CHANGES.rst

@@ -1,3 +1,8 @@
+0.2.4 (2025-09-02)
+-------------------
+- Release associated with diffskyopt+kdescent calibration by Alan Pearl
+
+
 0.2.3 (2025-07-30)
 -------------------
 - Fix bug in lightcone data loader (https://github.com/ArgonneCPAC/diffsky/pull/195)

{diffsky-0.2.3 → diffsky-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsky
-Version: 0.2.3
+Version: 0.2.4
 Summary: Library for differentiable generation of synthetic skies
 Author-email: Andrew Hearin <ahearin@anl.gov>
 License: BSD 3-Clause License

diffsky-0.2.4/diffsky/_version.py

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

{diffsky-0.2.3 → diffsky-0.2.4}/diffsky/data_loaders/hacc_utils/lightcone_utils.py

@@ -302,16 +302,23 @@ def get_diffsky_quantities_for_lc_patch(
 ):
     cf_first_row = cf_matrices["absolute_row_idx"][0]
     cf_last_row = cf_matrices["absolute_row_idx"][-1]
-    msk_olap = lc_patch_data["file_idx"] == cf_file_idx
-    msk_olap &= lc_patch_data["row_idx"] >= cf_first_row
-    msk_olap &= lc_patch_data["row_idx"] <= cf_last_row
+    msk_olap = lc_patch_data["coreforest_file_idx"] == cf_file_idx
+    msk_olap &= lc_patch_data["coreforest_row_idx"] >= cf_first_row
+    msk_olap &= lc_patch_data["coreforest_row_idx"] <= cf_last_row
     mpeak_history = np.maximum.accumulate(cf_matrices[DIFFMAH_MASS_COLNAME], axis=1)
     mp_obs = mpeak_history[:, timestep_idx]
     mp0 = mpeak_history[:, -1]
 
     n_olap = msk_olap.sum()
     if n_olap > 0:
-        olap_chunk_idx = lc_patch_data["row_idx"][msk_olap] - cf_first_row
+        olap_chunk_idx = lc_patch_data["coreforest_row_idx"][msk_olap] - cf_first_row
+
+        # Enforce exact match in core_tag
+        core_tags_cf = cf_matrices["core_tag"][:, -1][olap_chunk_idx]
+        core_tags_lc = lc_patch_data["core_tag"][msk_olap]
+        msg = "Mismatched core_tag between coreforest and lc-cores"
+        assert np.allclose(core_tags_cf, core_tags_lc), msg
+
         olap_chunk_ult_host_idx = cf_matrices["top_host_row"][:, timestep_idx][
             olap_chunk_idx
         ]

diffsky-0.2.4/diffsky/experimental/disk_bulge_modeling/disk_bulge_kernels.py

@@ -0,0 +1,314 @@
+"""
+Kernels for disk-bulge modeling
+"""
+from collections import OrderedDict
+
+import numpy as np
+from diffstar.utils import cumulative_mstar_formed
+from dsps.constants import SFR_MIN
+from dsps.sed.stellar_age_weights import calc_age_weights_from_sfh_table
+from jax import jit as jjit
+from jax import lax, vmap
+from jax import numpy as jnp
+
+from dsps.sfh.diffburst import (
+    _pureburst_age_weights_from_params as _burst_age_weights_from_params,
+)
+from .disk_knots import _disk_knot_kern, _disk_knot_vmap
+from ...utils.tw_utils import _tw_sigmoid
+
+FBULGE_MIN = 0.05
+FBULGE_MAX = 0.95
+
+BOUNDING_K = 0.1
+
+DEFAULT_FBULGE_EARLY = 0.75
+DEFAULT_FBULGE_LATE = 0.15
+
+DEFAULT_FBULGE_PDICT = OrderedDict(fbulge_tcrit=8.0, fbulge_early=0.5, fbulge_late=0.1)
+DEFAULT_FBULGE_PARAMS = np.array(list(DEFAULT_FBULGE_PDICT.values()))
+
+
+_linterp_vmap = jjit(vmap(jnp.interp, in_axes=(0, None, 0)))
+
+_A = (None, 0, 0)
+_burst_age_weights_from_params_vmap = jjit(
+    vmap(_burst_age_weights_from_params, in_axes=_A)
+)
+
+_D = (None, 0, None, 0)
+calc_age_weights_from_sfh_table_vmap = jjit(
+    vmap(calc_age_weights_from_sfh_table, in_axes=_D)
+)
+
+
+@jjit
+def _sigmoid_2d(x, x0, y, y0, kx, ky, zmin, zmax):
+    height_diff = zmax - zmin
+    return zmin + height_diff * lax.logistic(kx * (x - x0) - ky * (y - y0))
+
+
+@jjit
+def _bulge_fraction_kernel(t, thalf, frac_early, frac_late, dt):
+    """typical values of 10.0, 0.7, 0.1
+    frac_late < frac_late is needed to make bulges redder and so that
+    bulge fractions increase with stellar mass
+    """
+    tw_h = dt / 6.0
+    return _tw_sigmoid(t, thalf, tw_h, frac_early, frac_late)
+
+
+@jjit
+def calc_tform_kern(abscissa, xarr, tform_frac):
+    fracarr = xarr / xarr[-1]
+    return jnp.interp(tform_frac, fracarr, abscissa)
+
+
+_calc_tform_pop_kern = jjit(vmap(calc_tform_kern, in_axes=[None, 0, None]))
+
+
+@jjit
+def calc_tform_pop(tarr, smh_pop, tform_frac):
+    """Calculate the formation time of a population
+
+    Parameters
+    ----------
+    tarr : ndarray, shape(nt, )
+
+    smh_pop : ndarray, shape(npop, nt)
+
+    tform_frac : float
+        Fraction used in the formation time definition
+        tform_frac=0.5 corresponds to the half-mass time, for example
+
+    """
+    return _calc_tform_pop_kern(tarr, smh_pop, tform_frac)
+
+
+@jjit
+def _bulge_fraction_vs_tform(t, t10, t90, params):
+    fbulge_tcrit, fbulge_early, fbulge_late = params
+    dt = t90 - t10
+    fbulge = _bulge_fraction_kernel(t, fbulge_tcrit, fbulge_early, fbulge_late, dt)
+    return fbulge
+
+
+@jjit
+def _bulge_sfh(tarr, sfh, fbulge_params):
+    sfh = jnp.where(sfh < SFR_MIN, SFR_MIN, sfh)
+    smh = cumulative_mstar_formed(tarr, sfh)
+    fracmh = smh / smh[-1]
+    t10 = jnp.interp(0.1, fracmh, tarr)
+    t90 = jnp.interp(0.9, fracmh, tarr)
+    eff_bulge = _bulge_fraction_vs_tform(tarr, t10, t90, fbulge_params)
+    sfh_bulge = eff_bulge * sfh
+    smh_bulge = cumulative_mstar_formed(tarr, sfh_bulge)
+    bth = smh_bulge / smh
+    return smh, eff_bulge, sfh_bulge, smh_bulge, bth
+
+
+_B = (None, 0, 0)
+_bulge_sfh_vmap = jjit(vmap(_bulge_sfh, in_axes=_B))
+
+
+@jjit
+def _get_observed_quantity(t_obs, tarr, quantity):
+    q_obs = jnp.interp(t_obs, tarr, quantity)
+    return q_obs
+
+
+_C = (0, None, 0)  # map t_obs and sfh for each galaxy
+_get_observed_quantity_vmap = jjit(vmap(_get_observed_quantity, in_axes=_C))
+
+
+@jjit
+def get_observed_quantity_pop(t_obs, tarr, quantity):
+    """Calculate the observed quantity for population
+
+    Parameters
+    ----------
+    t_obs : times of observation, ndarray, shape(npop, )
+    tarr : times for histories, ndarray, shape(nt, )
+
+    quantity : history for quantity, ndarray, shape(npop, nt)
+    """
+    return _get_observed_quantity_vmap(t_obs, tarr, quantity)
+
+
+def decompose_sfhpop_into_bulge_disk_knots(
+    gal_fbulge_params,
+    gal_fknot,
+    gal_t_obs,
+    gal_t_table,
+    gal_sfh,
+    gal_fburst,
+    gal_burstshape_params,
+    ssp_lg_age_gyr,
+):
+    """Decompose the SFH of a Diffsky galaxy into three components:
+    bulges, diffuse disk, and star-forming knots in the disks
+
+    Parameters
+    ----------
+    gal_fbulge_params : ndarray, shape (n_gals, 3)
+        Bulge efficiency parameters (tcrit, fbulge_early, fbulge_late) for each galaxy
+
+    gal_fknot : ndarray, shape (n_gals, )
+        Fraction of the disk mass in bursty star-forming knots for each galaxy
+
+    gal_t_obs : ndarray, shape (n_gals, )
+        Age of the universe in Gyr at the redshift of each galaxy
+
+    gal_t_table : ndarray, shape (n_t, )
+        Grid in cosmic time t in Gyr at which SFH of the galaxy population is tabulated
+        gal_t_table should increase monotonically and it should span the
+        full range of gal_t_obs, including some padding of a few million years
+
+    gal_sfh : ndarray, shape (n_gals, n_t)
+        Grid in SFR in Msun/yr for each galaxy tabulated at the input gal_t_table
+
+    gal_fburst : ndarray, shape (n_gals, )
+        Fraction of stellar mass in the burst population of each galaxy
+
+    gal_burstshape_params : ndarray, shape (n_gals, 3)
+        Parameters controlling P(τ) for burst population in each galaxy
+        lgfburst = gal_burstshape_params[:, 0]
+        lgyr_peak = gal_burstshape_params[:, 1]
+        lgyr_max = gal_burstshape_params[:, 2]
+
+    ssp_lg_age_gyr : ndarray, shape (n_age, )
+        Grid in age τ at which the SSPs are computed, stored as log10(τ/Gyr)
+
+    Returns
+    -------
+    mbulge : ndarray, shape (n_gals, )
+        Total stellar mass in Msun formed in the bulge at time gal_t_obs
+
+    mdd : ndarray, shape (n_gals, )
+        Total stellar mass in Msun formed in the diffuse disk at time gal_t_obs
+
+    mknot : ndarray, shape (n_gals, )
+        Total stellar mass in Msun formed in star-forming knots at time gal_t_obs
+
+    mburst : ndarray, shape (n_gals, )
+        Total stellar mass in Msun in the burst population at time gal_t_obs
+
+    bulge_age_weights : ndarray, shape (n_gals, n_age)
+        Probability distribution P(τ_age) for bulge of each galaxy
+
+    dd_age_weights : ndarray, shape (n_gals, n_age)
+        Probability distribution P(τ_age) for diffuse disk of each galaxy
+
+    knot_age_weights : ndarray, shape (n_gals, n_age)
+        Probability distribution P(τ_age) for star-forming knots of each galaxy
+
+    bulge_sfh : ndarray, shape (n_gals, n_t)
+        Grid in SFR in Msun/yr for each galaxy bulge tabulated at the input gal_t_table
+
+    gal_frac_bulge_t_obs : ndarray, shape (n_gals, )
+        Bulge/total mass ratio at gal_t_obs for every galaxy
+
+    """
+    ssp_lg_age_yr = ssp_lg_age_gyr + 9.0
+    lgyr_peak = gal_burstshape_params[:, 1]
+    lgyr_max = gal_burstshape_params[:, 2]
+    gal_burst_age_weights = _burst_age_weights_from_params_vmap(
+        ssp_lg_age_yr, lgyr_peak, lgyr_max,
+    )
+    return _decompose_sfhpop_into_bulge_disk_knots(
+        gal_fbulge_params,
+        gal_fknot,
+        gal_t_obs,
+        gal_t_table,
+        gal_sfh,
+        gal_fburst,
+        gal_burst_age_weights,
+        ssp_lg_age_gyr,
+    )
+
+
+@jjit
+def _decompose_sfh_singlegal_into_bulge_disk_knots(
+    fbulge_params,
+    fknot,
+    t_obs,
+    t_table,
+    sfh_table,
+    fburst,
+    age_weights_burst,
+    ssp_lg_age_gyr,
+):
+    _res = _bulge_sfh(t_table, sfh_table, fbulge_params)
+    smh, eff_bulge, bulge_sfh, smh_bulge, bulge_to_total_history = _res
+
+    bulge_sfh = jnp.where(bulge_sfh < SFR_MIN, SFR_MIN, bulge_sfh)
+    frac_bulge_t_obs = jnp.interp(t_obs, t_table, bulge_to_total_history)
+
+    bulge_age_weights = calc_age_weights_from_sfh_table(
+        t_table, bulge_sfh, ssp_lg_age_gyr, t_obs
+    )
+    disk_sfh = sfh_table - bulge_sfh
+    disk_sfh = jnp.where(disk_sfh < SFR_MIN, SFR_MIN, disk_sfh)
+
+    args = (
+        t_table,
+        t_obs,
+        sfh_table,
+        disk_sfh,
+        fburst,
+        fknot,
+        age_weights_burst,
+        ssp_lg_age_gyr,
+    )
+    _knot_info = _disk_knot_kern(*args)
+    mstar_tot, mburst, mdd, mknot, dd_age_weights, knot_age_weights = _knot_info
+
+    mbulge = frac_bulge_t_obs * mstar_tot
+    masses = mbulge, mdd, mknot, mburst
+    age_weights = bulge_age_weights, dd_age_weights, knot_age_weights
+    ret = (*masses, *age_weights, bulge_sfh, frac_bulge_t_obs)
+    return ret
+
+
+@jjit
+def _decompose_sfhpop_into_bulge_disk_knots(
+    gal_fbulge_params,
+    gal_fknot,
+    gal_t_obs,
+    gal_t_table,
+    gal_sfh,
+    gal_fburst,
+    age_weights_burst,
+    ssp_lg_age_gyr,
+):
+    _res = _bulge_sfh_vmap(gal_t_table, gal_sfh, gal_fbulge_params)
+    smh, eff_bulge, bulge_sfh, smh_bulge, bulge_to_total_history = _res
+
+    bulge_sfh = jnp.where(bulge_sfh < SFR_MIN, SFR_MIN, bulge_sfh)
+    gal_frac_bulge_t_obs = _linterp_vmap(gal_t_obs, gal_t_table, bulge_to_total_history)
+
+    bulge_age_weights = calc_age_weights_from_sfh_table_vmap(
+        gal_t_table, bulge_sfh, ssp_lg_age_gyr, gal_t_obs
+    )
+
+    disk_sfh = gal_sfh - bulge_sfh
+    disk_sfh = jnp.where(disk_sfh < SFR_MIN, SFR_MIN, disk_sfh)
+
+    args = (
+        gal_t_table,
+        gal_t_obs,
+        gal_sfh,
+        disk_sfh,
+        gal_fburst,
+        gal_fknot,
+        age_weights_burst,
+        ssp_lg_age_gyr,
+    )
+    _knot_info = _disk_knot_vmap(*args)
+    mstar_tot, mburst, mdd, mknot, dd_age_weights, knot_age_weights = _knot_info
+
+    mbulge = gal_frac_bulge_t_obs * mstar_tot
+    masses = mbulge, mdd, mknot, mburst
+    age_weights = bulge_age_weights, dd_age_weights, knot_age_weights
+    ret = (*masses, *age_weights, bulge_sfh, gal_frac_bulge_t_obs)
+    return ret

diffsky-0.2.4/diffsky/experimental/disk_bulge_modeling/disk_knots.py

@@ -0,0 +1,57 @@
+"""
+Add knot component
+"""
+from diffstar.utils import cumulative_mstar_formed
+from dsps.constants import SFR_MIN
+from dsps.sed.stellar_age_weights import calc_age_weights_from_sfh_table
+from jax import jit as jjit
+from jax import numpy as jnp
+from jax import vmap
+
+FKNOT_MAX = 0.2
+
+
+@jjit
+def _disk_knot_kern(
+    tarr, tobs, sfh, sfh_disk, fburst, fknot, age_weights_burst, ssp_lg_age_gyr
+):
+    sfh = jnp.where(sfh < SFR_MIN, SFR_MIN, sfh)
+    sfh_disk = jnp.where(sfh_disk < SFR_MIN, SFR_MIN, sfh_disk)
+
+    sfh_knot = fknot * sfh_disk
+    sfh_diffuse_disk = sfh_disk * (1 - fknot)
+
+    smh = cumulative_mstar_formed(tarr, sfh)
+    smh_knot = cumulative_mstar_formed(tarr, sfh_knot)
+    smh_diffuse_disk = cumulative_mstar_formed(tarr, sfh_diffuse_disk)
+
+    age_weights_dd = calc_age_weights_from_sfh_table(
+        tarr, sfh_diffuse_disk, ssp_lg_age_gyr, tobs
+    )
+
+    lgt_table = jnp.log10(tarr)
+    mstar_tot = 10 ** jnp.interp(jnp.log10(tobs), lgt_table, jnp.log10(smh))
+    mknot = 10 ** jnp.interp(jnp.log10(tobs), lgt_table, jnp.log10(smh_knot))
+    mdd = 10 ** jnp.interp(jnp.log10(tobs), lgt_table, jnp.log10(smh_diffuse_disk))
+    mburst = fburst * mstar_tot
+
+    mburst_by_mknot = mburst / mknot
+    burst_knot_age_weights = (
+        mburst_by_mknot * age_weights_burst + (1 - mburst_by_mknot) * age_weights_dd
+    )
+    age_weights_knot = jnp.where(
+        mburst_by_mknot > 1, age_weights_burst, burst_knot_age_weights
+    )
+
+    mburst_dd = jnp.where(mburst_by_mknot > 1, mburst - mknot, 0.0)
+
+    mdd_tot = mdd + mburst_dd
+    age_weights_dd = (mdd / mdd_tot) * age_weights_dd + (
+        mburst_dd / mdd_tot
+    ) * age_weights_burst
+
+    return mstar_tot, mburst, mdd, mknot, age_weights_dd, age_weights_knot
+
+
+_V = (None, 0, 0, 0, 0, 0, 0, None)
+_disk_knot_vmap = jjit(vmap(_disk_knot_kern, in_axes=_V))

diffsky-0.2.4/diffsky/experimental/disk_bulge_modeling/generate_bulge_disk_sample.py

@@ -0,0 +1,80 @@
+import jax.numpy as jnp
+import numpy as np
+from ...mc_diffsky import mc_diffstar_cenpop, mc_diffstar_galpop
+from dsps.cosmology.defaults import DEFAULT_COSMOLOGY
+from dsps.cosmology.flat_wcdm import age_at_z
+from jax import random as jran
+
+from .mc_disk_bulge import (
+    DEFAULT_FBULGE_2dSIGMOID_PARAMS,
+    mc_disk_bulge,
+)
+
+ran_key = jran.key(0)
+halo_key, ran_key = jran.split(ran_key, 2)
+
+# Invert t_table for redshifts
+
+
+def get_redshifts_from_times(
+    t_table, cosmo_params, zmin=0.001, zmax=50, Ngrid=200,
+):
+    zgrid = np.logspace(np.log10(zmax), np.log10(zmin), Ngrid)
+    age_grid = age_at_z(zgrid, *cosmo_params)
+    redshifts = np.interp(t_table, age_grid, zgrid)
+
+    return redshifts
+
+
+# Generate subcat and SFH catalog
+
+
+def get_bulge_disk_test_sample(
+    ran_key,
+    lgmp_min=11.0,
+    redshift=0.05,
+    Lbox=100.0,
+    centrals=True,
+    cosmology=DEFAULT_COSMOLOGY,
+):
+    volume_com = Lbox**3
+    args = (ran_key, redshift, lgmp_min, volume_com)
+    if centrals:
+        diffstar = mc_diffstar_cenpop(*args, return_internal_quantities=True)
+    else:
+        diffstar = mc_diffstar_galpop(*args, return_internal_quantities=True)
+
+    print("Generated data shape = ", diffstar["sfh"].shape)
+
+    diffstar["sSFR"] = jnp.divide(diffstar["sfh"], diffstar["smh"])
+    diffstar["z_table"] = get_redshifts_from_times(diffstar["t_table"], cosmology)
+
+    return diffstar
+
+
+def get_bulge_disk_decomposition(
+    ran_key, diffstar, fbulge_2d_params=DEFAULT_FBULGE_2dSIGMOID_PARAMS,
+):
+    _res = mc_disk_bulge(
+        ran_key,
+        diffstar["t_table"],
+        diffstar["sfh"],
+        fbulge_2d_params=fbulge_2d_params,
+    )
+    fbulge_params, smh, eff_bulge, sfh_bulge, smh_bulge, bth = _res
+
+    diffstar["tcrit_bulge"] = fbulge_params[:, 0]
+    diffstar["fbulge_early"] = fbulge_params[:, 1]
+    diffstar["fbulge_late"] = fbulge_params[:, 2]
+    diffstar["bth"] = bth
+    diffstar["eff_bulge"] = eff_bulge
+
+    # Save or compute bulge and disk quantities
+    diffstar["sfh_bulge"] = sfh_bulge
+    diffstar["smh_bulge"] = smh_bulge
+    diffstar["sSFR_bulge"] = jnp.divide(sfh_bulge, smh_bulge)
+    diffstar["smh_disk"] = diffstar["smh"] - smh_bulge
+    diffstar["sfh_disk"] = diffstar["sfh"] - sfh_bulge
+    diffstar["sSFR_disk"] = jnp.divide(diffstar["sfh_disk"], diffstar["smh_disk"])
+
+    return diffstar

diffsky-0.2.4/diffsky/experimental/disk_bulge_modeling/mc_disk_bulge.py

@@ -0,0 +1,133 @@
+"""
+Generate disk-bulge decomposition
+"""
+from collections import OrderedDict, namedtuple
+
+import jax.numpy as jnp
+import numpy as np
+from diffstar.utils import cumulative_mstar_formed_galpop
+from dsps.constants import SFR_MIN
+
+from .disk_bulge_kernels import (
+    _bulge_sfh_vmap,
+    _sigmoid_2d,
+    calc_tform_pop,
+)
+
+DEFAULT_FBULGE_PDICT = OrderedDict(
+    early_logsm0_x0=10.0,
+    early_logssfr0_x0=-10.0,
+    early_logssfr0_k=0.2,
+    early_logsm0_k=0.2,
+    early_zmin=1.0,
+    early_zmax=0.2,
+    late_logsm0_x0=10.0,
+    late_logssfr0_x0=-10.0,
+    late_logssfr0_k=0.2,
+    late_logsm0_k=0.2,
+    late_zmax=0.1,
+    tcrit_logsm0_x0=10.0,
+    tcrit_logssfr0_x0=-10.0,
+    tcrit_logssfr0_k=0.5,
+    tcrit_logsm0_k=0.8,
+)
+Fbulge2dParams = namedtuple("Fbulge2dParams", DEFAULT_FBULGE_PDICT.keys())
+DEFAULT_FBULGE_2dSIGMOID_PARAMS = Fbulge2dParams(**DEFAULT_FBULGE_PDICT)
+
+
+def mc_disk_bulge(
+    ran_key, tarr, sfh_pop, fbulge_2d_params=DEFAULT_FBULGE_2dSIGMOID_PARAMS,
+):
+    """Decompose input SFHs into disk and bulge contributions
+
+    Parameters
+    ----------
+    ran_key : jax.random.PRNGKey
+
+    tarr : ndarray, shape (n_t, )
+
+    sfh_pop : ndarray, shape (n_gals, n_t)
+
+    fbulge_2d_params : named tuple of parameters for 2d-sigmoid
+
+    new_model : boolean flag to switch between new and old model
+
+    Returns
+    -------
+    fbulge_params : ndarray, shape (n_gals, 3)
+        tcrit_bulge = fbulge_params[:, 0]
+        fbulge_early = fbulge_params[:, 1]
+        fbulge_late = fbulge_params[:, 2]
+
+    smh : ndarray, shape (n_gals, n_t)
+        Stellar mass history of galaxy in units of Msun
+
+    eff_bulge : ndarray, shape (n_gals, n_t)
+        History of in-situ bulge growth efficiency for every galaxy
+
+    sfh_bulge : ndarray, shape (n_gals, n_t)
+        Star formation history of bulge in units of Msun/yr
+
+    smh_bulge : ndarray, shape (n_gals, n_t)
+        Stellar mass history of bulge in units of Msun
+
+    bth : ndarray, shape (n_gals, n_t)
+        History of bulge-to-total mass ratio of every galaxy
+
+    """
+    sfh_pop = np.where(sfh_pop < SFR_MIN, SFR_MIN, sfh_pop)
+    smh_pop = cumulative_mstar_formed_galpop(tarr, sfh_pop)
+    t10 = calc_tform_pop(tarr, smh_pop, 0.1)
+    t90 = calc_tform_pop(tarr, smh_pop, 0.9)
+    logsm0 = jnp.log10(smh_pop[:, -1])
+
+    ssfr = jnp.divide(sfh_pop, smh_pop)
+    logssfr0 = jnp.log10(ssfr[:, -1])
+    fbulge_params = generate_fbulge_parameters_2d_sigmoid(
+        ran_key, logsm0, logssfr0, t10, t90, fbulge_2d_params
+    )
+
+    _res = _bulge_sfh_vmap(tarr, sfh_pop, fbulge_params)
+    smh, eff_bulge, sfh_bulge, smh_bulge, bth = _res
+    return fbulge_params, smh, eff_bulge, sfh_bulge, smh_bulge, bth
+
+
+def generate_fbulge_parameters_2d_sigmoid(
+    ran_key, logsm0, logssfr0, t10, t90, f_bulge_params
+):
+    fbulge_early = _sigmoid_2d(
+        logssfr0,
+        f_bulge_params.early_logssfr0_x0,
+        logsm0,
+        f_bulge_params.early_logsm0_x0,
+        f_bulge_params.early_logssfr0_k,
+        f_bulge_params.early_logsm0_k,
+        f_bulge_params.early_zmin,
+        f_bulge_params.early_zmax,
+    )
+
+    fbulge_late = _sigmoid_2d(
+        logssfr0,
+        f_bulge_params.late_logssfr0_x0,
+        logsm0,
+        f_bulge_params.late_logsm0_x0,
+        f_bulge_params.late_logssfr0_k,
+        f_bulge_params.late_logsm0_k,
+        fbulge_early,
+        f_bulge_params.late_zmax,
+    )
+
+    fbulge_tcrit = _sigmoid_2d(
+        logssfr0,
+        f_bulge_params.tcrit_logssfr0_x0,
+        logsm0,
+        f_bulge_params.tcrit_logsm0_x0,
+        f_bulge_params.tcrit_logssfr0_k,
+        f_bulge_params.tcrit_logsm0_k,
+        t90,
+        t10,
+    )
+
+    fbulge_param_arr = np.asarray((fbulge_tcrit, fbulge_early, fbulge_late)).T
+
+    return fbulge_param_arr