# Solving the Discretised Neutron Diffusion Equations using Neural Networks

arXiv: https://arxiv.org/abs/2301.09939 · HTML: https://mlsanity.com/paper/2301.09939/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 77 references that are in the map.
- First posted on arXiv 24 Jan 2023.

## Builds on (most cited first, 13 of 13)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [MultiResUNet : Rethinking the U-Net Architecture for Multimodal Biomedical Image Segmentation](https://mlsanity.com/paper/1902.04049/) (2019)
- [Recurrent Residual Convolutional Neural Network based on U-Net (R2U-Net) for Medical Image Segmentation](https://mlsanity.com/paper/1802.06955/) (2018)
- [Deep Learning Methods for Reynolds-Averaged Navier-Stokes Simulations of Airfoil Flows](https://mlsanity.com/paper/1810.08217/) (2018)
- [Deep UQ: Learning deep neural network surrogate models for high dimensional uncertainty quantification](https://mlsanity.com/paper/1802.00850/) (2018)
- [A Perspective on Machine Learning Methods in Turbulence Modelling](https://mlsanity.com/paper/2010.12226/) (2020)
- [An autoencoder-based reduced-order model for eigenvalue problems with application to neutron diffusion](https://mlsanity.com/paper/2008.10532/) (2020)
- [Surrogate Modeling of Fluid Dynamics with a Multigrid Inspired Neural Network Architecture](https://mlsanity.com/paper/2105.03854/) (2021)
- [Accelerating MRI Reconstruction on TPUs](https://mlsanity.com/paper/2006.14080/) (2020)
- [Large Scale Distributed Linear Algebra With Tensor Processing Units](https://mlsanity.com/paper/2112.09017/) (2021)
- [Tensor Processing Units for Financial Monte Carlo](https://mlsanity.com/paper/1906.02818/) (2019)

## Built on by (most cited first, 5 of 5)

- [Solving the Discretised Multiphase Flow Equations with Interface Capturing on Structured Grids Using Machine Learning Libraries](https://mlsanity.com/paper/2401.06755/) (2024)
- [Integration of physics-informed operator learning and finite element method for parametric learning of partial differential equations](https://mlsanity.com/paper/2401.02363/) (2024)
- [Neural Physics: Using AI Libraries to Develop Physics-Based Solvers for Incompressible Computational Fluid Dynamics](https://mlsanity.com/paper/2402.17913/) (2024)
- [Parameter estimation for land-surface models using Neural Physics](https://mlsanity.com/paper/2505.02979/) (2025)
- [Quantum Neural Physics: Solving Partial Differential Equations on Quantum Simulators using Quantum Convolutional Neural Networks](https://mlsanity.com/paper/2603.24196/) (2026)

## Abstract (arXiv)

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.