Physics-Informed Neural Network Surrogates with Polynomial Chaos-Based Uncertainty Propagation for Stochastic Model Predictive Control
Stochastic partial differential equations (PDEs) govern critical engineering and geophysical systems but are challenging to use for real-time control under parametric uncertainty. We present a unified framework that couples Physics-Informed Neural Networks (PINNs) with Polynomial Chaos Expansion (PCE) to construct a fast and differentiable surrogate. The PCE representation enables analytical propagation of parametric uncertainty and computation of the corresponding moments without requiring Monte Carlo sampling. We provide an error decomposition for the PINN-PCE surrogate that separates PCE truncation, stochastic quadrature, and PINN approximation errors. Embedding this surrogate into a stochastic model predictive control (SMPC) scheme enables finite-horizon control updates based on analytic mean and covariance predictions. We further show how the surrogate approximation error can be incorporated into tightened probabilistic constraints. The approach is validated on three benchmarks: the Korteweg-de Vries equation, Burgers' equation, and the two-dimensional incompressible Navier-Stokes equations, representing dispersive, convective, and convective-diffusive dynamics. Across all cases, the surrogate enables real-time control updates while maintaining prescribed risk levels and closely matching the corresponding high-fidelity solvers at substantially lower computational cost.