# yuyangw/molclr architecture

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

- Interactive diagram: https://gitdiagram.com/yuyangw/molclr
- Repository: https://github.com/yuyangw/molclr

## Overview

MolCLR is a molecular graph representation-learning project. Its main workflow reads molecular structures, creates paired augmented graph views, trains a GCN or GIN encoder with a contrastive objective, and saves pretrained weights. A second workflow loads those weights to fine-tune molecular property predictors on benchmark datasets and report classification or regression results. The README also documents node, subgraph, and mixed augmentation and provides pretrained GCN/GIN checkpoints. Some wiring is inferred from documented entry points because excerpts are partial; uncertain internal links are omitted.

Main entry points: [`molclr.py`](https://github.com/yuyangw/molclr/blob/master/molclr.py) (Pretrain runner), [`finetune.py`](https://github.com/yuyangw/molclr/blob/master/finetune.py) (Fine-tune runner), [`dataset/dataset.py`](https://github.com/yuyangw/molclr/blob/master/dataset/dataset.py) (Pretrain loaders), [`dataset/dataset_subgraph.py`](https://github.com/yuyangw/molclr/blob/master/dataset/dataset_subgraph.py) (Subgraph views)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Molecular data"]
  node_molecule_data["Molecule data"]
  node_node_views["Node views<br/>[dataset.py]"]
  node_subgraph_views["Subgraph views"]
  node_mixed_views["Mixed views<br/>[dataset_mix.py]"]
  node_pretrain_loaders["Pretrain loaders<br/>[dataset.py]"]
  node_benchmark_data["Benchmark data"]
  node_benchmark_dataset["Benchmark loader<br/>[dataset_test.py]"]
end

subgraph group_pretrain["Contrastive pretraining"]
  node_pretrain_entry["Pretrain runner<br/>[molclr.py]"]
  node_gcn_encoder["GCN encoder<br/>[gcn_molclr.py]"]
  node_gin_encoder["GIN encoder<br/>[ginet_molclr.py]"]
  node_contrastive_loss["NT-Xent loss<br/>[nt_xent.py]"]
end

subgraph group_finetune["Property fine-tuning"]
  node_finetune_runner["Fine-tune runner<br/>[finetune.py]"]
  node_gcn_predictor["GCN predictor<br/>[gcn_finetune.py]"]
  node_gin_predictor["GIN predictor<br/>[ginet_finetune.py]"]
  node_property_results["Property metrics<br/>[finetune.py]"]
end

subgraph group_state["Model state"]
  node_pretrained_weights["Pretrained weights<br/>[model.pth]"]
  node_gin_weights["GIN weights<br/>[model.pth]"]
end

node_researcher(("Researcher"))

node_researcher -->|"starts pretraining"| node_pretrain_entry
node_researcher -->|"starts fine-tuning"| node_finetune_runner
node_molecule_data -->|"reads structures"| node_node_views
node_molecule_data -->|"reads structures"| node_subgraph_views
node_molecule_data -->|"reads structures"| node_mixed_views
node_node_views -->|"creates graph views"| node_pretrain_loaders
node_subgraph_views -->|"creates graph views"| node_pretrain_loaders
node_mixed_views -->|"creates graph views"| node_pretrain_loaders
node_pretrain_entry -->|"requests batches"| node_pretrain_loaders
node_pretrain_entry -->|"selects model"| node_gcn_encoder
node_pretrain_entry -->|"selects model"| node_gin_encoder
node_pretrain_entry -->|"computes objective"| node_contrastive_loss
node_pretrain_entry -->|"saves checkpoints"| node_pretrained_weights
node_pretrain_entry -->|"saves checkpoints"| node_gin_weights
node_benchmark_data -->|"reads benchmarks"| node_benchmark_dataset
node_finetune_runner -->|"builds dataset"| node_benchmark_dataset
node_finetune_runner -->|"selects model"| node_gcn_predictor
node_finetune_runner -->|"selects model"| node_gin_predictor
node_pretrained_weights -.->|"loads weights"| node_finetune_runner
node_gin_weights -.->|"loads weights"| node_finetune_runner
node_finetune_runner -->|"reports metrics"| node_property_results

click node_node_views "https://github.com/yuyangw/molclr/blob/master/dataset/dataset.py"
click node_subgraph_views "https://github.com/yuyangw/molclr/blob/master/dataset/dataset_subgraph.py"
click node_mixed_views "https://github.com/yuyangw/molclr/blob/master/dataset/dataset_mix.py"
click node_pretrain_loaders "https://github.com/yuyangw/molclr/blob/master/dataset/dataset.py"
click node_pretrain_entry "https://github.com/yuyangw/molclr/blob/master/molclr.py"
click node_gcn_encoder "https://github.com/yuyangw/molclr/blob/master/models/gcn_molclr.py"
click node_gin_encoder "https://github.com/yuyangw/molclr/blob/master/models/ginet_molclr.py"
click node_contrastive_loss "https://github.com/yuyangw/molclr/blob/master/utils/nt_xent.py"
click node_pretrained_weights "https://github.com/yuyangw/molclr/blob/master/ckpt/pretrained_gcn/checkpoints/model.pth"
click node_gin_weights "https://github.com/yuyangw/molclr/blob/master/ckpt/pretrained_gin/checkpoints/model.pth"
click node_benchmark_dataset "https://github.com/yuyangw/molclr/blob/master/dataset/dataset_test.py"
click node_finetune_runner "https://github.com/yuyangw/molclr/blob/master/finetune.py"
click node_gcn_predictor "https://github.com/yuyangw/molclr/blob/master/models/gcn_finetune.py"
click node_gin_predictor "https://github.com/yuyangw/molclr/blob/master/models/ginet_finetune.py"
click node_property_results "https://github.com/yuyangw/molclr/blob/master/finetune.py"

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_molecule_data,node_node_views,node_subgraph_views,node_mixed_views,node_pretrain_loaders,node_benchmark_data,node_benchmark_dataset toneBlue
class node_pretrain_entry,node_gcn_encoder,node_gin_encoder,node_contrastive_loss toneAmber
class node_finetune_runner,node_gcn_predictor,node_gin_predictor,node_property_results toneMint
class node_pretrained_weights,node_gin_weights toneRose
class node_researcher toneIndigo
```

## Components

### Molecular data

- **Molecule data**
- [Node views](https://github.com/yuyangw/molclr/blob/master/dataset/dataset.py): `dataset/dataset.py`
- [Subgraph views](https://github.com/yuyangw/molclr/blob/master/dataset/dataset_subgraph.py): `dataset/dataset_subgraph.py`
- [Mixed views](https://github.com/yuyangw/molclr/blob/master/dataset/dataset_mix.py): `dataset/dataset_mix.py`
- [Pretrain loaders](https://github.com/yuyangw/molclr/blob/master/dataset/dataset.py): `dataset/dataset.py`
- **Benchmark data**
- [Benchmark loader](https://github.com/yuyangw/molclr/blob/master/dataset/dataset_test.py): `dataset/dataset_test.py`

### Contrastive pretraining

- [Pretrain runner](https://github.com/yuyangw/molclr/blob/master/molclr.py): `molclr.py`
- [GCN encoder](https://github.com/yuyangw/molclr/blob/master/models/gcn_molclr.py): `models/gcn_molclr.py`
- [GIN encoder](https://github.com/yuyangw/molclr/blob/master/models/ginet_molclr.py): `models/ginet_molclr.py`
- [NT-Xent loss](https://github.com/yuyangw/molclr/blob/master/utils/nt_xent.py): `utils/nt_xent.py`

### Property fine-tuning

- [Fine-tune runner](https://github.com/yuyangw/molclr/blob/master/finetune.py): `finetune.py`
- [GCN predictor](https://github.com/yuyangw/molclr/blob/master/models/gcn_finetune.py): `models/gcn_finetune.py`
- [GIN predictor](https://github.com/yuyangw/molclr/blob/master/models/ginet_finetune.py): `models/ginet_finetune.py`
- [Property metrics](https://github.com/yuyangw/molclr/blob/master/finetune.py): `finetune.py`

### Model state

- [Pretrained weights](https://github.com/yuyangw/molclr/blob/master/ckpt/pretrained_gcn/checkpoints/model.pth): `ckpt/pretrained_gcn/checkpoints/model.pth`
- [GIN weights](https://github.com/yuyangw/molclr/blob/master/ckpt/pretrained_gin/checkpoints/model.pth): `ckpt/pretrained_gin/checkpoints/model.pth`

### Other components

- **Researcher**

## Connections

- Researcher → Pretrain runner: starts pretraining
- Researcher → Fine-tune runner: starts fine-tuning
- Molecule data → Node views: reads structures
- Molecule data → Subgraph views: reads structures
- Molecule data → Mixed views: reads structures
- Node views → Pretrain loaders: creates graph views
- Subgraph views → Pretrain loaders: creates graph views
- Mixed views → Pretrain loaders: creates graph views
- Pretrain runner → Pretrain loaders: requests batches
- Pretrain runner → GCN encoder: selects model
- Pretrain runner → GIN encoder: selects model
- Pretrain runner → NT-Xent loss: computes objective
- Pretrain runner → Pretrained weights: saves checkpoints
- Pretrain runner → GIN weights: saves checkpoints
- Benchmark data → Benchmark loader: reads benchmarks
- Fine-tune runner → Benchmark loader: builds dataset
- Fine-tune runner → GCN predictor: selects model
- Fine-tune runner → GIN predictor: selects model
- Pretrained weights → Fine-tune runner: loads weights
- GIN weights → Fine-tune runner: loads weights
- Fine-tune runner → Property metrics: reports metrics