Data-augmented few-shot neural stencil emulation for computer-model system identification

Abstract

Partial differential equations (PDEs) underpin the modeling of many natural and engineered systems. It can be convenient to express such models as neural PDEs (NPDEs) rather than using traditional numerical PDE solvers by replacing part or all of the PDE’s governing equations with a neural network representation. NPDEs are often easier to differentiate, linearize, reduce, or use for uncertainty quantification than the original numerical solver. They are usually trained on solution trajectories obtained by long-horizon rollout of the PDE solver. Here we propose a more sample-efficient data-augmentation strategy for generating NPDE training data from a computer model—specifically, one using an explicit, structured-grid finite-difference numerical solver—by space-filling sampling of local ‘stencil’ states. This approach removes a large degree of spatiotemporal redundancy present in trajectory data and oversamples states that may be rarely visited but help the NPDE generalize across the state space. We demonstrate that accurate NPDE stencil operators can be learned from augmented stencil training data generated by the computational equivalent of 10 timesteps’ worth of numerical simulation. Accuracy is further improved if a single full-trajectory simulation from the computer model is available, which is typically the case in practice. Across several PDE systems, we show that our data-augmented stencil data yield better trained neural stencil operators, with clear performance gains compared with naïvely sampled stencil data from simulation trajectories. Finally, using only 10 solver timesteps’ worth of augmented stencil data—with no full-trajectory simulation—our approach already matches or outperforms traditional ML emulators trained on hundreds to thousands of full trajectories in long-horizon rollout accuracy and stability. Reusing a single full trajectory within the same training budget further improves accuracy.