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

Resumen Técnico: Furax – Un Marco Modular JAX para Operadores Lineales en Análisis de Datos Astrofísicos y Cosmológicos

1. El Problema

Los experimentos actuales y futuros de Fondo Cósmico de Microondas (CMB), como el Simons Observatory, el South Pole Observatory, QUBIC y LiteBIRD, generarán flujos masivos de datos ordenados en el tiempo (TOD). El desafío central en el análisis de estos datos es la inversión de problemas lineales para recuperar la señal del cielo ($m$) a partir de observaciones ruidosas ($d$), modeladas mediante la ecuación:
$$d = Hm + Tt + n$$
Donde $H$ es el operador de adquisición (puntería, respuesta instrumental), $T$ son efectos aditivos no deseados y $n$ es el ruido. La solución óptima requiere un estimador de mínimos cuadrados generalizado que implica la aplicación eficiente de operadores y sus transpuestas, así como la inversión de grandes matrices de covarianza.

Históricamente, las cadenas de procesamiento de datos han estado atadas a experimentos específicos, utilizando tecnologías heredadas (legacy) que carecen de generalidad, no sobreviven a los experimentos y no aprovechan las arquitecturas modernas de hardware (GPUs/TPUs) ni la diferenciabilidad automática.

2. Metodología y Arquitectura

Furax (Framework for Unified and Robust data Analysis with JAX) es un marco de código abierto en Python construido sobre JAX. Su metodología se basa en los siguientes pilares:

• Álgebra de Operadores Lineales Composibles: Furax define operadores lineales como clases de datos de Python registradas como pytrees de JAX. Esto permite combinar operadores mediante notación matemática intuitiva (suma, composición @, multiplicación por escalar, transposición .T, e inversión .I).
• Diferenciabilidad Automática y Compilación JIT: Aprovecha las capacidades de compilación just-in-time y diferenciación automática de JAX para optimizar el rendimiento y habilitar la optimización basada en gradientes.
• Aceleración de Hardware: Está diseñado para ejecutarse eficientemente en GPUs y TPUs, permitiendo el procesamiento de conjuntos de datos a escala de terabytes.
• Estructura Modular:
  • Operadores Genéricos: Incluyen transformadas de Fourier (FFTOperator), operadores de máscara, reordenamiento de ejes y operadores de bloque (diagonales, filas/columnas) para manejar sistemas multi-observación.
  • Operadores Específicos del Dominio: Diseñados para la astrofísica y el CMB, incluyendo modelado de instrumentos (placas de media onda, polarizadores lineales), rotación de parámetros Stokes (QU), y separación de componentes espectrales (polvo, sincrotrón, CMB) con índices espectrales espacialmente variables.
  • Reducción Algebraica: Implementa un sistema basado en reglas para simplificar expresiones de operadores (ej. combinar rotaciones consecutivas o descomponer estructuras de bloque) antes de la compilación, complementando las optimizaciones de bajo nivel de XLA.

3. Contribuciones Clave

1. Marco Unificado y Agnóstico al Experimento: A diferencia de herramientas anteriores (como TOAST, que tiene un núcleo en C++ con soporte limitado de GPU, o Commander, que carece de aceleración GPU), Furax ofrece una base unificada, diferenciable y moderna en Python/JAX.
2. Herramientas para la Separación de Componentes y "Map-Making": Proporciona algoritmos avanzados para la creación de mapas (estimadores binned, máxima verosimilitud, deproyección de plantillas) y la separación de componentes astrofísicos, integrando modelos de ruido correlacionado mediante operadores de Toeplitz de banda simétrica.
3. Integración con Inferencia Probabilística: La diferenciabilidad nativa permite la integración con herramientas de programación probabilística basadas en JAX, habilitando métodos de inferencia estadística avanzada, como el modelado jerárquico bayesiano en espacios de alta dimensión.
4. Interoperabilidad: Se integra con bibliotecas existentes como jax-healpy, s2fft y PySM, y mantiene compatibilidad con cadenas de producción a través de la integración con TOAST.
5. Tipos de Datos Especializados: Define tipos JAX específicos para parámetros de Stokes ($I, QU, IQU, IQUV$) y tipos de paisaje (discretización del cielo, ej. HEALPix), facilitando operaciones matemáticas coherentes.

4. Resultados y Capacidades

El documento destaca que Furax permite:

• Prototipado Rápido: Los investigadores pueden validar pipelines de análisis con eficiencia computacional lista para producción.
• Rendimiento Competitivo: Logra un rendimiento superior gracias a la aceleración en GPU/TPU y la compilación JIT, superando a las implementaciones tradicionales en CPU.
• Flexibilidad en Modelado: Facilita la experimentación con modelos de instrumentos realistas y sistemas complejos de ruido sistemático.
• Innovación en Inferencia: Habilita enfoques híbridos que combinan redes neuronales con operadores lineales para inferencia basada en simulaciones.

5. Significado e Impacto

Furax representa un cambio de paradigma en el análisis de datos cosmológicos. Al proporcionar una infraestructura moderna, diferenciable y escalable, sienta las bases para el desarrollo de técnicas de análisis de próxima generación necesarias para los grandes experimentos del CMB (Simons Observatory, LiteBIRD, etc.).

Su capacidad para integrar la optimización basada en gradientes y el aprendizaje automático en pipelines de física tradicional abre nuevas posibilidades para la extracción de información cosmológica de alta precisión. Además, al ser agnóstico al experimento, promueve la reutilización de software y la colaboración entre diferentes grupos de investigación, rompiendo con el modelo histórico de pipelines aislados y obsoletos.