diffsky 0.2.4__tar.gz → 0.3.0__tar.gz

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

@@ -54,7 +54,6 @@ jobs:
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffmah.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffstar.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/dsps.git
-          pip install --no-deps git+https://${GITHUB_USER}:${GITHUB_TOKEN}@github.com/ArgonneCPAC/diffstarpop.git@main
           python -m pip install --no-build-isolation --no-deps -e .
 
       - name: test that no warnings are raised

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

@@ -53,7 +53,6 @@ jobs:
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffmah.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffstar.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/dsps.git
-          pip install --no-deps git+https://${GITHUB_USER}:${GITHUB_TOKEN}@github.com/ArgonneCPAC/diffstarpop.git@main
           python -m pip install --no-build-isolation --no-deps -e .
 
       - name: test

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

@@ -57,7 +57,6 @@ jobs:
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffmah.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/diffstar.git
           pip install --no-deps git+https://github.com/ArgonneCPAC/dsps.git
-          pip install --no-deps git+https://${GITHUB_USER}:${GITHUB_TOKEN}@github.com/ArgonneCPAC/diffstarpop.git@main
           python -m pip install --no-build-isolation --no-deps -e .
 
       - name: test

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

@@ -1,3 +1,13 @@
+0.3.0 (2025-10-03)
+-------------------
+- Update diffsky to be compatible with diffstar v1.0
+
+
+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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diffsky
-Version: 0.2.4
+Version: 0.3.0
 Summary: Library for differentiable generation of synthetic skies
 Author-email: Andrew Hearin <ahearin@anl.gov>
 License: BSD 3-Clause License
@@ -41,8 +41,7 @@ License-File: LICENSE.rst
 Requires-Dist: numpy
 Requires-Dist: jax
 Requires-Dist: diffmah>=0.7.2
-Requires-Dist: diffstar>=0.3.4
-Requires-Dist: diffstarpop>=0.0.2
+Requires-Dist: diffstar>=1.0.0
 Requires-Dist: dsps>=0.4.5
 Requires-Dist: h5py
 Dynamic: license-file
@@ -50,8 +49,8 @@ Dynamic: license-file
 diffsky
 ============
 
-Diffsky is a python library based on JAX for producing mock catalogs based on 
-`diffstar <https://diffstar.readthedocs.io/en/latest/>`_ 
+Diffsky is a python library based on JAX for producing mock catalogs based on
+`diffstar <https://diffstar.readthedocs.io/en/latest/>`_
 and `dsps <https://dsps.readthedocs.io/en/latest/>`_.
 
 Installation
@@ -77,18 +76,7 @@ For a typical development environment in conda-forge::
 Documentation
 -------------
 
-Online documentation for diffsky is available at 
+Online documentation for diffsky is available at
 `diffsky.readthedocs.io <https://diffsky.readthedocs.io/en/latest/>`_.
 
-Testing
--------
-To run the suite of unit tests::
-
-    $ cd /path/to/root/diffsky
-    $ pytest
-
-To build html of test coverage::
-
-    $ pytest -v --cov --cov-report html
-    $ open htmlcov/index.html
 

{diffsky-0.2.4 → diffsky-0.3.0}/README.rst

@@ -1,8 +1,8 @@
 diffsky
 ============
 
-Diffsky is a python library based on JAX for producing mock catalogs based on 
-`diffstar <https://diffstar.readthedocs.io/en/latest/>`_ 
+Diffsky is a python library based on JAX for producing mock catalogs based on
+`diffstar <https://diffstar.readthedocs.io/en/latest/>`_
 and `dsps <https://dsps.readthedocs.io/en/latest/>`_.
 
 Installation
@@ -28,18 +28,7 @@ For a typical development environment in conda-forge::
 Documentation
 -------------
 
-Online documentation for diffsky is available at 
+Online documentation for diffsky is available at
 `diffsky.readthedocs.io <https://diffsky.readthedocs.io/en/latest/>`_.
 
-Testing
--------
-To run the suite of unit tests::
-
-    $ cd /path/to/root/diffsky
-    $ pytest
-
-To build html of test coverage::
-
-    $ pytest -v --cov --cov-report html
-    $ open htmlcov/index.html
 

diffsky-0.3.0/diffsky/_version.py

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

{diffsky-0.2.4 → diffsky-0.3.0}/diffsky/data_loaders/hacc_utils/lc_mock_production.py

@@ -9,14 +9,16 @@ import h5py
 import numpy as np
 from diffmah import DEFAULT_MAH_PARAMS, logmh_at_t_obs
 from diffstar import DEFAULT_DIFFSTAR_PARAMS
-from diffstarpop import mc_diffstar_sfh_galpop
-from diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
-from diffstarpop.param_utils import mc_select_diffstar_params
+from diffstar.diffstarpop import mc_diffstar_sfh_galpop
+from diffstar.diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
+from diffstar.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.4 → diffsky-0.3.0}/diffsky/diagnostics/plot_delta_mag_burstiness.py

@@ -5,8 +5,8 @@ import os
 import numpy as np
 from diffmah.diffmahpop_kernels.bimod_censat_params import DEFAULT_DIFFMAHPOP_PARAMS
 from diffmah.diffmahpop_kernels.mc_bimod_cens import mc_cenpop
-from diffstarpop import mc_diffstar_sfh_galpop
-from diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
+from diffstar.diffstarpop import mc_diffstar_sfh_galpop
+from diffstar.diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
 from dsps.constants import T_TABLE_MIN
 from dsps.cosmology import DEFAULT_COSMOLOGY
 from dsps.cosmology.flat_wcdm import _age_at_z_kern, age_at_z0

{diffsky-0.2.4 → diffsky-0.3.0}/diffsky/diagnostics/plot_diffstar_fq.py

@@ -1,8 +1,7 @@
-"""
-"""
+""" """
 
 import numpy as np
-from diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
+from diffstar.diffstarpop.defaults import DEFAULT_DIFFSTARPOP_PARAMS
 from dsps.cosmology import flat_wcdm
 from dsps.cosmology.defaults import DEFAULT_COSMOLOGY
 from jax import random as jran

diffsky-0.3.0/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.3.0/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.3.0/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.3.0/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