Latent Residual-Closure Fourier Neural Operator for Robust Multi-Field Solving in Particle-in-Cell Simulations

Abstract: Particle-in-cell (PIC) simulations are widely used for kinetic plasma modeling in energy applications, but their efficiency is often limited by repeated field solves on dense meshes. This work proposes a Latent Residual-Closure Fourier Neural Operator (LRC-FNO) for robust surrogate multi-field solving in PIC simulations. Rather than treating field prediction as a purely data-driven regression task, LRC-FNO formulates PIC field solving as a two-level residual-closure problem involving source compression and source-to-field operator mapping. An autoencoder extracts compact representations of particle-deposited source fields, while a Latent Closure Refiner recovers unresolved residual structures lost during compression. A Coarse-FNO Solver captures the dominant field response, and a Residual-Closure FNO restores full-resolution corrections. The method is tested on three benchmarks with increasing complexity: 1D linear Landau damping (LLD), 2D two-stream instability (TSI), and a 2D scrape-off layer (SOL) fusion plasma model. In LLD and TSI, LRC-FNO better preserves charge-to-potential mapping, potential-mode evolution, residual charge structures, and particle-field energy exchange during closed-loop PIC integration. In the SOL case, LRC-FNO achieves relative L2 errors of 0.0447 for the self-consistent potential and 0.0251 for the magnetic vector potential in single-step prediction. More importantly, when used as a neural initial guess with 20 iterative corrections, LRC-FNO maintains strong physical consistency in extrapolated closed-loop simulations, preserving charge and current density structures over a time range close to twice the training horizon. These results demonstrate that LRC-FNO can serve as both a fast surrogate field solver and a high-quality initialization strategy for iterative PIC field solvers.

• The paper introduces a two-level neural operator combining latent and field-space residual closures to surrogate multi-field solving in PIC simulations.
• Experimental results show that the latent closure reduces reconstruction errors and improves closed-loop stability while accelerating simulations by up to 2.5×.
• LRC-FNO offers modular compatibility with various operator kernels and helps alleviate mesh-based field-solve bottlenecks in kinetic plasma modeling.

Latent Residual-Closure Fourier Neural Operator for Robust Multi-Field Solving in Particle-in-Cell Simulations

Introduction

This work introduces the Latent Residual-Closure Fourier Neural Operator (LRC-FNO) as a neural surrogate operator for multi-field solving in particle-in-cell (PIC) simulations. In kinetic plasma modeling, PIC frameworks are central to resolving nonlinear and multiscale plasma phenomena, but their efficiency is severely constrained by repeated mesh-based field solves, especially under fine resolution and complex boundaries. While recent ML-based PDE solvers and operator learning frameworks such as DeepONet and FNO have accelerated PDE solution, their application to PIC settings is hindered by the inherently noisy, locally oscillatory nature of particle-deposited fields and the need to guarantee closed-loop dynamical stability. The LRC-FNO is designed to address these limitations with a two-level residual-closure architecture, combining both reduced latent representations and high-fidelity field-space correction to enhance robustness, accuracy, and stability in multi-field PIC simulation.

Figure 1: Schematic of the particle-in-cell cycle outlining particle advancement, source deposition, field solve, and field interpolation.

Methodology

Problem Formulation and Challenges in PIC Surrogate Field Solving

In standard PIC algorithms, the field solve occurs upon mapping particle-deposited charge and current densities $\rho$, $\mathbf{J}$ onto a mesh, followed by the solution of elliptic field equations (e.g., $\rho \rightarrow \phi$, $J \rightarrow A$) and field interpolation for particle forces. Unlike continuum fluid regimes, PIC-deposited sources contain strong noise and local discontinuities, making standard regression or spectral operator-learning approaches, which assume smooth functional input spaces, inadequate in both accuracy and long-time dynamical stability.

Moreover, reduced-order modeling (ROM) methods such as POD, DMD, and autoencoder-based transforms can compress field representations, but inevitably lose unresolved local structures. The amplification of these unresolved components, especially near boundaries or in high-gradient regions, significantly affects closed-loop property retention when deployed in dynamical PIC loops.

LRC-FNO Architecture

LRC-FNO is explicitly formulated as a two-level residual-closure mapping, where source information loss from compression is compensated via latent-space and field-space correction modules. The architecture is as follows:

Figure 2: Overall architecture of LRC-FNO featuring sequential latent-space and field-space residual closures.

Latent-Space Residual Closure

A convolutional autoencoder encodes the particle-deposited source field ($\rho$, $J$) into a low-dimensional latent space. The residual error between the autoencoder reconstruction and the original source field, encapsulating high-frequency and unresolved structures, is then processed by a patch-wise Latent Closure Refiner (LCR). The LCR uses convolutional embedding, SE channel reweighting, and MLP-mixer style nonlinear transformations to represent local residuals as compact latent variables. This approach ensures that essential boundary and noise-driven features, often omitted during compression, are made available for downstream operator learning.

Figure 3: Detailed breakdown of the LRC-FNO latent and residual closure structure.

The concatenation of the autoencoder latent vector and LCR output forms a closure-enhanced latent feature, which is subsequently decoded to generate a denoised, task-oriented source proxy field.

Field-Space Residual Closure

The source proxy is provided as input to a two-tier FNO-based solver. The Coarse-FNO solver predicts the dominant component of the solution field ($\phi$, $A$) on a downsampled mesh using spectral convolutions, emphasizing nonlocal and low-rank field responses. The output is upsampled and passed, together with the source proxy, into a Residual-Closure FNO, which corrects for unresolved residuals, boundary responses, and local details lost at the coarse stage.

This sequential closure ensures that high-fidelity physical field features are robustly reconstructed—first by compressing and closing latent space errors, and second by correcting operator-induced field truncations—without directly regressing from the high-dimensional particle-deposited sources.

Numerical Experiments

1D Linear Landau Damping

In the 1D LLD testbed, the objective is robust surrogate solution of the electrostatic field ($\rho \rightarrow \phi$) under charge densities with significant mode structure and noise. The autoencoder alone yields a test error of 0.0969; augmenting with LCR reduces the reconstruction error to 0.0195, highlighting the importance of latent closure for accurate source representation.

Figure 4: Comparison of charge density reconstruction by AE and AE+LCR in the 1D LLD benchmark.

Ablations using L-FNO, L-FNO-CNN, and LRC-FNO architectures all incorporating the same latent compressions indicate that single-step field errors are similar for L-FNO-CNN and LRC-FNO (0.0177 vs. 0.0180), but long-time closed-loop PIC integration strongly favors LRC-FNO, which maintains both amplitude and phase accuracy in modal potential evolution and demonstrates the highest $R^2$ and normalized cross-correlation with the FFT ground truth, especially under out-of-distribution driving.

Figure 5: Comparison of the fourth electrostatic potential Fourier mode evolution using FFT and surrogate field solvers in the 1D LLD benchmark.

2D Two-Stream Instability

In the 2D two-stream instability (TSI) scenario, the latent closure mechanism is validated for capturing both global and fine-scale emergent field structures. Three variants—a 4-channel autoencoder, 12-channel autoencoder, and AE+LCR (4+8 channels)—were assessed for source-field reconstruction and downstream operator accuracy. The AE+LCR yields superior spatial structure preservation in reconstructed charge fields and provides lower $\mathbf{J}$0 error input to the FNO operator (0.0346, vs. 0.0390/0.0442 for standard AEs).

Figure 6: Comparison of charge density reconstruction using AE-4ch, AE-12ch, and AE+LCR in the 2D TSI benchmark.

Long-time closed-loop analysis demonstrates that the AE+LCR-enhanced operator robustly preserves charge density, 2D spatial pattern, and field/kinetic energy evolution during dynamical extrapolation, whereas standard latent encodings show error localization, phase drift, and dynamical distortion.

Figure 7: Charge density distributions and errors in closed-loop 2D TSI simulations using different latent representations.

Figure 8: Comparison of field-energy and kinetic-energy evolution using different latent representations in closed-loop 2D TSI simulations.

2D Scrape-Off Layer (SOL) Model

A 2D3V PIC model with complex boundary conditions, representing a scrape-off-layer (SOL) in a fusion context, was used to test multi-field surrogate performance under boundary-sensitive and highly localized field structures. LRC-FNO achieves single-step relative $\mathbf{J}$1 errors of 0.0447 for self-consistent $\mathbf{J}$2 and 0.0251 for $\mathbf{J}$3. Direct replacement of the iterative solver with LRC-FNO leads to rapid, globally accurate solutions but yields localized sheath-region error accumulation over time.

Figure 9: Schematic of the 2D SOL model with boundary conditions and oblique magnetic field.

When LRC-FNO is deployed as an initialization for the numerical solver (followed by 20 iterative corrections), field and density distributions maintain remarkable agreement with the reference for prolonged closed-loop integration, even into temporal extrapolation. Errors remain low until well past the training horizon, with degradation primarily observed only after multiple training-time windows.

Figure 10: Multi-field distributions and errors at step = 500 in the SOL-PIC benchmark using direct LRC-FNO and LRC-FNO-assisted iterative correction.

Figure 11: Multi-field distributions and errors at step = 1000 in the SOL-PIC benchmark using LRC-FNO-assisted iterative correction.

Figure 12: Evolution of relative L2 errors of electron and ion densities in extrapolated SOL-PIC simulations using LRC-FNO-assisted iterative correction.

Runtime measurements indicate dramatic field-solver acceleration: direct LRC-FNO and "LRC-FNO + 20 iter" reduce field-solve wallclock time by 33.33× and 17.74×, respectively, yielding overall PIC cycle speedups of approximately 2.5×. This establishes LRC-FNO as both a surrogate solver and as a high-quality initial guess provider for iterative solvers under realistic plasma conditions.

Implications and Future Directions

The LRC-FNO framework marks a methodological advance in surrogate operator learning for PIC environments. The structured, two-level closure design overcomes key challenges in field mapping from particle-deposited, noise-dominated sources under stiff boundary constraints, securing both instantaneous accuracy and closed-loop physical consistency. The approach generalizes readily to multi-field tasks—e.g., simultaneous solution of $\mathbf{J}$4 and $\mathbf{J}$5—and offers modular compatibility with various operator kernels and compression schemes.

Practically, LRC-FNO can be leveraged for:

• Fast parametric sweeps and digital-twin workflows, especially in applications constrained by field-solve bottlenecks.
• Providing high-quality initial guesses in iterative or implicit PIC solvers to reduce convergence time and iteration cost.
• Scenarios requiring robust extrapolation in time and parameter space, as in experimental control, uncertainty quantification, or optimization tasks.

Theoretically, the formulation demonstrates how targeted latent and operator residual compensation, when constructed for physically consistent surrogates, can both close ROM errors and maintain dynamical conservation under non-smooth, stochastic driving typical in kinetic plasma systems.

Future work can address direct adaptation to implicit PIC solvers, extension to particle–collision operators, porting to more generalized electromagnetic field configurations, and integration with end-to-end differentiable PIC stacks for in-the-loop optimization.

Conclusion

LRC-FNO offers a robust, modular framework for surrogate multi-field solving in PIC simulations, combining latent-space compression/closure and field-space operator correction. It achieves significant speed-ups, maintains closed-loop PIC consistency, and supports generalization and extrapolation under noisy, boundary-sensitive, dynamical regimes. Its dual role as a fast surrogate and high-quality initialization for iterative solvers positions it as a practical tool for large-scale plasma modeling, with methodological advances applicable to broader classes of operator learning tasks in noisy and multiscale PDE environments.