File size: 8,304 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 | """
Tests for the sigma-pepper reader (decode_sigma_pepper.SigmaPepper).
These build a tiny synthetic two-group dataset in the *exact* on-disk format (so the fixture
also serves as an executable spec of the layout) and exercise the decoder against it. No
dependency on the multi-hundred-MB real file. One opt-in test runs against the real
`sigma-pepper.hdf5` if it is present next to this file.
Run: pytest test_sigma_pepper.py -q
Requires: pytest, numpy, h5py.
"""
import os
import numpy as np
import h5py
import pytest
from decode_sigma_pepper import SigmaPepper
SCALE = 1e-4
def _i32(x):
return np.round(np.asarray(x, float) / SCALE).astype(np.int32)
def _write_group(f, name, *, n_atoms, znums, smiles, coords, shieldings, failed, attrs):
"""shieldings: dict name->per-atom array; failed: list of molecule indices."""
g = f.create_group(name)
g.attrs.update(attrs)
g.attrs["n_molecules"] = len(n_atoms)
g.attrs["n_atoms"] = int(np.sum(n_atoms))
assert len(znums) == int(np.sum(n_atoms))
g.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8))
g.create_dataset("n_atoms", data=np.asarray(n_atoms, np.uint8))
d = g.create_dataset("coordinates", data=_i32(coords)); d.attrs["scale"] = SCALE; d.attrs["units"] = "angstrom"
for sname, arr in shieldings.items():
d = g.create_dataset(sname, data=_i32(arr)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
g.create_dataset("failed_indices", data=np.asarray(failed, np.int32))
enc = [s.encode("utf-8") for s in smiles]
off = np.concatenate([[0], np.cumsum([len(e) for e in enc])]).astype(np.int64)
g.create_dataset("smiles_buffer", data=np.frombuffer(b"".join(enc), np.uint8))
g.create_dataset("smiles_offsets", data=off)
# group A: 3 molecules (2,1,3 atoms = 6), one shielding, molecule 1 failed
A_NATOMS = [2, 1, 3]
A_Z = [6, 1, 8, 6, 7, 1] # heteroatoms exercise the passthrough
A_SMILES = ["CC", "O", "CCC"]
A_COORDS = np.array([[0, 0, 0], [1.2345, 0, 0],
[5, 5, 5],
[3, 0, 0], [4.1, 0, 0], [3, 1.5, 0]], float)
A_PBE0 = np.array([180.1234, 31.0, 0.0, 150.0, 151.5, 29.9]) # mol 1 (atom idx 2) = sentinel 0
# group B: 2 molecules (3,2 atoms = 5), two shieldings, molecule 0 failed
B_NATOMS = [3, 2]
B_Z = [6, 8, 1, 6, 7]
B_SMILES = ["CCO", "CN"]
B_COORDS = np.array([[0, 0, 0], [1, 0, 0], [2, 0, 0],
[9, 9, 9], [9.5, 9, 9]], float)
B_WB97 = np.array([0.0, 0.0, 0.0, 110.0, 28.0]) # mol 0 failed (all zero)
B_WP04 = np.array([0.0, 0.0, 0.0, 98.0, 28.3])
@pytest.fixture(scope="module")
def mini_path(tmp_path_factory):
p = str(tmp_path_factory.mktemp("sp") / "mini.hdf5")
with h5py.File(p, "w") as f:
_write_group(f, "pbe0_pcSseg1", n_atoms=A_NATOMS, znums=A_Z, smiles=A_SMILES, coords=A_COORDS,
shieldings={"shielding_pbe0": A_PBE0}, failed=[1],
attrs=dict(level_of_theory="PBE0/pcSseg-1, gas",
geometry="AIMNet2 stationary (fmax=0.01)", gdb_range="1-10 heavy"))
_write_group(f, "wp04_wb97xd_pcSseg2", n_atoms=B_NATOMS, znums=B_Z, smiles=B_SMILES, coords=B_COORDS,
shieldings={"shielding_wb97xd": B_WB97, "shielding_wp04": B_WP04}, failed=[0],
attrs=dict(level_of_theory="WP04 + wB97X-D /pcSseg-2, gas",
geometry="AIMNet2 stationary (fmax=1e-4)", gdb_range="1-9 heavy"))
return p
@pytest.fixture
def ds(mini_path):
with SigmaPepper(mini_path) as d:
yield d
def test_group_names(ds):
assert set(ds.group_names) == {"pbe0_pcSseg1", "wp04_wb97xd_pcSseg2"}
def test_component_dimensions(ds):
a = ds["pbe0_pcSseg1"]
assert a.n_molecules == 3 and a.n_atoms == 6
b = ds["wp04_wb97xd_pcSseg2"]
assert b.n_molecules == 2 and b.n_atoms == 5
assert b.shieldings == ["shielding_wb97xd", "shielding_wp04"]
def test_offsets_slice_each_molecule(ds):
a = ds["pbe0_pcSseg1"]
assert len(a.molecule(0)["atomic_numbers"]) == 2
assert len(a.molecule(1)["atomic_numbers"]) == 1
assert len(a.molecule(2)["atomic_numbers"]) == 3
def test_values_roundtrip_and_smiles(ds):
m = ds["pbe0_pcSseg1"].molecule(2)
assert m["smiles"] == "CCC"
np.testing.assert_array_equal(m["atomic_numbers"], [6, 7, 1]) # heteroatom passthrough
np.testing.assert_allclose(m["coordinates"], A_COORDS[3:6], atol=5e-5)
np.testing.assert_allclose(m["shielding_pbe0"], A_PBE0[3:6], atol=5e-5)
assert not m["failed"]
def test_two_shieldings_present(ds):
m = ds["wp04_wb97xd_pcSseg2"].molecule(1)
assert m["smiles"] == "CN"
np.testing.assert_allclose(m["shielding_wp04"], B_WP04[3:5], atol=5e-5)
np.testing.assert_allclose(m["shielding_wb97xd"], B_WB97[3:5], atol=5e-5)
def test_failed_flag(ds):
assert ds["pbe0_pcSseg1"].molecule(1)["failed"] is True
assert ds["pbe0_pcSseg1"].molecule(0)["failed"] is False
b0 = ds["wp04_wb97xd_pcSseg2"].molecule(0)
assert b0["failed"] is True
# failed molecule's shieldings are sentinel zeros (the whole point of the flag)
assert np.all(b0["shielding_wp04"] == 0) and np.all(b0["shielding_wb97xd"] == 0)
# geometry/smiles still valid for a failed molecule
assert b0["smiles"] == "CCO"
np.testing.assert_allclose(b0["coordinates"], B_COORDS[0:3], atol=5e-5)
def test_empty_failed_indices(tmp_path):
p = str(tmp_path / "nofail.hdf5")
with h5py.File(p, "w") as f:
_write_group(f, "pbe0_pcSseg1", n_atoms=[2], znums=[6, 1], smiles=["CO"],
coords=np.array([[0, 0, 0], [1.0, 0, 0]], float),
shieldings={"shielding_pbe0": np.array([150.0, 30.0])}, failed=[],
attrs=dict(level_of_theory="x", geometry="x", gdb_range="x"))
with SigmaPepper(p) as d:
assert d["pbe0_pcSseg1"].molecule(0)["failed"] is False
def test_smiles_accessor(ds):
assert ds["pbe0_pcSseg1"].smiles(0) == "CC"
assert ds["wp04_wb97xd_pcSseg2"].smiles(1) == "CN"
def test_index_out_of_range_raises(ds):
with pytest.raises(IndexError):
ds["pbe0_pcSseg1"].molecule(-1)
with pytest.raises(IndexError):
ds["pbe0_pcSseg1"].molecule(3)
def test_context_manager_closes(mini_path):
d = SigmaPepper(mini_path)
d.close()
assert not d.f.id.valid
# ---------- opt-in: real file, if present ----------
REAL = os.path.join(os.path.dirname(__file__), "sigma-pepper.hdf5")
@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-pepper.hdf5 not present")
def test_real_file():
with SigmaPepper(REAL) as ds:
a = ds["pbe0_pcSseg1"]
b = ds["wp04_wb97xd_pcSseg2"]
assert a.n_molecules == 2_167_210 and a.n_atoms == 45_418_890
assert b.n_molecules == 320_105 and b.n_atoms == 6_113_101
assert len(a.g["failed_indices"]) == 1
assert len(b.g["failed_indices"]) == 472
# column identity: WP04 must be lower than wB97XD on carbon
Z = b.g["atomic_numbers"][:2_000_000]
wp = b.g["shielding_wp04"][:2_000_000].astype(float) * SCALE
wb = b.g["shielding_wb97xd"][:2_000_000].astype(float) * SCALE
C = (Z == 6) & (wp != 0) & (wb != 0)
assert wp[C].mean() < wb[C].mean()
@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-pepper.hdf5 not present")
def test_real_file_reproduces_table_s8_site_counts():
"""SI Table S8's 'MagNET-Zero (initial round)' row is this file's pbe0_pcSseg1 group exactly.
The 'MagNET-Zero (fine-tuning)' row is this group's wp04_wb97xd_pcSseg2 PLUS sigma-concentrate
(the 50,000 reused molecules) -- see test_sigma_concentrate.py's matching test for that half."""
with SigmaPepper(REAL) as ds:
a = ds["pbe0_pcSseg1"]
z_a = a.g["atomic_numbers"][:]
assert int((z_a == 1).sum()) == 24_047_999
assert int((z_a == 6).sum()) == 15_330_361
b = ds["wp04_wb97xd_pcSseg2"]
z_b = b.g["atomic_numbers"][:]
# this group alone is not a Table S8 row (fine-tuning adds sigma-concentrate's atoms too),
# but its counts are the fixed half of that sum, so lock them down here too.
assert int((z_b == 1).sum()) == 3_287_698
assert int((z_b == 6).sum()) == 2_028_642
|