Diffhalos: A Generative Model of Cosmological Lightcones of Dark Matter Halos

Resumen Técnico de "Diffhalos: Un Modelo Generativo de Conos de Luz de Materia Oscura"

Planteamiento del Problema
Los análisis cosmológicos modernos de los sondeos de galaxias requieren conos de luz sintéticos (catálogos simulados de galaxias que abarcan el tiempo cósmico) para probar modelos teóricos e interpretar datos observacionales. Mientras que los métodos tradicionales dependen de la ejecución de masivas simulaciones de N-cuerpos gravitatorios o de la interpolación entre instantáneas de salida para generar estos conos de luz, estos enfoques son computacionalmente costosos y a menudo carecen de diferenciabilidad. Además, generar poblaciones sintéticas físicamente realistas requiere incorporar Historias de Ensamblaje de Masa (MAH, por sus siglas en inglés), las cuales describen cómo los halos crecen mediante fusiones y acreción. Los modelos existentes para las MAH, como el formalismo de Press-Schechter extendido o las funciones de ajuste estándar, a menudo carecen de la flexibilidad para capturar la diversidad completa de las rutas de ensamblaje o la eficiencia computacional requerida para la inferencia basada en gradientes. Existe la necesidad de un marco generativo que pueda producir conos de luz de halos de materia oscura, subhalos y sus MAH con precisión estadística, mientras soporte la diferenciación automática para la optimación eficiente de parámetros.

Metodología
Los autores presentan Diffhalos, un modelo generativo implementado en JAX para aprovechar la diferenciación automática y la aceleración por GPU. El marco descompone la generación de un cono de luz cosmológico en tres etapas distintas y diferenciables:

1. Generación de Halos Anfitriones (Host Halos):
   El modelo genera una población de halos anfitriones dentro de un volumen de cono de luz muestreando de la Función de Masa de Halos (HMF), $n_{\text{halo}}(m_{\text{halo}}, z)$. Los autores exploran dos implementaciones diferenciables:

   • Basada en Halox: Una implementación en JAX del modelo analítico de halo (usando la fórmula de ajuste de Tinker et al. 2008) que calcula la HMF mediante el espectro de potencia de materia lineal. Esto permite el cálculo directo de gradientes con respecto a los parámetros cosmológicos.
   • Basada en Emuladores: Para reducir la sobrecarga de memoria asociada con la integración numérica en el enfoque analítico, los autores entrenan un Perceptrón Multicapa (MLP) para emular la HMF acumulativa, $N_{\text{halo}}(> m_{\text{halo}}, z)$. La HMF diferencial se deriva mediante jax.grad.
     El muestreo se realiza mediante el muestreo de transformación inversa en una rejilla de desplazamientos al rojo (redshifts) y masas. Los autores también implementan un generador de Cuasi-Monte Carlo (QMC) que utiliza rejillas ponderadas para aproximar la población, reduciendo significativamente la huella de memoria manteniendo la consistencia estadística.

2. Generación de Subhalos:
   Una vez establecidos los halos anfitriones, el modelo los puebla con subhalos utilizando una Función de Masa de Subhalos Condicional (CSHMF) diferenciable. La CSHMF está parametrizada como una distribución acumulativa $N_{\text{sub}}(> \mu | M_{\text{host}})$, donde $\mu = M_{\text{sub}}/M_{\text{host}}$. Los autores utilizan un núcleo "sig-slope" (una ley de potencia con una pendiente que varía mediante una sigmoide) calibrado contra las simulaciones Discovery (código HACC) y las funciones de ajuste de Jiang & van den-Bosch (2016). Al igual que en la generación de anfitriones, los métodos de Monte Carlo y QMC están disponibles para el muestreo de subhalos.

3. Asignación de la Historia de Ensamblaje de Masa (MAH):
   La etapa final asigna una MAH a cada halo anfitrión y subhalo. La MAH se aproxima utilizando la forma funcional Diffmah, parametrizada por un conjunto de parámetros $\theta_{\text{MAH}}$. Para mapear estos parámetros a la población simulada, los autores introducen DiffmahNet, una arquitectura de flujo de normalización entrenada con árboles de fusión de la simulación de N-cuerpos Small MultiDark Planck (SMDPL). DiffmahNet aprende la distribución de probabilidad $P(\theta_{\text{MAH}} | z, M_{\text{halo}}(z))$, permitiendo la generación de historias de ensamblaje diversas y físicamente plausibles condicionadas a la masa y el desplazamiento al rojo del halo en el momento de la observación.

Contribuciones Clave

• Marco Diffhalos: Una librería totalmente diferenciable basada en JAX para generar conos de luz cosmológicos de halos de materia oscura, subhalos y sus MAH.
• HMF y CSHMF Diferenciables: La implementación de funciones de masa de halo diferenciables (vía Halox y emuladores personalizados) y funciones de masa de subhalos condicionales, permitiendo la optimación basada en gradientes con respecto a parámetros cosmológicos y astrofísicos.
• DiffmahNet: Un modelo de flujo de normalización que generaliza los modelos paramétricos previos (Hearin et al. 2021) para generar MAHs para poblaciones enteras de halos y subhalos, capturando tanto las tendencias de crecimiento medio como la varianza de la población.
• Herramientas de Eficiencia: La introducción de generadores de Cuasi-Monte Carlo (QMC) para las poblaciones de anfitriones y subhalos, ofreciendo mejoras de órdenes de magnitud en la eficiencia de memoria en comparación con las realizaciones de Monte Carlo estándar.
• Código Público: La liberación del código de Diffhalos para apoyar el desarrollo y la calibración de modelos de conexión galaxia-halo.

Resultos
Los autores validan el marco Diffhalos contra datos de simulaciones de N-cuerpos:

• Precisión de la HMF: La implementación de Halox muestra un acuerdo a nivel de porcentaje con la librería Colossus en un amplio rango de masas y desplazamientos al rojo ($z=0$ a $5$) para diferentes cosmologías. El emulador basado en MLP logra una precisión $\le 10\%$, lo cual es comparable a la precisión intrínseca de la propia función de ajuste de Tinker et al. (2008).
• Estadísticas de Subhalos: El generador de subhalos reproduce la función de masa total de halos anfitriones y subhalos de las simulaciones Discovery con una fidelidad razonable para masas de anfitrión $10^{10} < M_{\text{host}}/M_\odot < 10^{14}$ y $0 < z < 4$.
• Fidelidad de la MAH: DiffmahNet reproduce con precisión el crecimiento de masa medio y la varianza de las MAH encontradas en la simulación SMDPL para halos identificados en diferentes épocas ($t_{\text{obs}} = 6.42$ y $13.8$ Gyr).
• Cálculo de Gradientes: El marco calcula con éxito los gradientes de la HMF con respecto a los parámetros cosmológicos ($\Omega_m, \sigma_8$), demostrando la utilidad del enfoque diferenciable.

Significancia y Reivindicaciones
El artículo posiciona a Diffhalos como una herramienta de "prueba de concepto" diseñada principalmente para apoyar el desarrollo y la calibración de modelos de conexión galaxia-halo. Los autores enfatizan que la naturaleza diferenciable del marco permite aplicaciones eficientes basadas en gradientes, tales como el Monte Carlo Hamiltoniano y el descenso de gradiente, que son críticos para optimizar los modelos semi-analíticos (SAM) o modelos empíricos frente a datos observacionales.

Los autores declaran explícitamente que la implementación actual es una prueba de concepto y señalan modestamente que la dependencia cosmológica de la función de masa de subhalos y las MAH aún no ha sido totalmente calibrada a través de una serie de cosmologías (actualmente entrenada solo con la cosmología SMDPL). Identifican que es necesario el trabajo futuro para integrar emuladores de vanguardia, calibrar la dependencia cosmológica para subhalos y MAH, y extender el marco para incluir el agrupamiento espacial (coordenadas celestes) y efectos bariónicos. El principal valor propuesto es proporcionar una alternativa computacionalmente eficiente y diferenciable a las simulaciones de N-cuerpos para generar poblaciones de halos sintéticos, acelerando así la optimación de los modelos de formación de galaxias.