How hysonlab/equihgnn works
1. What this repo is
- EquiHGNN is a training/evaluation codebase for molecular property prediction using hypergraph neural networks with optional rotationally equivariant geometric modules.
- The core design choice is: represent molecules as higher-order interactions, then optionally inject 3D geometry through an equivariant backbone.
2. Top-level orchestration
- main.py is the entry point for experiments.
- The scripts/ directory contains task-specific launchers for each dataset and for 2D vs 3D variants.
- Makefile, Dockerfile, pyproject.toml, and requirements.txt define environment/build/runtime setup.
- tests/ mirrors the training scripts for smoke/regression runs.
3. Experiment configuration and object construction
- equihgnn/common/registry.py and equihgnn/utils/create.py indicate a registry-based architecture for instantiating models, datasets, and related components from config/CLI flags.
- equihgnn/utils/data_split.py centralizes dataset splitting logic, so train/val/test partitioning is a shared boundary rather than duplicated per dataset.
4. Data subsystem
- equihgnn/data/ contains dataset adapters for:
- qm9.py
- opv3d.py
- pcqm4.py
- molecule3d.py
- md17.py
- md22.py
- equihgnn/data/utils.py and xyz2mol.py suggest preprocessing utilities for converting molecular representations and building graph/hypergraph inputs.
- The datasets/ directory is an external data mount/cache location, not source-controlled data.
5. Main data flow
- Raw dataset files are loaded through the dataset adapters.
- Preprocessing converts molecular inputs into graph/hypergraph structures and, for 3D runs, geometry-aware features.
- The prepared samples are split, batched, and fed into a selected model backbone.
- Model outputs are compared against task-specific targets such as QM9 properties, OPV properties, PCQM4Mv2 gap prediction, or Molecule3D regression targets.
6. Model families
- equihgnn/models/ contains the supported architectures:
- 2D baselines: baseline_2d.py, dimenet.py
- geometric/equivariant baselines: egnn.py, equiformer.py, mhnn.py
- EquiHGNN variants combining hypergraph structure with a geometric backbone: equihnn_egnn.py, equihnn_equiformer.py, equihnn_fa_former.py, equihnn_se3_transformer.py, equihnn_visnet.py
- This division makes the repo a comparison framework as much as a single model implementation.
7. Layer/submodule boundary
- equihgnn/models/layers/ hosts reusable building blocks: MLPs, convolution/message-passing layers, and specialized layers for each backbone.
- Separate subdirectories for equiformer and se3_transformer include copied/embedded support code, basis functions, reversible blocks, spherical harmonics, and cached constants like J_dense.pt/.npy.
- The presence of these modules implies the repo avoids treating those backbones as black-box dependencies; it integrates them directly into the model stack.
8. Architectural pattern inside EquiHGNN
- The repo appears to follow a shared hypergraph core + pluggable geometric encoder pattern.
- Hypergraph reasoning supplies higher-order molecular structure.
- The optional geometric branch injects 3D equivariant features, which are then fused in task-specific wrappers.
- This makes “2D only” vs “hypergraph + geometry” an explicit architectural boundary.
9. Task and dataset coverage
- The README and scripts show four primary benchmark families: OPV, QM9, PCQM4Mv2, Molecule3D.
- Each dataset has paired script variants for non-geometric and 3D/geometric training.
- Task IDs are used for multi-target datasets, so a single dataset adapter likely supports multiple prediction heads.
10. Infrastructure and execution
- The repo is container-ready: Dockerfile plus the README’s docker run workflow indicate training is expected to run in isolated environments with GPU access.
- logs/ is a mounted output location for experiment artifacts.
- The README references COMET_API_KEY, suggesting optional experiment tracking is part of the operational workflow.
11. External inspirations integrated into the codebase
- The implementation incorporates ideas/code from MHNN, EGNN, SE(3) Transformers, Equiformer, Frame Averaging Transformer, and VisNet.
- Architecturally, this means the repo is not a single monolithic model; it is a comparative framework with reusable equivariant components wrapped in a unified training interface.
12. How to draw the architecture
- Draw three main lanes: dataset/preprocessing, model family selection, and training/evaluation orchestration.
- Put main.py and scripts/ at the center as the control plane.
- Show equihgnn/data feeding standardized molecular tensors/features into equihgnn/models.
- Show equihgnn/common and equihgnn/utils as supporting config/factory/split infrastructure.
- Show logs/ and datasets/ as external runtime mounts rather than source components.
Main entry points: equihnn_egnn.py EquiHGNN variants · utils.py data utils · main.py main.py · egnn.py Geometric models
Components 25
Control plane
Entry points and experiment launchers.
- main.py · Primary experiment entry point.
main.py - scripts · Task- and dataset-specific training launch scripts.
scripts - registry · Central registry for constructing components from config.
equihgnn/common/registry.py - create utils · Config-driven creation helpers for models and datasets.
equihgnn/utils/create.py
Data subsystem
Dataset adapters, preprocessing, and splitting.
- data split · Shared train/validation/test partitioning logic.
equihgnn/utils/data_split.py - data utils · Common molecular preprocessing and input construction utilities.
equihgnn/data/utils.py - xyz2mol · Conversion helpers between coordinate and molecular representations.
equihgnn/data/xyz2mol.py - QM9 · QM9 benchmark loader and target preparation.
equihgnn/data/qm9.py - OPV3D · OPV benchmark loader with 3D support.
equihgnn/data/opv3d.py - PCQM4 · PCQM4Mv2-style benchmark loader.
equihgnn/data/pcqm4.py - Molecule3D · Molecule3D benchmark loader and regression targets.
equihgnn/data/molecule3d.py - MD17 · MD17 molecular dynamics benchmark loader.
equihgnn/data/md17.py - MD22 · MD22 molecular dynamics benchmark loader.
equihgnn/data/md22.py
Model stack
Baseline, geometric, and EquiHGNN model families.
- 2D baselines · Non-geometric baseline models for comparison.
equihgnn/models/baseline_2d.py - Geometric models · Equivariant and geometric backbone implementations.
equihgnn/models/egnn.py - EquiHGNN variants · Hypergraph plus geometric backbone model wrappers.
equihgnn/models/equihnn_egnn.py - MHNN · Hypergraph neural network baseline and inspiration source.
equihgnn/models/mhnn.py
Shared layers
Reusable building blocks and integrated equivariant code.
- Core layers · MLP, convolution, and message-passing building blocks.
equihgnn/models/layers/mlp.py - Equiformer layer · Integrated Equiformer-specific layer implementation.
equihgnn/models/layers/equiformer_layer.py - SE3 layer · Integrated SE(3) Transformer-specific layer implementation.
equihgnn/models/layers/se3_transformer_layer.py - FA/VisNet layers · Frame averaging and VisNet-specific layer implementations.
equihgnn/models/layers/fa_former_layer.py - Equiformer code · Copied Equiformer support code, basis functions, and constants.
equihgnn/models/layers/equiformer - SE3 code · Copied SE(3) Transformer support code, basis functions, and constants.
equihgnn/models/layers/se3_transformer
Connections 30
- main.py → registry: loads · Main uses the registry to resolve components.
- main.py → create utils: builds · Main constructs models and datasets from config.
- scripts → main.py: invoke · Launcher scripts run the main experiment entry point.
- scripts → datasets/: read data · Scripts assume mounted datasets are available.
- main.py → data split: splits · Experiment flow uses shared data splitting logic.
- QM9 → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- OPV3D → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- PCQM4 → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- Molecule3D → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- MD17 → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- MD22 → data utils: preprocess · Dataset adapter uses common molecular preprocessing.
- data utils → xyz2mol: convert · Preprocessing may convert coordinate inputs to molecular form.
- QM9 → 2D baselines: feeds · Prepared samples are consumed by selected model families.
- OPV3D → Geometric models: feeds · Prepared samples are consumed by selected model families.
- PCQM4 → EquiHGNN variants: feeds · Prepared samples are consumed by selected model families.
- Molecule3D → EquiHGNN variants: feeds · Prepared samples are consumed by selected model families.
- MD17 → Geometric models: feeds · Prepared samples are consumed by selected model families.
- MD22 → Geometric models: feeds · Prepared samples are consumed by selected model families.
- 2D baselines → Core layers: uses · 2D baselines depend on shared core layers.
- Geometric models → Equiformer layer: uses · Geometric models rely on specialized backbone layers.
- Geometric models → SE3 layer: uses · Geometric models rely on specialized backbone layers.
- EquiHGNN variants → Core layers: uses · Hybrid models use shared message-passing and MLP layers.
- EquiHGNN variants → Equiformer layer: wraps · Some variants integrate an Equiformer backbone.
- EquiHGNN variants → SE3 layer: wraps · Some variants integrate an SE(3) Transformer backbone.
- EquiHGNN variants → FA/VisNet layers: wraps · Some variants integrate frame averaging or VisNet layers.
- EquiHGNN variants → MHNN: extends · Hybrid models build on hypergraph reasoning ideas.
- Equiformer layer → Equiformer code: depends on · Layer implementation depends on embedded Equiformer support code.
- SE3 layer → SE3 code: depends on · Layer implementation depends on embedded SE(3) support code.
- main.py → logs/: writes · Training and evaluation outputs are written to logs.
- main.py → datasets/: reads · Main training flow reads from mounted datasets.