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

TL;DR

This paper introduces a methodology for integrating machine learning models into particle-in-cell simulations, enabling efficient, scalable, and accurate physics extensions with improved computational performance.

Contribution

It presents a novel approach for embedding neural networks in PIC codes using Python, demonstrated with a proof-of-concept in OSIRIS for Compton scattering.

Findings

ML models accurately reproduce conventional results

ML integration improves computational efficiency

Method enables scalable physics extensions in PIC simulations

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 an ML-based approach can provide a higher performance increase.