Machine Learning for neutron source distributions

Este artículo propone y evalúa un enfoque novedoso que utiliza modelos generativos probabilísticos para estimar distribuciones de fuentes de neutrones a partir de listas de partículas de Monte Carlo, permitiendo un muestreo eficiente e independiente de la memoria una vez que los modelos están entrenados.

Lo esencial

Imagina que estás intentando recrear la receta perfecta para un pastel complejo y de múltiples capas. En el mundo de la ciencia de neutrones, este "pastel" es un flujo de neutrones (partículas diminutas) que salen disparados de una fuente, cada uno con su propia velocidad, dirección, energía y temporización específicas.

Tradicionalmente, los científicos han intentado recrear este flujo de dos maneras:

1. El Método "Copiar y Pegar": Ejecutan una simulación informática masiva y lenta para generar una lista gigante de cada neutrón individual. Guardan esta lista (llamada archivo MCPL) e intentan usarla una y otra vez. ¿El problema? Si necesitas más neutrones de los que tiene la lista, simplemente copias y pegas los mismos una y otra vez. Esto crea "fallos" o "puntos calientes" en la simulación, como ver el mismo patrón de migas repetido infinitamente.
2. El Método "Regla Empírica": Intentan adivinar la receta observando los ingredientes por separado (por ejemplo, "¿cuántos son rápidos?", "¿cuántos son lentos?"). ¿El problema? Esto ignora cómo se mezclan los ingredientes entre sí. En la realidad, un neutrón rápido podría estar siempre moviéndose en una dirección específica, pero este método los trata como si no tuvieran relación, perdiendo el "sabor" de los datos reales.

El Nuevo Enfoque: El "Chef de IA"
Este artículo presenta una nueva forma de resolver este problema utilizando Aprendizaje Automático. En lugar de copiar la lista o adivinar las reglas, los autores entrenaron cuatro tipos diferentes de "Chefs de IA" (Modelos Generativos) para aprender la esencia de la receta de neutrones.

Así es como el artículo lo desglosa:

1. La Fase de Entrenamiento (Aprendiendo la Receta)

Los chefs de IA se alimentan de una muestra de la simulación informática original y lenta (los "datos de entrenamiento"). No solo memorizan la lista; aprenden las relaciones complejas entre todas las variables.

• La Analogía: Imagina mostrarle a un chef mil fotos de un tipo específico de nube. No solo memoriza las fotos; aprende qué hace que una nube se vea como esa nube: la forma en que se curvan los bordes, la densidad y cómo incide la luz sobre ella. Una vez que aprende esto, puede pintar una nueva nube que nunca ha existido antes pero que se ve exactamente correcta.

2. Los Cuatro Chefs de IA

Los autores probaron cuatro tipos diferentes de modelos de IA para ver cuál aprendía mejor la receta:

• Flujos Normalizadores (NF): Piensa en esto como un chef que puede estirar y apretar perfectamente una pieza de masa. Comienzan con una bola de masa simple y uniforme (ruido aleatorio) y la estiran hasta la forma compleja exacta de la nube de neutrones. El artículo encontró que este era el mejor chef, creando los "nuevos" neutrones más precisos que coincidían perfectamente con los datos originales.
• Autoencoders Variacionales (VAE): Este chef intenta comprimir la receta en un resumen y luego reconstruirla. Es rápido y bueno con formas complejas, pero a veces el pastel reconstruido sale un poco "borroso" o menos nítido que el original.
• Redes Generativas Antagónicas (GAN): Esto es un "tira y afloja" entre dos chefs. Uno intenta hornear un pastel falso y el otro intenta detectar el engaño. Siguen compitiendo hasta que el pastel falso es indistinguible del real. Este artículo los encontró un poco difíciles de entrenar y propensos a "hacer trampa" (repetir los mismos pocos patrones).
• Modelos de Difusión (DM): Este chef comienza con un pastel ruidoso y desordenado y lo limpia lentamente paso a paso hasta que es perfecto. Funciona bien pero es muy lento y costoso computacionalmente, como intentar limpiar una habitación recogiendo un grano de polvo a la vez.

3. Los Resultados: Por Qué Importa

El artículo probó a estos chefs de IA en dos escenarios del mundo real:

• Escenario A (El Conjunto de Datos TDR): Una fuente de neutrones compleja y de alta energía. Los chefs de IA aprendieron la receta tan bien que pudieron generar millones de nuevos neutrones que se veían estadísticamente idénticos a la simulación original, pero sin los fallos de "copiar y pegar".
• Escenario B (El Conjunto de Datos de Referencia): Un experimento del mundo real donde compararon los neutrones generados por IA contra mediciones reales tomadas en un laboratorio. La IA (específicamente el Flujo Normalizador) coincidió con los datos del mundo real casi perfectamente.

La Ventaja Clave:
Una vez que el chef de IA aprende la receta, ya no se necesita la lista gigante y pesada de neutrones originales. El modelo de IA es diminuto (como unos pocos kilobytes) y puede generar instantáneamente neutrones nuevos ilimitados que son estadísticamente perfectos. Esto ahorra cantidades masivas de tiempo y memoria informática.

Lo Que el Artículo No Dice

Los autores tienen cuidado de señalar que estos modelos son impulsados por datos. Aprenden estrictamente de los datos que se les proporcionan.

• Si la simulación original carecía de cierto tipo de neutrón, la IA no lo inventará (a menos que el modelo se ajuste específicamente para adivinar fuera de los datos, lo cual el artículo señala como una característica específica de otros métodos, no el objetivo principal aquí).
• El artículo no afirma que estos modelos puedan predecir nueva física o corregir datos defectuosos; son herramientas para recrear eficientemente patrones de datos existentes para su uso en el diseño de instrumentos de neutrones.

En Resumen:
El artículo demuestra que podemos reemplazar las listas pesadas y propensas a fallos de datos de neutrones con modelos de IA pequeños e inteligentes. Estos modelos aprenden el "ADN" del flujo de neutrones y pueden generar neutrones frescos y realistas bajo demanda, haciendo que el diseño de futuros experimentos de neutrones sea más rápido, más barato y más preciso. De los cuatro modelos probados, el Flujo Normalizador fue el claro ganador.

Resumen técnico

Resumen Técnico: Aprendizaje Automático para Distribuciones de Fuentes de Neutrones

Planteamiento del Problema
En óptica de neutrones y diseño de instrumentos, la caracterización precisa de las distribuciones del espacio de fases de los neutrones (posición, dirección, tiempo, energía, peso y polarización) en superficies específicas es crítica. Tradicionalmente, estas distribuciones se derivan de simulaciones de Monte Carlo (MC) (por ejemplo, MCNP, PHITS, OpenMC) y se almacenan como archivos de Lista de Partículas de Monte Carlo (MCPL). Aunque los archivos MCPL sirven como entradas estándar para software de trazado de rayos (por ejemplo, Vitess, McStas), su reutilización directa presenta limitaciones significativas:

1. Artefactos Estadísticos: El muestreo insuficiente y el ruido estadístico en el archivo MCPL original se propagan a las simulaciones aguas abajo, creando estructuras artificiales (por ejemplo, "puntos calientes") y una convergencia deficiente.
2. Pérdida de Correlación: Enfoques alternativos que ajustan modelos analíticos o semianalíticos a variables individuales a menudo fallan en preservar correlaciones complejas y de alta dimensión entre las variables del espacio de fases, las cuales son esenciales para un transporte realista, especialmente en fuentes de alto brillo.
3. Costo Computacional: Volver a ejecutar simulaciones de MC de alta estadística para generar nuevas listas de fuentes para diferentes configuraciones de instrumentos es lento y computacionalmente costoso.

Las soluciones existentes de aprendizaje automático, como la Estimación de Densidad de Kernel (KDSource), abordan algunos problemas pero a menudo requieren un ajuste manual de los pesos de importancia y luchan con la estimación completa del espacio de fases multivariado, particularmente cuando los pesos de las partículas varían.

Metodología
Los autores proponen un marco que utiliza Modelos Generativos Probabilísticos (PGM) para aprender la distribución de probabilidad conjunta de las variables del espacio de fases de los neutrones directamente desde archivos MCPL. Una vez entrenados, estos modelos pueden generar nuevas muestras de neutrones estadísticamente independientes que preservan las correlaciones de la distribución original sin requerir la lista de partículas original.

• Conjuntos de Datos: Se utilizaron dos conjuntos de datos distintos para evaluar los modelos:

  1. Conjunto de Datos TDR: Una simulación PHITS del moderador de hidrógeno para-hidrógeno de la Fuente de Alto Brillo (HBS) (676.558 neutrones). Este conjunto de datos incluye pesos de partículas y abarca un amplio espectro de energía (frío, térmico, rápido), presentando un desafío de modelado complejo y de alta dimensión (7 variables).
  2. Conjunto de Datos de Referencia: Una simulación PHITS de la Plataforma de Neutrones JULIC (118.219 neutrones), contrastada con mediciones experimentales de 2023. Este conjunto de datos tiene pesos de partículas constantes (la dimensionalidad se reduce a 6 variables) y permite una validación directa contra espectros experimentales.
  3. Preprocesamiento: Todas las variables se normalizaron utilizando escalado MinMax. Se aplicaron transformaciones logarítmicas a las variables de tiempo y energía para manejar su amplio rango dinámico.

• Modelos Generativos: Se compararon cuatro arquitecturas PGM distintas:

  1. Flujos de Normalización (NF): Específicamente Flujos de Acoplamiento y Flujos Autoregresivos Enmascarados (MAF). Estos modelos aprenden transformaciones invertibles entre una distribución de referencia simple (Gaussiana) y la distribución de datos compleja, permitiendo una evaluación exacta de la verosimilitud y el muestreo.
  2. Autoencoders Variacionales (VAE): Una arquitectura codificador-decodificador que aprende una representación latente de los datos, optimizada mediante una pérdida de reconstrucción y un regularizador de divergencia de Kullback–Leibler (KL).
  3. Redes Generativas Antagónicas (GAN): Una red generadora entrenada contra una red discriminadora para producir muestras indistinguibles de los datos reales.
  4. Modelos de Difusión (DM): Modelos que aprenden a revertir un proceso de adición de ruido hacia adelante para reconstruir datos a partir de ruido aleatorio.

• Implementación: Los modelos se entrenaron utilizando PyTorch en el clúster booster JUWELS. Se desarrolló un módulo personalizado en C++, source AI, para el software MC Vitess para cargar modelos entrenados de PyTorch (mediante TorchScript) y realizar muestreo en tiempo real, reemplazando efectivamente la entrada de archivos MCPL.

• Métricas de Evaluación:

  • Discrepancia Media Máxima (MMD): Utilizada para cuantificar la distancia entre la distribución aprendida y la distribución MCPL original en un Espacio de Hilbert de Núcleo Reproductivo (RKHS).
  • Velocidad de Muestreo: Tiempo requerido para generar 20.000 neutrones.
  • Validación Física: Comparación de los espectros de longitud de onda simulados contra datos experimentales de la plataforma JULIC.

Resultados Clave

• Fidelidad de la Distribución: Todos los modelos aprendieron con éxito las distribuciones subyacentes, pero los Flujos de Normalización (NF) demostraron un rendimiento superior. Los NF lograron las puntuaciones MMD más bajas en ambos conjuntos de datos (por ejemplo, $1,19 \times 10^{-4}$ para TDR y $1,20 \times 10^{-4}$ para Referencia), superando a los VAE, GAN, Modelos de Difusión y la línea base KDSource.
• Preservación de Correlaciones: El análisis visual de histogramas 2D confirmó que los NF preservaron con éxito correlaciones complejas entre variables (por ejemplo, tiempo-energía, posición-ángulo) que se perdieron en métodos de ajuste más simples.
• Eficiencia de Muestreo: Aunque KDSource ofreció los tiempos de muestreo más rápidos, los NF proporcionaron un equilibrio favorable entre velocidad y precisión. Los VAE y GAN también fueron rápidos, mientras que los Modelos de Difusión fueron significativamente más lentos para muestrear debido a su naturaleza iterativa de eliminación de ruido.
• Validación Experimental: En el conjunto de datos de referencia, las fuentes generadas por NF reprodujeron el espectro de longitud de onda experimental con alta fidelidad, igualando el rendimiento de KDSource. Sin embargo, a diferencia de KDSource, el modelo NF se adhirió estrictamente a los límites de los datos de entrenamiento (debido a la capa de activación Sigmoid), evitando la generación de extrapolaciones no físicas (por ejemplo, neutrones fuera de la guía o con longitudes de onda no observadas) a menos que estuviera diseñado explícitamente para ello.
• Robustez: Las mismas arquitecturas de modelos requirieron un ajuste mínimo al aplicarse a los dos conjuntos de datos diferentes, demostrando robustez a través de diversas dimensionalidades y complejidades de distribución.

Significado y Afirmaciones
El artículo afirma presentar el primer estudio comparativo que integra modelos generativos modernos de aprendizaje automático en tuberías de Monte Carlo de neutrones y los valida contra datos experimentales. El significado principal radica en demostrar que:

1. Los PGM son sustitutos viables: Los modelos generativos probabilísticos pueden reemplazar archivos MCPL crudos, ofreciendo un método eficiente en memoria (modelos de tamaño kilobyte frente a archivos de gigabytes) y estadísticamente robusto para la representación de fuentes.
2. Los Flujos de Normalización son óptimos: Entre las arquitecturas probadas, los NF ofrecen el mejor compromiso para esta aplicación específica, proporcionando un aprendizaje de distribución de alta fidelidad con la capacidad de incorporar restricciones físicas estrictas (por ejemplo, espacio de fases acotado) directamente en la arquitectura.
3. Integración del Flujo de Trabajo: El módulo source AI desarrollado en Vitess demuestra que estos modelos pueden integrarse sin problemas en flujos de trabajo de simulación de neutrones existentes, permitiendo un muestreo rápido e independiente sin necesidad de los datos de entrenamiento originales.

Los autores concluyen que, aunque los modelos generativos no reemplazan la necesidad de simulaciones MC iniciales de alta fidelidad, proporcionan una herramienta eficiente, flexible y físicamente significativa para el diseño y optimización de instrumentos aguas abajo, particularmente en escenarios que requieren simulaciones repetidas o muestreo de alta estadística donde la reutilización directa de MCPL es impráctica.