| """ |
| decode_sigma_concentrate.py - reader/decoder for the sigma-concentrate HDF5 dataset. |
| |
| sigma-concentrate holds B3LYP/pcSseg-2 NMR shieldings computed both in the gas phase and with |
| an implicit chloroform solvent (PCM), on AIMNet2 stationary geometries for 50,000 molecules |
| sampled from sigma-shake. It is the data behind the MagNET-PCM implicit-solvent model, which |
| learns the solvent correction shielding_b3lyp_pcm - shielding_b3lyp_gas. |
| |
| The file is flat (a single molecule set, one geometry per molecule): |
| |
| shielding_b3lyp_gas isotropic sigma, B3LYP/pcSseg-2, gas phase |
| shielding_b3lyp_pcm isotropic sigma, B3LYP/pcSseg-2, PCM, solvent = chloroform |
| |
| Storage format: |
| - Per-atom arrays are concatenated across all molecules. Each molecule's atom count is in |
| `n_atoms`; the row range for molecule i is [start[i], start[i+1]) where start = 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). Coordinates are mean-centered per |
| molecule (centroid at the origin), a translation only; interatomic distances are unchanged. |
| - The molecule's array index IS its identifier (row order is fixed to the source). `sigma_shake_id` |
| additionally gives each molecule's id in the sigma-shake dataset (same molecule there, at a |
| different level of theory), so the two datasets can be joined. |
| - `failed_indices` lists molecules whose shielding calculation failed (none in this release): |
| their shieldings would be sentinel zeros and must be ignored. There are no SMILES. |
| |
| Requires: Python >= 3.7, numpy, h5py. No other dependencies. |
| """ |
| from __future__ import annotations |
| import numpy as np |
| import h5py |
|
|
|
|
| class SigmaConcentrate: |
| """Reader for the sigma-concentrate HDF5 file: gas/PCM shieldings on 50,000 sigma-shake molecules.""" |
|
|
| def __init__(self, path: str): |
| self.f = h5py.File(path, "r") |
| self.n_molecules = int(self.f.attrs["n_molecules"]) |
| self.n_atoms_total = int(self.f.attrs["n_atoms"]) |
| self.level_of_theory = str(self.f.attrs["level_of_theory"]) |
| self.geometry = str(self.f.attrs["geometry"]) |
| |
| na = self.f["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._sid = self.f["sigma_shake_id"][:] |
| self._failed = set(int(x) for x in self.f["failed_indices"][:]) |
| self.shieldings = sorted(k for k in self.f.keys() if k.startswith("shielding_")) |
|
|
| 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 _scaled(self, name: str, sl: slice) -> np.ndarray: |
| d = self.f[name] |
| return d[sl].astype(np.float64) * float(d.attrs["scale"]) |
|
|
| def sigma_shake_id(self, index: int) -> int: |
| """Return molecule `index`'s id in the sigma-shake dataset.""" |
| self._check(index) |
| return int(self._sid[index]) |
|
|
| def molecule(self, index: int) -> dict: |
| """Return one molecule's decoded data (physical units) by array index. |
| |
| Keys: index, sigma_shake_id, atomic_numbers, coordinates (n,3 Angstrom), |
| shielding_b3lyp_gas (n, ppm), shielding_b3lyp_pcm (n, ppm), and `failed`. If `failed` |
| is True the shieldings are sentinel zeros and should be ignored; geometry is still valid. |
| """ |
| self._check(index) |
| a, b = int(self._start[index]), int(self._start[index + 1]) |
| sl = slice(a, b) |
| out = { |
| "index": index, |
| "sigma_shake_id": int(self._sid[index]), |
| "atomic_numbers": self.f["atomic_numbers"][sl], |
| "coordinates": self._scaled("coordinates", sl), |
| "failed": index in self._failed, |
| } |
| for s in self.shieldings: |
| out[s] = self._scaled(s, sl) |
| return out |
|
|
| def __len__(self): |
| return self.n_molecules |
|
|
| 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"<SigmaConcentrate: {self.n_molecules:,} molecules, " |
| f"{self.n_atoms_total:,} atoms, {self.level_of_theory}>") |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| path = sys.argv[1] if len(sys.argv) > 1 else "sigma-concentrate.hdf5" |
| with SigmaConcentrate(path) as ds: |
| print(ds) |
| m = ds.molecule(0) |
| print(f" mol 0: sigma_shake_id={m['sigma_shake_id']} Z={list(m['atomic_numbers'][:6])} " |
| f"failed={m['failed']}") |
| print(f" gas={np.round(m['shielding_b3lyp_gas'][:3], 4)} " |
| f"pcm={np.round(m['shielding_b3lyp_pcm'][:3], 4)}") |
|
|
