# tum-pbs/pbfm architecture

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

- Interactive diagram: https://gitdiagram.com/tum-pbs/pbfm
- Repository: https://github.com/tum-pbs/pbfm

## Overview

PBFM is a physics-constrained flow-matching implementation for generating fields in three PDE benchmarks: Darcy flow, dynamic stall, and Kolmogorov flow. The main workflow loads benchmark data, trains a time-conditioned model with flow-matching and physics-residual objectives, saves checkpoints, then samples fields from a selected checkpoint for evaluation. The README also highlights unrolled training, conflict-free gradient updates, and stochastic sampling; source excerpts confirm some training and sampling details, but not all benchmark-specific wiring, so uncertain links are omitted.

Main entry points: [`darcy_flow/train_ddp.py`](https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/train_ddp.py) (Darcy training), [`darcy_flow/logs/PBFM/checkpoints`](https://github.com/tum-pbs/pbfm/tree/main/darcy_flow/logs/PBFM/checkpoints) (Model checkpoints), [`dynamic_stall/train_ddp.py`](https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/train_ddp.py) (Stall training), [`kolmogorov_flow/train_ddp.py`](https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/train_ddp.py) (Kolmogorov training)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Benchmark data"]
  node_darcy_loader["Darcy loader<br/>[dataset_loader.py]"]
  node_stall_loader["Stall loader"]
  node_kolmogorov_loader["Kolmogorov loader"]
end

subgraph group_training["Training"]
  node_darcy_train["Darcy training<br/>[train_ddp.py]"]
  node_stall_train["Stall training<br/>[train_ddp.py]"]
  node_kolmogorov_train["Kolmogorov training<br/>[train_ddp.py]"]
  node_darcy_flow["Darcy flow matching"]
  node_stall_flow["Stall flow matching<br/>[flow_matching.py]"]
  node_kolmogorov_flow["Kolmogorov flow matching<br/>[flow_matching.py]"]
  node_grad_helper["Gradient constraints<br/>[grad_utils.py]"]
end

subgraph group_generation["Sampling and evaluation"]
  node_darcy_sample["Darcy sampling<br/>[sample.py]"]
  node_stall_sample["Stall sampling<br/>[sample.py]"]
  node_kolmogorov_sample["Kolmogorov sampling<br/>[sample.py]"]
  node_checkpoints[("Model checkpoints")]
  node_generated["Generated fields"]
end

subgraph group_models["Neural models"]
  node_darcy_model["Darcy DiT / U-Net"]
  node_stall_model["Stall DiT<br/>[dit.py]"]
  node_kolmogorov_model["Kolmogorov DiT<br/>[dit.py]"]
end

node_researcher(("Researcher"))
node_hf_data[("HF datasets")]
node_pidm_data[("PIDM Darcy data")]

node_researcher -->|"starts training"| node_darcy_train
node_researcher -->|"starts training"| node_stall_train
node_researcher -->|"starts training"| node_kolmogorov_train
node_researcher -->|"requests samples"| node_darcy_sample
node_researcher -->|"requests samples"| node_stall_sample
node_researcher -->|"requests samples"| node_kolmogorov_sample
node_pidm_data -.->|"provides data"| node_darcy_loader
node_hf_data -.->|"provides data"| node_stall_loader
node_hf_data -.->|"provides data"| node_kolmogorov_loader
node_darcy_train -->|"loads batches"| node_darcy_loader
node_stall_train -->|"loads batches"| node_stall_loader
node_kolmogorov_train -->|"loads batches"| node_kolmogorov_loader
node_darcy_train -->|"computes objective"| node_darcy_flow
node_stall_train -->|"computes objective"| node_stall_flow
node_kolmogorov_train -->|"computes objective"| node_kolmogorov_flow
node_darcy_train -->|"instantiates model"| node_darcy_model
node_stall_train -->|"instantiates model"| node_stall_model
node_kolmogorov_train -->|"instantiates model"| node_kolmogorov_model
node_darcy_flow -->|"queries model"| node_darcy_model
node_stall_flow -->|"queries model"| node_stall_model
node_kolmogorov_flow -->|"queries model"| node_kolmogorov_model
node_darcy_train -->|"uses gradients"| node_grad_helper
node_darcy_train -->|"saves weights"| node_checkpoints
node_stall_train -->|"saves weights"| node_checkpoints
node_kolmogorov_train -->|"saves weights"| node_checkpoints
node_darcy_sample -->|"loads test data"| node_darcy_loader
node_stall_sample -->|"loads test data"| node_stall_loader
node_kolmogorov_sample -->|"loads test data"| node_kolmogorov_loader
node_darcy_sample -->|"loads weights"| node_checkpoints
node_stall_sample -->|"loads weights"| node_checkpoints
node_kolmogorov_sample -->|"loads weights"| node_checkpoints
node_darcy_sample -->|"runs sampler"| node_darcy_flow
node_stall_sample -->|"runs sampler"| node_stall_flow
node_kolmogorov_sample -->|"runs sampler"| node_kolmogorov_flow
node_darcy_sample -->|"produces fields"| node_generated
node_stall_sample -->|"produces fields"| node_generated
node_kolmogorov_sample -->|"produces fields"| node_generated

click node_darcy_loader "https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/dataset_loader.py"
click node_stall_loader "https://github.com/tum-pbs/pbfm/tree/main/dynamic_stall"
click node_kolmogorov_loader "https://github.com/tum-pbs/pbfm/tree/main/kolmogorov_flow"
click node_darcy_train "https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/train_ddp.py"
click node_stall_train "https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/train_ddp.py"
click node_kolmogorov_train "https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/train_ddp.py"
click node_darcy_flow "https://github.com/tum-pbs/pbfm/tree/main/darcy_flow"
click node_stall_flow "https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/flow_matching.py"
click node_kolmogorov_flow "https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/flow_matching.py"
click node_grad_helper "https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/grad_utils.py"
click node_darcy_model "https://github.com/tum-pbs/pbfm/tree/main/darcy_flow"
click node_stall_model "https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/dit.py"
click node_kolmogorov_model "https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/dit.py"
click node_darcy_sample "https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/sample.py"
click node_stall_sample "https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/sample.py"
click node_kolmogorov_sample "https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/sample.py"
click node_checkpoints "https://github.com/tum-pbs/pbfm/tree/main/darcy_flow/logs/PBFM/checkpoints"

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_darcy_loader,node_stall_loader,node_kolmogorov_loader toneBlue
class node_darcy_train,node_stall_train,node_kolmogorov_train,node_darcy_flow,node_stall_flow,node_kolmogorov_flow,node_grad_helper,node_hf_data,node_pidm_data toneAmber
class node_darcy_sample,node_stall_sample,node_kolmogorov_sample,node_checkpoints,node_generated toneMint
class node_darcy_model,node_stall_model,node_kolmogorov_model toneRose
class node_researcher toneIndigo
```

## Components

### Benchmark data

- [Darcy loader](https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/dataset_loader.py): `darcy_flow/dataset_loader.py`
- [Stall loader](https://github.com/tum-pbs/pbfm/tree/main/dynamic_stall): `dynamic_stall`
- [Kolmogorov loader](https://github.com/tum-pbs/pbfm/tree/main/kolmogorov_flow): `kolmogorov_flow`

### Training

- [Darcy training](https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/train_ddp.py): `darcy_flow/train_ddp.py`
- [Stall training](https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/train_ddp.py): `dynamic_stall/train_ddp.py`
- [Kolmogorov training](https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/train_ddp.py): `kolmogorov_flow/train_ddp.py`
- [Darcy flow matching](https://github.com/tum-pbs/pbfm/tree/main/darcy_flow): `darcy_flow`
- [Stall flow matching](https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/flow_matching.py): `dynamic_stall/flow_matching.py`
- [Kolmogorov flow matching](https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/flow_matching.py): `kolmogorov_flow/flow_matching.py`
- [Gradient constraints](https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/grad_utils.py): `darcy_flow/grad_utils.py`

### Sampling and evaluation

- [Darcy sampling](https://github.com/tum-pbs/pbfm/blob/main/darcy_flow/sample.py): `darcy_flow/sample.py`
- [Stall sampling](https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/sample.py): `dynamic_stall/sample.py`
- [Kolmogorov sampling](https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/sample.py): `kolmogorov_flow/sample.py`
- [Model checkpoints](https://github.com/tum-pbs/pbfm/tree/main/darcy_flow/logs/PBFM/checkpoints): `darcy_flow/logs/PBFM/checkpoints`
- **Generated fields**

### Neural models

- [Darcy DiT / U-Net](https://github.com/tum-pbs/pbfm/tree/main/darcy_flow): `darcy_flow`
- [Stall DiT](https://github.com/tum-pbs/pbfm/blob/main/dynamic_stall/dit.py): `dynamic_stall/dit.py`
- [Kolmogorov DiT](https://github.com/tum-pbs/pbfm/blob/main/kolmogorov_flow/dit.py): `kolmogorov_flow/dit.py`

### Other components

- **Researcher**
- **HF datasets**
- **PIDM Darcy data**

## Connections

- Researcher → Darcy training: starts training
- Researcher → Stall training: starts training
- Researcher → Kolmogorov training: starts training
- Researcher → Darcy sampling: requests samples
- Researcher → Stall sampling: requests samples
- Researcher → Kolmogorov sampling: requests samples
- PIDM Darcy data → Darcy loader: provides data
- HF datasets → Stall loader: provides data
- HF datasets → Kolmogorov loader: provides data
- Darcy training → Darcy loader: loads batches
- Stall training → Stall loader: loads batches
- Kolmogorov training → Kolmogorov loader: loads batches
- Darcy training → Darcy flow matching: computes objective
- Stall training → Stall flow matching: computes objective
- Kolmogorov training → Kolmogorov flow matching: computes objective
- Darcy training → Darcy DiT / U-Net: instantiates model
- Stall training → Stall DiT: instantiates model
- Kolmogorov training → Kolmogorov DiT: instantiates model
- Darcy flow matching → Darcy DiT / U-Net: queries model
- Stall flow matching → Stall DiT: queries model
- Kolmogorov flow matching → Kolmogorov DiT: queries model
- Darcy training → Gradient constraints: uses gradients
- Darcy training → Model checkpoints: saves weights
- Stall training → Model checkpoints: saves weights
- Kolmogorov training → Model checkpoints: saves weights
- Darcy sampling → Darcy loader: loads test data
- Stall sampling → Stall loader: loads test data
- Kolmogorov sampling → Kolmogorov loader: loads test data
- Darcy sampling → Model checkpoints: loads weights
- Stall sampling → Model checkpoints: loads weights
- Kolmogorov sampling → Model checkpoints: loads weights
- Darcy sampling → Darcy flow matching: runs sampler
- Stall sampling → Stall flow matching: runs sampler
- Kolmogorov sampling → Kolmogorov flow matching: runs sampler
- Darcy sampling → Generated fields: produces fields
- Stall sampling → Generated fields: produces fields
- Kolmogorov sampling → Generated fields: produces fields