bace-gnn

Imported from SakshamKapoor2911/bace-gnn (AGENTS.md). Install upstream with npx skills add SakshamKapoor2911/bace-gnn. Copyright stays with the author.

AGENTS.md — bace-gnn (demo workspace)

Task

Port notebooks/VibecodingExample_Part1.ipynb into a testable Python package at src/bace_gnn/ that is functionally identical to the notebook's results.

Dataset

• BACE (MoleculeNet), cached at data/BACE (offline-safe; first load would download bace.csv ~10 MB from deepchemdata S3).
• 1513 molecules, 9 integer node features, binary labels: class 0 = 822 (54.33%), class 1 = 691 (45.67%). BACE is classification, NOT pIC50 regression.

Models

• GAT: 2x GATConv(9→64 heads=4, 256→64 heads=4) + GATConv(256→1 heads=1, concat=False) + global_mean_pool + Linear(1,1). ELU + dropout(0.5) between.
• GCN baseline: GCNConv(9→64) + GCNConv(64→64) + global_mean_pool + dropout(0.5) + Linear(64,1). torch.manual_seed(42) inside __init__.

Hyperparameters (must stay identical to the notebook; centralized in config.py)

• lr=0.005, weight_decay=5e-4, batch_size=32, dropout=0.5, epochs=100, early-stop patience=15 (on val loss), ReduceLROnPlateau(mode=min, factor=0.7, patience=5), scaffold split 60/20/20, split_seed=0, model seed=42.

Headline numbers (notebook executed outputs)

• GAT: test ROC-AUC 0.6594, PR-AUC 0.6418 (stopped at epoch 43).
• GCN: test ROC-AUC 0.6436, PR-AUC 0.6023 (stopped at epoch 47).
• Run-to-run variance exists (notebook seeds GAT poorly); verification target is within-noise (|delta| <= 0.02) vs GCN.

Known pitfalls — guardrails that bite (see PLAN.md)

1. Integer node features must be float-cast (x.float()) before the GNN conv. GAT casts inside forward; GCN in the notebook casts at the call site.
2. Flatten outputs and labels to [batch_size] (.view(-1)) before BCEWithLogitsLoss — avoids silent 2D loss-matrix broadcasting.
3. Balanced scaffold split: greedy 60/20/20 whole-scaffold groups with a retry guard (both classes per split, seed increments) and zero overlap between splits — a single-class test set yields NaN ROC-AUC / 0.0 PR-AUC.
4. Dropout must be off in eval (F.dropout(..., training=self.training) + model.eval() + torch.no_grad()).

Commands (Windows)

• venv python: .venv/Scripts/python.exe
• tests: .venv/Scripts/python.exe -m pytest
• train: .venv/Scripts/python.exe -m bace_gnn.train --model gcn --epochs 30