MolecularDiffusion.runmodes.data.graph3d_import¶

Convert MiDi’s raw QM9 / GEOM-Drugs data into the platform’s 3D-graph ASE db.

The platform’s existing dbs (e.g. data/qm9.db) store SMILES but no mol_block, so they carry no bond ground truth, and SMILES atom order does not match the xyz order. Rather than perceive bonds from geometry, this imports from MiDi’s sources, which have real bond orders.

Storage schema, one row per molecule/conformer:

Atoms  -> explicit H, RAW uncentered positions, RDKit atom order
data   -> smiles          canonical
          bond_index      (2, E) int32, UPPER-TRIANGULAR (i < j), real bonds only
          bond_type       (E,)  int8, 1=SINGLE 2=DOUBLE 3=TRIPLE 4=AROMATIC
          formal_charge   (N,)  int8, RAW SIGNED
          source          "qm9" | "geom"
          split           MiDi's split label, or ""
          <extra>         optional scalar per-molecule fields (QM9 property
                          targets + jodo_idx under ``--targets``); absent
                          unless asked for, so existing dbs are unchanged

Bond class 0 (“no bond”) is never stored: absence from bond_index means no bond, and the no-bond count is derived at statistics/materialization time.

## Aromaticity

MiDi reads gdb9 with sanitize=False, so RDKit performs no aromatic perception and its QM9 bond class 4 is never populated (verified: benzene yields bond orders {1.0, 2.0} unsanitized vs {1.0, 1.5} sanitized). We deliberately sanitize instead, so the db keeps the richer aromatic form, and reproduce MiDi’s Kekule distribution on demand via the dataset’s kekulize=True load flag rather than by storing a second lossy copy. --no-sanitize restores exact MiDi parity if you need it.

Attributes¶

`QM9_RAW_ARTEFACTS`
`logger`

Functions¶

`delete_raw`(→ None)	Remove ONLY the named download artefacts under `raw_dir`.
`ensure_qm9_splits`(→ None)	Materialize MiDi's published QM9 split csvs if they are not already there.
`iter_geom`(→ Iterator[Tuple[object, str]])	Yield `(mol, split)` from MiDi's GEOM-Drugs pickles.
`iter_qm9`(→ Iterator[Tuple])	Yield `(mol, split)` -- or `(mol, split, extra)` -- from gdb9.
`main`(→ int)
`mol_to_row`(mol[, sanitize_failed_ok, extra])	RDKit mol -> `(Atoms, data_dict)` for `ase.db.core.Database.write()`.
`qm9_targets`(→ Tuple[list, numpy.ndarray])	`gdb9.sdf.csv` -> `(names, values)` in PyG/JODO order and units.
`verify_db`(→ None)	Re-open the db and round-trip a sample back through RDKit.
`write_graph3d_db`(→ dict)	Write rows to an ASE db. Returns a counts dict.

Module Contents¶

MolecularDiffusion.runmodes.data.graph3d_import.delete_raw(raw_dir: str, artefacts) → None¶

Remove ONLY the named download artefacts under raw_dir.

Never touches the produced db and never leaves raw_dir. Prints every path before unlinking, because this is irreversible.

MolecularDiffusion.runmodes.data.graph3d_import.ensure_qm9_splits(raw_dir: str) → None¶

Materialize MiDi’s published QM9 split csvs if they are not already there.

Byte-for-byte the procedure in QM9Dataset.download (others/midi/midi/datasets/qm9_dataset.py:121-134): shuffle gdb9.sdf.csv with random_state=42 and cut at 100k train / remainder val / 10% test. Reproduced here rather than imported so the converter does not depend on MiDi being importable – but it must stay in sync, or published numbers stop being comparable.

MolecularDiffusion.runmodes.data.graph3d_import.iter_geom(pickle_dir: str, max_conformers: int = 5, sanitize: bool = True) → Iterator[Tuple[object, str]]¶

Yield (mol, split) from MiDi’s GEOM-Drugs pickles.

Each pickle is a list of (smiles, [conformers]); MiDi caps at 5 conformers per molecule (geom_dataset.py:90-91).

MolecularDiffusion.runmodes.data.graph3d_import.iter_qm9(raw_dir: str, sanitize: bool = True, targets: bool = False) → Iterator[Tuple]¶

Yield (mol, split) – or (mol, split, extra) – from gdb9.

Skips the indices listed in uncharacterized.txt, exactly as others/midi/midi/datasets/qm9_dataset.py:141-143 does.

With targets=True each molecule additionally carries the 19 QM9 property targets under their csv names, plus jodo_idx – the compacted index of upstream’s Data list (the SDF enumerated with only the uncharacterized entries removed), so a published split file that indexes into it is joinable.

MolecularDiffusion.runmodes.data.graph3d_import.main(argv=None) → int¶

MolecularDiffusion.runmodes.data.graph3d_import.mol_to_row(mol, sanitize_failed_ok: bool = False, extra: dict | None = None)¶

RDKit mol -> (Atoms, data_dict) for ase.db.core.Database.write().

Returns None when the molecule carries a bond type outside the canonical vocabulary (DATIVE, UNSPECIFIED, …). Such molecules are counted and skipped rather than silently coerced to SINGLE.

extra merges caller-supplied scalar per-molecule fields (property targets, an upstream row index, …) into the written data dict. They land in the row metadata, where the dataset reads them through target_fields. Default None -> the row schema is unchanged.

MolecularDiffusion.runmodes.data.graph3d_import.qm9_targets(raw_dir: str) → Tuple[list, numpy.ndarray]¶

gdb9.sdf.csv -> (names, values) in PyG/JODO order and units.

values is indexed by raw SDF index, so it is joinable with the enumerate in iter_qm9() before the uncharacterized skip. Names come from the csv header rather than a hardcoded list, so they cannot drift out of step with the columns they label.

MolecularDiffusion.runmodes.data.graph3d_import.verify_db(out_path: str, expected: int, n_check: int = 25) → None¶

Re-open the db and round-trip a sample back through RDKit.

Raises if the db does not contain what we think it does. This gates the --delete-raw step, so it must be strict.

MolecularDiffusion.runmodes.data.graph3d_import.write_graph3d_db(mol_iter, out_path: str, source: str, limit: int | None = None, verbose: bool = True) → dict¶: Write rows to an ASE db. Returns a counts dict.

MolecularDiffusion.runmodes.data.graph3d_import.QM9_RAW_ARTEFACTS = ('gdb9.sdf', 'gdb9.sdf.csv', 'uncharacterized.txt')¶

MolecularDiffusion.runmodes.data.graph3d_import.logger¶