README.md

---
license: apache-2.0
library_name: pytorch
tags:
- precipitation-nowcasting
- weather-forecasting
- video-transformer
- space-time-attention
- satellite-imagery
pipeline_tag: video-classification
base_model:
- facebook/timesformer-base-finetuned-k400
---

# SaTformer: A Space-Time Transformer for Precipitation Nowcasting

**Authors:** Levi Harris, Tianlong Chen — *The University of North Carolina at Chapel Hill*

[![arXiv](https://img.shields.io/badge/arXiv-2511.11090-b31b1b.svg)](https://arxiv.org/abs/2511.11090)
[![NeurIPS](https://img.shields.io/badge/NeurIPS_2025-1st_Place_CUMSUM-4b44ce.svg)](https://neurips.cc/virtual/2025/loc/san-diego/135896)
[![GitHub](https://img.shields.io/badge/GitHub-satformer-181717.svg?logo=github)](https://github.com/leharris3/satformer)

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from src.model.SaTformer.SaTformer import SaTformer

model = SaTformer(
    dim=512,
    num_frames=4,
    num_classes=64,
    image_size=32,
    patch_size=4,
    channels=11,
    depth=12,
    heads=8,
    dim_head=64,
    attn_dropout=0.1,
    ff_dropout=0.1,
    rotary_emb=False,
    attn="ST^2"
)

weights = hf_hub_download(repo_id="leharris3/satformer", filename="sf-64-cls.pt")
model.load_state_dict(torch.load(weights, weights_only=True), strict=False)
model.eval()

with torch.no_grad():
    x = torch.rand(1, 4, 11, 32, 32)  # (batch, frames, channels, H, W)
    logits = model(x)                  # -> [1, 64]
```

## Citation

```bibtex
@article{harris2025satformer,
  title={A Space-Time Transformer for Precipitation Forecasting},
  author={Harris, Levi and Chen, Tianlong},
  journal={arXiv preprint arXiv:2511.11090},
  year={2025}
}
```