Machine Learning for neutron source distributions

This paper proposes and evaluates a novel approach using probabilistic generative models to estimate neutron source distributions from Monte Carlo particle lists, enabling efficient, memory-independent sampling once the models are trained.

The Gist

Imagine you are trying to recreate the perfect recipe for a complex, multi-layered cake. In the world of neutron science, this "cake" is a stream of neutrons (tiny particles) shooting out of a source, each with its own specific speed, direction, energy, and timing.

Traditionally, scientists have tried to recreate this stream in two ways:

1. The "Copy-Paste" Method: They run a massive, slow computer simulation to generate a giant list of every single neutron. They save this list (called an MCPL file) and try to use it again and again. The problem? If you need more neutrons than the list has, you just copy and paste the same ones over and over. This creates "glitches" or "hot spots" in the simulation, like seeing the same crumb pattern repeated endlessly.
2. The "Rule-of-Thumb" Method: They try to guess the recipe by looking at the ingredients separately (e.g., "how many are fast?" "how many are slow?"). The problem? This ignores how the ingredients mix together. In reality, a fast neutron might always be moving in a specific direction, but this method treats them as if they are unrelated, losing the "flavor" of the real data.

The New Approach: The "AI Chef"
This paper introduces a new way to solve this problem using Machine Learning. Instead of copying the list or guessing the rules, the authors trained four different types of "AI Chefs" (Generative Models) to learn the essence of the neutron recipe.

Here is how the paper breaks it down:

1. The Training Phase (Learning the Recipe)

The AI chefs are fed a sample of the original, slow computer simulation (the "training data"). They don't just memorize the list; they learn the complex relationships between all the variables.

• The Analogy: Imagine showing a chef a thousand photos of a specific type of cloud. They don't just memorize the photos; they learn what makes a cloud look like that cloud—the way the edges curl, the density, and how the light hits it. Once they learn this, they can paint a new cloud that has never existed before but looks exactly right.

2. The Four AI Chefs

The authors tested four different types of AI models to see which one learned the recipe best:

• Normalizing Flows (NF): Think of this as a chef who can perfectly stretch and squeeze a piece of dough. They start with a simple, uniform ball of dough (random noise) and stretch it into the exact complex shape of the neutron cloud. The paper found this was the best chef, creating the most accurate "new" neutrons that perfectly matched the original data.
• Variational Autoencoders (VAE): This chef tries to compress the recipe into a summary and then rebuild it. It's fast and good at complex shapes, but sometimes the rebuilt cake comes out a little "blurry" or less sharp than the original.
• Generative Adversarial Networks (GAN): This is a "tug-of-war" between two chefs. One tries to bake a fake cake, and the other tries to spot the fake. They keep competing until the fake cake is indistinguishable from the real one. This paper found them a bit difficult to train and prone to "cheating" (repeating the same few patterns).
• Diffusion Models (DM): This chef starts with a noisy, messy cake and slowly cleans it up step-by-step until it's perfect. It works well but is very slow and computationally expensive, like trying to clean a room by picking up one grain of dust at a time.

3. The Results: Why It Matters

The paper tested these AI chefs on two real-world scenarios:

• Scenario A (The TDR Dataset): A complex, high-energy neutron source. The AI chefs learned the recipe so well that they could generate millions of new neutrons that looked statistically identical to the original simulation, but without the "copy-paste" glitches.
• Scenario B (The Benchmark Dataset): A real-world experiment where they compared the AI-generated neutrons against actual measurements taken in a lab. The AI (specifically the Normalizing Flow) matched the real-world data almost perfectly.

The Key Advantage:
Once the AI chef learns the recipe, the giant, heavy list of original neutrons is no longer needed. The AI model is tiny (like a few kilobytes) and can instantly generate unlimited new neutrons that are statistically perfect. This saves massive amounts of computer time and memory.

What the Paper Doesn't Say

The authors are careful to state that these models are data-driven. They learn strictly from the data they are given.

• If the original simulation was missing a certain type of neutron, the AI won't invent it (unless the model is specifically tweaked to guess outside the data, which the paper notes is a specific feature of other methods, not the primary goal here).
• The paper does not claim these models can predict new physics or fix bad data; they are tools to efficiently recreate existing data patterns for use in designing neutron instruments.

In Summary:
The paper demonstrates that we can replace heavy, glitch-prone lists of neutron data with tiny, smart AI models. These models learn the "DNA" of the neutron stream and can generate fresh, realistic neutrons on demand, making the design of future neutron experiments faster, cheaper, and more accurate. Among the four models tested, the Normalizing Flow was the clear winner.

Technical Summary

Technical Summary: Machine Learning for Neutron Source Distributions

Problem Statement
In neutron optics and instrument design, accurate characterization of neutron phase-space distributions (position, direction, time, energy, weight, and polarization) at specific surfaces is critical. Traditionally, these distributions are derived from Monte Carlo (MC) simulations (e.g., MCNP, PHITS, OpenMC) and stored as Monte Carlo Particle List (MCPL) files. While MCPL files serve as standard inputs for ray-tracing software (e.g., Vitess, McStas), their direct reuse presents significant limitations:

1. Statistical Artifacts: Undersampling and statistical noise in the original MCPL file propagate into downstream simulations, creating artificial structures (e.g., "hot spots") and poor convergence.
2. Correlation Loss: Alternative approaches that fit analytical or semi-analytical models to individual variables often fail to preserve complex, high-dimensional correlations between phase-space variables, which are essential for realistic transport, especially in high-brilliance sources.
3. Computational Cost: Re-running high-statistics MC simulations to generate new source lists for different instrument configurations is time-consuming and computationally expensive.

Existing machine learning solutions, such as Kernel Density Estimation (KDSource), address some issues but often require manual tuning of importance weights and struggle with full multivariate phase-space estimation, particularly when particle weights vary.

Methodology
The authors propose a framework utilizing Probabilistic Generative Models (PGMs) to learn the joint probability distribution of neutron phase-space variables directly from MCPL files. Once trained, these models can generate new, statistically independent neutron samples that preserve the original distribution's correlations without requiring the original particle list.

• Datasets: Two distinct datasets were used to evaluate the models:

  1. TDR Dataset: A PHITS simulation of the High Brilliance Source (HBS) para-hydrogen moderator (676,558 neutrons). This dataset includes particle weights and covers a broad energy spectrum (cold, thermal, fast), presenting a complex, high-dimensional (7 variables) modeling challenge.
  2. Benchmark Dataset: A PHITS simulation of the JULIC Neutron Platform (118,219 neutrons), benchmarked against 2023 experimental measurements. This dataset has constant particle weights (dimensionality reduced to 6 variables) and allows for direct validation against experimental spectra.
  3. Preprocessing: All variables were normalized using MinMax scaling. Logarithmic transformations were applied to time and energy variables to handle their wide dynamic range.

• Generative Models: Four distinct PGM architectures were compared:

  1. Normalizing Flows (NF): Specifically Coupling Flows and Masked Autoregressive Flows (MAF). These models learn invertible transformations between a simple reference distribution (Gaussian) and the complex data distribution, allowing for exact likelihood evaluation and sampling.
  2. Variational Autoencoders (VAE): An encoder-decoder architecture learning a latent representation of the data, optimized via a reconstruction loss and a Kullback–Leibler (KL) divergence regularizer.
  3. Generative Adversarial Networks (GAN): A generator network trained against a discriminator network to produce samples indistinguishable from the real data.
  4. Diffusion Models (DM): Models that learn to reverse a forward noising process to reconstruct data from random noise.

• Implementation: Models were trained using PyTorch on the JUWELS booster cluster. A custom C++ module, source AI, was developed for the Vitess MC software to load trained PyTorch models (via TorchScript) and perform on-the-fly sampling, effectively replacing the MCPL file input.

• Evaluation Metrics:

  • Maximum Mean Discrepancy (MMD): Used to quantify the distance between the learned distribution and the original MCPL distribution in a Reproducing Kernel Hilbert Space (RKHS).
  • Sampling Speed: Time required to generate 20,000 neutrons.
  • Physical Validation: Comparison of simulated wavelength spectra against experimental data from the JULIC platform.

Key Results

• Distribution Fidelity: All models successfully learned the underlying distributions, but Normalizing Flows (NF) demonstrated superior performance. NFs achieved the lowest MMD scores on both datasets (e.g., $1.19 \times 10^{-4}$ for TDR and $1.20 \times 10^{-4}$ for Benchmark), outperforming VAEs, GANs, Diffusion Models, and the KDSource baseline.
• Correlation Preservation: Visual analysis of 2D histograms confirmed that NFs successfully preserved complex correlations between variables (e.g., time-energy, angular-position) that were lost in simpler fitting methods.
• Sampling Efficiency: While KDSource offered the fastest sampling times, NFs provided a favorable balance between speed and accuracy. VAEs and GANs were also fast, whereas Diffusion Models were significantly slower to sample due to their iterative denoising nature.
• Experimental Validation: In the Benchmark dataset, NF-generated sources reproduced the experimental wavelength spectrum with high fidelity, matching the performance of KDSource. However, unlike KDSource, the NF model strictly adhered to the bounds of the training data (due to the Sigmoid activation layer), preventing the generation of unphysical extrapolations (e.g., neutrons outside the guide or with unobserved wavelengths) unless explicitly designed to do so.
• Robustness: The same model architectures required minimal retuning when applied to the two different datasets, demonstrating robustness across varying dimensionalities and distribution complexities.

Significance and Claims
The paper claims to present the first comparative study integrating modern generative ML models into neutron Monte Carlo pipelines and validating them against experimental data. The primary significance lies in demonstrating that:

1. PGMs are viable surrogates: Probabilistic generative models can replace raw MCPL files, offering a memory-efficient (kilobyte-sized models vs. gigabyte files) and statistically robust method for source representation.
2. Normalizing Flows are optimal: Among the tested architectures, NFs offer the best trade-off for this specific application, providing high-fidelity distribution learning with the ability to incorporate hard physical constraints (e.g., bounded phase-space) directly into the architecture.
3. Workflow Integration: The developed source AI module in Vitess proves that these models can be seamlessly integrated into existing neutron simulation workflows, allowing for rapid, independent sampling without the need for the original training data.

The authors conclude that while generative models do not replace the need for initial high-fidelity MC simulations, they provide an efficient, flexible, and physically meaningful tool for downstream instrument design and optimization, particularly in scenarios requiring repeated simulations or high-statistics sampling where direct MCPL reuse is impractical.