Levi Harris, Tianlong Chen
The University of North Carolina at Chapel Hill
This project uses uv for dependency management.
# install uv (if you don't have it)
curl -LsSf https://astral.sh/uv/install.sh | sh
# sync dependencies
uv syncDownload pretrained weights from HuggingFace:
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="leharris3/satformer", filename="sf-64-cls.pt", local_dir="weights")import torch
import warnings
from huggingface_hub import hf_hub_download
from src.model.SaTformer.SaTformer import SaTformer
# quiet some annoying UserWarnings thrown by xarray
# when opening datasets with phony_dims=None
warnings.simplefilter("ignore")
model = SaTformer(
dim=512,
num_frames=4, # number HRIT input frames
num_classes=64, # number precipitation bins to use
image_size=32, # HRIT input spatial dimensions
patch_size=4,
channels=11, # number HRIT radiance channels
depth=12, # number transformer encoder blocks
heads=8,
dim_head=64,
attn_dropout=0.1,
ff_dropout=0.1,
rotary_emb=False, # i.e., use postitional embeds
attn="ST^2"
)
WEIGHTS_FP = hf_hub_download(repo_id="leharris3/satformer", filename="sf-64-cls.pt")
model.load_state_dict(torch.load(WEIGHTS_FP, weights_only=True), strict=False);
model.eval()
with torch.no_grad():
inputs = torch.rand(1, 4, 11, 32, 32) # randomly generated HRIT input
logits = model(inputs) # call model forward pass
print(logits.shape) # -> [1, 64]; raw model probs over output classes
Model predicted cumulative mass function (CMF) for a random input.
satformer/
├── train.py # training entrypoint
├── test.py # inference entrypoint
├── demo.ipynb # interactive demo notebook
├── configs/ # training & test configs
├── scripts/ # launcher scripts
└── src/
├── model/
│ └── SaTformer/
│ ├── SaTformer.py # model architecture
│ └── rotary.py # rotary positional embeddings
├── dataloader/ # dataset & preprocessing
└── util/ # losses, metrics, logging, plotting
If you use this code in your research, please cite:
@article{harris2025satformer,
title={A Space-Time Transformer for Precipitation Forecasting},
author={Harris, Levi and Chen, Tianlong},
journal={arXiv preprint arXiv:2511.11090},
year={2025}
}