Same model, new framework

ICEBERG predicts a molecule's tandem mass spectrum from its fragmentation graph. I ported the full pipeline from DGL to PyTorch Geometric — 19 files, three GNN families, zero change in behaviour.

The pipeline

molecule → RDKit → PyG Data → MAGMa fragment DAG
   → Batch.from_data_list → GNN message passing (GGNN · PNA · GINE)
   → per-fragment intensities → predicted MS/MS spectrum

A molecule becomes a fragment DAG, each fragment a graph; a two-stage model generates the fragments, then predicts their intensities into a binned spectrum.

Precursor molecule through a collision cell to a learned neural simulator predicting the mass spectrum — Precursor → fragmentation → learned simulator → predicted spectrum.

The port

Custom layers were rebuilt PyG-style with torch_scatter in place of DGL's update_all: dgl.batch → Batch.from_data_list, node/edge stores → data.x / edge_attr, pooling → global_mean_pool. Result: zero remaining DGL imports, same numbers.

GitHub ↗