The inference of Fokker-Planck equations via transport maps

TL;DR

This work addresses inferring the Fokker-Planck equation and the driving potential from trajectory data by learning a time-resolved Knothe-Rosenblatt transport map $S$ that pulls a reference density $\eta$ to the evolving density $\rho$, yielding a velocity field $v(x,t) = - (\nabla_x S)^{-1} \partial_t S$. By matching the FP flux through a PDE-constrained loss $L_{\text{PDE}} = \int |\rho v + \rho \nabla \Psi + D \nabla \rho |^2 \mathrm{d}x$ and exploiting the triangular structure of $S$, the framework jointly identifies $\Psi$ and $D$, while a Bayesian variational formulation quantifies uncertainty. The method demonstrates accurate density and parameter recovery in 2D and 5D problems with substantially fewer data points than competing grid-based or neural PDE approaches, and provides both structured (physics-informed) and neural-network representations of the potential. This transport-map–driven, flux-based inference offers a scalable, grid-free route to high-dimensional FP dynamics with practical uncertainty quantification and improved data efficiency.

Abstract

We present a framework, which, from the trajectories detailing the spatiotemporal dynamics of a population, simultaneously reconstructs a transport map as well as the Fokker-Planck equation governing the coarse-grained probability distribution. Leveraging the Knothe-Rosenblatt rearrangement, we model the transport map from a fixed reference distribution to the target distribution, and derive the velocity fields of the flows from the trajectory of transport maps. Exploiting the velocity fields, we circumvent spatial gradients to infer the Fokker-Planck equation's potential and diffusivity. The sparsity of trajectories injects uncertainty, which we treat in a Bayesian setting using variational inference. The approach is applied to inferring the Fokker-Planck dynamics in spaces of up to five dimensions, demonstrating both accurate identification of the system and efficiency with respect to data size.

Figures (17)

• Figure 1: Evolution of a data distribution governed by an Itô stochastic differential equation, with snapshots at time instances $t=t_1, t_2, t_3$.
• Figure 2: Illustration of the transport maps $S(x,t_1)$, $S(x,t_2)$, $S(x,t_3)$ pulling back a fixed reference distribution $\eta(y)$ to target distributions $\rho(x,t_1)$, $\rho(x,t_2)$, $\rho(x,t_3)$ at different time instants. Thus, $S(x,t)$ captures the temporal evolution of the system through a sequence of diffeomorphic mappings.
• Figure 3: Overview of the model architecture. The pullback map $S$ and the the Fokker-Planck equation quantities $\Psi,D$ are parametrized by $\bm{\Theta}_1$ and $\bm{\Theta}_2$, respectively, and optimized based on data observed at discrete time points.
• Figure 4: Contour plots of the potential function for the two-dimensional isotropic, quadratic potential problem. The lower row is the zoomed-in version of the upper row images (see axis limits). Each column shows: (left) the ground truth potential, (middle) reconstructed potential using the structured model, and (right) the reconstructed potential using the neural network model.
• Figure 5: Slices of the potential function along each axis: (left) at $x_2=0$ and (right) at $x_1=0$ for the two-dimensional isotropic, quadratic potential problem.