FURAX: a modular JAX toolbox for solving inverse problems in science
JAX GPU Compilation Chain
From the Python code to the GPU-native code
Dead Code Elimination (DCE)
XLA Common Subexpression Elimination (CSE)
import jax
import jax.numpy as jnp
@jax.jit
def func_cse(theta):
a = jnp.sin(theta)
b = jnp.sin(theta) + 1
return a + bXLA Common Subexpression Elimination (CSE)
Definition: CSE identifies and eliminates duplicate computations within a function to optimize performance.
Example in Code:
- Without CSE:
jnp.sin(theta)computed twice. - After CSE: Shared computation across
aandb.
- Without CSE:
Benefits:
- Reduces redundant computation.
- Enhances runtime efficiency and memory usage.
JAX GPU Compilation Chain with FURAX
From the Python code to the GPU-native code
FURAX Algebraic Reductions: Composition of Rotations
import jax
import jax.numpy as jnp
def rot(x, y, theta):
rotated_x = x * jnp.cos(theta) - y * jnp.sin(theta)
rotated_y = x * jnp.sin(theta) + y * jnp.cos(theta)
return rotated_x, rotated_y
@jax.jit
def func(x, y, theta1, theta2):
return rot(x, y, theta=theta1 + theta2)
FURAX Algebraic Reductions: Composition of Rotations
import jax
import jax.numpy as jnp
def rot(x, y, theta):
rotated_x = x * jnp.cos(theta) - y * jnp.sin(theta)
rotated_y = x * jnp.sin(theta) + y * jnp.cos(theta)
return rotated_x, rotated_y
@jax.jit
def func(x, y, theta1, theta2):
x, y = rot(x, y, theta=theta1)
return rot(x, y, theta=theta2)
Maximum-Likelihood Map-Making
Classic data model
\[ d = \mathbf{P}s + n \]
- \(d\): time-ordered data
- \(\mathbf{P}\): pointing matrix ( telescope scanning)
- \(s\): discretized sky signal
- \(n\): stochastic contribution (noise)
Optimal (GLS) solution:
\[ \widehat{s} = (\mathbf{P}^\top \mathbf{N}^{-1} \mathbf{P})^{-1} \mathbf{P}^\top \mathbf{N}^{-1} d \]
Generalized parametric data model
\[ d_{\nu, i, t} = \int_{\textrm{BP}_\nu} d\nu' \mathbf{M}^{(\gamma)}_{\nu', i, t, p} \mathbf{A}^{(\beta)}_{\nu', t, c, p} s_{c, p} + n_{\nu, i, t} \]
- \(\mathbf{M}\) includes HWP parameters, band passes, beam properties, gains, …
- \(\mathbf{A}\) includes the frequency modeling of CMB, astrophysical foregrounds, atmosphere, ground
- \(\mathbf{H} = \mathbf{MA}\) is the generalized pointing operator
Credit: Simon Biquart
Maximum-Likelihood Map-Making
Furax example
h = bandpass @ pol @ rot @ hwp @ proj @ mixing
tod = h(sky)
solution = ((h.T @ invN @ h).I @ h.T @ invN)(gap_fill(key, tod))Credit: Simon Biquart
Writing this with FURAX tools
polarizer = LinearPolarizerOperator.create(detector_angles)
hwp = HWPOperator.create(hwp_angles)
rotation = QURotationOperator.create(telescope_angle)
sampling = SamplingOperator(pixel_indices)
h = polarizer @ hwp @ rotation @ sampling
invN = SymmetricBandToeplitzOperator(noise_correlations)
L = (h.T @ invN @ h).I @ h.T @ invN
estimate = L(data)
FURAX Map-Making
Can be extended & complexified easily!. - non-ideal components - parametric data model
Credit: Simon Biquart
Quantifying Biases from Non-Ideal Half Wave Plates
Non-ideal Half Wave Plate
- Made of several stacked layers
- Takes into account the transmission and reflection of the incident electromagnetic field at the layers boundaries
Furax modeling
Credit: Ema Tsang King Sang
Component Separation
- The Cosmic Microwave Background (CMB) signal is obscured by various foregrounds, making it challenging to detect the true cosmological information.
- Dust: Emission from galactic dust adds significant noise to the CMB, particularly affecting polarization measurements.
- Synchrotron Radiation: Electrons spiraling in the galaxy’s magnetic fields produce synchrotron radiation, another major contaminant.
Component separation methods
- Blind Methods: Like SMICA (Spectral Matching Independent Component Analysis)
- Parametric Methods: Like FGbuster (Foreground Buster)
Credit: Wassim Kabalan
On-going development of parametric component separation within the FURAX framework
Does everything fgbuster does but “better”
- Uses FURAX linear algebra operators to efficiently represent the mixing matrix
- Is written in JAX, is hardware accelerated
- Provides easy access to gradients
Beyond fgbuster
- Automatic cluster detection for spectral index parameters
- Flexible likelihood model (can be extended to include more components and different objective functions)
Credit: Wassim Kabalan
Atmosphere Decontamination Using Time-Domain Component Separation
Science Goal
For the Simons Observatory, characterize the observed atmospheric template from the recorded time-ordered data to separate the atmosphere from the sky signal we are after.
Data Model
\[ d_{\text{atm}} = \mathbf{A}(\text{PWV}) \mathbf{P}(w_x, w_y) s_{\text{atm}} + n \]
Model Parameters
- Wind velocity: \(\vec{w} = (w_x, w_y)\)
- Precipitable Water Vapour (PWV): ~Amplitude of atmospheric fluctuations.
- Estimate parameters by minimizing the spectral likelihood.
Atmospheric time-ordered data recorded by our telescope at two frequencies: \(d_{\text{atm}}\)
Mixing matrix \(\mathbf{A} = \begin{bmatrix} a(\text{PWV}_{1}) & 0 \\ 0 & a(\text{PWV}_{2}) \end{bmatrix}\)
Pointing matrix \(\mathbf{P} = \begin{bmatrix} P_{1} \\ P_{2} \end{bmatrix}\)
Atmospheric template: \(s_{\text{atm}}\)
Noise: \(n\), of covariance matrix \(\mathbf{N} = \begin{bmatrix} \sigma_{1}^{2} & 0 \\ 0 & \sigma_{2}^{2} \end{bmatrix}\)
Spectral Likelihood
\[ \boxed{ \langle \delta \mathcal{S}_\text{spec}(w_x, w_y, \text{PWV} \mid \vec{d}_{\text{atm}}) \rangle = \vec{d}_{\text{atm}}^\top \cdot \mathbf{AP} \Big[ (\mathbf{AP})^\top \mathbf{N}^{-1} (\mathbf{AP}) \Big]^{-1} (\mathbf{AP})^\top \mathbf{N}^{-1} d_{\text{atm}} } \]
Credit: Amalia Villarrubia Aguilar
Atmosphere Decontamination Using Time-Domain Component Separation
Spectral likelihood minimization
Gridding process
➜ This spectral likelihood is minimised through gridding: we compute \(\langle \delta \mathcal{S}_\text{spec}(w_x, w_y \mid \text{PWV}_{\text{sim}}) \rangle\)
for 22,500 different combinations of \((w_x, w_y)\).
Credit: Amalia Villarrubia Aguilar