Furax: A Modular JAX Framework for Linear Operators in Astrophysical and Cosmological Data Analysis
Pierre Chanial, Simon Biquard, Wassim Kabalan, Wuhyun Sohn, Artem Basyrov, Benjamin Beringue, Alexandre Boucaud, Andréa Landais, Magdy Morshed, Radek Stompor, Ema Tsang King Sang, Amalia Villarrubia-Aguilar, Josquin Errard
Abstract
The Framework for Unified and Robust data Analysis with JAX (Furax) is an open-source Python framework for modeling data acquisition systems and solving inverse problems in astrophysics and cosmology. Built on JAX, Furax provides composable building blocks in the form of general-purpose and domain-specific linear operators, along with preconditioners and solvers for their numerical inversion. Domain-specific tools are provided for astrophysical and cosmic microwave background (CMB) data analysis$-$including map-making, instrument modeling, and astrophysical component separation$-$with a modular architecture designed to extend to other fields.