ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers

Paper · Project Page · Datasets · Pretrained Models

Public code for ICML Oral paper ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers.

MHD velocity

MHD magnetic

TCF velocity

Installation

Option A — Quick Install (pip, library use only)

Best if you just want to import revit in your own project:

conda create -n revit python=3.10 -y
conda activate revit

# Step 1: Install PyTorch for YOUR CUDA version (check nvidia-smi)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

# Step 2: Install ReViT
pip install revit-pde[all]         # or: revit-pde, revit-pde[datasets], revit-pde[train]

Option B — Development Install

git clone https://github.com/Howw-Way/ReViT.git
cd ReViT
conda env create -f environment.yml   # creates 'revit' env with Python 3.10+
conda activate revit

# Step 1: Install PyTorch for YOUR CUDA version (check nvidia-smi)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

# Step 2: Install ReViT in editable mode
pip install -e ".[all]"

Option C — HPC Clusters

The setup script auto-detects the platform and handles PyTorch CUDA selection:

conda create -n revit python=3.10 -y && conda activate revit
bash scripts/setup_env.sh                      # default: PyTorch + cu124
TORCH_CUDA=cu121 bash scripts/setup_env.sh     # override for CUDA 12.1
TORCH_CUDA=cpu bash scripts/setup_env.sh       # CPU-only
bash scripts/setup_env.sh --help               # all options

Quick Use

import torch
from revit import ReViT2DConfig, build_revit_2d

model = build_revit_2d(
    ReViT2DConfig(
        img_size=64,
        patch_size=4,
        depth=[1],
        embed_dim=48,
        num_heads=4,
    )
)
x = torch.randn(2, 2, 64, 64)
y = model(x)
print(y.shape)

The public API exposes:

• build_revit_2d: hierarchical 2D ReViT PDE model
• build_revit_3d: hierarchical 3D ReViT PDE model
• build_revit_classifier: ReViT image classifier for rotated image benchmarks
• revit.datasets: APEBench, MHD_64/The Well, and TCF data adapters

Testing

ReViT's core property is rotational equivariance. Verify it directly on small 2D and 3D models (CPU, a few seconds) with the standalone check:

python scripts/test_equivariance.py

It builds tiny ReViT2D/ReViT3D models and reports the worst-case error of model(R . x) vs R . model(x) over the square (C4) and cube (24-rotation) groups — near machine-zero means equivariance holds. The script also runs under pytest if preferred (pytest scripts/test_equivariance.py -v).

Supported Models

Model	Use case
ReViT2D	2D PDE rollout (KF2D, ADV_2D)
ReViT3D	3D PDE rollout (MHD_64, TCF)
EqViT	Image classification (rotMNIST)

The classifier's Python builder is build_revit_classifier (class ReViT4Class); EqViT is the name used on the training CLI (--model_type EqViT).

Datasets

All datasets are hosted on Hugging Face: huggingface.co/datasets/thuerey-group/ReViT

Dataset	Task	Grid	Channels	Size
KF2D	2D Kolmogorov Flow	160²	2 (velocity)	~3.3 GB
MHD_64	3D Magnetohydrodynamics	64³	7 (B + v + ρ)	~24.5 GB
TCF	3D Turbulent Channel Flow	96³	4 (v + p)	~6.3 GB

Download

# Install the downloader dependency
pip install huggingface_hub

# Download all datasets
python scripts/download_dataset.py

# Download specific dataset(s)
python scripts/download_dataset.py --dataset MHD_64 TCF

# Custom output directory
python scripts/download_dataset.py --output_dir /path/to/Dataset

# Verify existing download
python scripts/download_dataset.py --verify

Or download directly via the Python API:

from huggingface_hub import snapshot_download
snapshot_download(repo_id="thuerey-group/ReViT", repo_type="dataset", local_dir="./Dataset")

Dataset Layout

Set REVIT_DATASET_ROOT to the directory containing dataset folders:

export REVIT_DATASET_ROOT=/path/to/Dataset
# or create a symlink:
ln -s /path/to/Dataset ./Dataset

Expected directory structure:

Dataset/
├── apebench-scraped/data/          # KF2D
│   ├── *_train_V.npy
│   ├── *_test_V.npy
│   └── *_test_V_{angle}.npy        # rotated test data
├── Well/MHD_64/                    # MHD_64
│   ├── data/train/*.hdf5
│   ├── data/valid/*.hdf5
│   ├── data/valid/rotated/*.npy    # rotated validation data
│   └── stats.yaml
├── TCF/cropped/                    # TCF (turbulent channel flow)
│   ├── sim0_data_{train,test}.h5
│   └── rotated/sim0_data_test_rot_*.h5
└── ADV_2D/                         # Advection (coming soon)
    └── adv_*_{vector,scalar}_uniform_grid.pkl

The legacy environment variable EQUIVAR_DATASET_ROOT is still accepted for backward compatibility.

Pretrained Models

Two trained ReViT3D checkpoints are hosted on Hugging Face (huggingface.co/thuerey-group/ReViT, model repo) and load in one line:

from revit import load_pretrained, list_pretrained

list_pretrained()                          # {'TCF': ..., 'MHD_64': ...}
model = load_pretrained("TCF")             # downloads weights + config, returns an eval() model

import torch
y = model(torch.randn(1, 4, 96, 96, 96))   # TCF: 4 channels on a 96³ grid

Model	Task	Grid	Channels	Patch
TCF	3D turbulent channel flow	96³	4 (v + p)	3
MHD_64	3D magnetohydrodynamics	64³	7 (B + v + ρ)	2

load_pretrained needs huggingface_hub (bundled in revit-pde[datasets], or pip install huggingface_hub). To pre-fetch the raw files for offline use:

python scripts/download_models.py                  # all models -> ./Pretrained
python scripts/download_models.py --model TCF       # just one

Rollout inference with saved figures

load_pretrained returns a bare model for programmatic use. To reproduce the full evaluation of a local test run — autoregressive rollout over every test rotation, MSE per angle, and the same saved figures + .npy files — point the rollout trainer at a pretrained checkpoint with --pretrained instead of a local --model_id:

python scripts/Train_rollout.py --mode test --task TCF --pretrained TCF

python scripts/Train_rollout.py --mode test --task MHD_64 --pretrained MHD_64

Figures land in figures/test/angle_<name>/ and rollout .npy in test/angle_<name>/, under Results/<task>/pretrained_<name>/<timestamp>_step<N>/ — the same layout a --model_id test run produces. TensorBoard is optional (logging is skipped if it isn't installed; figure/.npy saving is unaffected).

Maintainers can (re)publish checkpoints from local training runs with scripts/upload_models.py (use --dry-run to stage without uploading).

Training Entry Point

The rollout trainer supports all ReViT models and the valid benchmark cases:

# 3D tasks
python scripts/Train_rollout.py --task MHD_64 --model_type ReViT3D --epochs 100
python scripts/Train_rollout.py --task TCF --model_type ReViT3D --epochs 100

# 2D tasks
python scripts/Train_rollout.py --task KF2D --model_type ReViT2D --epochs 100
python scripts/Train_rollout.py --task ADV_2D --model_type ReViT2D --epochs 100

Citation

If you use this code, please cite the ReViT paper.

@inproceedings{ReViT2026,
  title     = {{ReViT}: Rotational-equivariant Vision Transformers for Neural {PDE} Solvers},
  author    = {Hao Wei and Bjoern List and Nils Thuerey},
  booktitle = {Forty-Third International Conference on Machine Learning},
  year      = {2026},
}