arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2606.01547astro-ph.IMastro-ph.CO
收藏

JAXtronomy: lenstronomy 的 JAX 移植

JAXtronomy: A JAX port of lenstronomy

Stony Brook University(石溪大学);Laboratoire Univers et Particules de Montpellier, CNRS and Université de Montpellier (UMR-5299)(蒙彼利埃宇宙与粒子实验室,法国国家科学研究中心和蒙彼利埃大学(UMR-5299));Max-Planck-Institut für Astrophysik(马克斯·普朗克天体物理学研究所);Technical University of Munich(慕尼黑工业大学);University of Chicago(芝加哥大学);Kavli Institute for Cosmological Physics(卡弗里宇宙物理学研究所);Independent University, Bangladesh(孟加拉独立大学);University of Illinois(伊利诺伊大学)

Alan Huang, Simon Birrer, Daniel Gilman, Natalie B. Hogg, Anowar J. Shajib, Aymeric Galan, Nan Zhang

AI总结:

本文介绍 JAXtronomy,一个基于 JAX 重新实现的引力透镜软件包 lenstronomy,通过即时编译和自动微分等技术提升性能并支持 GPU 加速。

摘要:

引力透镜是一种光线在巨大物体周围弯曲的现象,导致观测者看到扭曲的图像。研究引力透镜系统为宇宙学和天体物理学提供了见解,包括对宇宙膨胀速率和暗物质分布的限制。因此,我们引入了 JAXtronomy,这是使用 JAX (Bradbury et al., 2018) 对引力透镜软件包 lenstronomy (Birrer, 2021; Birrer & Amara, 2018) 的重新实现。JAX 是一个 Python 库,它使用加速线性代数 (XLA) 编译器来提高计算软件的性能。JAXtronomy 的核心设计原则是保持与 lenstronomy 相同的 API。JAXtronomy 中利用的主要 JAX 特性是即时编译,这可以显著减少执行时间,以及自动微分,这允许实现以前不可能的基于梯度的算法。此外,JAX 允许代码在 GPU 上运行或在 CPU 核心上并行化,进一步提升了 JAXtronomy 的性能。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑