yellowshippo/fluxgnn-icml2024

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.