PDE-NHF: Neural Hamiltonian Flow

• PDE-NHF is a machine learning framework that parameterizes and learns symplectic, invertible maps to simulate the time evolution of kinetic Hamiltonian PDEs while preserving physical invariants.
• It leverages a neural Hamiltonian flow architecture with learned Leapfrog integration steps, enabling efficient and accurate evolution of particle distributions under varied initial conditions.
• Empirical results show strong performance improvements, with lower Wasserstein-1 errors and reduced computational cost compared to traditional Particle-In-Cell methods.

The PDE-NHF (Partial Differential Equation–Neural Hamiltonian Flow) model is a machine learning framework that parameterizes and learns symplectic, invertible maps representing the time evolution of particle distributions governed by kinetic Hamiltonian PDEs. In particular, it has been introduced as an efficient, interpretable surrogate for the 1D Vlasov–Poisson system, enabling fast sampling, generalization to arbitrary initial conditions, and preservation of physical invariants through its construction (Souveton et al., 7 May 2025).

1. Mathematical Foundation: The Vlasov–Poisson System

PDE-NHF targets the 1D Vlasov–Poisson system, which describes the evolution of a phase-space probability density $f(x, v, t)$ for collisionless particles in either plasma physics or cosmology. This system is given by

$$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$

where the self-consistent force $F$ derives from a potential $\Phi$ that solves the Poisson equation: $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$ with mass density $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$. The system conserves total mass, total momentum, and total energy. It is subject to periodic boundary conditions in $x$ and decay of $f$ as $|v|\to\infty$ (Souveton et al., 7 May 2025).

2. Hamiltonian Structure and Symplecticity

In phase-space variables $(q, p)\equiv(x,v)$, the system can be derived from the Hamiltonian

$$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$0

with canonical equations

$$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$1

Liouville’s theorem guarantees preservation of phase-space volume, i.e., the transformation $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$2 is symplectic and volume-preserving ($$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$3) (Souveton et al., 7 May 2025).

3. Neural Hamiltonian Flow Architecture

PDE-NHF replaces explicit particle-in-cell (PIC) time integration with a symplectic normalizing flow. The total evolution time $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$4 is split into $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$5 equal sub-steps of size $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$6, each realized by an invertible, symplectic map $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$7 based on a learned Hamiltonian $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$8:

• Kinetic term: $$\partial_t f + v \,\partial_x f + F(x, t) \,\partial_v f = 0,$$9 with learnable scalar $F$0.
• Potential: $F$1, a DeepSet network imposing permutation and translation invariance ($F$2).

Each Leapfrog step computes

$F$3

with manifest invertibility and symplecticity ($F$4). The full flow over $F$5 is $F$6 (Souveton et al., 7 May 2025).

4. Training Objective and Data Generation

The model is trained to map an initial phase-space Gaussian $F$7 (mean and covariance $F$8) to a final empirical distribution $F$9 (from PIC) at $\Phi$0. The objective is forward Kullback–Leibler divergence,

$\Phi$1

with $\Phi$2 the model's pushforward of $\Phi$3 by the learned flow. Due to exact volume preservation, the model density at $\Phi$4 is given by the density of the corresponding inverse-mapped point under the initial Gaussian.

Training uses a dataset of $\Phi$5 PIC-generated trajectories, each with $\Phi$6 macroparticles, box length $\Phi$7, grid $\Phi$8, initial Gaussian parameters over $\Phi$9, evolved by Leapfrog with $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$0, $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$1 steps ($$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$2) (Souveton et al., 7 May 2025).

5. Inference, Generalization, and Sampling

After training, inference for a new initial Gaussian $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$3 at arbitrary time $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$4 proceeds by:

• Sampling $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$5 from $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$6.
• Integrating the learned symplectic flow for $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$7 steps using the trained $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$8 and scalar $$\partial_{x}^2 \Phi(x, t) = \begin{cases} -\rho(x, t)/\varepsilon_0 & \text{(plasma)} \ 4\pi G\,\rho(x, t) & \text{(self-gravitating)} \end{cases}$$9.
• Outputting the transformed $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$0.

As the flow is continuous-time and equivariant, evaluating the model at any intermediate $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$1 yields the corresponding approximate density without retraining. Only $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$2 network gradient evaluations (for $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$3) and sum updates per sample are needed (Souveton et al., 7 May 2025).

6. Empirical Performance and Physical Interpretation

On a held-out test set, measured 1D Wasserstein-1 distances ($\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$4) between true and generated samples at $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$5 are:

• PDE-NHF (25 Leapfrog steps): $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$6, $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$7.
• Baseline direct MLP: $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$8.

For intermediate $\rho(x,t) = \int_{\mathbb{R}} f(x,v,t) dv$9, PDE-NHF achieves $x$0 errors below $x$1 (position) and $x$2 (velocity), signaling strong generalization to unseen temporal slices. The learned kinetic coefficient $x$3 converges to $x$4, closely recovering the physical Hamiltonian.

The architecture is robust to coarse discretization: reducing $x$5 (e.g., to 5, $x$6) yields errors within $x$7 of the full-scheme, and per-sample runtime halves relative to standard PIC ($x$8 s per sample on A100 GPU vs. $x$9 s for PIC). Exact volume preservation ($f$0) rules out spurious diffusion or density collapse (Souveton et al., 7 May 2025).

7. Significance and Scope

PDE-NHF constitutes an overview of kinetic PDE theory, Hamiltonian mechanics, and deep generative modeling. Its main features are:

• Preservation of symplectic invariants by construction.
• Fast, parallelizable sampling for arbitrary initial data and times.
• Accurate learning of self-consistent physical potentials.
• Orders-of-magnitude improvements in repeated-inference throughput versus simulation-based approaches.
• Seamless interpolation across unseen (time, initial condition) domains.

These attributes distinguish PDE-NHF from direct regression architectures and enable its use as a surrogate, interpretive, and generative tool for high-dimensional kinetic systems (Souveton et al., 7 May 2025).