MolecularDiffusion.data.component.graph3d_dataset¶

3D molecular graph dataset with EXPLICIT BONDS.

The platform’s other two data shapes (pointcloud, pyg) are point clouds: coordinates + atom types, with at most a radius_graph/k-NN edge_index as a message-passing convenience and no edge features. This module adds a third shape, graph3d, carrying real bond orders and formal charges so bond-generating models become expressible – FlowMol’s e modality (currently dropped, see modules/tasks/diffusion_flowmol.py:10) and MiDi.

It is strictly additive. Nothing here is imported by the existing loaders; the shared helpers flow the other way (this file imports from .dataset), so no existing path can regress.

## The one storage rule

FlowMol and MiDi look different but agree on the underlying object:

FlowMol disk: sparse upper-triangular real bonds model: [upper||lower] directed MiDi disk: symmetric edge_index + edge_attr model: dense (B,N,N) one-hot

Both are recoverable from the smaller of the two. So:

Store upper-triangular real bonds only. Materialize per model in collate.

bond_type uses 5 classes – 0=none, 1=SINGLE, 2=DOUBLE, 3=TRIPLE, 4=AROMATIC – matching data/component/pharmacophore.py:186, which already uses this vocabulary. Class 0 is never stored: absence from bond_index means “no bond”, and the no-bond count is derived, not counted.

## Why the fields are named bond_index / bond_type

Not edge_index/edge_attr, for two independent reasons:

PyG’s Batch.from_data_list auto-increments any key whose name contains index. bond_index therefore batches correctly for free; a key named anything else is silently concatenated without the node offset, which produces wrong graphs rather than an error.
A sparse real-bonds-only tensor called edge_index would be misread as “the message-passing graph” by every existing pyg-consuming module.

## Charges are stored raw and signed

FlowMol hardcodes [-2,+3] (6 classes), MiDi QM9 uses 3 classes offset +1, MiDi GEOM 6 offset +2. Storing the raw signed integer means a model changing its range never forces a dataset reconversion – the offset/one-hot is applied model-side.

Attributes¶

`BOND_ORDER_TO_CLASS`
`BOND_VOCAB`
`Chem`
`N_BOND_CLASSES`
`logger`

Classes¶

`Graph3DDataset`	3D molecular graphs with explicit bonds, loaded from an ASE db.
`Graph3DStats`	Preprocessing-time statistics. Both target models refuse to build without

Functions¶

`build_rdkit_mol`(atomic_numbers, bond_index, bond_type)	Rebuild an RDKit mol from the stored fields.
`graph3d_dense_collate`(→ Dict[str, Any])	Pad a batch into MiDi's dense shapes.
`kekulize_bonds`(atomic_numbers, bond_index, bond_type)	Aromatic (class 4) bonds -> alternating SINGLE/DOUBLE.
`rdkit_bond_types`()	Index -> `Chem.BondType`, for rebuilding molecules. Requires RDKit.
`remap_bonds_after_atom_removal`(bond_index, bond_type, ...)	Drop bonds touching removed atoms and renumber the survivors.

Module Contents¶

class MolecularDiffusion.data.component.graph3d_dataset.Graph3DDataset(root: str = '', ase_db_path: str | None = None, dataset_name: str = 'graph3d', max_atom: int = 200, load_only: bool = False, **kwargs: Any)¶

Bases: torch.utils.data.Dataset

3D molecular graphs with explicit bonds, loaded from an ASE db.

Exposes the same public surface as GraphDataset/LazyChunkedGraphDataset (smiles_list, num_atoms, get_property, atom_types, targets) so the DataModule tail at runmodes/train/data.py:320-357 drives it with no changes.

atom_types() → List[int]¶

get_item(index: int) → Dict[str, Any]¶

get_property(task: str) → torch.Tensor | None¶

load_db(db_path: str, atom_vocab: List[str] | None = None, target_fields: List[str] | None = None, max_atom: int = 200, with_hydrogen: bool = True, forbidden_atoms: List[str] | None = None, allow_unknown: bool = False, use_ohe_feature: bool = True, center_coords: bool = False, kekulize: bool = False, compute_stats: bool = True, transform: Callable | None = None, verbose: int = 0, chunk_size: int | None = None, chunk_dir: str | None = None, **kwargs: Any)¶

Build the dataset from ASE db rows written by the graph3d converter.

Each row must carry bond_index (2,E upper-triangular), bond_type (E,) and formal_charge (N,) in row.data; rows without them are counted and skipped rather than silently treated as bond-free.

load_pickle(pkl_file: str, verbose: int = 0)¶

save_pickle(pkl_file: str, verbose: int = 0)¶

dataset_name = 'graph3d'¶

max_atom = 200¶

property num_atoms: torch.Tensor¶

root = ''¶

property tasks: List[str]¶

transform = None¶

class MolecularDiffusion.data.component.graph3d_dataset.Graph3DStats¶

Preprocessing-time statistics. Both target models refuse to build without some of these, and both compute the same things under different names.

bond_type_counts is in upper-triangular pair units: [0] = sum_mols(n(n-1)/2 - n_real_bonds), [1..4] = real bond counts. FlowMol’s p_e uses this directly. MiDi’s edge_counts wants directed counts, i.e. exactly 2x every entry – one multiply, stated once here so it is never re-derived incorrectly downstream.

atom_type_marginal() → numpy.ndarray¶

bond_type_marginal(directed: bool = False) → numpy.ndarray¶: Normalized bond-class distribution. directed=True gives MiDi’s convention; the scaling cancels, so the result is identical – the flag exists for callers that also want the raw counts to match.

n_atoms_histogram()¶: (sizes, counts) sorted by size, as FlowMol’s sampler expects.

atom_type_counts: numpy.ndarray¶

bond_type_counts: numpy.ndarray¶

charge_counts: Dict[int, Dict[int, int]]¶

property charge_range¶: Observed (min, max) raw signed formal charge, or None.

n_atoms_hist: Dict[int, int]¶

n_molecules: int = 0¶

valencies: Dict[int, Dict[int, Dict[float, int]]]¶

MolecularDiffusion.data.component.graph3d_dataset.build_rdkit_mol(atomic_numbers, bond_index, bond_type, formal_charge=None, coords=None, sanitize: bool = True, assign_stereo: bool = True)¶

Rebuild an RDKit mol from the stored fields.

This is the inverse of what the converter stores, and is what makes the round-trip testable: if a molecule survives db -> dataset -> here with its canonical SMILES intact, the bond pipeline is lossless.

Pass coords to recover stereochemistry, which lives in the geometry rather than in the bond table.

Measured on 2000 real QM9 molecules: 99.9% canonical-SMILES round-trip. The residue is carbenes and similar radicals ([CH]), where (z, bonds, formal_charge) alone cannot distinguish a radical from a missing hydrogen without a radical-electron count. MiDi’s build_molecule has the same ambiguity, and neither target model consumes such a field, so it is not stored.

MolecularDiffusion.data.component.graph3d_dataset.graph3d_dense_collate(batch: List[Dict[str, Any]]) → Dict[str, Any]¶

Pad a batch into MiDi’s dense shapes.

bond_type comes back as integer class ids (B,N,N), not a (B,N,N,5) one-hot: one-hotting is F.one_hot(E, 5) on the model side (exactly where MiDi does it, in to_one_hot), and at GEOM sizes the int8-equivalent dense matrix is ~20x smaller than the float one-hot.

MolecularDiffusion.data.component.graph3d_dataset.kekulize_bonds(atomic_numbers, bond_index, bond_type, formal_charge=None)¶

Aromatic (class 4) bonds -> alternating SINGLE/DOUBLE.

Kekulization is a ring-level operation, not a per-bond relabel, so this round-trips through RDKit. It runs once at dataset-build time (the result is cached in the processed pickle / chunks), never per __getitem__.

Storage always keeps the aromatic form; this exists so a model trained on a Kekule distribution – MiDi’s published QM9 run reads gdb9 with sanitize=False, so its bond class 4 is never populated – can be reproduced without a second copy of the dataset.

Returns the bond_type array unchanged if there is nothing aromatic to do, or if RDKit cannot kekulize the molecule.

MolecularDiffusion.data.component.graph3d_dataset.rdkit_bond_types()¶: Index -> Chem.BondType, for rebuilding molecules. Requires RDKit.

MolecularDiffusion.data.component.graph3d_dataset.remap_bonds_after_atom_removal(bond_index, bond_type, keep_mask)¶

Drop bonds touching removed atoms and renumber the survivors.

keep_mask is the boolean array over the original atom order returned by _drop_hydrogens(), whose docstring states it exists precisely so parallel per-atom data can be filtered identically.

i < j is preserved because the remap is monotone on kept indices.

MolecularDiffusion.data.component.graph3d_dataset.BOND_ORDER_TO_CLASS: Dict[float, int]¶

MolecularDiffusion.data.component.graph3d_dataset.BOND_VOCAB: List[str | None] = [None, 'SINGLE', 'DOUBLE', 'TRIPLE', 'AROMATIC']¶

MolecularDiffusion.data.component.graph3d_dataset.Chem = None¶

MolecularDiffusion.data.component.graph3d_dataset.N_BOND_CLASSES = 5¶

MolecularDiffusion.data.component.graph3d_dataset.logger¶