Application of PDE emulators for a multi-material, multi-physics simulation problem
Authors: Diane Oyen, Morgan Jones, Adithan Kathirgamaraju, Abu Bucker Siddik, Alexander Most, David Schodt, Christine Sweeney, Kyle Hickmann, Michael McCann, Jemma Stachelek, Alexius Wadell, Michal Kucer, Mahindra Rautela, Siddharth Mansingh, James Amarel, Gerd J. Kunde, Emily Casleton Taylor, Soumi De, Ayan Biswas, Aimee Hungerford, Nathan Debardeleben, Earl Lawrence
Organizations: Los Alamos National Laboratory, Los Alamos, New Mexico 87545, USA
Machine learning (ML) architectures are being developed for fast emulation of numerical simulators to accelerate both scientific discovery and engineering design. While these approaches are well-tested on popular PDE (partial differential equation) benchmarks, it is necessary to evaluate their performance on datasets that more accurately reflect industrial and engineering applications. We evaluate several computer vision-based architectures trained as surrogate models on a challenging computational physics task involving multiple materials, revealing significant implementation challenges for adapting certain pre-trained models, including transformer-based architectures. For such multi-physics multi-material datasets, our results underscore the importance of learning algorithms designed for sparse and imbalanced data, techniques for adapting spatial encodings and physical consistency metrics, in addition to methods for reducing the cost of generating sufficient training data.
Figures & tables
Fig. 1 : Area-average density of the initial state for a single PLI trajectory. The labeled materials consist of: (1) copper throw , ρ=8.94 ; (2) aluminum, ρ=2.71g/cm3 ; (3) polymer, ρ=1.04g/cm3 ; (4) high explosive, ρ=1.83g/cm3 ; (5) steel case , ρ=7.91g/cm3 ; and (6) air, ρ=0.120×10−2g/cm3 . Materials 1-3 form the liner, which varies in geometry defined by splines, across samples. Materials and physics of the simulation are detailed in the published dataset [ 1 ] .
Fig. 2 : Comparison of material evolution and deformation between (top) the ground truth, and (bottom) an emulation from the EDM-L model. Both consist of 100 steps with a physical timestep of Δtp=0.25μs , for a total duration of 25μs . This figure displays the area-average density across the 6 channels of density fields. RMSE values are computed as the sum of density-normalized RMSEs for each material, each weighted by the initial total mass.
Fig. 3 : Comparison of models by pixelwise error (lower is better) across test trajectories. Results are shown for (a) steel case density fraction and (b) copper throw density fraction. Vertical bars denote interquartile range. The vertical dashed line indicates the timestep at which deformation typically begins across all trajectories.
Fig. 4 : Comparison of models by Pearson correlation coefficient between predictions and ground truth (higher is better) across test trajectories. Results are shown for (a) steel case density fraction and (b) copper throw density fraction. Vertical bars denote interquartile range.
Fig. 5 : Comparison of models for copper throw density after correction for the optimal offset. Panel (a) shows the offset that maximizes agreement with the ground truth, while panel (b) reports the corresponding maximum correlation after alignment, averaged across test trajectories. Vertical bars denote interquartile range.
Fig. 6 : Comparison of models by mass conservation. The plotted quantity compares the total predicted mass relative to the total mass in the ground truth field. Results are shown for (a) steel case density and (b) copper throw density. Markers indicate the median across test trajectories and vertical bars denote the interquartile range.
Training/Finetuning
Inference
Model
# Traj
CPU node-hours
GPU-hours (h)
GPU-seconds per trajectory (s)
GPU-hours for test set (h)
EDM-S
3834
5751
432
1494
41.50
EDM-L
3834
5751
910
1968
54.67
SPUS finetuning
38
57
184
1.10
0.03
TABLE I : Computational cost of generating data, training the model, and inferring 100 trajectories for the test set.
Machine learning offers a promising avenue to accelerate physical simulations by replacing computationally expensive traditional Partial Differential Equation (PDE) solvers with fast, differentiable surrogate models. However, standard auto-regressive ML emulators often suffer from error accumulation over long horizons and struggle to capture the stochasticity of complex physical systems. In this paper, we propose Kastor, a comprehensive methodology to adapt a deterministic physics foundation model into a highly efficient and accurate generative surrogate. First, we introduce a two-stage inference scheme that combines a large-stride causal auto-regressive model with a non-causal temporal super-resolution network, significantly reducing error accumulation while minimizing computational cost. Second, we present Mean prediction regularization (MPR), a novel training objective that constrains the generative model to predict the deterministic distribution mean under null noise conditioning. This regularization dramatically improves the performance and stability of both Functional Generative Networks (FGN) and diffusion-based emulators. Finally, we demonstrate that incorporating spatial gradient matching improves the accuracy and physical fidelity of the simulations as measured by power spectrum density. Extensive evaluations on diverse simulation datasets of the benchmark The Well show that with these components, our model outperforms competing methods in forecasting accuracy, spectral consistency, and computational efficiency. Our model achieves a 42.9% average reduction in forecasting compared to our reference based on the Walrus finetuning methodology, and outperforms Walrus for 8 out of 10 datasets on variance-normalized RMSE (VRMSE).
Simulation is central to modern engineering and science, but the cost of numerical solvers for partial differential equations (PDEs) remains a bottleneck whenever fast or many-query evaluations are required. Neural emulators trained on solver-generated data promise significant speedups, yet they are usually framed as opaque alternatives to the very methods that produce their training signal. This thesis argues the two paradigms are more alike than different: neural architectures mirror classical discretizations, their errors are amenable to the same spectral analysis, and insight flows profitably in both directions. We approach the relationship by disentangling the multiple roles a solver plays in the emulator learning pipeline. Mode-wise Fourier analysis then provides a common language in which solver errors, architectural inductive biases, and training objectives can all be read off simultaneously. Taken together, this allows synthesizing three contributions. (1) APEBench, a comprehensive benchmarking suite for autoregressive neural emulators of PDEs that uses fast differentiable pseudo-spectral solvers in JAX. (2) Progressively Refined Differentiable Physics, an investigation of the effect of unconverged solvers on surrogate training. (3) Neural Emulator Superiority, an analysis of the influence of numerical errors and architectural inductive biases.
Felix Koehler
Technische Universität München TUM School of Computation, Information and Technology
Neural surrogate models for Partial Differential Equations (PDEs) on unstructured 3D geometries are often limited by poor generalization and the high cost of generating large-scale training datasets. Consequently, pre-training on massive datasets of related PDE dynamics has emerged as a critical alternative to enhance the robustness and scalability of these models. However, this strategy is neither compute- nor data-efficient, as it relies on massive pre-computed data that is very costly to generate. In this work, we introduce a disk-data-free pre-training framework tailored to both steady-state and transient regimes. For steady-state problems, we propose a geometry-driven strategy that leverages intrinsic shape descriptors to learn representations of complex 3D domains. For transient problems, we introduce a physics-driven approach based on online generation of synthetic PDE data, enabling scalable pre-training without reliance on expensive datasets. Across multiple experiments, our approach achieves faster convergence, greater data efficiency, and higher accuracy during fine-tuning, particularly under realistic low-data regimes. This methodology provides a practical pathway toward data-efficient neural emulators for large-scale simulations. Code and models are available at https://github.com/tum-pbs/GMP.
Luis Medrano-Navarro, Giacomo Baldan, Qiang Liu +4
Technical University of Munich · University of Stuttgart