# yellowshippo/fluxgnn-icml2024 architecture

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

- Interactive diagram: https://gitdiagram.com/yellowshippo/fluxgnn-icml2024
- Repository: https://github.com/yellowshippo/fluxgnn-icml2024

## Overview

1. This repository is a research codebase for PDE solving with graph neural networks. It is organized around two problem families: convection–diffusion and mixture. The main architectural idea is config-driven training/evaluation of multiple solver families against shared datasets.

2. The core model stack lives under `fluxgnn/`. This package contains the FluxGNN implementation itself (`fluxgnn.py`), PDE/problem-specific wrappers (`navier_stokes.py`, `poisson.py`, `mixture.py`), finite-volume support (`torch_fv.py`, `fv_data.py`), and shared infrastructure for datasets, settings, and utility code. In practice, this is the “engine” that loads data, builds the solver, and runs inference/training.

3. `fluxgnn/dataset.py` and `fluxgnn/setting.py` define the boundary between data/configuration and the model code. The repo is heavily YAML-driven: experiment behavior, dataset paths, and model hyperparameters are supplied through `inputs/**.yml`, and these settings flow into the FluxGNN experiment runner.

4. The `experiments/` directory is for postprocessing and experiment orchestration rather than model definition. `run.py` appears to be the execution entrypoint, `stats.py` computes aggregate metrics/results, and `generate_vtu.py` exports outputs to VTU for visualization in mesh-based tools. This separates simulation runtime from analysis and visualization.

5. Baseline methods are isolated from FluxGNN in their own top-level package `penn/`, which contains `train.py` and `eval.py`. That separation suggests a distinct training/evaluation pipeline for the PENN baseline, even though it shares datasets and experiment inputs with the main model family.

6. The `inputs/` tree is the experiment matrix. It contains many YAML files grouped by task and solver family (`convection_diffusion/fluxgnn`, `mixture/fluxgnn`, `mixture/fv`, `mixture/penn`). The architecture is therefore not hardcoded around one experiment; it is parameterized by input files that select the solver, replication settings, and finite-volume/baseline variants.

7. Data is externally supplied and stored under `data/`. The README shows three distinct roles there: raw/primary training data, intermediate/preprocessed forms for baseline methods, and pretrained models. For the mixture task, the split into `raw`, `interim`, `preprocessed`, `transformed`, and `taller` indicates that different solver families consume different representations of the same underlying simulation data.

8. Pretrained weights are treated as first-class runtime assets, not as source code. The repository expects them in `data/pretrained/**`, which means evaluation can run without retraining and architecture diagrams should show an explicit model-artifact store feeding inference.

9. The repository includes a vendored or linked external component path via `lib/siml` and a git submodule `MP-Neural-PDE-Solvers`. That strongly suggests some baseline/utility code is maintained outside this repo and integrated as an external dependency boundary rather than being rewritten here.

10. `Makefile` is the main user-facing command layer. The README’s train/eval commands map to make targets such as `cd_fluxgnn_train`, `mixture_penn_eval`, and `mixture_fvm_eval`, so the operational workflow is: choose a task/solver target, load the matching YAML config, then dispatch into the corresponding Python package.

11. The test layout mirrors the two main problem domains: `tests/cd` and `tests/mixture`. This reinforces that the repository’s architecture is split primarily by PDE family, with shared solver infrastructure reused across both.

12. At a high level, the data flow is: external datasets enter `data/`, config files under `inputs/` select a solver and task, Python packages under `fluxgnn/` or `penn/` run training/evaluation, and `experiments/` plus pretrained assets support analysis and reproducible evaluation. The main boundaries are between shared infrastructure, solver-specific code, baseline code, and data/model artifacts.

Main entry points: [`fluxgnn/fluxgnn.py`](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/fluxgnn.py) (FluxGNN), [`fluxgnn/experiment.py`](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/experiment.py) (Experiment), [`penn`](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/penn) (PENN), [`Makefile`](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/Makefile) (Makefile)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_core["Core solver stack"]
  node_experiment_runner["Experiment<br/>runner<br/>[experiment.py]"]
  node_fluxgnn_core["FluxGNN<br/>model engine<br/>[fluxgnn.py]"]
  node_dataset_layer["Dataset<br/>data boundary<br/>[dataset.py]"]
  node_setting_layer["Setting<br/>config boundary<br/>[setting.py]"]
  node_fv_support["FV support<br/>finite volume<br/>[torch_fv.py]"]
  node_fv_data["FV data<br/>[fv_data.py]"]
  node_navier_poisson["PDE wrappers<br/>problem wrappers<br/>[navier_stokes.py]"]
  node_mixture_wrapper["Mixture<br/>problem wrapper<br/>[mixture.py]"]
end

subgraph group_orchestration["Run orchestration"]
  node_makefile["Makefile<br/>command layer"]
  node_experiments_dir["Experiments<br/>postprocess"]
  node_penn_pkg["PENN<br/>baseline pipeline"]
end

subgraph group_configs["Inputs"]
  node_inputs_dir["Inputs<br/>yaml matrix"]
  node_cd_inputs["CD configs<br/>yaml set"]
  node_mixture_inputs["Mixture configs<br/>yaml set"]
end

subgraph group_assets["Data and artifacts"]
  node_data_dir[("Data<br/>artifacts store")]
  node_sample_data[("Sample data<br/>training sample")]
  node_pretrained_store[("Pretrained<br/>model store")]
end

subgraph group_external["External deps"]
  node_siml_lib["siml<br/>vendored lib"]
  node_mp_neural_pde["MP-Neural<br/>submodule"]
end

subgraph group_validation["Tests"]
  node_tests_cd["CD tests<br/>domain tests"]
  node_tests_mixture["Mixture tests<br/>domain tests"]
end

node_makefile -->|"dispatches"| node_experiment_runner
node_makefile -->|"targets"| node_penn_pkg
node_makefile -->|"postprocess"| node_experiments_dir
node_inputs_dir -->|"configures"| node_experiment_runner
node_cd_inputs -->|"selects"| node_fluxgnn_core
node_mixture_inputs -->|"selects"| node_fluxgnn_core
node_mixture_inputs -->|"selects"| node_penn_pkg
node_dataset_layer -->|"loads"| node_data_dir
node_setting_layer -->|"parses"| node_inputs_dir
node_experiment_runner -->|"uses"| node_dataset_layer
node_experiment_runner -->|"uses"| node_setting_layer
node_experiment_runner -->|"runs"| node_fluxgnn_core
node_fluxgnn_core -->|"uses"| node_fv_support
node_fv_support -->|"reads"| node_fv_data
node_navier_poisson -->|"wraps"| node_fluxgnn_core
node_mixture_wrapper -->|"wraps"| node_fluxgnn_core
node_experiments_dir -->|"analyzes"| node_data_dir
node_experiments_dir -->|"exports"| node_pretrained_store
node_penn_pkg -->|"trains on"| node_data_dir
node_penn_pkg -->|"depends on"| node_siml_lib
node_penn_pkg -->|"integrates"| node_mp_neural_pde
node_tests_cd -->|"covers"| node_cd_inputs
node_tests_mixture -->|"covers"| node_mixture_inputs

click node_makefile "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/Makefile"
click node_experiment_runner "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/experiment.py"
click node_fluxgnn_core "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/fluxgnn.py"
click node_dataset_layer "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/dataset.py"
click node_setting_layer "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/setting.py"
click node_fv_support "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/torch_fv.py"
click node_fv_data "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/fv_data.py"
click node_navier_poisson "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/navier_stokes.py"
click node_mixture_wrapper "https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/mixture.py"
click node_experiments_dir "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/experiments"
click node_penn_pkg "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/penn"
click node_inputs_dir "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs"
click node_cd_inputs "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs/convection_diffusion"
click node_mixture_inputs "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs/mixture"
click node_data_dir "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data"
click node_sample_data "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data/sample_data"
click node_pretrained_store "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data"
click node_siml_lib "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/lib/siml"
click node_mp_neural_pde "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/MP-Neural-PDE-Solvers"
click node_tests_cd "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/tests/cd"
click node_tests_mixture "https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/tests/mixture"

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_experiment_runner,node_fluxgnn_core,node_dataset_layer,node_setting_layer,node_fv_support,node_fv_data,node_navier_poisson,node_mixture_wrapper toneBlue
class node_makefile,node_experiments_dir,node_penn_pkg toneAmber
class node_inputs_dir,node_cd_inputs,node_mixture_inputs toneMint
class node_data_dir,node_sample_data,node_pretrained_store toneRose
class node_siml_lib,node_mp_neural_pde toneIndigo
class node_tests_cd,node_tests_mixture toneTeal
```

## Components

### Core solver stack

Shared FluxGNN engine and PDE-specific wrappers.

- [Experiment](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/experiment.py): runner · Central experiment orchestration and config plumbing. · `fluxgnn/experiment.py`
- [FluxGNN](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/fluxgnn.py): model engine · Main graph neural PDE solver implementation. · `fluxgnn/fluxgnn.py`
- [Dataset](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/dataset.py): data boundary · Dataset loading and preparation boundary for shared tasks. · `fluxgnn/dataset.py`
- [Setting](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/setting.py): config boundary · Typed settings and experiment configuration handling. · `fluxgnn/setting.py`
- [FV support](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/torch_fv.py): finite volume · Finite-volume data and operator helpers used by solver variants. · `fluxgnn/torch_fv.py`
- [FV data](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/fv_data.py): fv data · Finite-volume specific data structures and transforms. · `fluxgnn/fv_data.py`
- [PDE wrappers](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/navier_stokes.py): problem wrappers · Problem-specific wrappers for Navier-Stokes and Poisson style tasks. · `fluxgnn/navier_stokes.py`
- [Mixture](https://github.com/yellowshippo/fluxgnn-icml2024/blob/main/fluxgnn/mixture.py): problem wrapper · Mixture-task model wrapper and task-specific wiring. · `fluxgnn/mixture.py`

### Run orchestration

Entry points and experiment execution layer.

- [Makefile](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/Makefile): command layer · User-facing targets for training, evaluation, and task dispatch. · `Makefile`
- [Experiments](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/experiments): postprocess · Run scripts, stats, and export utilities for analysis and visualization. · `experiments`
- [PENN](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/penn): baseline pipeline · Separate baseline training and evaluation pipeline. · `penn`

### Inputs

YAML experiment matrix and task-specific settings.

- [Inputs](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs): yaml matrix · Task and solver YAML configurations selecting experiments and hyperparameters. · `inputs`
- [CD configs](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs/convection_diffusion): yaml set · Convection–diffusion experiment definitions. · `inputs/convection_diffusion`
- [Mixture configs](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/inputs/mixture): yaml set · Mixture-task experiment definitions for FluxGNN, FV, and PENN. · `inputs/mixture`

### Data and artifacts

External datasets, preprocessing outputs, and pretrained weights.

- [Data](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data): artifacts store · External datasets and runtime artifacts. · `data`
- [Sample data](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data/sample_data): training sample · Small bundled dataset for local runs and tests. · `data/sample_data`
- [Pretrained](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/data): model store · First-class runtime asset store for pretrained weights. · `data`

### External deps

Vendored or submodule-based dependency boundaries.

- [siml](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/lib/siml): vendored lib · Linked external library boundary used by the repo. · `lib/siml`
- [MP-Neural](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/MP-Neural-PDE-Solvers): submodule · External submodule dependency boundary. · `MP-Neural-PDE-Solvers`

### Tests

Domain-focused smoke and regression coverage.

- [CD tests](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/tests/cd): domain tests · Convection–diffusion validation coverage. · `tests/cd`
- [Mixture tests](https://github.com/yellowshippo/fluxgnn-icml2024/tree/main/tests/mixture): domain tests · Mixture-task validation coverage. · `tests/mixture`

## Connections

- Makefile → Experiment: dispatches · Invokes the main training/evaluation runner with a chosen target.
- Makefile → PENN: targets · Provides baseline train/eval targets.
- Makefile → Experiments: postprocess · Runs stats and export utilities through command targets.
- Inputs → Experiment: configures · YAML experiment definitions feed the shared runner.
- CD configs → FluxGNN: selects · Convection–diffusion configs select the FluxGNN solver family.
- Mixture configs → FluxGNN: selects · Mixture FluxGNN configs parameterize the main solver.
- Mixture configs → PENN: selects · Mixture PENN configs drive the baseline pipeline.
- Dataset → Data: loads · Dataset code reads the shared data store and task assets.
- Setting → Inputs: parses · Settings are populated from YAML experiment files.
- Experiment → Dataset: uses · Coordinates dataset loading before solver execution.
- Experiment → Setting: uses · Consumes settings to build and run experiments.
- Experiment → FluxGNN: runs · Constructs and executes the FluxGNN solver stack.
- FluxGNN → FV support: uses · Solver relies on finite-volume operations for PDE handling.
- FV support → FV data: reads · Finite-volume support is paired with FV-specific data structures.
- PDE wrappers → FluxGNN: wraps · PDE-specific wrappers extend the core solver.
- Mixture → FluxGNN: wraps · Mixture task wiring sits on top of the shared engine.
- Experiments → Data: analyzes · Postprocessing reads run outputs and produced artifacts.
- Experiments → Pretrained: exports · Can emit or consume model artifacts for evaluation workflows.
- PENN → Data: trains on · Baseline pipeline shares the external data and transformed assets.
- PENN → siml: depends on · Baseline code relies on the vendored external library boundary.
- PENN → MP-Neural: integrates · Uses the external PDE solver submodule as a dependency.
- CD tests → CD configs: covers · Domain tests mirror the convection–diffusion configuration set.
- Mixture tests → Mixture configs: covers · Domain tests mirror the mixture configuration set.