pysersic: A Python package for determining galaxy structural properties via Bayesian inference, accelerated with jax

A standard practice in extragalactic population studies is the fitting of parametric models to galaxy images. From such fits, key structural parameters of galaxies such as total flux and effective radius (size) can be extracted. One of the most popular parametric forms is that of the S\'ersic profile, which is flexible enough to reasonably fit the light distribution of nearly all galaxies. Here we present pysersic, a Bayesian framework created to facilitate the inference of structural parameters from galaxy images. Pysersic is written in pure Python, and is built using the jax framework, allowing for just-in-time compilation, auto-differentiation and seamless execution on CPUs, GPUs or TPUs. Inference is performed with the numpyro package using gradient based methods, e.g., No U-Turn Sampling, for efficient and robust posterior estimation in only a few minutes on a modern laptop. Pysersic is designed to have a user-friendly interface, allowing users to fit single or multiple sources in a few lines of code, while also being flexible enough for integration into current and future analysis pipelines. In addition to sampling, pysersic can produce point estimates of the best model via optimization in several seconds, and approximate the posterior via stochastic variational inference. The use of the numpyro probabilistic language provides future extensibility to arbitrary models beyond the S\'ersic.

where the total flux,  total , half-light radius,   and Sérsic index,  are the parameters of interest to be fit and subsequently used to characterize a galaxy's morphology.
Here we present pysersic, a Bayesian framework created to facilitate the inference of structural parameters from galaxy images.It is written in pure Python, and built using the jax framework (Bradbury et al., 2018) allowing for just-in-time (JIT) compilation, auto-differentiation and seamless execution on CPUs, GPUs or TPUs.Inference is performed with the numpyro (Bingham et al., 2019;Phan et al., 2019) package utilizing gradient based methods, e.g., No U-Turn Sampling (NUTS) (Hoffman et al., 2014), for efficient and robust posterior estimation.pysersic was designed to have a user-friendly interface, allowing users to fit single or multiple sources in a few lines of code.It was also designed to scale to many images, such that it can be seamlessly integrated into current and future analysis pipelines.

Statement of need
Parametric profile fitting has become a ubiquitous and essential tool for numerous applications including measuring the photometry -or total flux -of galaxies, as well as the investigation of the structural evolution of galaxies over cosmic time (Kawinwanichakij et al., 2021;Lange et al., 2015;Mowla et al., 2019).This approach allows one to both extrapolate galaxy surface brightness profiles beyond the noise limit of images, as well as account for the PSF to accurately measure the structure of galaxies near the resolution limit of those images.The empirically derived Sérsic profile is the most common parametric form for the surface-brightness profile as it provides a reasonable approximation to nearly all galaxies, given the additional freedom of the Sérsic index, , over fixed-index profiles.
Given the long history of Sérsic fitting codes with many available tools, the development of pysersic was largely motivated by two related factors, first and foremost of which was the desire to implement Sérsic fitting in a fully Bayesian context at speed.The ability to place the typical Sérsic fitting problem into a Bayesian context with runtimes that are not prohibitive (the traditional drawback of MCMC methods) has recently been unlocked by the second motivation: to leverage the jax library.jax utilizes JIT compilation to decrease computational runtimes, provides seamless integration with hardware accelerators such as GPUs and TPUs for further improvements in performance, and enables automatic differentiation, facilitating gradient based optimization and sampling methods.Together, these features greatly increase speed and efficiency, especially when sampling or optimizing a large number of parameters.
Inference in pysersic is implemented using the numpyro probabilistic programming language (PPL).This allows for total control over the priors and methods used for inference.The numpyro package utilizes jax's auto-differentiation capabilities for gradient based samplers such as Hamiltonian Monte Carlo (HMC) and No-U-Turn-Sampling (NUTS).In addition, there are recently-developed techniques for posterior estimation, including variational inference (Ranganath et al., 2014) utilizing normalizing flows (De Cao et al., 2020).These techniques dramatically reduce the number of likelihood calls required to provide accurate estimates of the posterior relative to gradient-free methods.Combined with the jax's JIT compilation, posteriors can now be generated in a few minutes or less on modern laptops.

Code Description
pysersic was designed to have a user-friendly API with sensible defaults.Tools are provided to automatically generate priors for all free parameters based on an initial characterization of a given image -but can also easily be set manually.We provide default inference routines for NUTS MCMC and variational inference using neural flows.Users can access the underlying numpyro model if desired, to perform inference using any tools available within the numpyro ecosystem.The goal for pysersic is to provide reasonable defaults for new users interested in a handful of galaxies, yet maintain the ability for advanced users to tweak options as necessary to perform inference for entire surveys.
A crucial component of any Sérsic fitting code is an efficient and accurate rendering algorithm.Sérsic profiles with high index,  ≳ 3 are notoriously difficult to render accurately given the steep increase in brightness as  → 0. In pysersic, the rendering module is kept separate from the frontend API and inference modules, such that different algorithms can be interchanged and therefore easily tested (and hopefully encourage innovation as well).In this initial release, we provide three algorithms.The first is a traditional rendering algorithm in which the intrinsic profile is rendered in real space, with oversampling in the center to ensure accurate results for high index profiles.The second and third methods render the profiles in Fourier space, providing accurate results even for strongly peaked profiles and avoiding artifacts due to pixelization.In pysersic, this is achieved by representing the profiles using a series of Gaussian following the algorithm presented in Shajib (2019).We include one algorithm that is fully based in Fourier space, along with a version of the hybrid real-Fourier algorithm introduced in Lang (2020) which helps avoid some of the aliasing present when rendering solely in Fourier space.

Related Software
There is a long history and many software tools designed for Sérsic profile fitting.Some of the most popular libraries are listed below.