Solving the Discretised Neutron Diffusion Equations using Neural Networks

TL;DR

The paper presents a novel method to solve discretised PDEs by encoding finite-volume and finite-element discretisations as pre-determined convolutional neural network weights, eliminating the need for training. It employs Jacobi iterations within a diffusion-convolution framework and implements a sawtooth multigrid solver using a U‑Net architecture, extended to multi-group energy treatment. Demonstrations on KAIST reactor benchmarks show that neural-network solvers reproduce Fortran Gauss-Seidel solutions with machine-precision tolerances and offer comparable or faster performance on AI hardware, illustrating cross‑architecture efficiency. This approach enables exact discretisation reproduction while leveraging AI libraries and hardware accelerators, with potential for integration into digital twins, sensitivities, and uncertainty quantification in reactor physics.

Abstract

This paper presents a new approach which uses the tools within Artificial Intelligence (AI) software libraries as an alternative way of solving partial differential equations (PDEs) that have been discretised using standard numerical methods. In particular, we describe how to represent numerical discretisations arising from the finite volume and finite element methods by pre-determining the weights of convolutional layers within a neural network. As the weights are defined by the discretisation scheme, no training of the network is required and the solutions obtained are identical (accounting for solver tolerances) to those obtained with standard codes often written in Fortran or C++. We also explain how to implement the Jacobi method and a multigrid solver using the functions available in AI libraries. For the latter, we use a U-Net architecture which is able to represent a sawtooth multigrid method. A benefit of using AI libraries in this way is that one can exploit their power and their built-in technologies. For example, their executions are already optimised for different computer architectures, whether it be CPUs, GPUs or new-generation AI processors. In this article, we apply the proposed approach to eigenvalue problems in reactor physics where neutron transport is described by diffusion theory. For a fuel assembly benchmark, we demonstrate that the solution obtained from our new approach is the same (accounting for solver tolerances) as that obtained from the same discretisation coded in a standard way using Fortran. We then proceed to solve a reactor core benchmark using the new approach.

Figures (23)

• Figure 1: A 3 by 3 convolutional filter which applies the discretised diffusion operator (a five-point finite volume stencil) in 2D to 9 cells. The filter is first applied to all the cells in the blue block on the left; the result of which can be seen in the blue block following the equals sign. The 9 values are then summed to give the value in the red block on the right which is the output value, which represents the value of the diffusion operator acting on the input approximated at the central cell.
• Figure 2: Schematic of the neural network used for a single Jacobi iteration written as $\text{J}(\,\cdot\,)$ in Equation \ref{['eq:matrix_full']}. This network performs a single Jacobi iteration on the flux of a single energy group. The inputs are the flux ($\bm{\Phi}_g^{(k)}$), representative source ($\bm{s}_g$), diffusion coefficients ($\bm{D}_g$) and the strictly diagonal coefficients ($\bm{A_g^{0,0}})^{\odot -1}$) (green boxes). A number of layer operations are performed, mathematical operations are shown in orange and convolutional operations are in cyan. The output is the flux of the next Jacobi iteration ($\bm{\Phi}_g^{(k+1)}$). Arrows originate from which layer the data originated and the end of an arrow indicates which layer takes that data as input. The dimensions of the layers are given on the second line of each box.
• Figure 3: A schematic diagram showing the U-Net architecture (left) that is used to form a single multigrid sawtooth cycle. On the right, we can see how multiple cycles are brought together to form the overall solution method.
• Figure 4: Multigrid iteration, with the subscript indicating the resolution and superscript representing the multigrid iteration and $\text{J}(\,\cdot\,)$ representing a Jacobi iteration. The residual is calculated and restricted twice, indicated by the cyan nodes. These residuals are used with Jacobi smoothing, indicated by yellow nodes. After smoothing, prolongation is performed using UpSampling layers, indicated by teal nodes. After the finest level is reached, the flux is updated and the process repeats.
• Figure 5: Multigrid network, MG$(\,\cdot\,)$, representing a single multigrid iteration. This network performs a single multigrid iteration on the flux of a single energy group. Takes the flux ($\bm{\Phi}_{\bm{1}}^{(k)}$), representative source ($\bm{s_1}$), diffusion coefficients ($\bm{D_1}$) and the strictly diagonal coefficients ($\bm{A}_{\bm{1}}^{0,0}$), along with the coarser resolution coefficients, as inputs (green boxes). A number of layer operations are performed, mathematical operations in orange, convolutional passes in cyan, sub-model operations in yellow and upsampling in teal. The sub-models can be iterated on multiple times. Finally it outputs the flux of the next multigrid iteration flux ($\bm{\Phi}_{\bm{1}}^{(k+1)}$). Arrow origins show which layer the data originated and the end of the arrow shows which layer takes that data as input. Dimensions of layers are given on the second line of each box.