File size: 6,704 Bytes
fefb9a6
 
 
 
23c6f41
 
 
 
 
 
 
 
 
 
 
fefb9a6
 
23c6f41
 
 
 
fefb9a6
 
 
 
23c6f41
 
 
fefb9a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23c6f41
 
 
 
fefb9a6
 
 
 
 
 
 
 
 
 
23c6f41
 
 
fefb9a6
 
23c6f41
 
 
 
 
 
 
 
 
 
 
 
 
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""
decode_delta50.py - reader/decoder for the delta50 HDF5 dataset.

DELTA50 is a small benchmark of 50 common organic molecules (reference 47 in the paper). This
file holds, per atom on each molecule's AIMNet2 geometry: MagNET's predicted shieldings, the DFT
reference shieldings MagNET-Zero is trained to reproduce, and the published experimental shifts.

  nn_magnet_zero          MagNET-Zero (WP04 for 1H, wB97X-D for 13C; pcSseg-2; 20-pass), the
                          gas-phase triple-zeta-quality prediction
  nn_b3lyp                the gas-phase component of the MagNET-PCM correction (B3LYP/pcSseg-2)
  nn_b3lyp_pcm            the chloroform component of the MagNET-PCM correction (B3LYP/pcSseg-2 + PCM)
  shielding_wp04_pcSseg2  DFT reference, WP04/pcSseg-2 gas (the level MagNET-Zero targets for 1H)
  shielding_wb97xd_pcSseg2 DFT reference, wB97X-D/pcSseg-2 gas (the level MagNET-Zero targets for 13C)
  experimental_shift      experimental 1H/13C shift (ppm, CDCl3, TMS ref) from DELTA50 (Cohen et
                          al., Molecules 2023, 28, 2449; CC BY 4.0), mapped onto this atom order

The MagNET-PCM implicit-solvent correction for a molecule is nn_b3lyp_pcm minus nn_b3lyp. The
MagNET-Zero-vs-DFT residual (nn_magnet_zero minus shielding_wp04_pcSseg2 at 1H, minus
shielding_wb97xd_pcSseg2 at 13C) is what the paper uses to probe outliers (nitromethane, nitroethane,
and 2-methyl-2-nitropropane are the notable ones). experimental_shift is NaN at atoms with no
reported shift (heteroatoms).

Storage format:
  - Per-atom arrays are concatenated across molecules. Molecule i has `n_atoms[i]` atoms and owns
    the rows [start[i], start[i+1]) where start is the cumulative sum of `n_atoms`.
  - coordinates, shieldings, and shifts are int32 fixed point: physical = stored * 1e-4, so
    reconstruction error is <= 5e-5 (Angstrom or ppm). A missing value is the marker
    -2147483648 -> NaN.
  - molecules are identified by name (`molecule_names`).

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 Delta50:
    """Reader for delta50.hdf5. Use as a context manager; index molecules by position or by name."""

    def __init__(self, path: str):
        self.f = h5py.File(path, "r")
        self.molecule_names = [str(n) for n in self.f["molecule_names"].asstr()[:]]
        self.n_molecules = int(self.f.attrs["n_molecules"])
        self._n_atoms = self.f["n_atoms"][:].astype(np.int64)
        self._start = np.empty(self.n_molecules + 1, np.int64)
        self._start[0] = 0
        np.cumsum(self._n_atoms, out=self._start[1:])
        self._name_to_index = {n: i for i, n in enumerate(self.molecule_names)}

    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, name: str) -> int:
        """Position of the molecule named `name`; raises KeyError if absent."""
        if name not in self._name_to_index:
            raise KeyError(f"molecule {name!r} not in delta50")
        return self._name_to_index[name]

    def molecule(self, index: int) -> dict:
        """Return one molecule's data (physical units) by position.

        Keys: name; atomic_numbers (n,); coordinates (n, 3 Angstrom); the three predicted shieldings
        (nn_magnet_zero, nn_b3lyp, nn_b3lyp_pcm); the two DFT reference shieldings
        (shielding_wp04_pcSseg2, shielding_wb97xd_pcSseg2); and experimental_shift. Every per-atom
        array is (n,) ppm; experimental_shift is NaN at atoms with no reported shift.
        """
        self._check(index)
        sl = slice(int(self._start[index]), int(self._start[index + 1]))
        return {
            "name": self.molecule_names[index],
            "atomic_numbers": self.f["atomic_numbers"][sl],
            "coordinates": _decode(self.f["coordinates"][sl]),
            "nn_magnet_zero": _decode(self.f["nn_magnet_zero"][sl]),
            "nn_b3lyp": _decode(self.f["nn_b3lyp"][sl]),
            "nn_b3lyp_pcm": _decode(self.f["nn_b3lyp_pcm"][sl]),
            "shielding_wp04_pcSseg2": _decode(self.f["shielding_wp04_pcSseg2"][sl]),
            "shielding_wb97xd_pcSseg2": _decode(self.f["shielding_wb97xd_pcSseg2"][sl]),
            "experimental_shift": _decode(self.f["experimental_shift"][sl]),
        }

    def all_atoms(self) -> dict:
        """Every atom in the dataset as flat arrays (physical units), for whole-dataset analyses
        such as the MagNET-Zero-vs-DFT residual or DFT-vs-experiment comparisons.

        Keys: atomic_numbers (n_total,); nn_magnet_zero, nn_b3lyp, nn_b3lyp_pcm,
        shielding_wp04_pcSseg2, shielding_wb97xd_pcSseg2, experimental_shift, each (n_total,) ppm.
        """
        keys = ("nn_magnet_zero", "nn_b3lyp", "nn_b3lyp_pcm", "shielding_wp04_pcSseg2",
                "shielding_wb97xd_pcSseg2", "experimental_shift")
        out = {"atomic_numbers": self.f["atomic_numbers"][:]}
        out.update({k: _decode(self.f[k][:]) for k in keys})
        return out

    def molecule_by_name(self, name: str) -> dict:
        """Return one molecule's data (see `molecule`) by name."""
        return self.molecule(self.index_of(name))

    def pcm_correction(self, index: int) -> np.ndarray:
        """The MagNET-PCM implicit-solvent correction (nn_b3lyp_pcm minus nn_b3lyp), (n,) ppm."""
        m = self.molecule(index)
        return m["nn_b3lyp_pcm"] - m["nn_b3lyp"]

    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"<Delta50: {self.n_molecules} molecules>"


if __name__ == "__main__":
    import sys
    path = sys.argv[1] if len(sys.argv) > 1 else "delta50.hdf5"
    with Delta50(path) as ds:
        print(ds)
        m = ds.molecule(ds.index_of("nitromethane") if "nitromethane" in ds.molecule_names else 0)
        c = m["atomic_numbers"] == 6
        print(f"  {m['name']}: {len(m['atomic_numbers'])} atoms")
        print(f"    13C MagNET-Zero[:3] = {np.round(m['nn_magnet_zero'][c][:3], 3)} ppm")