Physics-Informed Neural Compression of High-Dimensional Plasma Data

High-fidelity scientific simulations are now producing unprecedented amounts of data, creating a storage and analysis bottleneck. A single simulation can generate tremendous data volumes, often forcing researchers to discard valuable information. A prime example of this is plasma turbulence described by the gyrokinetic equations: nonlinear, multiscale, and 5D in phase space. It constitutes one of the most computationally demanding frontiers of modern science, with runs taking weeks and yielding tens of terabytes of data dumps. The increasing storage demands underscore the importance of compression. However, reconstructed snapshots do not necessarily preserve essential physical quantities. We present a spatiotemporal evaluation pipeline, accounting for structural phenomena and multi-scale transient fluctuations to assess the degree of physical fidelity. Indeed, we find that various compression techniques lack preservation of both spatial mode structure and temporal turbulence characteristics. Therefore, we explore Physics-Informed Neural Compression (PINC), which incorporates physics-informed losses tailored to gyrokinetics and enables extreme compressions ratios of over 70,000x. Entropy coding on top of PINC further pushes it to 120,000x. This direction provides a viable and scalable solution to the prohibitive storage demands of gyrokinetics, enabling post-hoc analyses that were previously infeasible.

💡 Research Summary

**
The paper tackles the growing storage bottleneck caused by high‑fidelity gyrokinetic plasma simulations, which generate multi‑terabyte, five‑dimensional (5D) datasets describing the distribution function f(v∥, µ, s, x, y). While these simulations are essential for understanding turbulence in magnetically confined fusion devices, researchers typically retain only a few diagnostics (e.g., electrostatic potential ϕ, heat flux Q) and discard the full 5D fields, limiting post‑hoc analysis. Existing compression schemes either achieve modest compression ratios or, when pushed to higher ratios, severely distort the physical quantities that are crucial for scientific interpretation.

To address this, the authors introduce two major contributions: (1) a spatiotemporal evaluation pipeline that quantifies how well compressed snapshots preserve both spatial mode structures and temporal turbulence dynamics, and (2) Physics‑Informed Neural Compression (PINC), a family of neural compression models that incorporate physics‑aware loss terms tailored to gyrokinetics.

Evaluation Pipeline
The pipeline separates spatial and temporal fidelity assessments. Spatial fidelity is measured using nonlinear integrals of the distribution function—specifically the electrostatic potential ϕ and the heat flux Q—as well as spectral diagnostics (k‑spec and Q‑spec), which capture the distribution of energy across perpendicular wave numbers. Temporal fidelity is evaluated through two complementary metrics: (i) the Energy Cascade (EC) error, computed as the sum of Wasserstein distances between the true and reconstructed spectral sequences during the transition from linear growth to statistically steady turbulence, and (ii) the End‑Point Error (EPE) of optical flow fields, a standard video‑analysis measure that captures the consistency of motion between reconstructed and reference sequences.

Neural Compression Architectures
Two families of learned compressors are explored:

1. 5D Autoencoders – Built on n‑dimensional Swin‑Transformer blocks (window‑based self‑attention) to handle the high‑dimensional data efficiently. The encoder partitions the 5D field into patches, applies hierarchical attention and down‑sampling, and projects to a latent space whose dimensionality controls the compression ratio. Both standard autoencoders (AE) and Vector‑Quantized Variational Autoencoders (VQ‑VAE) are implemented. These models share parameters across all time steps, enabling fast inference after a costly pre‑training phase.

2. Neural Implicit Fields – Each snapshot is fitted independently by a coordinate‑based multilayer perceptron (MLP) that maps the 5‑tuple (v∥, µ, s, x, y) to the complex distribution value. Positional embeddings are realized via learnable hash maps, and activations such as SiLU, sine, or Gabor are explored. Training a single field takes 1–2 minutes on an NVIDIA H100; because each snapshot is independent, training can be fully parallelized. The resulting network weights constitute the compressed representation, offering resolution‑agnostic reconstruction at the cost of higher encoding time.

Physics‑Informed Loss (PINC)
Pure reconstruction loss (complex MSE) is insufficient to guarantee conservation of key plasma quantities. Therefore, the authors augment training with a suite of global, physics‑driven losses:

• Integral losses: L_Q = |Q_pred – Q_GT| and L_ϕ = L1(ϕ_pred, ϕ_GT).
• Spectral losses: L_k = L1(k_spec_pred, k_spec_GT) and L_Qspec = L1(Q_spec_pred, Q_spec_GT).
• Monotonicity (isotonic) loss: Enforces that after the dominant wave number k_peak, the spectra decay monotonically. Implemented as a log‑transformed isotonic penalty L_iso(s) that penalizes negative slopes.

The total PINC loss is a weighted sum of these terms, guiding the network to reproduce not only pointwise values but also the global energy distribution and cascade behavior intrinsic to gyrokinetic turbulence. Notably, the losses are applied to nonlinear integrals of f rather than to PDE residuals, distinguishing this approach from traditional Physics‑Informed Neural Networks (PINNs).

Experimental Results
A 50 GB validation set of gyrokinetic simulations is released, and baseline results for conventional compressors are provided. At a compression ratio of ~70 000×, PINC (both VQ‑VAE‑based and implicit‑field variants) achieves reconstruction errors comparable to or better than standard methods, while dramatically improving physics fidelity: spatial mode reconstruction error drops by a factor of 2–3, EC error is reduced by ~60 %, and EPE improves by ~45 %. Adding entropy coding on top of PINC pushes the effective compression to ~120 000× with less than 1 % deviation in the key diagnostics. Moreover, a clear rate‑distortion curve emerges, allowing practitioners to predict the trade‑off between storage savings and physical accuracy before training.

Additional experiments demonstrate latent‑space temporal interpolation (producing intermediate time steps) and hybrid schemes that combine autoencoder latent codes with implicit‑field refinements, showcasing the flexibility of the framework.

Impact and Future Directions
The work establishes a concrete methodology for compressing ultra‑high‑dimensional scientific data without sacrificing the physical integrity required for downstream analysis. By embedding domain‑specific constraints directly into the loss function, the authors bridge the gap between data compression and scientific fidelity—a gap that has limited the reuse of massive simulation outputs. The proposed evaluation pipeline and physics‑informed loss design are readily transferable to other domains with multi‑scale, multi‑dimensional data, such as climate modeling, astrophysical simulations, and high‑resolution fluid dynamics. Future research avenues include real‑time (in‑transit) compression pipelines, extension to electromagnetic gyrokinetics, and leveraging the compressed latent space for accelerated surrogate modeling or uncertainty quantification.