JAXtronomy: A JAX port of lenstronomy

Abstract

Gravitational lensing is a phenomenon where light bends around massive objects, resulting in distorted images seen by an observer. Studying gravitationally lensed systems provides insights into cosmology and astrophysics, including constraints of the expansion rate of the Universe and the distribution of dark matter. Thus, we introduce JAXtronomy, a re-implementation of the gravitational lensing software package lenstronomy (Birrer, 2021; Birrer & Amara, 2018) using JAX (Bradbury et al., 2018). JAX is a Python library that uses an accelerated linear algebra (XLA) compiler to improve the performance of computing software. Our core design principle of JAXtronomy is to maintain an identical API to that of lenstronomy. The main JAX features utilized in JAXtronomy are just-in-time compilation, which can lead to significant reductions in execution time, and automatic differentiation, which allows for the implementation of gradient-based algorithms that were previously impossible. Additionally, JAX allows code to be run on GPUs or parallelized across CPU cores, further boosting the performance of JAXtronomy.

Published by JOSS, software available at https://github.com/lenstronomy/JAXtronomy

References in corpus (16)

• SciPy 1.0--Fundamental Algorithms for Scientific Computing in Python
• Array Programming with NumPy
• TDCOSMO IV: Hierarchical time-delay cosmography -- joint inference of the Hubble constant and galaxy density profiles
• Warm dark matter chills out: constraints on the halo mass function and the free-streaming length of dark matter with 8 quadruple-image strong gravitational lenses
• TDCOSMO. I. An exploration of systematic uncertainties in the inference of $H_0$ from time-delay cosmography
• The mass relations between supermassive black holes and their host galaxies at 1<z<2 with HST-WFC3
• Dark matter halos of massive elliptical galaxies at $z\sim 0.2$ are well described by the Navarro-Frenk-White profile
• TDCOSMO V: strategies for precise and accurate measurements of the Hubble constant with strong lensing
• TDCOSMO. IX. Systematic comparison between lens modelling software programs: time delay prediction for WGD 2038-4008
• Project Dinos I: A joint lensing-dynamics constraint on the deviation from the power law in the mass profile of massive ellipticals
• TDCOSMO III: Dark matter substructure meets dark energy -- the effects of (sub)halos on strong-lensing measurements of $H_0$
• From Images to Dark Matter: End-To-End Inference of Substructure From Hundreds of Strong Gravitational Lenses
• TDCOSMO 2025: Cosmological constraints from strong lensing time delays
• GIGA-Lens: Fast Bayesian Inference for Strong Gravitational Lens Modeling
• Project Dinos II: Redshift evolution of dark and luminous matter density profiles in strong-lensing elliptical galaxies across 0.1 < z < 0.9
• TDCOSMO: XX. WFI2033--4723, the First Quadruply-Imaged Quasar Modeled with JWST Imaging