Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$α$ Forest Power Spectrum

Resumen Técnico: Uso de Emuladores Neuronales y Monte Carlo Hamiltoniano para Restringir la Historia de la Época de Reionización

Planteamiento del Problema
La Época de Reionización (EoR) representa una transición crítica en la historia cósmica, aunque su cronología detallada y su estructura inhomogénea permanecen insuficientemente restringidas. Las mediciones del Fondo Cósmico de Microondas (CMB) proporcionan una restricción de la profundidad óptica integrada, pero carecen de sensibilidad respecto a la duración y topología específicas de la reionización. Las restricciones actuales provenientes de los espectros de cuásares de alto desplazamiento al rojo (por ejemplo, alas de amortiguamiento, fracciones de píxeles oscuros) ofrecen una precisión limitada y a menudo dependen de interpretaciones basadas en modelos. El bosque Lyman-$\alpha$ (Ly$\alpha$) en $z \sim 5$ ofrece una promesa como sonda de los estados térmicos e de ionización del medio intergaláctico (IGM), pero extraer historias de reionización detalladas requiere aprovechar la distribución completa de las opacidades de Ly$\alpha$.

Una barrera significativa para la inferencia cuantitativa es el costo computacional. Generar simulaciones hidrodinámicas con la resolución suficiente para capturar las estructuras de pequeña escala relevantes para la EoR requiere miles de horas de GPU por modelo. En consecuencia, estudios previos se han visto limitados a comparaciones cualitativas o espacios de parámetros restringidos. Además, la estimación robusta de parámetros requiere matrices de covarianza precisas, las cuales son difíciles de derivar empíricamente a partir de las escasas muestras de cuásares de alto $z$ disponibles. Los enfoques tradicionales suelen asumir estructuras de covarianza fijas o dependen de conjuntos de simulaciones prohibitivamente costosos para computar covarianzas dependientes del modelo.

Metodología
Este trabajo presenta un nuevo marco de inferencia basado en JAX que integra observables del bosque Ly$\alpha$ modelados hacia adelante con emuladores neuronales diferenciables y Monte Carlo Hamiltoniano (HMC). La metodología procede a través de los siguientes pasos:

1. Construcción del Conjunto de Datos de Simulación: Los autores generaron un conjunto de datos de 501 simulaciones hidrodinámicas de baja resolución ($256^3$ vóxeles, $L_{box} = 20 \, h^{-1}$ Mpc) utilizando el código Nyx. Estas simulaciones fueron impulsadas por historias de reionización generadas por AMBER, un código seminumérico que mapea parámetros definidos por el usuario a campos de ionización. El espacio de parámetros incluye el punto medio de la reionización ($z_{mid}$), la duración ($\Delta z$), la asimetría ($A_z$), la temperatura máxima de calentamiento ($\Delta T_{re}$) y el flujo medio ($\langle F \rangle$).
2. Modelado hacia Adelante (Forward Modeling): A partir de estas simulaciones, los autores generaron trayectorias (skewers) de flujo 1D de Ly$\alpha$, aplicando efectos observacionales como la resolución del espectrógrafo ($R=30,000$) y el ruido ($SNR=30$). Computaron el espectro de potencia medio 1D ($\langle P_{Ly\alpha} \rangle$) y las matrices de covarianza dependientes del modelo ($\Sigma_{P_{Ly\alpha}}$) para cada modelo.
3. Emulación Neuronal: Se entrenaron dos emuladores neuronales independientes:
   • Emulador del Espectro de Potencia: Una red neuronal feed-forward totalmente conectada que predice $\log_{10}\langle P_{Ly\alpha} \rangle$ a partir de los parámetros de entrada.
   • Emulador de la Matriz de Covarianza: Para asegurar que el resultado sea simétrico y definido positivo, los autores emplearon la descomposición de Cholesky ($\Sigma = LL^T$). Una red neuronal predice el factor de Cholesky inferior aplanado ($L_{flat}$), aplicando una activación exponencial a los elementos diagonales para forzar la positividad.
4. Pipeline de Inferencia: Los emuladores se combinaron con Monte Carlo Hamiltoniano (HMC) mediante el paquete Numpyro. La diferenciabilidad de los emuladores basados en JAX permite el cálculo de gradientes requeridos por HMC, acelerando significativamente la exploración del espacio de parámetros en comparación con los algoritmos estándar de Metropolis-Hastings.

Contribuciones Clave

• Emulación de Covarianza Diferenciable: Este trabajo presenta el primer intento de emular matrices de covarianza del bosque Ly$\alpha$ utilizando redes neuronales. Al predecir el factor de Cholesky, el método garantiza la validez física (simetría y definición positiva) mientras captura correlaciones dependientes de los parámetros.
• Marco Completamente Diferenciable (End-to-End): La integración del modelado hacia adelante, la emulación neuronal y el HMC dentro del ecosistema JAX crea un flujo de trabajo completamente diferenciable para la inferencia de parámetros cosmológicos.
• Validación Cuantitativa: El marco se valida contra un conjunto de 100 observaciones simuladas (mock), demostrando la capacidad de recuperar los parámetros reales y evaluando la validez estadística de las distribuciones posteriores resultantes.

Resultos

• Precisión del Emulador: El emulador del espectro de potencia alcanzó un Error Porcentual Absoluto Medio (MAPE) mediano de 0.15% en el conjunto de prueba, con el 95% de los errores por debajo del 2%. El emulador de la matriz de covarianza logró un error mediano de 0.42% (relativo al producto de los elementos diagonales), con errores que generalmente se mantuvieron por debajo del 5% en todas las escalas. Ambos emuladores funcionaron con una precisión muy por debajo de los niveles de ruido observacional esperados asumidos en el modelado hacia adelante.
• Recuperación de Parámetros: Al aplicarse a observaciones simuladas, el marco HMC recuperó con éxito los parámetros reales de entrada. Las distribuciones posteriores fueron visualmente consistentes con la verdad fundamental (ground truth).
• Validez Estadística: Una prueba de cobertura realizada en 100 observaciones simuladas reveló que las distribuciones posteriores son ligeramente sobreconfiadas (la probabilidad de cobertura estimada se sitúa marginalmente por debajo de la línea ideal). Los autores atribuyen esto principalmente al supuesto de una función de verosimilitud gaussiana multivariante, señalando que el bosque de Ly$\alpha$ es inherentemente un campo aleatorio no gaussiano.
• Eficiencia de Datos: Los autores probaron el método con un conjunto de entrenamiento reducido de solo 100 simulaciones. Aunque los errores de emulación aumentaron (el MAPE del espectro de potencia subió a ~1.2%, los errores de covarianza también aumentaron), las distribuciones posteriores resultantes y la cobertura estadística se mantuvieron comparables a las derivadas con el conjunto de datos completo. Esto sugiere que el método es robusto incluso con datos de entrenamiento limitados, un factor crucial para futuras simulaciones de alta resolución.

Significado y Limitaciones
El artículo afirma que este marco proporciona una vía viable hacia la restricción cuantitativa de la historia de la EoR utilizando los espectros de potencia del bosque de Ly$\alpha$. La principal significancia radica en demostrar que los emuladores neuronales pueden reemplazar eficientemente los modelos hacia adelante costosos y que las matrices de covarianza dependientes del modelo pueden emularse, permitiendo una inferencia bayesiana más rápida y rigurosa.

Sin embargo, los autores declaran explícitamente que los resultados actuales aún no pueden aplicarse a datos observacionales reales debido a la baja resolución de las simulaciones utilizadas, la cual suprime la potencia en las escalas pequeñas. Además, el modelo de reionización depende de prescripciones de inyección de calor simplificadas y carece de fluctuaciones espaciales en el fondo ultravioleta (UVB). Los autores concluyen que el trabajo futuro debe incorporar simulaciones de mayor resolución, tratamientos térmicos más motivados físicamente y modelos de UVB fluctuante. También sugieren que la ligera sobreconfianza observada en las posteriores puede abordarse mediante métodos de inferencia basada en simulaciones (SBI) que no dependan de supuestos de verosimilitud gaussianos.