⚛️ general relativity

HIcosmo: a differentiable JAX-based framework for cosmology inference

Este artículo presenta HIcosmo, un marco de trabajo de código abierto basado en JAX para la inferencia cosmológica eficiente que aprovecha la diferenciación automática y la aceleración por GPU para lograr velocidades de muestreo significativamente más rápidas que las herramientas existentes como Cobaya, al tiempo que soporta una amplia gama de verosimilitudes de sondeos de la Etapa IV y modelos cosmológicos.

Autores originales: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Publicado 2026-06-29

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Jing-Zhao Qi, Jing-Fei Zhang, Xin Zhang

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina el universo como un rompecabezas gigante y complejo que los científicos están tratando de resolver. Durante décadas, han estado utilizando un conjunto estándar de reglas (llamado modelo $\Lambda$ CDM) para comprender cómo comenzó el universo, cómo se está expandiendo y de qué está hecho. Pero recientemente, las piezas del rompecabezas no han encajado perfectamente. Las mediciones del universo temprano (como el Fondo Cósmico de Microondas) no concuerdan con las mediciones del universo cercano (como las estrellas que explotan llamadas Supernovas). Esto se conoce como la "Tensión de Hubble".

Para resolver esto, se están construyendo nuevos y masivos telescopios que recolectarán miles de millones de puntos de datos. ¿El problema? Los programas informáticos que utilizan los científicos para resolver estos acertijos son demasiado lentos y torpes para manejar este nuevo aluvión de información. Son como intentar contar granos de arena en una playa usando una sola y lenta calculadora.

Este artículo presenta HIcosmo, una herramienta informática nueva y superrápida diseñada para resolver estos rompecabezas cósmicos de manera eficiente. Así es como funciona, explicado de forma sencilla:

1. La "Calculadora Inteligente" (JAX y Diferenciabilidad)

Las herramientas tradicionales son como un chef que tiene que medir manualmente cada ingrediente, anotar el resultado y luego calcular manualmente el siguiente paso. Si necesita saber cómo el cambio de un ingrediente afecta el sabor final, tiene que adivinar y probar, lo que lleva una eternidad.

HIcosmo es diferente. Está construido sobre una tecnología llamada JAX, que actúa como una "calculadora inteligente" que comprende la receta completa de una sola vez.

Diferenciación Automática: En lugar de adivinar, HIcosmo sabe exactamente cómo cambia el resultado si se ajusta un solo número. Es como tener un GPS que te indica instantáneamente la mejor ruta a seguir sin necesidad de recorrer todos los caminos posibles para ver cuál es el más rápido.
Compilación Just-in-Time: La primera vez que le pides a HIcosmo que realice un cálculo, tarda una fracción de segundo en "calentarse" y traducir las instrucciones en un código de máquina superoptimizado. Después de eso, funciona a la velocidad de la luz, evitando el lento "traductor" (el intérprete de Python) que ralentiza a las herramientas más antiguas.

2. El Pastel de Tres Capas

Los autores construyeron HIcosmo como un pastel de tres capas, donde cada capa tiene un trabajo específico:

La Capa Inferior (El Modelo): Esta es la parte de la "teoría". Toma las reglas del universo (como cuánta materia oscura o energía oscura hay) y calcula cómo debería verse el universo.
La Capa Media (La Verosimilitud): Esta es la parte de la "comparación". Toma la teoría de la capa inferior y la compara con los datos reales de los telescopios (como los datos de supernovas de Pantheon+ o los mapas de galaxias de DESI). Calcula qué tan "equivocada" está la teoría.
La Capa Superior (La Inferencia): Esta es la parte del "solucionador". Utiliza la información de la capa media para encontrar la mejor respuesta posible. Utiliza un método inteligente llamado NUTS (No-U-Turn Sampler) que utiliza los gradientes de la "calculadora inteligente" para dirigirse directamente a la respuesta correcta, en lugar de vagar sin rumbo como lo hacen las herramientas más antiguas.

3. Lo que Puede Hacer (y lo que No Puede Hacer)

Lo que hace:
HIcosmo es actualmente el maestro de la cosmología de fondo. Piensa en esto como medir el "panorama general" del universo: qué tan rápido se está expandiendo, qué tan lejos están las cosas y la forma general del espacio. Maneja datos de:

Estrellas que explotan (Supernovas).
Agrupamiento de galaxias (BAO).
El resplandor remanente del Big Bang (prior de distancia del CMB).
Retrasos temporales en lentes gravitacionales.

Ha sido probado contra el software que es el estándar de oro actual (llamado Cobaya) y se ha encontrado que es estadísticamente idéntico en sus respuestas pero mucho más rápido.

En un procesador de computadora estándar (CPU), es aproximadamente 8.7 veces más rápido que la herramienta antigua.
En una potente tarjeta gráfica (GPU), que es excelente para manejar cantidades masivas de datos a la vez, puede ser hasta 20 veces más rápido a medida que los datos crecen.

Lo que no hace:
El artículo es muy claro en que HIcosmo no maneja actualmente los "detalles finos" o las "perturbaciones" del universo, como las ondulaciones complejas de la red cósmica o los espectros de potencia detallados del universo temprano. Esos requieren un motor diferente, muy pesado (llamado un resolvedor de Boltzmann), que HIcosmo deja a otras herramientas. Piensa en HIcosmo como el experto en la forma del globo, mientras que otras herramientas son necesarias para analizar la textura del caucho.

4. Resultados del Mundo Real

Los autores no solo lo construyeron; lo usaron para resolver problemas reales:

Restricciones Conjuntas: Combinaron datos de supernovas, sondeos de galaxias y el CMB para fijar la tasa de expansión del universo ( $H_0$ ) y la densidad de la materia ( $\Omega_m$ ) con una precisión increíble.
Pronóstico Futuro: Usaron HIcosmo para predecir qué tan bien los futuros telescopios de radio (como SKA, MeerKAT y CHIME) podrán medir la "ecuación de estado" de la energía oscura (una forma elegante de preguntar: "¿Es la energía oscura constante o cambia con el tiempo?"). Encontraron que estos nuevos telescopios podrían medir esto con una precisión asombrosa.

La Conclusión

HIcosmo es un motor moderno de alta velocidad para la cosmología. Reemplaza los métodos lentos y manuales del pasado con un sistema optimizado y automático que puede manejar las inundaciones masivas de datos provenientes de la próxima generación de telescopios. No cambia las leyes de la física, pero cambia la rapidez y la eficiencia con la que podemos probarlas frente al universo.

Resumen Técnico: HIcosmo – Un Marco de Inferencia Cosmológica Diferenciable Basado en JAX

1. Planteamiento del Problema

Los próximos sondeosos de la etapa IV de la cosmología (por ejemplo, Euclid, LSST, DESI, SKA) generarán datos observacionales de un volumen sin precedentes. Mientras que la inferencia cosmológica a nivel de fondo (que involucra distancias, razones de distancias y módulos de distancia) opera típicamente en espacios de parámetros de baja dimensionalidad ( $d \sim 5\text{--}7$ ), el costo computacional de la inferencia sigue siendo un cuello de botella debido al enorme número de evaluaciones de verosimilitud requeridas por los métodos de Monte Carlo por Cadenas de Markov (MCMC).

Los procesos tradicionales (por ejemplo, CosmoMC, Cobaya, MontePython) dependen de muestreadores de Metropolis–Hastings (MH) de paseo aleatorio. Estos métodos sufren de tiempos de mezcla lentos que escalan linealmente con la dimensionalidad ( $d$ ) y requieren aproximaciones de diferencias finitas para los gradientes, las cuales son computacionalmente costosas ( $O(d)$ evaluaciones adicionales) y sensibles a la elección del tamaño del paso. Además, los códigos maduros existentes no están implementados como programas diferenciables, lo que impide el uso de muestreadores basados en gradientes como el muestreador No-U-Turn (NUTS) y la diferenciación automática (AD) para el pronóstico de Fisher. Aunque han surgido herramientas basadas en JAX para análisis a nivel de perturbación (por ejemplo, jax-cosmo, CosmoPower), no existía un marco sistemáticamente validado y de extremo a extremo para la cosmología de nivel de fondo que integrara verosimilitudes, muestreo NUTS y pronósticos de Fisher basados en AD dentro de una sola arquitectura.

2. Metodología y Arquitectura

El artículo presenta HIcosmo (Inferencia de Alto Rendimiento para la Cosmología), un marco de código abierto construido enteramente sobre JAX. El marco está diseñado en torno al principio de diferenciabilidad de extremo a extremo, asegurando que cada operación, desde los parámetros cosmológicos hasta la log-verosimilitud escalar, sea trazable por el sistema de diferenciación automática de JAX.

2.1 Arquitectura por Capas

HIcosmo consta de tres capas principales:

Capa de Modelo: Calcula las predicciones teóricas. Implementa cuatro modelos de fondo: $\Lambda$ CDM, $w$ CDM, $w_0w_a$ CDM (parametrización CPL) y un modelo de energía oscura interactuante ( $I\Lambda$ CDM). Crucialmente, reemplaza los integradores adaptativos (que no son diferenciables) por cuadratura trapezoidal de rejilla fija en rejillas de desplazamiento al rojo densas. Esto permite el cálculo exacto de gradientes mediante la regla de la cadena sin diferencias finitas.
Capa de Verosimilitud: Compara las predicciones del modelo con los datos observacionales. Soporta:
- Supernovas de Tipo Ia: Pantheon+, DES-SN5YR y Union3.
- Oscilaciones Acústicas de Bariones (BAO): DESI DR1/DR2, SDSS, BOSS y 6dFGS.
- CMB: Priores de distancia de Planck 2018 (comprimidos en $R$ , $l_A$ , $\omega_b$ ).
- $H_0$ Local: Escala de distancia SH0ES.
- Lente Fuerte: Retrasos temporales de H0LiCOW.
- Todas las verosimilitudes están implementadas en JAX puro, soportando compilación Just-In-Time (JIT) y vectorización.
Capa de Inferencia: Consume la log-verosimilitud y sus derivadas. Proporciona:
- Muestreo NUTS: Implementado vía NumPyro, utilizando gradientes directamente de JAX para mejorar los tiempos de mezcla (escalando como $d^{1/4}$ ).
- Pronóstico de Fisher: Calcula la matriz Hessiana mediante diferenciación automática anidada, eliminando la necesidad de selección de tamaño de paso y errores de diferencias finitas.

2.2 Detalles de Implementación

Diferenciabilidad: Al utilizar integración de paso fijo y precomputar cantidades independientes de los parámetros (por ejemplo, matrices de covarianza inversas), el marco asegura que los gradientes y las Hesianas se obtengan directamente mediante AD de modo inverso.
Aceleración de Hardware: El marco aprovecha la ejecución unificada CPU/GPU/TPU de JAX. La compilación JIT optimiza la verosimilitud en código de máquina, evitando la sobrecarga del intérprete de Python durante las ejecciones de MCMC.
Calibración: Para asegurar la precisión física, el marco calibra las cantidades derivadas (como el horizonte de sonido $r_d$ ) contra el resolvedor de Boltzmann CAMB, reduciendo los residuos a $\sim 0.01\%$ .

3. Resultados Clave

3.1 Validación y Precisión

HIcosmo fue validado sistemáticamente contra implementaciones de referencia (Cobaya) y códigos de verosimilitud oficiales a través de los modelos $\Lambda$ CDM, $w$ CDM y $w_0w_a$ CDM.

Precisión de la Verosimilitud: Los valores de $\chi^2$ concuerdan con los códigos de referencia con diferencias absolutas de $10^{-6}$ a $10^{-2}$ .
Restricciones de Parámetros: Las restricciones marginalizadas de HIcosmo y Cobaya difieren en menos de 0.2 $\sigma$ en todos los análisis probados. Para los análisis conjuntos de SN+BAO+CMB, el acuerdo es dentro de 0.04 $\sigma$ .
Consistencia: El marco reproduce las restricciones establecidas, tales como $H_0 = 68.24 \pm 0.38$ km/s/Mpc y $\Omega_m = 0.307 \pm 0.005$ para el análisis conjunto, confirmando que las aproximaciones diferenciables no introducen un sesgo significativo.

3.2 Benchmarks de Rendimiento

Mejora de Velocidad sobre Scipy: Para cálculos de distancia, HIcosmo logra una aceleración de 140 $\times$ sobre la cuadratura adaptativa de scipy tras la compilación JIT. Para el flujo completo de la verosimilitud de Pantheon+, es 16 $\times$ más rápido en CPU.
Rendimiento de Muestreo: Comparado con Cobaya (usando muestreadores MH) en un CPU de 8 núcleos, HIcosmo (usando NUTS) logra un rendimiento de Tamaño de Muestra Efectiva (ESS) 8.7 $\times$ mayor (400 ESS/s frente a 45.8 ESS/s). Esto se debe tanto a la evaluación de verosimilitud más rápida como a la mejor mezcla de NUTS ( $\tau \approx 1.2$ frente a $\tau \approx 11.4$ ).
Aceleración por GPU: Aunque las evaluaciones individuales en conjuntos de datos pequeños (Pantheon+) muestran ganancias modestas en GPU debido a la sobrecarga, la aceleración por GPU escala con el tamaño del conjunto de datos. Para conjuntos de datos sintéticos con $N_{SN} = 20,000$ , la GPU logra una aceleración de 20.8 $\times$ sobre la CPU.
Matriz de Fisher: El cálculo de la matriz de Fisher basado en AD es estable y no requiere ajuste de tamaño de paso, evitando los errores de redondeo de punto flotante que plagan los métodos de diferencias finitas a tamaños de paso pequeños.

3.3 Aplicaciones

El artículo demuestra la utilidad del marco a través de:

Restricciones Conjuntas de Multi-sonda: Combinando Pantheon+, DESI BAO y priores de Planck para estrechar las restricciones sobre $H_0$ y $\Omega_m$ .
Restricciones de Energía Oscura: Derivando restricciones sobre $w_0$ y $w_a$ consistentes con los resultados de DESI 2024.
Pronósticos de Mapeo de Intensidad de 21 cm: Realizando pronósticos de Fisher para seis sondeos (SKA1, MeerKAT, BINGO, Tianlai, CHIME). El marco predice que SKA1 Banda 1 podría restringir la ecuación de estado de la energía oscura a $\sigma(w) \approx 0.046$ .

4. Significado y Reivindicaciones

Los autores afirman que HIcosmo llena un vacío crítico en el panorama de la inferencia cosmológica al proporcionar el primer pipeline de cosmología de fondo nativo de JAX, de extremo a extremo y sistemáticamente validado.

Eficiencia: Demuestra que el muestreo basado en gradientes (NUTS) combinado con verosimilitudes compiladas con JIT puede superar significativamente al MCMC tradicional de paseo aleatorio tanto en velocidad computacional como en eficiencia estadística (ESS).
Escalabilidad: El marco está posicionado de manera única para manejar los masivos conjuntos de datos de futuros sondeos (Etapa IV) aprovechando la aceleración por GPU, donde las mejoras de velocidad pueden alcanzar las 20 $\times$ sobre la ejecución en CPU.
Modularidad y Extensibilidad: Al separar las capas de modelo, verosimilitud e inferencia, el marco permite la adición de nuevos modelos o conjuntos de datos sin reescribir el código de derivadas.
Base para la Inferencia a Nivel de Perturbación: Aunque actualmente está restringido a la cosmología de fondo (excluyendo resolvedores de Boltzmann completos), el artículo argumenta que la validación exitosa de extremo a extremo en el nivel de fondo establece una base robusta. La arquitectura está diseñada para extenderse a análisis de nivel de perturbación reemplazando la capa de modelo con emuladores de Boltzmann diferenciables (por ejemplo, CosmoPower-JAX), permitiendo así la inferencia basada en gradientes para problemas de alta dimensionalidad de la Etapa IV.

En resumen, HIcosmo ofrece una alternativa reproducible y de alto rendimiento a las herramientas existentes, demostrando que la diferenciación automática y la aceleración de hardware moderna pueden aplicarse eficazmente a las tareas estándar de inferencia cosmológica.