# thuml/corrformer architecture

> Architecture diagram of the GitHub repository thuml/corrformer, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-26). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/thuml/corrformer
- Repository: https://github.com/thuml/corrformer

## Overview

1. **What this repo is**
   - This is a PyTorch implementation of **Corrformer**, a unified deep model for large-scale weather forecasting.
   - The code is organized around a single training/inference pipeline that can run on different station-scale datasets, especially the **Global Temp** and **Global Wind** benchmarks.

2. **Top-level execution path**
   - `run.py` is the main entry point for launching experiments.
   - The actual experiment logic is delegated into the `exp/` package, which handles the lifecycle of **train, validation, test, and demo/evaluation** runs.

3. **Experiment orchestration boundary**
   - `exp/exp_basic.py` defines shared experiment behavior and common setup.
   - `exp/exp_main.py` contains the concrete experiment implementation used for Corrformer.
   - This layer is the control plane: it wires together data loading, model creation, optimization, evaluation, and checkpoint usage.

4. **Data subsystem**
   - `data_provider/` encapsulates dataset loading and batching.
   - `data_factory.py` selects the appropriate dataset/dataloader setup for each experiment.
   - `data_loader.py` likely defines dataset classes and split handling for train/val/test.
   - The repo expects external data under `./dataset`, so raw data is not stored here.

5. **Model subsystem**
   - `models/Corrformer.py` is the top-level model definition.
   - It assembles the architecture from reusable components in `layers/`.
   - The model is designed around **spatiotemporal forecasting across many stations**, not a generic sequence classifier.

6. **Core architectural building blocks**
   - `layers/Embed.py` handles input embedding and feature projection.
   - `layers/Corrformer_EncDec.py` defines the encoder-decoder structure used in the paper.
   - `layers/Causal_Conv.py` supports causal convolution for cross-correlation modeling.
   - `layers/Multi_Correlation.py` contains the Multi-Correlation mechanism that unifies temporal and spatial dependencies.
   - `layers/` is therefore the main research/novelty layer of the repo.

7. **Utility layer**
   - `utils/` holds shared helpers that are not model-specific.
   - `metrics.py` provides evaluation metrics.
   - `timefeatures.py` builds time-based covariates.
   - `masking.py` likely supports attention or causal masking.
   - `tools.py` contains general training utilities such as logging, seeding, or checkpoint helpers.

8. **Training and evaluation outputs**
   - The repo writes learned weights and pretrained artifacts to `./checkpoints`.
   - The README indicates both full training runs and demo runs can load pretrained checkpoints.
   - This creates a clear boundary between **offline model training** and **later evaluation/demo inference**.

9. **Scripted experiment entry points**
   - `scripts/Global_Temp/` and `scripts/Global_Wind/` contain the production training scripts for the two main benchmarks.
   - `scripts/Demo/` contains shorter runs for testing pretrained models.
   - These scripts are the operational interface for reproducing paper results.

10. **End-to-end data flow**
   - External dataset files are placed in `./dataset`.
   - `data_provider/` loads and prepares them into model-ready batches.
   - `exp/` drives the run loop and feeds batches into `models/Corrformer.py`.
   - The model uses the specialized `layers/` modules to compute forecasts.
   - `utils/metrics.py` evaluates predictions, and results/checkpoints are persisted for later reuse.

11. **Technology and runtime assumptions**
   - The project is built for **Python + PyTorch** and was reproduced on a **single GPU** environment.
   - The README notes the intended setup is Linux with a CUDA-capable device.
   - This matters architecturally because the code and scripts are geared toward batch GPU training rather than distributed serving.

12. **Architectural shape to diagram**
   - Draw this as a pipeline with four main zones: **data ingestion**, **experiment orchestration**, **Corrformer model internals**, and **evaluation/checkpointing**.
   - Treat `layers/` as the internal model backbone, `exp/` as the control layer, and `data_provider/` + `utils/` as supporting infrastructure around them.

Main entry points: [`exp/exp_main.py`](https://github.com/thuml/corrformer/blob/main/exp/exp_main.py) (Main exp), [`run.py`](https://github.com/thuml/corrformer/blob/main/run.py) (run.py), [`models/Corrformer.py`](https://github.com/thuml/corrformer/blob/main/models/Corrformer.py) (Corrformer), [`layers/Corrformer_EncDec.py`](https://github.com/thuml/corrformer/blob/main/layers/Corrformer_EncDec.py) (Enc-Dec)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Entry & scripts"]
  node_run_py["run.py<br/>entrypoint<br/>[run.py]"]
  node_global_temp_script["Temp script<br/>benchmark script<br/>[Corrformer.sh]"]
  node_global_wind_script["Wind script<br/>benchmark script<br/>[Corrformer.sh]"]
  node_demo_temp_script["Temp demo<br/>demo script"]
  node_demo_wind_script["Wind demo<br/>demo script"]
end

subgraph group_data["Data ingestion"]
  node_data_factory{{"Data factory<br/>dataset selector<br/>[data_factory.py]"}}
  node_data_loader[("Data loader<br/>dataset/batching<br/>[data_loader.py]")]
end

subgraph group_control["Experiment control"]
  node_exp_basic["Base exp<br/>experiment base<br/>[exp_basic.py]"]
  node_exp_main["Main exp<br/>corrformer experiment<br/>[exp_main.py]"]
end

subgraph group_model["Corrformer model"]
  node_corrformer_model["Corrformer<br/>forecast model<br/>[Corrformer.py]"]
  node_embed["Embedding<br/>[Embed.py]"]
  node_encdec["Enc-Dec<br/>model backbone"]
  node_causal_conv["Causal conv<br/>temporal conv<br/>[Causal_Conv.py]"]
  node_multi_corr["Multi-corr<br/>correlation module"]
end

subgraph group_eval["Evaluation & utilities"]
  node_metrics["Metrics<br/>evaluation<br/>[metrics.py]"]
  node_timefeatures["Time feats<br/>feature utility<br/>[timefeatures.py]"]
  node_masking["Masking<br/>attention utility<br/>[masking.py]"]
  node_tools["Tools<br/>training utils<br/>[tools.py]"]
end

node_run_py -->|"launches"| node_exp_basic
node_run_py -.->|"used by"| node_global_temp_script
node_run_py -.->|"used by"| node_global_wind_script
node_run_py -.->|"used by"| node_demo_temp_script
node_run_py -.->|"used by"| node_demo_wind_script
node_exp_basic -->|"extends"| node_exp_main
node_exp_main -->|"loads data"| node_data_factory
node_data_factory -->|"builds"| node_data_loader
node_exp_main -->|"builds model"| node_corrformer_model
node_exp_main -->|"evaluates"| node_metrics
node_exp_main -->|"uses"| node_tools
node_exp_main -->|"uses"| node_timefeatures
node_corrformer_model -->|"assembles"| node_embed
node_corrformer_model -->|"assembles"| node_encdec
node_encdec -->|"uses"| node_causal_conv
node_encdec -->|"uses"| node_multi_corr
node_multi_corr -->|"uses"| node_masking
node_data_loader -->|"augments"| node_timefeatures
node_exp_main -->|"trains"| node_corrformer_model
node_metrics -.->|"reports via"| node_tools

click node_run_py "https://github.com/thuml/corrformer/blob/main/run.py"
click node_global_temp_script "https://github.com/thuml/corrformer/blob/main/scripts/Global_Temp/Corrformer.sh"
click node_global_wind_script "https://github.com/thuml/corrformer/blob/main/scripts/Global_Wind/Corrformer.sh"
click node_demo_temp_script "https://github.com/thuml/corrformer/blob/main/scripts/Demo/Global_Temp_demo.sh"
click node_demo_wind_script "https://github.com/thuml/corrformer/blob/main/scripts/Demo/Global_Wind_demo.sh"
click node_data_factory "https://github.com/thuml/corrformer/blob/main/data_provider/data_factory.py"
click node_data_loader "https://github.com/thuml/corrformer/blob/main/data_provider/data_loader.py"
click node_exp_basic "https://github.com/thuml/corrformer/blob/main/exp/exp_basic.py"
click node_exp_main "https://github.com/thuml/corrformer/blob/main/exp/exp_main.py"
click node_corrformer_model "https://github.com/thuml/corrformer/blob/main/models/Corrformer.py"
click node_embed "https://github.com/thuml/corrformer/blob/main/layers/Embed.py"
click node_encdec "https://github.com/thuml/corrformer/blob/main/layers/Corrformer_EncDec.py"
click node_causal_conv "https://github.com/thuml/corrformer/blob/main/layers/Causal_Conv.py"
click node_multi_corr "https://github.com/thuml/corrformer/blob/main/layers/Multi_Correlation.py"
click node_metrics "https://github.com/thuml/corrformer/blob/main/utils/metrics.py"
click node_timefeatures "https://github.com/thuml/corrformer/blob/main/utils/timefeatures.py"
click node_masking "https://github.com/thuml/corrformer/blob/main/utils/masking.py"
click node_tools "https://github.com/thuml/corrformer/blob/main/utils/tools.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_run_py,node_global_temp_script,node_global_wind_script,node_demo_temp_script,node_demo_wind_script toneBlue
class node_data_factory,node_data_loader toneAmber
class node_exp_basic,node_exp_main toneMint
class node_corrformer_model,node_embed,node_encdec,node_causal_conv,node_multi_corr toneRose
class node_metrics,node_timefeatures,node_masking,node_tools toneIndigo
```

## Components

### Entry & scripts

Command-line entry points and benchmark launch scripts.

- [run.py](https://github.com/thuml/corrformer/blob/main/run.py): entrypoint · Main launcher for training, validation, testing, and demo runs. · `run.py`
- [Temp script](https://github.com/thuml/corrformer/blob/main/scripts/Global_Temp/Corrformer.sh): benchmark script · Training script for the Global Temp benchmark. · `scripts/Global_Temp/Corrformer.sh`
- [Wind script](https://github.com/thuml/corrformer/blob/main/scripts/Global_Wind/Corrformer.sh): benchmark script · Training script for the Global Wind benchmark. · `scripts/Global_Wind/Corrformer.sh`
- [Temp demo](https://github.com/thuml/corrformer/blob/main/scripts/Demo/Global_Temp_demo.sh): demo script · Short demo run for pretrained evaluation on Global Temp. · `scripts/Demo/Global_Temp_demo.sh`
- [Wind demo](https://github.com/thuml/corrformer/blob/main/scripts/Demo/Global_Wind_demo.sh): demo script · Short demo run for pretrained evaluation on Global Wind. · `scripts/Demo/Global_Wind_demo.sh`

### Data ingestion

Dataset selection, loading, and batch preparation.

- [Data factory](https://github.com/thuml/corrformer/blob/main/data_provider/data_factory.py): dataset selector · Chooses the dataset and dataloader configuration for each experiment. · `data_provider/data_factory.py`
- [Data loader](https://github.com/thuml/corrformer/blob/main/data_provider/data_loader.py): dataset/batching · Defines dataset splits and batch construction for train, validation, and test. · `data_provider/data_loader.py`

### Experiment control

Shared experiment lifecycle and Corrformer-specific run logic.

- [Base exp](https://github.com/thuml/corrformer/blob/main/exp/exp_basic.py): experiment base · Shared experiment setup and lifecycle behavior. · `exp/exp_basic.py`
- [Main exp](https://github.com/thuml/corrformer/blob/main/exp/exp_main.py): corrformer experiment · Concrete Corrformer experiment implementation and run loop. · `exp/exp_main.py`

### Corrformer model

Model assembly and core architectural blocks.

- [Corrformer](https://github.com/thuml/corrformer/blob/main/models/Corrformer.py): forecast model · Top-level model that assembles the full forecasting architecture. · `models/Corrformer.py`
- [Embedding](https://github.com/thuml/corrformer/blob/main/layers/Embed.py): input embedding · Projects raw station and time features into model space. · `layers/Embed.py`
- [Enc-Dec](https://github.com/thuml/corrformer/blob/main/layers/Corrformer_EncDec.py): model backbone · Encoder-decoder structure used by Corrformer. · `layers/Corrformer_EncDec.py`
- [Causal conv](https://github.com/thuml/corrformer/blob/main/layers/Causal_Conv.py): temporal conv · Causal convolution block for cross-correlation modeling. · `layers/Causal_Conv.py`
- [Multi-corr](https://github.com/thuml/corrformer/blob/main/layers/Multi_Correlation.py): correlation module · Core mechanism for unifying temporal and spatial dependencies. · `layers/Multi_Correlation.py`

### Evaluation & utilities

Metrics, time features, masking, and training helpers.

- [Metrics](https://github.com/thuml/corrformer/blob/main/utils/metrics.py): evaluation · Forecast error metrics and result scoring. · `utils/metrics.py`
- [Time feats](https://github.com/thuml/corrformer/blob/main/utils/timefeatures.py): feature utility · Builds time-based covariates for the input pipeline. · `utils/timefeatures.py`
- [Masking](https://github.com/thuml/corrformer/blob/main/utils/masking.py): attention utility · Mask construction for causal or attention-based computation. · `utils/masking.py`
- [Tools](https://github.com/thuml/corrformer/blob/main/utils/tools.py): training utils · General helpers such as logging, seeding, and checkpoint support. · `utils/tools.py`

## Connections

- run.py → Base exp: launches · CLI entry point instantiates the experiment layer.
- run.py → Temp script: used by · Benchmark script drives the main launcher for Global Temp.
- run.py → Wind script: used by · Benchmark script drives the main launcher for Global Wind.
- run.py → Temp demo: used by · Demo script drives pretrained evaluation for Global Temp.
- run.py → Wind demo: used by · Demo script drives pretrained evaluation for Global Wind.
- Base exp → Main exp: extends · Concrete Corrformer experiment inherits shared experiment behavior.
- Main exp → Data factory: loads data · Experiment selects the appropriate dataset pipeline.
- Data factory → Data loader: builds · Factory instantiates the dataset and batching logic.
- Main exp → Corrformer: builds model · Experiment creates the Corrformer network for training or inference.
- Main exp → Metrics: evaluates · Experiment computes forecast quality during validation and test.
- Main exp → Tools: uses · Experiment relies on shared helpers for runtime support.
- Main exp → Time feats: uses · Experiment consumes time covariates from the utility layer.
- Corrformer → Embedding: assembles · Top-level model starts with input embedding.
- Corrformer → Enc-Dec: assembles · Model backbone uses the encoder-decoder structure.
- Enc-Dec → Causal conv: uses · Backbone incorporates causal convolution for temporal correlation.
- Enc-Dec → Multi-corr: uses · Backbone centers on the multi-correlation mechanism.
- Multi-corr → Masking: uses · Correlation module depends on masking for constrained computation.
- Data loader → Time feats: augments · Dataset processing enriches inputs with temporal features.
- Main exp → Corrformer: trains · Control plane feeds batches into the model for optimization and inference.
- Metrics → Tools: reports via · Metric results are handled alongside shared training utilities.