File size: 5,056 Bytes
fefb9a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """
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"])
# atom row offsets: start[i]..start[i+1] for molecule i (one cumsum, cached)
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), # (n,3) Angstrom
"failed": index in self._failed,
}
for s in self.shieldings:
out[s] = self._scaled(s, sl) # (n,) ppm
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)}")
|