File size: 10,733 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 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 | """
Tests for the dft8k reader (dft8k_reader.py).
Most tests build a tiny synthetic dft8k.hdf5 in the exact on-disk format (so the fixture
doubles as a spec of what the reader expects) and need no large data. One opt-in test runs
against the real dft8k.hdf5 if it is present (skipped otherwise, e.g. in CI).
Run: pytest test_dft8k.py -q
Requires: pytest, numpy, h5py.
"""
import os
import numpy as np
import h5py
import pytest
from dft8k_reader import DFT8k, _decode, _MARKER
SCALE = 1e4
def _enc(arr):
"""Encode floats the way the build script does: int32 = round(value * 1e4), NaN -> marker."""
a = np.asarray(arr, dtype=np.float64)
nan = np.isnan(a)
out = np.rint(np.where(nan, 0.0, a) * SCALE).astype(np.int32)
out[nan] = _MARKER
return out
# ---------------------------------------------------------------------------
# fixed-point decoding
# ---------------------------------------------------------------------------
def test_decode_round_trip_and_marker():
values = np.array([0.0, 207.07, -355.5, np.nan, 875.0], dtype=np.float64)
decoded = _decode(_enc(values))
assert np.isnan(decoded[3])
np.testing.assert_allclose(decoded[~np.isnan(decoded)],
values[~np.isnan(values)], atol=5e-5)
def test_decode_passes_through_float():
# a plain-decimal array must read back unchanged (decode is a no-op on floats)
values = np.array([1.0, 2.5, -7.25], dtype=np.float64)
np.testing.assert_array_equal(_decode(values), values)
def test_decode_handles_int64():
# any integer dtype decodes the same way as int32
vals = np.array([10000, -3550000, _MARKER], dtype=np.int64)
out = _decode(vals)
np.testing.assert_allclose(out[:2], [1.0, -355.0])
assert np.isnan(out[2])
# ---------------------------------------------------------------------------
# synthetic dft8k.hdf5 in the real two-group layout
# ---------------------------------------------------------------------------
# two molecules: ids 100 and 200, with 2 and 3 atoms (atoms are H, C, ...)
A_IDS = [100, 200]
A_COUNTS = [2, 3]
A_Z = [1, 6, 6, 1, 7] # mol100: H C ; mol200: C H N
# per-atom DFT shieldings (concatenated): pbe0, wp04, wb97xd
A_PBE0 = [30.0, 150.0, 151.0, 31.0, 200.0]
A_WP04 = [30.5, 150.5, 151.5, 31.5, 200.5] # the 1H reference
A_WB97 = [30.7, 150.7, 151.7, 31.7, 200.7] # the 13C reference
# MagNET predicts H with the WP04 model and C with the wB97X-D model; other atoms blank
A_NN_WP04 = [30.6, np.nan, np.nan, 31.6, np.nan] # only at H atoms (index 0, 3)
A_NN_WB97 = [np.nan, 150.6, 151.6, np.nan, np.nan] # only at C atoms (index 1, 2)
def _make_synthetic_dft8k(path):
with h5py.File(path, "w") as f:
g = f.create_group("aimnet2_wp04_wb97xd_pcsseg2")
g.attrs["geometry"] = "AIMNet2"
g.attrs["level_of_theory"] = "PBE0/pcSseg-1; WP04/pcSseg-2; wB97X-D/pcSseg-2"
g.attrs["n_molecules"] = len(A_IDS)
g.attrs["n_atoms"] = sum(A_COUNTS)
g.attrs["scale"] = 1e-4
g.create_dataset("molecule_ids", data=np.array(A_IDS, np.uint32))
g.create_dataset("n_atoms", data=np.array(A_COUNTS, np.uint8))
# one SMILES per molecule; mol100 has a real one, mol200 carries the literal "none"
g.create_dataset("smiles", data=np.array(["CC", "none"], dtype=object),
dtype=h5py.string_dtype(encoding="utf-8"))
g.create_dataset("atomic_numbers", data=np.array(A_Z, np.int8))
coords = np.arange(sum(A_COUNTS) * 3, dtype=np.float64).reshape(-1, 3) * 0.1
g.create_dataset("coordinates", data=_enc(coords))
g.create_dataset("shielding_pbe0_pcSseg1", data=_enc(A_PBE0))
g.create_dataset("shielding_wp04_pcSseg2", data=_enc(A_WP04))
g.create_dataset("shielding_wb97xd_pcSseg2", data=_enc(A_WB97))
g.create_dataset("nn_shielding_wp04_pcSseg2", data=_enc(A_NN_WP04))
g.create_dataset("nn_shielding_wb97xd_pcSseg2", data=_enc(A_NN_WB97))
# second group, SAME ids but reversed row order, three geometries per molecule
b = f.create_group("b3lyp_pbe0_pcsseg1")
b.attrs["geometry"] = "B3LYP/pcSseg-1"
b.attrs["level_of_theory"] = "PBE0/pcSseg-1 on B3LYP/pcSseg-1 geometries"
b.attrs["n_molecules"] = 2
b.attrs["n_atoms"] = 5
b.attrs["n_geometries"] = 3
b.attrs["scale"] = 1e-4
b.create_dataset("molecule_ids", data=np.array([200, 100], np.uint32)) # reversed
b.create_dataset("n_atoms", data=np.array([3, 2], np.uint8))
b.create_dataset("atomic_numbers", data=np.array([6, 1, 7, 1, 6], np.int8))
bcoords = np.arange(3 * 5 * 3, dtype=np.float64).reshape(3, 5, 3) * 0.01
b.create_dataset("coordinates", data=_enc(bcoords))
bshield = np.arange(3 * 5, dtype=np.float64).reshape(3, 5) + 100.0
b.create_dataset("shielding", data=_enc(bshield))
b.create_dataset("weights", data=np.array([1, 1, 1, 0.5, 0.5], np.float32))
# symmetry groups are 0-based atom indices, local to each molecule.
# mol200 (first row, 3 atoms): two groups [[1, 2]] and [[0]];
# mol100 (second row, 2 atoms): one group [[0, 1]].
b.create_dataset("symmetric_atoms_n_groups", data=np.array([2, 1], np.int32))
b.create_dataset("symmetric_atoms_group_lengths", data=np.array([2, 1, 2], np.int32))
b.create_dataset("symmetric_atoms_members", data=np.array([1, 2, 0, 0, 1], np.int32))
def test_aimnet2_group(tmp_path):
p = tmp_path / "dft8k.hdf5"
_make_synthetic_dft8k(str(p))
with DFT8k(str(p)) as ds:
assert len(ds.aimnet2) == 2
m = ds.aimnet2.molecule(0)
assert m["id"] == 100
assert m["smiles"] == "CC"
assert ds.aimnet2.smiles(1) == "none" # mol200 has no SMILES
np.testing.assert_array_equal(m["atomic_numbers"], [1, 6])
np.testing.assert_allclose(m["wp04_pcSseg2"], [30.5, 150.5], atol=5e-5)
# H atom (index 0) has a WP04 prediction; C atom (index 1) does not
assert not np.isnan(m["nn_wp04_pcSseg2"][0])
assert np.isnan(m["nn_wp04_pcSseg2"][1])
# C atom has a wB97X-D prediction; H atom does not
assert np.isnan(m["nn_wb97xd_pcSseg2"][0])
assert not np.isnan(m["nn_wb97xd_pcSseg2"][1])
np.testing.assert_allclose(m["nn_wb97xd_pcSseg2"][1], 150.6, atol=5e-5)
def test_all_shieldings_flat(tmp_path):
p = tmp_path / "dft8k.hdf5"
_make_synthetic_dft8k(str(p))
with DFT8k(str(p)) as ds:
flat = ds.aimnet2.all_shieldings()
# one entry per atom across both molecules, in stored order
np.testing.assert_array_equal(flat["atomic_numbers"], [1, 6, 6, 1, 7])
np.testing.assert_allclose(flat["wp04_pcSseg2"], A_WP04, atol=5e-5)
# MagNET prediction is present only at the atoms its model covers
assert np.isnan(flat["nn_wp04_pcSseg2"]).tolist() == [False, True, True, False, True]
# the 1H residual (DFT minus MagNET) at the two H atoms is -0.1
Z = flat["atomic_numbers"]
h = (Z == 1) & np.isfinite(flat["nn_wp04_pcSseg2"])
np.testing.assert_allclose(flat["wp04_pcSseg2"][h] - flat["nn_wp04_pcSseg2"][h],
[-0.1, -0.1], atol=5e-5)
def test_b3lyp_group_and_symmetry(tmp_path):
p = tmp_path / "dft8k.hdf5"
_make_synthetic_dft8k(str(p))
with DFT8k(str(p)) as ds:
b = ds.b3lyp
assert b.n_geometries == 3
m = b.molecule(0) # first row is id 200
assert m["id"] == 200
assert m["coordinates"].shape == (3, 3, 3) # (geometries, atoms, xyz)
assert m["shielding"].shape == (3, 3)
assert m["symmetric_atoms"] == [[1, 2], [0]]
m2 = b.molecule(1) # id 100
assert m2["symmetric_atoms"] == [[0, 1]]
np.testing.assert_allclose(m2["weights"], [0.5, 0.5])
# the symmetry indices are valid 0-based positions into this molecule's atoms
for group in m["symmetric_atoms"] + m2["symmetric_atoms"]:
for atom_index in group:
assert 0 <= atom_index < len(m["atomic_numbers"]) or \
0 <= atom_index < len(m2["atomic_numbers"])
def test_match_across_groups_by_id(tmp_path):
p = tmp_path / "dft8k.hdf5"
_make_synthetic_dft8k(str(p))
with DFT8k(str(p)) as ds:
# rows are ordered differently in the two groups; matching by id must still work
a = ds.aimnet2.molecule_by_id(200)
b = ds.b3lyp.molecule_by_id(200)
assert a["id"] == b["id"] == 200
np.testing.assert_array_equal(a["atomic_numbers"], b["atomic_numbers"])
# ---------------------------------------------------------------------------
# opt-in test against the real file (skipped when absent, e.g. in CI)
# ---------------------------------------------------------------------------
_HERE = os.path.dirname(os.path.abspath(__file__))
_REAL = os.path.join(_HERE, "dft8k.hdf5")
@pytest.mark.skipif(not os.path.exists(_REAL), reason="real dft8k.hdf5 not present")
def test_real_dft8k():
with DFT8k(_REAL) as ds:
assert len(ds.aimnet2) == 7111
assert len(ds.b3lyp) == 7111
assert ds.b3lyp.n_geometries == 3
# the two groups carry the same id set
assert set(int(x) for x in ds.aimnet2.molecule_ids) == \
set(int(x) for x in ds.b3lyp.molecule_ids)
m = ds.aimnet2.molecule(0)
# SMILES: 6780 of the 7111 molecules have one; the rest carry the literal "none"
n_real_smiles = sum(1 for i in range(len(ds.aimnet2)) if ds.aimnet2.smiles(i) != "none")
assert n_real_smiles == 6780
assert isinstance(m["smiles"], str) and m["smiles"] != "none"
z = m["atomic_numbers"]
h = z == 1
c = z == 6
other = ~(h | c)
# the WP04 model predicts every hydrogen and nothing else; the wB97X-D model
# predicts every carbon and nothing else; non-target atoms read back blank
assert not np.isnan(m["nn_wp04_pcSseg2"][h]).any()
assert np.isnan(m["nn_wp04_pcSseg2"][~h]).all()
assert not np.isnan(m["nn_wb97xd_pcSseg2"][c]).any()
assert np.isnan(m["nn_wb97xd_pcSseg2"][~c]).all()
assert other.sum() == 0 or np.isnan(m["nn_wp04_pcSseg2"][other]).all()
# the same molecule resolves in the other group by id
b = ds.b3lyp.molecule_by_id(m["id"])
np.testing.assert_array_equal(m["atomic_numbers"], b["atomic_numbers"])
# every symmetry index is a valid 0-based position into this molecule's atoms
for group in b["symmetric_atoms"]:
for atom_index in group:
assert 0 <= atom_index < len(b["atomic_numbers"])
|