ReViT

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

Installation

1. Create a Conda Environment

ReViT requires Python ≥ 3.10. Create and activate a dedicated conda environment before installing:

# Option A: from the provided environment.yml
conda env create -f environment.yml
conda activate revit

# Option B: manually
conda create -n revit python=3.10 -y
conda activate revit

Note — System Python will not work. Running pip install outside a virtual/conda environment on modern Debian/Ubuntu systems will fail with error: externally-managed-environment (see PEP 668). Always install inside a conda env or python -m venv.

2. Install the Package

Install from source (recommended):

git clone https://github.com/Howw-Way/ReViT.git
cd ReViT
pip install -e ".[all]"

Install from PyPI (once published):

pip install revit-pde              # core only
pip install revit-pde[datasets]    # + h5py, the-well
pip install revit-pde[train]       # + TensorBoard
pip install revit-pde[all]         # everything

Note: PyTorch must be installed with the correct CUDA version for your GPU. If the default pip install torch doesn't match your driver, install it first:

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

HPC Clusters

The setup script handles PyTorch CUDA selection automatically:

# Installs PyTorch (cu121) + ReViT into your active env
bash scripts/setup_env.sh

# CPU-only or different CUDA:
TORCH_CUDA=cpu bash scripts/setup_env.sh
TORCH_CUDA=cu124 bash scripts/setup_env.sh

# See all options:
bash scripts/setup_env.sh --help

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 P3D/TCF data adapters

Supported Models

Model	Params (approx.)	Use case
EqPDETiny	~1M	2D PDE, fast experiments
EqPDESmall	~5M	2D PDE, paper results
EqPDEBase	~15M	2D PDE, larger experiments
EqPDELarge	~50M	2D PDE, scaling study
EqPDE3DTiny	~1M	3D PDE (MHD_64, TCF)
EqPDE3DSmall	~5M	3D PDE, paper results
EqPDE3DBase	~15M	3D PDE, larger experiments
EqPDE3DLarge	~50M	3D PDE, scaling study
EqViT	varies	Image classification (rotMNIST, PCam)

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 (ρ + v + B)	~24.5 GB
P3D	3D Periodic 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 P3D

# Skip rotated test data (smaller download)
python scripts/download_dataset.py --no-rotated

# 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 / TD2D
│   ├── *_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
├── P3D/cropped/                    # P3D / TCF
│   ├── 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.

Training Entry Point

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

# 2D tasks (default decoder: local_query)
python scripts/Train_rollout.py --task KF2D --model_type EqPDETiny --epochs 100
python scripts/Train_rollout.py --task TD2D --model_type EqPDESmall --epochs 100
python scripts/Train_rollout.py --task ADV_2D --model_type EqPDETiny --epochs 100

# 3D tasks (default decoder: trilinear)
python scripts/Train_rollout.py --task MHD_64 --model_type EqPDE3DTiny --epochs 100
python scripts/Train_rollout.py --task P3D --model_type EqPDE3DSmall --epochs 100

Testing

pip install revit-pde[dev]
python -m pytest tests/ -v

Publishing to PyPI

For maintainers — one-time setup and release workflow:

# 1. Install build tools
pip install build twine

# 2. Build source distribution + wheel
python -m build

# 3. (Optional) Test on TestPyPI first
twine upload --repository testpypi dist/*
pip install --index-url https://test.pypi.org/simple/ revit-pde

# 4. Upload to PyPI
twine upload dist/*

You'll need a PyPI account and an API token. Store the token in ~/.pypirc:

[pypi]
username = __token__
password = pypi-AgEI...  # your API token

After uploading, pip install revit-pde will work for everyone.

Citation

If you use this code, please cite the ReViT paper. See Revit.pdf for details.