| """ |
| decode_gdb_qcd.py - reader/decoder for the gdb_qcd HDF5 dataset. |
| |
| gdb_qcd is a quasiclassical-dynamics benchmark: replicate molecular-dynamics trajectories for |
| 2461 small molecules drawn from GDB (the MagNET / sigma-shake test set). It is used to benchmark |
| rovibrational corrections, which are computed by comparing NMR shieldings on the moving |
| (displaced) geometries against the resting (stationary) geometry. |
| |
| For each molecule there are several trajectories; each trajectory has 33 frames; each frame gives |
| every atom's position and its PBE0/pcSseg-1 NMR shielding. Trajectory counts vary by molecule |
| (about 23 on average); the frame count is always 33. The resting (stationary) reference shieldings |
| are not included here: they are the test-set entries of the MagNET training data (sigma-shake), |
| matched by the same molecule id. |
| |
| Storage format: |
| - Per-atom arrays are concatenated across molecules. Molecule i has `n_atoms[i]` atoms and |
| `n_trajectories[i]` trajectories. Its atomic numbers occupy `n_atoms[i]` rows; its coordinates |
| and shieldings occupy `n_trajectories[i] * 33 * n_atoms[i]` rows, which reshape to |
| (n_trajectories, 33, n_atoms, ...). Offsets are reconstructed once, on open. |
| - coordinates and shieldings are int32 fixed point: physical = stored * 1e-4, so reconstruction |
| error is <= 5e-5 (Angstrom or ppm). A missing value would be the marker -2147483648 -> NaN |
| (none occur in this dataset). |
| - `molecule_ids` are integers and match the sigma-shake test-set ids. |
| |
| 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 GDBQCD: |
| """Reader for gdb_qcd.hdf5. Use as a context manager; index molecules by position or by id.""" |
|
|
| def __init__(self, path: str): |
| self.f = h5py.File(path, "r") |
| self.molecule_ids = self.f["molecule_ids"][:] |
| self.n_molecules = int(self.f.attrs["n_molecules"]) |
| self.n_frames = int(self.f.attrs["n_frames"]) |
| self.nmr_level_of_theory = str(self.f.attrs["nmr_level_of_theory"]) |
| self._n_atoms = self.f["n_atoms"][:].astype(np.int64) |
| self._n_traj = self.f["n_trajectories"][:].astype(np.int64) |
| |
| self._astart = np.empty(self.n_molecules + 1, np.int64) |
| self._astart[0] = 0 |
| np.cumsum(self._n_atoms, out=self._astart[1:]) |
| block = self._n_traj * self.n_frames * self._n_atoms |
| self._cstart = np.empty(self.n_molecules + 1, np.int64) |
| self._cstart[0] = 0 |
| np.cumsum(block, out=self._cstart[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 index_of(self, molecule_id: int) -> int: |
| """Position of the molecule with this id; raises KeyError if absent.""" |
| if int(molecule_id) not in self._id_to_index: |
| raise KeyError(f"molecule id {molecule_id} not in gdb_qcd") |
| return self._id_to_index[int(molecule_id)] |
|
|
| def n_trajectories(self, index: int) -> int: |
| """How many trajectories molecule i has.""" |
| self._check(index) |
| return int(self._n_traj[index]) |
|
|
| def molecule(self, index: int) -> dict: |
| """Return one molecule's trajectory data (physical units) by position. |
| |
| Keys: |
| id - integer molecule id (matches the sigma-shake test set) |
| atomic_numbers - (n_atoms,) uint8 |
| coordinates - (n_trajectories, 33, n_atoms, 3) Angstrom |
| shieldings - (n_trajectories, 33, n_atoms) ppm, PBE0/pcSseg-1 |
| """ |
| self._check(index) |
| na = int(self._n_atoms[index]) |
| nt = int(self._n_traj[index]) |
| asl = slice(int(self._astart[index]), int(self._astart[index + 1])) |
| csl = slice(int(self._cstart[index]), int(self._cstart[index + 1])) |
| coords = _decode(self.f["coordinates"][csl]).reshape(nt, self.n_frames, na, 3) |
| shields = _decode(self.f["shieldings"][csl]).reshape(nt, self.n_frames, na) |
| return { |
| "id": int(self.molecule_ids[index]), |
| "atomic_numbers": self.f["atomic_numbers"][asl], |
| "coordinates": coords, |
| "shieldings": shields, |
| } |
|
|
| def molecule_by_id(self, molecule_id: int) -> dict: |
| """Return one molecule's trajectory data (see `molecule`) by id.""" |
| return self.molecule(self.index_of(molecule_id)) |
|
|
| def close(self): |
| """Close the underlying HDF5 file.""" |
| self.f.close() |
|
|
| def __enter__(self): |
| return self |
|
|
| def __exit__(self, *a): |
| self.close() |
|
|
| def __len__(self): |
| return self.n_molecules |
|
|
| def __repr__(self): |
| return (f"<GDBQCD: {self.n_molecules:,} molecules, " |
| f"{int(self.f.attrs['n_trajectories']):,} trajectories, " |
| f"{self.n_frames} frames each, {self.nmr_level_of_theory}>") |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| path = sys.argv[1] if len(sys.argv) > 1 else "gdb_qcd.hdf5" |
| with GDBQCD(path) as ds: |
| print(ds) |
| m = ds.molecule(0) |
| print(f" molecule id={m['id']}: {len(m['atomic_numbers'])} atoms, " |
| f"coordinates {m['coordinates'].shape}, shieldings {m['shieldings'].shape}") |
| print(f" trajectory 0, frame 0, atom 0: xyz={np.round(m['coordinates'][0,0,0], 4)} " |
| f"shielding={m['shieldings'][0,0,0]:.4f} ppm") |
|
|
