diffsky 0.2.2__tar.gz → 0.2.4__tar.gz

{diffsky-0.2.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → diffsky-0.2.4}/CHANGES.rst

@@ -1,4 +1,14 @@
-0.2.2 (unreleased)
+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)
+
+
+0.2.2 (2025-07-17)
 -------------------
 - Change little h convention of Monte Carlo halo generators (https://github.com/ArgonneCPAC/diffsky/pull/184)
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsky
-Version: 0.2.2
+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.2 → diffsky-0.2.4}/diffsky/data_loaders/hacc_utils/lc_mock_production.py

@@ -80,8 +80,9 @@ def write_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
         hdf_out["data/logssfr_obs"] = diffsky_data["logssfr_obs"]
 
         hdf_out["data/redshift_true"] = lc_data["redshift_true"]
-        hdf_out["data/ra"] = lc_data["phi"]
-        hdf_out["data/dec"] = np.pi / 2.0 - lc_data["theta"]
+        ra, dec = hlu._get_lon_lat_from_theta_phi(lc_data["theta"], lc_data["phi"])
+        hdf_out["data/ra"] = ra
+        hdf_out["data/dec"] = dec
         hdf_out["data/snapnum"] = lc_data["snapnum"]
 
         for key in LC_DATA_KEYS_OUT:

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

@@ -132,6 +132,12 @@ def get_ra_dec(x, y, z):
     return ra, dec
 
 
+@jjit
+def get_ra_dec_from_theta_phi(theta, phi):
+    """Change sky coordinates from {theta, phi} (radians) to {ra, dec} (degrees)"""
+    return _get_lon_lat_from_theta_phi(theta, phi)
+
+
 @jjit
 def _get_lon_lat_from_theta_phi(theta, phi):
     lon = jnp.degrees(phi)
@@ -166,6 +172,27 @@ def get_redshift_from_xyz(x_mpch, y_mpch, z_mpch, cosmo_params):
 
 
 def read_lc_ra_dec_patch_decomposition(fn):
+    """Read and parse 'lc_cores-decomposition.txt'
+
+    Parameters
+    ----------
+    fname : string
+
+    Returns
+    -------
+    patch_decomposition : array shape (n_patches, 5)
+        Each row stores (patchnum, theta_low, theta_high, phi_low, phi_high)
+
+        The relationship between {theta, phi} and {ra, dec} is given by the
+        lightcone_utils.get_ra_dec_from_theta_phi function
+
+    sky_frac : array shape (n_patches, )
+        Size of the sky patch as a fraction of 4π steradians
+
+    solid_angles : array shape (n_patches, )
+        Size of the sky patch in deg^2
+
+    """
     collector = []
     with open(fn, "r") as f:
         next(f)
@@ -275,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.2 → diffsky-0.2.4}/diffsky/data_loaders/hacc_utils/load_lc_mock.py

@@ -29,7 +29,7 @@ def load_diffsky_lightcone(drn, sim_name, z_min, z_max, patch_list):
 
     data_collector = []
     for fn in fn_collector:
-        lc_diffsky_data = load_flat_hdf5(fn)
+        lc_diffsky_data = load_flat_hdf5(fn, dataset="data")
         data_collector.append(lc_diffsky_data)
 
     diffsky_data = dict()

{diffsky-0.2.2 → diffsky-0.2.4}/diffsky/data_loaders/io_utils.py

@@ -31,11 +31,16 @@ def load_flat_hdf5(fn, istart=0, iend=None, keys=None, dataset=None):
     data = dict()
     with h5py.File(fn, "r") as hdf:
 
-        if keys is None:
-            if dataset is None:
+        if dataset is None:
+            if keys is None:
                 keys_in = list(hdf.keys())
             else:
+                keys_in = keys
+        else:
+            if keys is None:
                 keys_in = [dataset + "/" + key for key in list(hdf[dataset].keys())]
+            else:
+                keys_in = [dataset + "/" + key for key in keys]
 
         for key_in in keys_in:
             if dataset is None:

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