# hysonlab/equihgnn architecture

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

- Interactive diagram: https://gitdiagram.com/hysonlab/equihgnn
- Repository: https://github.com/hysonlab/equihgnn

## Overview

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: [`equihgnn/models/equihnn_egnn.py`](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/equihnn_egnn.py) (EquiHGNN variants), [`equihgnn/data/utils.py`](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/utils.py) (data utils), [`main.py`](https://github.com/hysonlab/equihgnn/blob/main/main.py) (main.py), [`equihgnn/models/egnn.py`](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/egnn.py) (Geometric models)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_control["Control plane"]
  node_main["main.py<br/>entrypoint<br/>[main.py]"]
  node_scripts["scripts<br/>launchers"]
  node_registry["registry<br/>[registry.py]"]
  node_create["create utils<br/>object factory<br/>[create.py]"]
end

subgraph group_data["Data subsystem"]
  node_split["data split<br/>split policy<br/>[data_split.py]"]
  node_data_utils["data utils<br/>preprocessing<br/>[utils.py]"]
  node_xyz2mol["xyz2mol<br/>structure conversion<br/>[xyz2mol.py]"]
  node_qm9["QM9<br/>dataset adapter<br/>[qm9.py]"]
  node_opv3d["OPV3D<br/>dataset adapter<br/>[opv3d.py]"]
  node_pcqm4["PCQM4<br/>dataset adapter<br/>[pcqm4.py]"]
  node_molecule3d["Molecule3D<br/>dataset adapter<br/>[molecule3d.py]"]
  node_md17["MD17<br/>dataset adapter<br/>[md17.py]"]
  node_md22["MD22<br/>dataset adapter<br/>[md22.py]"]
end

subgraph group_models["Model stack"]
  node_baseline2d["2D baselines<br/>model family<br/>[baseline_2d.py]"]
  node_geometric["Geometric models<br/>model family<br/>[egnn.py]"]
  node_equihgnn_variants["EquiHGNN variants<br/>hybrid models<br/>[equihnn_egnn.py]"]
  node_mhnn["MHNN<br/>model family<br/>[mhnn.py]"]
end

subgraph group_layers["Shared layers"]
  node_layers_core["Core layers<br/>shared layers<br/>[mlp.py]"]
  node_equiformer_layer["Equiformer layer<br/>equivariant layer"]
  node_se3_layer["SE3 layer<br/>equivariant layer"]
  node_fa_vis_layers["FA/VisNet layers<br/>specialized layers<br/>[fa_former_layer.py]"]
  node_equiformer_pkg["Equiformer code<br/>embedded dependency"]
  node_se3_pkg["SE3 code<br/>embedded dependency"]
end

subgraph group_runtime["Runtime assets"]
  node_datasets_mount[("datasets/<br/>data mount")]
  node_logs_mount[("logs/<br/>output mount")]
end

node_main -->|"loads"| node_registry
node_main -->|"builds"| node_create
node_scripts -->|"invoke"| node_main
node_scripts -.->|"read data"| node_datasets_mount
node_main -->|"splits"| node_split
node_qm9 -->|"preprocess"| node_data_utils
node_opv3d -->|"preprocess"| node_data_utils
node_pcqm4 -->|"preprocess"| node_data_utils
node_molecule3d -->|"preprocess"| node_data_utils
node_md17 -->|"preprocess"| node_data_utils
node_md22 -->|"preprocess"| node_data_utils
node_data_utils -->|"convert"| node_xyz2mol
node_qm9 -->|"feeds"| node_baseline2d
node_opv3d -->|"feeds"| node_geometric
node_pcqm4 -->|"feeds"| node_equihgnn_variants
node_molecule3d -->|"feeds"| node_equihgnn_variants
node_md17 -->|"feeds"| node_geometric
node_md22 -->|"feeds"| node_geometric
node_baseline2d -->|"uses"| node_layers_core
node_geometric -->|"uses"| node_equiformer_layer
node_geometric -->|"uses"| node_se3_layer
node_equihgnn_variants -->|"uses"| node_layers_core
node_equihgnn_variants -->|"wraps"| node_equiformer_layer
node_equihgnn_variants -->|"wraps"| node_se3_layer
node_equihgnn_variants -->|"wraps"| node_fa_vis_layers
node_equihgnn_variants -.->|"extends"| node_mhnn
node_equiformer_layer -->|"depends on"| node_equiformer_pkg
node_se3_layer -->|"depends on"| node_se3_pkg
node_main -.->|"writes"| node_logs_mount
node_main -.->|"reads"| node_datasets_mount

click node_main "https://github.com/hysonlab/equihgnn/blob/main/main.py"
click node_scripts "https://github.com/hysonlab/equihgnn/tree/main/scripts"
click node_registry "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/common/registry.py"
click node_create "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/utils/create.py"
click node_split "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/utils/data_split.py"
click node_data_utils "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/utils.py"
click node_xyz2mol "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/xyz2mol.py"
click node_qm9 "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/qm9.py"
click node_opv3d "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/opv3d.py"
click node_pcqm4 "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/pcqm4.py"
click node_molecule3d "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/molecule3d.py"
click node_md17 "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/md17.py"
click node_md22 "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/md22.py"
click node_baseline2d "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/baseline_2d.py"
click node_geometric "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/egnn.py"
click node_equihgnn_variants "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/equihnn_egnn.py"
click node_mhnn "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/mhnn.py"
click node_layers_core "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/mlp.py"
click node_equiformer_layer "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/equiformer_layer.py"
click node_se3_layer "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/se3_transformer_layer.py"
click node_fa_vis_layers "https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/fa_former_layer.py"
click node_equiformer_pkg "https://github.com/hysonlab/equihgnn/tree/main/equihgnn/models/layers/equiformer"
click node_se3_pkg "https://github.com/hysonlab/equihgnn/tree/main/equihgnn/models/layers/se3_transformer"
click node_datasets_mount "https://github.com/hysonlab/equihgnn/tree/main/datasets"
click node_logs_mount "https://github.com/hysonlab/equihgnn/tree/main/logs"

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_main,node_scripts,node_registry,node_create toneBlue
class node_split,node_data_utils,node_xyz2mol,node_qm9,node_opv3d,node_pcqm4,node_molecule3d,node_md17,node_md22 toneAmber
class node_baseline2d,node_geometric,node_equihgnn_variants,node_mhnn toneMint
class node_layers_core,node_equiformer_layer,node_se3_layer,node_fa_vis_layers,node_equiformer_pkg,node_se3_pkg toneRose
class node_datasets_mount,node_logs_mount toneIndigo
```

## Components

### Control plane

Entry points and experiment launchers.

- [main.py](https://github.com/hysonlab/equihgnn/blob/main/main.py): entrypoint · Primary experiment entry point. · `main.py`
- [scripts](https://github.com/hysonlab/equihgnn/tree/main/scripts): launchers · Task- and dataset-specific training launch scripts. · `scripts`
- [registry](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/common/registry.py): factory registry · Central registry for constructing components from config. · `equihgnn/common/registry.py`
- [create utils](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/utils/create.py): object factory · Config-driven creation helpers for models and datasets. · `equihgnn/utils/create.py`

### Data subsystem

Dataset adapters, preprocessing, and splitting.

- [data split](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/utils/data_split.py): split policy · Shared train/validation/test partitioning logic. · `equihgnn/utils/data_split.py`
- [data utils](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/utils.py): preprocessing · Common molecular preprocessing and input construction utilities. · `equihgnn/data/utils.py`
- [xyz2mol](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/xyz2mol.py): structure conversion · Conversion helpers between coordinate and molecular representations. · `equihgnn/data/xyz2mol.py`
- [QM9](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/qm9.py): dataset adapter · QM9 benchmark loader and target preparation. · `equihgnn/data/qm9.py`
- [OPV3D](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/opv3d.py): dataset adapter · OPV benchmark loader with 3D support. · `equihgnn/data/opv3d.py`
- [PCQM4](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/pcqm4.py): dataset adapter · PCQM4Mv2-style benchmark loader. · `equihgnn/data/pcqm4.py`
- [Molecule3D](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/molecule3d.py): dataset adapter · Molecule3D benchmark loader and regression targets. · `equihgnn/data/molecule3d.py`
- [MD17](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/md17.py): dataset adapter · MD17 molecular dynamics benchmark loader. · `equihgnn/data/md17.py`
- [MD22](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/data/md22.py): dataset adapter · MD22 molecular dynamics benchmark loader. · `equihgnn/data/md22.py`

### Model stack

Baseline, geometric, and EquiHGNN model families.

- [2D baselines](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/baseline_2d.py): model family · Non-geometric baseline models for comparison. · `equihgnn/models/baseline_2d.py`
- [Geometric models](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/egnn.py): model family · Equivariant and geometric backbone implementations. · `equihgnn/models/egnn.py`
- [EquiHGNN variants](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/equihnn_egnn.py): hybrid models · Hypergraph plus geometric backbone model wrappers. · `equihgnn/models/equihnn_egnn.py`
- [MHNN](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/mhnn.py): model family · Hypergraph neural network baseline and inspiration source. · `equihgnn/models/mhnn.py`

### Shared layers

Reusable building blocks and integrated equivariant code.

- [Core layers](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/mlp.py): shared layers · MLP, convolution, and message-passing building blocks. · `equihgnn/models/layers/mlp.py`
- [Equiformer layer](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/equiformer_layer.py): equivariant layer · Integrated Equiformer-specific layer implementation. · `equihgnn/models/layers/equiformer_layer.py`
- [SE3 layer](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/se3_transformer_layer.py): equivariant layer · Integrated SE(3) Transformer-specific layer implementation. · `equihgnn/models/layers/se3_transformer_layer.py`
- [FA/VisNet layers](https://github.com/hysonlab/equihgnn/blob/main/equihgnn/models/layers/fa_former_layer.py): specialized layers · Frame averaging and VisNet-specific layer implementations. · `equihgnn/models/layers/fa_former_layer.py`
- [Equiformer code](https://github.com/hysonlab/equihgnn/tree/main/equihgnn/models/layers/equiformer): embedded dependency · Copied Equiformer support code, basis functions, and constants. · `equihgnn/models/layers/equiformer`
- [SE3 code](https://github.com/hysonlab/equihgnn/tree/main/equihgnn/models/layers/se3_transformer): embedded dependency · Copied SE(3) Transformer support code, basis functions, and constants. · `equihgnn/models/layers/se3_transformer`

### Runtime assets

External mounts and outputs used during training.

- [datasets/](https://github.com/hysonlab/equihgnn/tree/main/datasets): data mount · External dataset cache and input mount. · `datasets`
- [logs/](https://github.com/hysonlab/equihgnn/tree/main/logs): output mount · Experiment outputs and run artifacts. · `logs`

## Connections

- 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.