diffsky 0.2.3__tar.gz → 0.2.5__tar.gz

{diffsky-0.2.3 → diffsky-0.2.5}/.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.5}/.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.5}/.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.5}/.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.5}/.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.5}/CHANGES.rst

@@ -1,3 +1,13 @@
+0.2.5 (2025-10-01)
+-------------------
+- Minor changes before updating to diffstar v1.0
+
+
+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.5}/PKG-INFO

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

diffsky-0.2.5/diffsky/_version.py

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

{diffsky-0.2.3 → diffsky-0.2.5}/diffsky/data_loaders/hacc_utils/lc_mock_production.py

@@ -15,8 +15,10 @@ from diffstarpop.param_utils import mc_select_diffstar_params
 from dsps.cosmology import flat_wcdm
 from jax import random as jran
 
+from ...experimental import mc_diffsky_seds
 from ...fake_sats import halo_boundary_functions as hbf
 from ...fake_sats import nfw_config_space as nfwcs
+from ...param_utils import diffsky_param_wrapper as dpw
 from ...utils.sfh_utils import get_logsm_logssfr_at_t_obs
 from . import lightcone_utils as hlu
 from . import load_lc_cf
@@ -94,6 +96,10 @@ def write_lc_sfh_mock_to_disk(fnout, lc_data, diffsky_data):
             hdf_out[key_out] = diffsky_data[key]
 
 
+def write_lc_sed_mock_to_disk(fnout, lc_data, diffsky_data):
+    raise NotImplementedError("Need to add SED quantities to write_lc_sed_mock_to_disk")
+
+
 def add_sfh_quantities_to_mock(sim_info, lc_data, diffsky_data, ran_key):
     lc_data["t_obs"] = flat_wcdm.age_at_z(
         lc_data["redshift_true"], *sim_info.cosmo_params
@@ -105,6 +111,7 @@ def add_sfh_quantities_to_mock(sim_info, lc_data, diffsky_data, ran_key):
     for pname, pval in zip(mah_params._fields, mah_params):
         diffsky_data[pname] = pval
     diffsky_data["has_diffmah_fit"] = msk_has_diffmah_fit
+    diffsky_data["mah_params"] = mah_params
 
     logmp0 = logmh_at_t_obs(
         mah_params, np.zeros(mah_params.logm0.size) + 10**sim_info.lgt0, sim_info.lgt0
@@ -121,7 +128,7 @@ def add_sfh_quantities_to_mock(sim_info, lc_data, diffsky_data, ran_key):
     gyr_since_infall = np.zeros_like(logmp0)
     upids = np.where(lc_data["central"] == 1, -1, 0)
 
-    t_table = np.linspace(0.1, 10**sim_info.lgt0, 100)
+    diffsky_data["t_table"] = np.linspace(0.1, 10**sim_info.lgt0, 100)
 
     args = (
         DEFAULT_DIFFSTARPOP_PARAMS,
@@ -132,27 +139,102 @@ def add_sfh_quantities_to_mock(sim_info, lc_data, diffsky_data, ran_key):
         logmhost_infall,
         gyr_since_infall,
         ran_key,
-        t_table,
+        diffsky_data["t_table"],
     )
 
     _res = mc_diffstar_sfh_galpop(*args)
     sfh_params_ms, sfh_params_q, sfh_ms, sfh_q, frac_q, mc_is_q = _res
 
-    sfh = np.where(mc_is_q.reshape((-1, 1)), sfh_q, sfh_ms)
+    diffsky_data["sfh_table"] = np.where(mc_is_q.reshape((-1, 1)), sfh_q, sfh_ms)
     sfh_params = mc_select_diffstar_params(sfh_params_q, sfh_params_ms, mc_is_q)
 
+    diffsky_data["sfh_table_ms"] = sfh_ms
+    diffsky_data["sfh_table_q"] = sfh_q
+
+    diffsky_data["mc_is_q"] = mc_is_q
+    # Need to add SFH table after MC selection
+
     for key in sfh_params.ms_params._fields:
         diffsky_data[key] = getattr(sfh_params.ms_params, key)
     for key in sfh_params.q_params._fields:
         diffsky_data[key] = getattr(sfh_params.q_params, key)
 
-    logsm_obs, logssfr_obs = get_logsm_logssfr_at_t_obs(lc_data["t_obs"], t_table, sfh)
+    logsm_obs, logssfr_obs = get_logsm_logssfr_at_t_obs(
+        lc_data["t_obs"], diffsky_data["t_table"], diffsky_data["sfh_table"]
+    )
     diffsky_data["logsm_obs"] = logsm_obs
     diffsky_data["logssfr_obs"] = logssfr_obs
 
+    logsm_obs_ms, logssfr_obs_ms = get_logsm_logssfr_at_t_obs(
+        lc_data["t_obs"], diffsky_data["t_table"], diffsky_data["sfh_table_ms"]
+    )
+    diffsky_data["logsm_obs_ms"] = logsm_obs_ms
+    diffsky_data["logssfr_obs_ms"] = logssfr_obs_ms
+
+    logsm_obs_q, logssfr_obs_q = get_logsm_logssfr_at_t_obs(
+        lc_data["t_obs"], diffsky_data["t_table"], diffsky_data["sfh_table_q"]
+    )
+    diffsky_data["logsm_obs_q"] = logsm_obs_q
+    diffsky_data["logssfr_obs_q"] = logssfr_obs_q
+
     return lc_data, diffsky_data
 
 
+def add_sed_quantities_to_mock(
+    sim_info,
+    lc_data,
+    diffsky_data,
+    ssp_data,
+    u_param_arr,
+    precomputed_ssp_mag_table,
+    z_phot_table,
+    wave_eff_table,
+    ran_key,
+):
+    u_param_collection = dpw.get_u_param_collection_from_u_param_array(u_param_arr)
+    param_collection = dpw.get_param_collection_from_u_param_collection(
+        *u_param_collection
+    )
+
+    (
+        diffstarpop_params,
+        mzr_params,
+        spspop_params,
+        scatter_params,
+        ssp_err_pop_params,
+    ) = param_collection
+
+    n_z_table, n_bands, n_met, n_age = precomputed_ssp_mag_table.shape
+
+    ran_key, sfh_key = jran.split(ran_key, 2)
+    lc_data, diffsky_data = add_sfh_quantities_to_mock(
+        sim_info, lc_data, diffsky_data, sfh_key
+    )
+    n_gals = diffsky_data["logsm_obs"].size
+
+    ran_key, sed_key = jran.split(ran_key, 2)
+    args = (
+        sed_key,
+        lc_data["redshift_true"],
+        lc_data["t_obs"],
+        diffsky_data["mah_params"],
+        diffsky_data["logmp0"],
+        diffsky_data["t_table"],
+        ssp_data,
+        precomputed_ssp_mag_table,
+        z_phot_table,
+        wave_eff_table,
+        diffstarpop_params,
+        mzr_params,
+        spspop_params,
+        scatter_params,
+        ssp_err_pop_params,
+        sim_info.cosmo_params,
+    )
+    phot_info = mc_diffsky_seds._mc_diffsky_phot_kern(*args)
+    return phot_info, lc_data, diffsky_data
+
+
 def reposition_satellites(sim_info, lc_data, diffsky_data, ran_key, fixed_conc=5.0):
 
     pos = np.array((lc_data["x"], lc_data["y"], lc_data["z"])).T

{diffsky-0.2.3 → diffsky-0.2.5}/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.5/diffsky/ellipsoidal_shapes/bulge_opt.py

@@ -0,0 +1,157 @@
+"""Functions and script for optimizing parameters of bulge_shapes model"""
+
+from collections import namedtuple
+
+import numpy as np
+from jax import jit as jjit
+from jax import numpy as jnp
+from jax import random as jran
+from jax import value_and_grad
+
+from .. import diffndhist
+from . import bulge_shapes as eshape
+from . import ellipse_proj_kernels as eproj
+
+
+@jjit
+def _mae(pred, target):
+    diff = pred - target
+    return jnp.mean(jnp.abs(diff))
+
+
+@jjit
+def get_all_params_from_varied(varied_params, default_params):
+    all_params = default_params._replace(**varied_params._asdict())
+    return all_params
+
+
+@jjit
+def _nondiff_loss_kern(bulge_params, loss_data):
+    x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins = loss_data
+    pred_data = ran_key, mu_ran, phi_ran, ba_bins, x_target
+
+    ba_pdf_pred = _nondiff_pred_kern(bulge_params, pred_data)
+    loss = _mae(ba_pdf_pred, ba_pdf_target)
+    return loss
+
+
+@jjit
+def loss_kern(varied_params, loss_data):
+    x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins, default_params = (
+        loss_data
+    )
+    bulge_params = get_all_params_from_varied(varied_params, default_params)
+    loss = _loss_kern(
+        bulge_params, x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins
+    )
+    return loss
+
+
+@jjit
+def _loss_kern(
+    bulge_params, x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins
+):
+    pred_data = x_target, ran_key, mu_ran, phi_ran, ba_bins
+
+    ba_pdf_pred = _pred_kern(bulge_params, pred_data)
+    loss = _mae(ba_pdf_pred, ba_pdf_target)
+    return loss
+
+
+@jjit
+def _pred_kern(bulge_params, pred_data):
+    x_target, ran_key, mu_ran, phi_ran, ba_bins = pred_data
+
+    ran_key, bulge_key = jran.split(ran_key, 2)
+    ellipse2d = _pred_ellipse_samples_kern(bulge_params, bulge_key, mu_ran, phi_ran)
+
+    nddata = jnp.array(ellipse2d.alpha / ellipse2d.beta).reshape((-1, 1))
+    ndsig = jnp.zeros_like(nddata) + 0.05
+
+    ndbins_lo = ba_bins[:-1].reshape((-1, 1))
+    ndbins_hi = ba_bins[1:].reshape((-1, 1))
+
+    wcounts = diffndhist.tw_ndhist(nddata, ndsig, ndbins_lo, ndbins_hi)
+    ba_pdf_pred_table = (wcounts / wcounts.sum()) / jnp.diff(ba_bins)
+
+    ba_binmids = 0.5 * (ba_bins[:-1] + ba_bins[1:])
+    ba_pdf_pred = jnp.interp(x_target, ba_binmids, ba_pdf_pred_table)
+
+    return ba_pdf_pred
+
+
+@jjit
+def _pred_ellipse_samples_kern(bulge_params, bulge_key, mu_ran, phi_ran):
+    ngals = mu_ran.shape[0]
+    axis_ratios = eshape.sample_bulge_axis_ratios(bulge_key, ngals, bulge_params)
+    a = jnp.ones_like(mu_ran)
+    b = a * axis_ratios.b_over_a
+    c = a * axis_ratios.c_over_a
+    ellipse2d = eproj.compute_ellipse2d(a, b, c, mu_ran, phi_ran)
+    return ellipse2d
+
+
+@jjit
+def _nondiff_pred_kern(bulge_params, pred_data):
+    ran_key, mu_ran, phi_ran, ba_bins, x_target = pred_data
+
+    ran_key, bulge_key = jran.split(ran_key, 2)
+    ellipse2d = _pred_ellipse_samples_kern(bulge_params, bulge_key, mu_ran, phi_ran)
+
+    alpha_over_beta = ellipse2d.alpha / ellipse2d.beta
+    ba_pdf_table, __ = jnp.histogram(alpha_over_beta, ba_bins, density=True)
+    ba_binmids = 0.5 * (ba_bins[:-1] + ba_bins[1:])
+    ba_pdf_pred = jnp.interp(x_target, ba_binmids, ba_pdf_table)
+
+    return ba_pdf_pred
+
+
+loss_and_grad_kern = jjit(value_and_grad(loss_kern, argnums=0))
+
+
+if __name__ == "__main__":
+    import argparse
+
+    from diffsky.mass_functions.fitting_utils.fitting_helpers import jax_adam_wrapper
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-seed", help="Integer input to jax.random.key(seed)", default=0, type=int
+    )
+    parser.add_argument(
+        "-ngals", help="Number of galaxies to estimate pdf", default=50_000, type=int
+    )
+    parser.add_argument(
+        "-tdata_fname",
+        help="Number of galaxies to estimate pdf",
+        default="ellipsoid_b_over_a_pdf_rodriguez_padilla_2013.txt",
+    )
+    args = parser.parse_args()
+
+    VariedParams = namedtuple("VariedParams", ("ba_peak", "ba_sigma", "c_min"))
+    varied_params = VariedParams._make(
+        [getattr(eshape.DEFAULT_BULGE_PARAMS, key) for key in VariedParams._fields]
+    )
+
+    ran_key = jran.key(args.seed)
+    ran_key, loss_key = jran.split(ran_key, 2)
+
+    target_data = np.loadtxt(args.tdata_fname, delimiter=",")
+    X = target_data[:, 0]
+    Y = target_data[:, 1]
+
+    ba_bins = np.linspace(0.01, 0.99, 50)
+
+    ran_key, mu_key, phi_key = jran.split(ran_key, 3)
+    mu_ran = jran.uniform(mu_key, minval=-1, maxval=1, shape=(args.ngals,))
+    phi_ran = jran.uniform(phi_key, minval=0, maxval=2 * np.pi, shape=(args.ngals,))
+
+    default_params = eshape.DEFAULT_BULGE_PARAMS
+    loss_data = X, Y, loss_key, mu_ran, phi_ran, ba_bins, default_params
+    args = loss_and_grad_kern, varied_params, loss_data, 100
+    _res = jax_adam_wrapper(*args)
+    p_best, loss, loss_hist = _res[:3]
+    bulge_params_best = default_params._replace(**p_best._asdict())
+    print(f"\nInitial parameters:\n{default_params}")
+
+    print(f"\nBest-fitting parameters:\n{bulge_params_best}")

diffsky-0.2.5/diffsky/ellipsoidal_shapes/bulge_shapes.py

@@ -0,0 +1,43 @@
+"""Ellipsoidal model for PDF of bulge shapes"""
+
+from collections import namedtuple
+from functools import partial
+
+from jax import jit as jjit
+from jax import numpy as jnp
+from jax import random as jran
+
+from .pdf_model_utils import truncated_normal_sample
+
+AxisRatios = namedtuple("AxisRatios", ("b_over_a", "c_over_a"))
+BulgeAxisRatioParams = namedtuple(
+    "BulgeAxisRatioParams", ("ba_peak", "ba_sigma", "ba_min", "ba_max", "c_min")
+)
+DEFAULT_BULGE_PARAMS = BulgeAxisRatioParams(
+    ba_peak=1.3, ba_sigma=0.38, ba_min=0.3, ba_max=1.0, c_min=0.61
+)
+
+
+@partial(jjit, static_argnames=["n_samples"])
+def sample_bulge_axis_ratios(ran_key, n_samples, bulge_params):
+    """Draw samples of axis ratios for bulges"""
+    ran_key_b, ran_key_c = jran.split(ran_key, 2)
+
+    b_over_a = truncated_normal_sample(
+        ran_key_b,
+        (n_samples,),
+        bulge_params.ba_peak,
+        bulge_params.ba_sigma,
+        bulge_params.ba_min,
+        bulge_params.ba_max,
+    )
+
+    c_over_b = jran.uniform(
+        ran_key_c, (n_samples,), minval=bulge_params.c_min, maxval=1.0
+    )
+
+    c_over_a = c_over_b * b_over_a
+    c_over_a = jnp.minimum(c_over_a, b_over_a * 0.99)  # Ensure c <= b for safety
+
+    axis_ratios = AxisRatios(b_over_a, c_over_a)
+    return axis_ratios

diffsky-0.2.5/diffsky/ellipsoidal_shapes/disk_opt.py

@@ -0,0 +1,130 @@
+"""Functions and script for optimizing parameters of disk_shapes model"""
+
+from collections import namedtuple
+
+import numpy as np
+from jax import jit as jjit
+from jax import numpy as jnp
+from jax import random as jran
+from jax import value_and_grad
+
+from .. import diffndhist
+from . import disk_shapes as dshape
+from . import ellipse_proj_kernels as eproj
+
+
+@jjit
+def _mae(pred, target):
+    diff = pred - target
+    return jnp.mean(jnp.abs(diff))
+
+
+@jjit
+def get_all_params_from_varied(varied_params, default_params):
+    all_params = default_params._replace(**varied_params._asdict())
+    return all_params
+
+
+@jjit
+def loss_kern(varied_params, loss_data):
+    x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins, default_params = (
+        loss_data
+    )
+    disk_params = get_all_params_from_varied(varied_params, default_params)
+    loss = _loss_kern(
+        disk_params, x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins
+    )
+    return loss
+
+
+@jjit
+def _loss_kern(disk_params, x_target, ba_pdf_target, ran_key, mu_ran, phi_ran, ba_bins):
+    pred_data = x_target, ran_key, mu_ran, phi_ran, ba_bins
+
+    ba_pdf_pred = _pred_kern(disk_params, pred_data)
+    loss = _mae(ba_pdf_pred, ba_pdf_target)
+    return loss
+
+
+@jjit
+def _pred_kern(disk_params, pred_data):
+    x_target, ran_key, mu_ran, phi_ran, ba_bins = pred_data
+
+    ran_key, disk_key = jran.split(ran_key, 2)
+    ellipse2d = _pred_ellipse_samples_kern(disk_params, disk_key, mu_ran, phi_ran)
+
+    nddata = jnp.array(ellipse2d.alpha / ellipse2d.beta).reshape((-1, 1))
+    ndsig = jnp.zeros_like(nddata) + 0.05
+
+    ndbins_lo = ba_bins[:-1].reshape((-1, 1))
+    ndbins_hi = ba_bins[1:].reshape((-1, 1))
+
+    wcounts = diffndhist.tw_ndhist(nddata, ndsig, ndbins_lo, ndbins_hi)
+    ba_pdf_pred_table = (wcounts / wcounts.sum()) / jnp.diff(ba_bins)
+
+    ba_binmids = 0.5 * (ba_bins[:-1] + ba_bins[1:])
+    ba_pdf_pred = jnp.interp(x_target, ba_binmids, ba_pdf_pred_table)
+
+    return ba_pdf_pred
+
+
+loss_and_grad_kern = jjit(value_and_grad(loss_kern, argnums=0))
+
+
+@jjit
+def _pred_ellipse_samples_kern(disk_params, disk_key, mu_ran, phi_ran):
+    ngals = mu_ran.shape[0]
+    axis_ratios = dshape.sample_disk_axis_ratios(disk_key, ngals, disk_params)
+    a = jnp.ones_like(mu_ran)
+    b = a * axis_ratios.b_over_a
+    c = a * axis_ratios.c_over_a
+    ellipse2d = eproj.compute_ellipse2d(a, b, c, mu_ran, phi_ran)
+    return ellipse2d
+
+
+if __name__ == "__main__":
+    import argparse
+
+    from diffsky.mass_functions.fitting_utils.fitting_helpers import jax_adam_wrapper
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-seed", help="Integer input to jax.random.key(seed)", default=0, type=int
+    )
+    parser.add_argument(
+        "-ngals", help="Number of galaxies to estimate pdf", default=50_000, type=int
+    )
+    parser.add_argument(
+        "-tdata_fname",
+        help="Number of galaxies to estimate pdf",
+        default="spiral_b_over_a_pdf_rodriguez_padilla_2013.txt",
+    )
+    args = parser.parse_args()
+
+    VariedParams = namedtuple("VariedParams", ("c_min", "c_max"))
+    varied_params = VariedParams._make(
+        [getattr(dshape.DEFAULT_DISK_PARAMS, key) for key in VariedParams._fields]
+    )
+
+    ran_key = jran.key(args.seed)
+    ran_key, loss_key = jran.split(ran_key, 2)
+
+    target_data = np.loadtxt(args.tdata_fname, delimiter=",")
+    X = target_data[:, 0]
+    Y = target_data[:, 1]
+
+    ba_bins = np.linspace(0.01, 0.99, 50)
+
+    ran_key, mu_key, phi_key = jran.split(ran_key, 3)
+    mu_ran = jran.uniform(mu_key, minval=-1, maxval=1, shape=(args.ngals,))
+    phi_ran = jran.uniform(phi_key, minval=0, maxval=2 * np.pi, shape=(args.ngals,))
+
+    default_params = dshape.DEFAULT_DISK_PARAMS
+    loss_data = X, Y, loss_key, mu_ran, phi_ran, ba_bins, default_params
+    args = loss_and_grad_kern, varied_params, loss_data, 100
+    _res = jax_adam_wrapper(*args)
+    p_best, loss, loss_hist = _res[:3]
+    disk_params_best = default_params._replace(**p_best._asdict())
+    print(f"\nInitial parameters:\n{default_params}")
+
+    print(f"\nBest-fitting parameters:\n{disk_params_best}")

diffsky-0.2.5/diffsky/ellipsoidal_shapes/disk_shapes.py

@@ -0,0 +1,34 @@
+"""Ellipsoidal model for PDF of disk shapes"""
+
+from collections import namedtuple
+from functools import partial
+
+from jax import jit as jjit
+from jax import numpy as jnp
+from jax import random as jran
+
+AxisRatios = namedtuple("AxisRatios", ("b_over_a", "c_over_a"))
+DiskAxisRatioParams = namedtuple(
+    "DiskAxisRatioParams", ("ba_min", "ba_max", "c_min", "c_max")
+)
+DEFAULT_DISK_PARAMS = DiskAxisRatioParams(ba_min=0.8, ba_max=1.0, c_min=0.2, c_max=0.5)
+
+
+@partial(jjit, static_argnames=["n_samples"])
+def sample_disk_axis_ratios(ran_key, n_samples, disk_params):
+    """Draw samples of axis ratios for disks"""
+    ran_key_b, ran_key_c = jran.split(ran_key, 2)
+
+    b_over_a = jran.uniform(
+        ran_key_b, (n_samples,), minval=disk_params.ba_min, maxval=disk_params.ba_max
+    )
+
+    c_over_b = jran.uniform(
+        ran_key_c, (n_samples,), minval=disk_params.c_min, maxval=disk_params.c_max
+    )
+
+    c_over_a = c_over_b * b_over_a
+    c_over_a = jnp.minimum(c_over_a, b_over_a * 0.99)  # Ensure c <= b for safety
+
+    axis_ratios = AxisRatios(b_over_a, c_over_a)
+    return axis_ratios