Machine-learning-based models in particle-in-cell codes for advanced physics extensions

Abstract

In this paper we propose a methodology for the efficient implementation of machine learning (ML)-based methods in particle-in-cell (PIC) codes, with a focus on Monte Carlo or statistical extensions to the PIC algorithm. The presented approach allows for neural networks to be developed in a Python environment, where advanced ML tools are readily available to proficiently train and test them. Those models are then efficiently deployed within highly scalable and fully parallelized PIC simulations during runtime. We demonstrate this methodology with a proof-of-concept implementation within the PIC code OSIRIS, where a fully connected neural network is used to replace a section of a Compton scattering module. We demonstrate that the ML-based method reproduces the results obtained with the conventional method and achieves better computational performance. These results offer a promising avenue for future applications of ML-based methods in PIC, particularly for physics extensions where a ML-based approach can provide a higher performance increase.