""" dft8k_reader.py - reader/decoder for the dft8k HDF5 dataset. DFT8k is an external NMR shielding benchmark of 7111 organic molecules (Guan, Paton and co-workers; reference 14 in the paper), used to evaluate the MagNET models out of sample. The file has TWO groups, one per level of theory. Both groups cover the SAME 7111 molecules and use the SAME integer molecule ids, but in different row order, so match a molecule across the two groups by id (use `molecule_by_id`), not by row position. /aimnet2_wp04_wb97xd_pcsseg2 AIMNet2 geometries; per atom, three DFT shieldings plus the two MagNET predictions. This is the file behind the paper's DFT8k panels. /b3lyp_pbe0_pcsseg1 B3LYP/pcSseg-1 geometries; three geometries per molecule with one PBE0/pcSseg-1 shielding each, plus symmetry-equivalence groups and per-atom averaging weights. This is the format the model's own training/eval loader consumes. Storage format: - Per-atom arrays are concatenated across the molecules in a group. Each molecule's atom count is in `n_atoms`; molecule i owns the rows [start[i], start[i+1]) where start is the cumulative sum of `n_atoms` (reconstructed once, on open). - coordinates and shieldings are int32 fixed-point: physical = stored * scale (scale = 1e-4), so reconstruction error is <= 5e-5 (Angstrom or ppm). A missing value (an atom a model does not predict, or an absent shielding) is the marker -2147483648 and decodes to NaN. - The molecule's integer `molecule_ids` entry is its identifier. Requires: Python >= 3.7, numpy, h5py. No other dependencies. """ from __future__ import annotations import numpy as np import h5py _SCALE = 1e-4 _MARKER = -2147483648 def _decode(values: np.ndarray) -> np.ndarray: """int32 fixed-point -> float64 physical units; marker -> NaN. No-op on float input.""" values = np.asarray(values) if not np.issubdtype(values.dtype, np.integer): return np.asarray(values, dtype=np.float64) out = values.astype(np.float64) * _SCALE out[values == _MARKER] = np.nan return out class _Group: """Common molecule-indexing logic shared by both groups.""" def __init__(self, grp: h5py.Group): self.g = grp self.name = grp.name.lstrip("/") self.molecule_ids = grp["molecule_ids"][:] self.n_molecules = int(grp.attrs["n_molecules"]) self.n_atoms = int(grp.attrs["n_atoms"]) self.level_of_theory = str(grp.attrs["level_of_theory"]) self.geometry = str(grp.attrs["geometry"]) na = grp["n_atoms"][:].astype(np.int64) self._start = np.empty(self.n_molecules + 1, np.int64) self._start[0] = 0 np.cumsum(na, out=self._start[1:]) self._id_to_index = {int(m): i for i, m in enumerate(self.molecule_ids)} def _check(self, index: int): if not 0 <= index < self.n_molecules: raise IndexError(f"molecule index {index} out of range [0, {self.n_molecules})") def _rows(self, index: int) -> slice: return slice(int(self._start[index]), int(self._start[index + 1])) def index_of(self, molecule_id: int) -> int: """Return the row position of a molecule id within this group.""" if int(molecule_id) not in self._id_to_index: raise KeyError(f"molecule id {molecule_id} not in group {self.name}") return self._id_to_index[int(molecule_id)] def molecule_by_id(self, molecule_id: int) -> dict: """Look up a molecule by its id instead of its row position.""" return self.molecule(self.index_of(molecule_id)) def __len__(self): return self.n_molecules class AIMNet2Group(_Group): """/aimnet2_wp04_wb97xd_pcsseg2 - AIMNet2 geometries, DFT and MagNET shieldings per atom.""" def __init__(self, grp: h5py.Group): super().__init__(grp) # one SMILES per molecule. 6780 of 7111 are real; the rest are the literal "none". # Older files may lack the field, so guard for it. self._smiles = grp["smiles"].asstr()[:] if "smiles" in grp else None def smiles(self, index: int) -> str: """Return molecule i's SMILES string (the literal "none" if it was not provided).""" self._check(index) if self._smiles is None: return "none" return str(self._smiles[index]) def molecule(self, index: int) -> dict: """Return one molecule's atom data (physical units) by position. Keys: id, smiles; atomic_numbers, coordinates (n,3 Angstrom); the three DFT shieldings (pbe0_pcSseg1, wp04_pcSseg2, wb97xd_pcSseg2; each (n,) ppm); and the two MagNET predictions (nn_wp04_pcSseg2, nn_wb97xd_pcSseg2; (n,) ppm). A MagNET prediction is NaN at atoms that model does not predict (WP04 is the 1H model, wB97X-D the 13C). """ self._check(index) sl = self._rows(index) return { "id": int(self.molecule_ids[index]), "smiles": self.smiles(index), "atomic_numbers": self.g["atomic_numbers"][sl], "coordinates": _decode(self.g["coordinates"][sl]), "pbe0_pcSseg1": _decode(self.g["shielding_pbe0_pcSseg1"][sl]), "wp04_pcSseg2": _decode(self.g["shielding_wp04_pcSseg2"][sl]), "wb97xd_pcSseg2": _decode(self.g["shielding_wb97xd_pcSseg2"][sl]), "nn_wp04_pcSseg2": _decode(self.g["nn_shielding_wp04_pcSseg2"][sl]), "nn_wb97xd_pcSseg2": _decode(self.g["nn_shielding_wb97xd_pcSseg2"][sl]), } def all_shieldings(self) -> dict: """Every atom in the dataset as flat arrays (physical units), for whole-dataset analyses like the MagNET-vs-DFT residual histograms. Keys: atomic_numbers (n_atoms_total,); the three DFT shieldings (pbe0_pcSseg1, wp04_pcSseg2, wb97xd_pcSseg2) and the two MagNET predictions (nn_wp04_pcSseg2, nn_wb97xd_pcSseg2), each (n_atoms_total,) ppm. A MagNET prediction is NaN at atoms that model does not predict (WP04 is the 1H model, wB97X-D the 13C). To get one nucleus' residual, difference the DFT and MagNET column for that model at that element and drop the NaNs. """ return { "atomic_numbers": self.g["atomic_numbers"][:], "pbe0_pcSseg1": _decode(self.g["shielding_pbe0_pcSseg1"][:]), "wp04_pcSseg2": _decode(self.g["shielding_wp04_pcSseg2"][:]), "wb97xd_pcSseg2": _decode(self.g["shielding_wb97xd_pcSseg2"][:]), "nn_wp04_pcSseg2": _decode(self.g["nn_shielding_wp04_pcSseg2"][:]), "nn_wb97xd_pcSseg2": _decode(self.g["nn_shielding_wb97xd_pcSseg2"][:]), } def __repr__(self): return (f"") class B3LYPGroup(_Group): """/b3lyp_pbe0_pcsseg1 - three geometries per molecule, symmetry groups and weights.""" def __init__(self, grp: h5py.Group): super().__init__(grp) self.n_geometries = int(grp.attrs["n_geometries"]) # symmetry groups are stored flat: n_groups per molecule, each group's length, then # all member atom indices. Reconstruct the per-molecule offsets once, on open. self._sym_n = grp["symmetric_atoms_n_groups"][:].astype(np.int64) self._sym_len = grp["symmetric_atoms_group_lengths"][:].astype(np.int64) self._sym_mem = grp["symmetric_atoms_members"][:] self._group_start = np.empty(self.n_molecules + 1, np.int64) self._group_start[0] = 0 np.cumsum(self._sym_n, out=self._group_start[1:]) self._member_start = np.empty(self._sym_len.shape[0] + 1, np.int64) self._member_start[0] = 0 np.cumsum(self._sym_len, out=self._member_start[1:]) def symmetric_atoms(self, index: int) -> list: """Return molecule i's symmetry-equivalence groups as a list of lists of atom indices. The indices are 0-based positions into this molecule's atom arrays (atomic_numbers, coordinates), so a group can be used to index those arrays directly. """ self._check(index) groups = [] for gi in range(int(self._group_start[index]), int(self._group_start[index + 1])): a, b = int(self._member_start[gi]), int(self._member_start[gi + 1]) groups.append([int(x) for x in self._sym_mem[a:b]]) return groups def molecule(self, index: int) -> dict: """Return one molecule's data (physical units) by position. Keys: id, atomic_numbers; coordinates (n_geometries, n, 3 Angstrom); shielding (n_geometries, n, ppm), where geometry 0 is a reference and 1, 2 are perturbations; weights (n,) per-atom symmetry-averaging weight; symmetric_atoms (list of lists of equivalent atoms, each a 0-based index into this molecule's atom arrays). """ self._check(index) sl = self._rows(index) return { "id": int(self.molecule_ids[index]), "atomic_numbers": self.g["atomic_numbers"][sl], "coordinates": _decode(self.g["coordinates"][:, sl, :]), "shielding": _decode(self.g["shielding"][:, sl]), "weights": self.g["weights"][sl].astype(np.float64), "symmetric_atoms": self.symmetric_atoms(index), } def __repr__(self): return (f"") class DFT8k: """Reader for dft8k.hdf5. Use as a context manager; access the two groups by attribute.""" def __init__(self, path: str): self.f = h5py.File(path, "r") self.aimnet2 = AIMNet2Group(self.f["aimnet2_wp04_wb97xd_pcsseg2"]) self.b3lyp = B3LYPGroup(self.f["b3lyp_pbe0_pcsseg1"]) def close(self): """Close the underlying HDF5 file.""" self.f.close() def __enter__(self): return self def __exit__(self, *a): self.close() def __repr__(self): return f"" if __name__ == "__main__": import sys path = sys.argv[1] if len(sys.argv) > 1 else "dft8k.hdf5" with DFT8k(path) as ds: print(ds) m = ds.aimnet2.molecule(0) h = m["atomic_numbers"] == 1 print(f" aimnet2 mol id={m['id']}: {len(m['atomic_numbers'])} atoms smiles={m['smiles']}") print(f" 1H WP04 DFT[:3]={np.round(m['wp04_pcSseg2'][h][:3], 3)} " f"MagNET[:3]={np.round(m['nn_wp04_pcSseg2'][h][:3], 3)}") b = ds.b3lyp.molecule(0) print(f" b3lyp mol id={b['id']}: coordinates {b['coordinates'].shape}, " f"{len(b['symmetric_atoms'])} symmetry groups")