🔭 astrophysics

Furax: A Modular JAX Framework for Linear Operators in Astrophysical and Cosmological Data Analysis

Furax is an open-source, JAX-based Python framework that offers a modular architecture of composable linear operators, preconditioners, and solvers to model data acquisition systems and solve inverse problems specifically for astrophysical and cosmological applications like CMB analysis.

Original authors: 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-A

Published 2026-03-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: 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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to listen to a faint, beautiful symphony (the Universe) being played in a very noisy, crowded room (the Earth's atmosphere and our instruments). You have thousands of microphones recording the sound, but the recordings are a mess of static, overlapping voices, and the microphones themselves are slightly out of tune.

Your goal is to clean up the recording, figure out exactly where each instrument is in the room, and separate the violins from the drums. This is what cosmologists do with data from telescopes looking at the Cosmic Microwave Background (CMB)—the oldest light in the universe.

Enter Furax.

What is Furax?

Think of Furax as a super-smart, modular Lego kit for building data analysis tools. Before Furax, building these tools was like trying to build a house by hand-carving every single brick from a different type of wood, often using old, rusty tools. If you wanted to change the design, you had to start over.

Furax changes the game by providing:

Universal Bricks (Linear Operators): Instead of hard-coding every step, Furax gives you pre-made, high-quality Lego bricks. You can snap them together to build complex systems.
A Magic Blueprint (JAX): It's built on a modern technology called JAX, which is like having a blueprint that automatically figures out the most efficient way to build your house, even if you have a massive construction crew (GPUs/TPUs) ready to work.

How Does It Work? (The Analogy)

1. The "Lego" Approach (Modular Operators)

In the past, if a scientist wanted to model how a telescope sees the sky, they had to write thousands of lines of code from scratch. With Furax, they just grab a "Pointing Brick" (to know where the telescope is looking), a "Noise Brick" (to model the static), and a "Polarization Brick" (to measure the light's direction).

They snap them together like this:

Telescope View + Noise Filter + Sky Map = The Solution

Because these are "composable," scientists can swap out a "Noise Brick" for a more complex one without breaking the whole machine. It's like swapping a standard lightbulb for a smart bulb without rewiring the whole house.

2. The "Super-Speed" Engine (JAX & GPUs)

Processing data from modern telescopes is like trying to drink from a firehose. The data is massive (Terabytes!).

Old Way: A single worker (CPU) trying to drink the hose, taking days or weeks.
Furax Way: It uses JAX, which is like having a team of thousands of workers (GPUs/TPUs) working in perfect sync. Furax automatically organizes the work so the "firehose" is emptied in minutes.

3. The "Mathematical Translator" (Differentiability)

This is the coolest part. Furax doesn't just solve the puzzle; it understands how it solved it.
Imagine you are trying to find the perfect recipe for a cake.

Old Way: You bake a cake, taste it, guess what to change, and bake again.
Furax Way: It's like having a recipe that tells you, "If you add 1% more sugar, the cake will be 2% sweeter." Because Furax is "differentiable," it can use Neural Networks (AI) to automatically tweak the analysis parameters to find the perfect "cake" (the most accurate map of the universe) instantly.

Why Do We Need This?

We are about to launch new, incredibly powerful telescopes (like the Simons Observatory and LiteBIRD). They will generate so much data that old software will drown in it.

Furax is the lifeboat.
It is experiment-agnostic, meaning it doesn't care if you are using a telescope in Chile or Antarctica; it works for everyone.
It allows scientists to prototype new ideas quickly. Instead of spending a year building a tool, they can build a prototype in a day, test it, and if it works, scale it up to handle the massive data.

The Bottom Line

Furax is a modern, flexible, and incredibly fast toolkit that helps scientists turn the chaotic noise of the universe into a clear, beautiful picture of our cosmic origins. It replaces the "duct tape and baling wire" of old data analysis with a sleek, automated, and intelligent assembly line, ensuring we don't miss any of the universe's secrets when the next generation of telescopes comes online.

Based on the provided paper, here is a detailed technical summary of Furax: A Modular JAX Framework for Linear Operators in Astrophysical and Cosmological Data Analysis.

1. Problem Statement

Modern and future Cosmic Microwave Background (CMB) experiments (e.g., Simons Observatory, South Pole Observatory, QUBIC, LiteBIRD) generate massive Time-Ordered Data (TOD) streams. The core challenge in analyzing this data is map-making: recovering the sky signal ( $m$ ) from noisy observations ( $d$ ) while separating astrophysical components and removing systematic effects.

This is mathematically modeled as a linear inverse problem:
$d = Hm + Tt + n$
Where:

$H$ : The data acquisition operator (pointing, instrument response).
$T$ : Unwanted additive effects (systematics) to be deprojected.
$n$ : Noise.

Solving for the sky signal requires computing a generalized least-squares estimator:
$\hat{m} = (H^\top F_T^{-1} H)^{-1} H^\top F_T^{-1} d$
Current limitations in the field include:

Lack of Generality: Existing pipelines are often tied to specific experiments and do not survive hardware or paradigm shifts.
Performance Bottlenecks: Legacy frameworks (e.g., C++ based TOAST) often lack full support for differentiability and GPU acceleration.
Rigidity: Many tools are either end-to-end black boxes or limited to specific pipeline steps, hindering the exploration of novel statistical methods (e.g., Bayesian hierarchical modeling) and hybrid neural network approaches.

2. Methodology

Furax addresses these challenges by building a modular, experiment-agnostic framework entirely in Python and JAX (Just-in-Time compilation library).

Core Architecture

Composable Linear Operators: The framework treats data processing steps as composable linear operators implemented as Python dataclasses registered as JAX pytrees. This allows for intuitive mathematical composition (e.g., H = detector @ band_pass @ pointing).
Operator Algebra: Furax provides a rich algebraic interface supporting addition, composition (@), scalar multiplication, transposition (.T), and inversion (.I).
Automatic Differentiation & JIT: By leveraging JAX, Furax enables automatic differentiation and Just-In-Time compilation, allowing for gradient-based optimization and massive speedups on GPUs and TPUs.

Key Technical Components

Generic Operators: A suite of mathematical operators including Identity, Diagonal, FFT, Reshape, and specialized block operators (BlockDiagonal, BlockRow/Column) for handling multi-observation systems.
Domain-Specific Operators: Specialized tools for CMB physics, including:
- Instrument Modeling: Half-wave plate (HWPOperator), polarizers, and on-the-fly pointing matrices (PointingOperator) that expand time-dependent data dynamically to save memory.
- Astrophysical Components: Parametric spectral energy distribution (SED) operators for CMB, Dust, and Synchrotron, supporting spatially varying spectral indices.
- Polarization Handling: Dedicated JAX pytree types for Stokes parameters ( $I, Q, U, V$ ) supporting arithmetic and broadcasting.
Algebraic Reductions: A rule-based system simplifies operator chains before execution (e.g., combining consecutive rotations or decomposing block structures like $P^\top N^{-1} P$ ) to reduce computational cost, complementing XLA's low-level optimizations.
Solvers and Preconditioners: The framework includes numerical solvers for inverting large linear systems, accelerated by preconditioners like the Block Jacobi (BJPreconditioner).
Map-Making Algorithms: Implementation of binned estimators, maximum-likelihood solvers, and template deprojection methods.

3. Key Contributions

Unified Differentiable Framework: Furax is the first framework to unify low-level JAX libraries (like jax-healpy, s2fft) with high-level astrophysical operators, enabling end-to-end differentiability.
GPU/TPU Acceleration: It fully exploits modern hardware accelerators, offering competitive performance for terabyte-scale datasets, a significant leap over CPU-bound legacy tools.
Modular Extensibility: The architecture is designed to be experiment-agnostic, allowing researchers to rapidly prototype new instrument models and noise systematics without rewriting core solvers.
Integration with Probabilistic Programming: The differentiability of the framework unlocks advanced statistical inference, such as Bayesian hierarchical modeling in high-dimensional spaces and integration with neural networks for simulation-based inference.
Hybrid Workflows: It supports combining neural networks with linear operators, facilitating novel "simulation-based inference" approaches.

4. Results and Capabilities

While the paper is a software description rather than a specific scientific discovery, it demonstrates the framework's capabilities through:

Performance: Achieves production-ready computational efficiency via GPU acceleration and JIT compilation.
Flexibility: Successfully models complex scenarios including correlated noise (via SymmetricBandToeplitzOperator), gap-filling, and multi-observation reconstruction (MultiObservationMapMaker) without recompiling kernels.
Interoperability: Integrates with existing tools like PySM for sky simulations and TOAST for production pipeline compatibility, bridging the gap between research prototyping and operational data reduction.
Validation: The framework has been developed within the SciPol project specifically to enable gradient-based optimization in CMB pipelines, proving its utility for next-generation experiments.

5. Significance

Furax represents a paradigm shift in astrophysical data analysis software:

Sustainability: By moving away from legacy, experiment-specific codebases to a modern, Python/JAX foundation, it ensures that analysis tools can evolve with hardware and scientific needs.
Scientific Discovery: The ability to perform gradient-based optimization on the entire data pipeline allows for the discovery of optimal analysis strategies and the rigorous application of Bayesian methods that were previously computationally prohibitive.
Future-Proofing: It provides the essential infrastructure required to handle the data deluge from upcoming CMB missions (Simons Observatory, LiteBIRD), ensuring that the community can extract maximum cosmological information from noisy, complex datasets.

In summary, Furax is not just a library but a foundational infrastructure that modernizes the CMB data analysis workflow, making it faster, more flexible, and mathematically rigorous through the power of differentiable programming.