ReViT - GitLite

KO | EN

gitlite — search

#tailwindcss #cli #swift #terminal #animation #c #open-source #web #proxy #shell

ReViT

(★ 9)

No description available.

파일 탐색기

ZIP 다운로드

- - logo_mark_previous.svg
  - revit_3d_mhd_magnetic.gif
  - revit_3d_mhd_velocity.gif
  - revit_3d_tcf_velocity.gif
- __init__.py
- config_helper.py
- config_revit.py
- config_revit3d.py
- config_ViT.py
- __init__.py
- Adv.py
- Apebench.py
- mnist_rot.py
- ROT_dataset.py
- TCF.py
- Well.py
- - __init__.py
  - Data_helper.py
  - Local_basis.py
- - __init__.py
  - argument.py
  - logger.py
  - logging.py
  - plotting.py
  - profiling.py
- __init__.py
- - __init__.py
  - revit_blocks.py
  - revit_blocks_3d.py
  - revit_classifier.py
  - revit_layers.py
  - revit_layers_3d.py
  - revit_transformer.py
  - revit_transformer_3d.py
- __init__.py
- Attentions.py
- compat.py
- Decoder.py
- Encoder.py
- eq_swin_layers.py
- equivariant_ops.py
- Positional_encoder.py
- __init__.py
- bases.py
- configs.py
- datasets.py
- models.py
- pretrained.py
- download_dataset.py
- download_models.py
- run_train_3d.sh
- setup_env.sh
- test_equivariance.py
- Train_rollout.py
- upload_models.py
- rollout_utils.py
- train_helper.py
.gitignore
environment.yml
LICENSE
pyproject.toml
README.md
requirements.txt

# CDN으로 사용하기

jsDelivr

jsDelivr는 공개 GitHub 리포지토리를 별도 설정 없이 CDN으로 즉시 서빙합니다. 버전과 파일을 고르면 웹페이지에 바로 붙일 수 있는 링크와 예시 코드가 만들어집니다.

버전

파일

링크

예시

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<table align="center">
  <tr>
    <td><img src=".github/assets/logo_mark_previous.svg" alt="ReViT logo" width="110"></td>
    <td><h1>ReViT: Rotational-equivariant Vision Transformers for Neural PDE Solvers</h1></td>
  </tr>
</table>

<p align="center">
  <a href="https://ge.in.tum.de/download/wei-icml2026-revit.pdf">Paper</a> ·
  <a href="https://tum-pbs.github.io/revit-web/">Project Page</a> ·
  <a href="https://huggingface.co/datasets/thuerey-group/ReViT">Datasets</a> ·
  <a href="https://huggingface.co/thuerey-group/ReViT">Pretrained Models</a>
</p>

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

<p align="center">
  <table>
    <tr>
      <td align="center" width="33%">
        <img src=".github/assets/revit_3d_mhd_velocity.gif" width="100%" alt="Animated ReViT MHD velocity prediction across 24 rotations">
        <br><sub><b>MHD velocity</b></sub>
      </td>
      <td align="center" width="33%">
        <img src=".github/assets/revit_3d_mhd_magnetic.gif" width="100%" alt="Animated ReViT MHD magnetic-field prediction across 24 rotations">
        <br><sub><b>MHD magnetic</b></sub>
      </td>
      <td align="center" width="33%">
        <img src=".github/assets/revit_3d_tcf_velocity.gif" width="100%" alt="Animated ReViT turbulent-channel-flow velocity prediction across 24 rotations">
        <br><sub><b>TCF velocity</b></sub>
      </td>
    </tr>
  </table>
</p>

## Installation

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

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

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

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

```bash
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:

```bash
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

```python
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:

```bash
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](https://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

```bash
# 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:

```python
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:

```bash
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](https://huggingface.co/thuerey-group/ReViT)**,
model repo) and load in one line:

```python
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:

```bash
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`:

```bash
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:

```bash

# 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},
}
```