sedpy_jax

JAX-accelerated SED fitting tools, filter handling, and broadband photometry projection for astronomy.

Built as a jaxified, GPU-ready reimplementation of sedpy by Ben Johnson, to support differentiable, scalable spectral energy distribution (SED) fitting — ideal for simulation-based inference, gradient-based samplers (e.g. HMC, NUTS), and high-throughput photometric pipelines.

---

Features

• Fully JAX-compatible filter operations and broadband projection
• Clean, modular implementation of the Filter and FilterSet classes
• Easy visualization of filter transmission curves with colormap scaling
• Accurate AB and Vega magnitude computation
• Utility functions for wavelength conversion, rebinning, and interpolation
• Built-in effective wavelength, pivot wavelength, width, and zero-point computation

---

Installation

From GitHub (recommended for users)

pip install git+https://github.com/Espe13/sedpy_jax.git

For development (recommended for contributors)

git clone https://github.com/Espe13/sedpy_jax.git
cd sedpy_jax
pip install -e .

---

Example: Using FilterSet to Project SEDs

from sedpy_jax.observate import FilterSet
import jax.numpy as jnp
import matplotlib.pyplot as plt

Load a set of filters

fset = FilterSet(['sdss_u0', 'sdss_g0', 'sdss_r0', 'sdss_i0', 'sdss_z0'])

Plot transmission curves

fset.display(normalize=False, colormap='coolwarm')
plt.show()

Project an SED onto filters

# --- Generate toy spectrum ---
wave = jnp.linspace(3000, 10000, 500)
wave_mask = (wave >= fset.lam.min()) & (wave <= fset.lam.max())
wave = wave[wave_mask]
flux = jnp.exp(-0.5 * ((wave - 6000) / 800)**2)
maggies = fset.get_sed_maggies(flux * 10**-8.5, sourcewave=wave)

# --- Create figure and use shared axis for filters + spectrum ---
fig, ax = plt.subplots()

# Plot transmission curves in background
fset.display(ax=ax)

# Plot toy SED and photometry points
ax.plot(wave, flux, color="black", linewidth=1, label="Toy SED")
ax.plot([f.wave_effective for f in fset.filters], maggies, markersize = 10,
         color="red", marker="s", linestyle="", label="Observed Photometry")

# --- Final touches ---
ax.set_xlabel("Wavelength [Å]")
ax.set_ylabel("Flux & Transmission")
ax.set_title("Toy SED with Filter Transmission and Photometry")
ax.legend()
plt.show()

---

Table: Filter Properties

print(fset)

yields

                  FilterSet Summary                  
┏━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┓
┃ Index ┃ Name    ┃ λ_eff [Å] ┃ Width [Å] ┃ AB→Vega ┃
┡━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━┩
│ 0     │ sdss_u0 │ 3546.0    │ 558.6     │ -0.9305 │
│ 1     │ sdss_g0 │ 4669.5    │ 1159.9    │ 0.1003  │
│ 2     │ sdss_r0 │ 6155.6    │ 1112.1    │ -0.1423 │
│ 3     │ sdss_i0 │ 7471.5    │ 1070.4    │ -0.3560 │
│ 4     │ sdss_z0 │ 8917.2    │ 1126.8    │ -0.5176 │
└───────┴─────────┴───────────┴───────────┴─────────┘

while printing a single filter gives a little more information:

print(fset.filters[0])

                               Filter: sdss_u0                                
┏━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━┓
┃ Name    ┃ Nickname ┃ λ_eff [Å] ┃ Pivot λ [Å] ┃ Width [Å] ┃ AB→Vega ┃ N pts ┃
┡━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━┩
│ sdss_u0 │ sdss_u0  │ 3546.0    │ 3556.5      │ 558.6     │ -0.9305 │ 47    │
└─────────┴──────────┴───────────┴─────────────┴───────────┴─────────┴───────┘

---

Included Tools

•	Filter: load and manipulate individual transmission curves
•	FilterSet: project SEDs across many filters efficiently
•	air2vac / vac2air: wavelength conversion (Allen/Morton standards)
•	rebin: flux-conserving rebinning of filters
•	getSED(): compute AB mags or maggies from source spectra

---

Tests

To run tests:

pytest tests/

---

Citation

If you use this package, please cite:

@misc{stoffers2025sedpyjax,
  author       = {Amanda Stoffers, Benjamin D. Johnson},
  title        = {sedpy\_jax: A JAX-compatible filter and SED utility library},
  year         = {2025},
  publisher    = {GitHub},
  journal      = {GitHub repository},
  howpublished = {\url{https://github.com/Espe13/sedpy_jax}},
  note         = {Version 1.0},
}

Feedback & Contributions

Feel free to open Issues or pull requests.

Created and maintained by Amanda Stoffers (@Espe13)

---

License

MIT License. See LICENSE.