P3D: Highly Scalable 3D Neural Surrogates for Physics Simulations

Authors: Benjamin Holzschuh, Georg Kohl, Florian Redinger, Nils Thuerey

P3D is a scalable framework for learning neural surrogates for high-resolution 3D physics simulations. The P3D model series combines convolutional blocks for fast local feature processing with windowed self-attention for deep representation learning in a hierarchical U-shaped architecture.

Quick Installation

# Install from PyPI
pip install p3d-surrogate

# Or install from source
git clone https://github.com/tum-pbs/P3D.git
cd P3D
pip install -e .

Usage

P3D is designed to efficiently process and predict the evolution of physical systems described by partial differential equations (PDEs) on 3D volumetric grids.

from p3d_surrogate import P3D_S
import torch

# Load pre-trained model
model = P3D_S(channel_size=4, channel_size_out=4).cuda()

# Run a forward pass on a 3D physics field
# Shape: (batch, channels, depth, height, width)
x_128 = torch.randn((1, 4, 128, 128, 128), dtype=torch.float32).cuda()
x_64 = torch.randn((1, 4, 64, 64, 64), dtype=torch.float32).cuda()

prediction_128 = model(x_128).sample
prediction_64 = model(x_64).sample

Examples

• Isotropic Turbulence Rollout: Autoregressive rollout of isotropic turbulence.
• Probabilistic Channel Flow Generation: Generate turbulent channel flow using P3D-L and visualize the results.
• More to come!

Datasets

• APE3D: Modified version of APEBench. For dataset generation, see data/simulations_apebench.
• Isotropic Turbulence: Available on JHTDB. See data/download_jhtdb for download instructions. We provide a single snapshot at t=420 on Hugging Face.
• Turbulence Channel Flow: Data generation based on PICT. Data is provided on Hugging Face.

Citation

@inproceedings{holzschuh2026p3d,
  title     = {P3D: Highly Scalable 3D Neural Surrogates for Physics
               Simulations with Global Context},
  author    = {Benjamin Holzschuh and Georg Kohl and
               Florian Redinger and Nils Thuerey},
  booktitle = {The Fourteenth International Conference on
               Learning Representations},
  year      = {2026},
  url       = {https://openreview.net/forum?id=8UdCE5nhFl}
}