[论文解读] pysersic: A Python package for determining galaxy structural properties via Bayesian inference, accelerated with jax

Imad Pasha, Tim B. Miller|arXiv (Cornell University)|Jun 8, 2023

Galaxies: Formation, Evolution, Phenomena被引用 4

一句话总结

pysersic 是一个用于使用贝叶斯推断对星系图像拟合 Sérsic 分布的 Python 包，通过 JAX 加速，支持 GPU/TPU，实现快速、可扩展且精确的结构参数估计。它利用 NUTS 采样和变分推断实现稳健的后验估计，同时提供多种高精度渲染方法，包括实空间和基于傅里叶的算法。

ABSTRACT

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érsic 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érsic.

研究动机与目标

• 开发一个完整的贝叶斯框架，用于以高计算效率对星系图像拟合 Sérsic 分布。
• 利用 JAX 的 JIT 编译和硬件加速（CPU/GPU/TPU）加速星系结构参数的贝叶斯推断。
• 提供一个用户友好且可扩展的 API，支持初学者用户和大规模调查中的高级工作流。
• 实现准确且灵活的 Sérsic 分布渲染算法，尤其适用于高 Sérsic 指数（n ≳ 3）的情况。
• 与概率编程工具（如 NumPyro）无缝集成，实现对先验和推断方法的完全控制。

提出的方法

• 使用 JAX 实现 JIT 编译、自动微分和硬件加速，以加快似然评估和采样速度。
• 采用 NumPyro 进行基于梯度的推断，包括无偏随机抽样（NUTS）和基于归一化流的变分推断。
• 提供三种不同的渲染算法：实空间过采样、全傅里叶空间方法和混合实空间-傅里叶高斯混合方法。
• 根据图像特征自动生成先验，同时允许专家手动覆盖以实现控制。
• 将渲染模块与推断流程分离，支持新算法的即插即用式测试。
• 支持单源和多源拟合，提供简洁直观的 API，便于快速部署。

实验结果

研究问题

• RQ1 能否使 Sérsic 分布拟合的贝叶斯推断在准确性和效率上均满足大规模星系调查的需求？
• RQ2 JAX 加速对星系测光中 MCMC 和变分推断的运行时间性能有何影响？
• RQ3 针对高 Sérsic 指数分布（n ≳ 3）的最准确且稳定的渲染技术是什么？
• RQ4 能否构建一个统一且用户友好的接口，同时支持快速分析和复杂的大规模调查工作流？
• RQ5 不同推断方法（NUTS 与归一化流）在收敛速度和后验准确性方面有何比较？

主要发现

• 得益于 JAX 的 JIT 编译和高效的基于梯度的采样，pysersic 可在标准笔记本电脑上于几分钟内完成后验估计。
• 使用 NUTS 和归一化流可显著减少所需的似然评估次数，大幅提升采样效率。
• 混合实空间-傅里叶渲染方法能有效减轻混叠伪影，同时保持对尖锐 Sérsic 分布的精度。
• 通过专用数值算法，该软件包可实现对高指数 Sérsic 分布（n ≳ 3）的高精度拟合。
• 通过与 NumPyro 生态系统的集成，用户可完全控制先验和采样方法，实现稳健且可重现的推断。
• 模块化设计使得渲染算法的基准测试和替换变得容易，促进了方法论的创新。