File size: 6,361 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 | """
Tests for the sigma-concentrate reader (decode_sigma_concentrate.SigmaConcentrate).
These build a tiny synthetic 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 real
file. One opt-in test runs against the real `sigma-concentrate.hdf5` if it is present next to this
file.
Run: pytest test_sigma_concentrate.py -q
Requires: pytest, numpy, h5py.
"""
import os
import numpy as np
import h5py
import pytest
from decode_sigma_concentrate import SigmaConcentrate
SCALE = 1e-4
def _i32(x):
return np.round(np.asarray(x, float) / SCALE).astype(np.int32)
def _write(f, *, n_atoms, znums, sids, coords, gas, pcm, failed):
assert len(znums) == int(np.sum(n_atoms))
f.attrs["dataset"] = "sigma-concentrate"
f.attrs["level_of_theory"] = "B3LYP/pcSseg-2"
f.attrs["geometry"] = "AIMNet2 stationary"
f.attrs["solvent"] = "gas + PCM chloroform"
f.attrs["n_molecules"] = len(n_atoms)
f.attrs["n_atoms"] = int(np.sum(n_atoms))
f.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8))
f.create_dataset("n_atoms", data=np.asarray(n_atoms, np.uint8))
f.create_dataset("sigma_shake_id", data=np.asarray(sids, np.int32))
d = f.create_dataset("coordinates", data=_i32(coords)); d.attrs["scale"] = SCALE; d.attrs["units"] = "angstrom"
d = f.create_dataset("shielding_b3lyp_gas", data=_i32(gas)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
d = f.create_dataset("shielding_b3lyp_pcm", data=_i32(pcm)); d.attrs["scale"] = SCALE; d.attrs["units"] = "ppm"
f.create_dataset("failed_indices", data=np.asarray(failed, np.int32))
# 3 molecules (2,1,3 atoms = 6); heteroatoms exercise the passthrough; molecule 1 failed
NATOMS = [2, 1, 3]
Z = [6, 1, 8, 6, 7, 1]
SIDS = [1000265, 42, 999808]
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)
GAS = np.array([180.1234, 31.0, 0.0, 150.0, 151.5, 29.9]) # mol 1 (atom 2) = sentinel 0
PCM = np.array([180.4567, 30.8, 0.0, 149.6, 151.2, 29.7])
@pytest.fixture(scope="module")
def mini_path(tmp_path_factory):
p = str(tmp_path_factory.mktemp("sc") / "mini.hdf5")
with h5py.File(p, "w") as f:
_write(f, n_atoms=NATOMS, znums=Z, sids=SIDS, coords=COORDS, gas=GAS, pcm=PCM, failed=[1])
return p
@pytest.fixture
def ds(mini_path):
with SigmaConcentrate(mini_path) as d:
yield d
def test_dimensions(ds):
assert ds.n_molecules == 3 and ds.n_atoms_total == 6
assert ds.shieldings == ["shielding_b3lyp_gas", "shielding_b3lyp_pcm"]
assert len(ds) == 3
def test_offsets_slice_each_molecule(ds):
assert len(ds.molecule(0)["atomic_numbers"]) == 2
assert len(ds.molecule(1)["atomic_numbers"]) == 1
assert len(ds.molecule(2)["atomic_numbers"]) == 3
def test_values_roundtrip_and_id(ds):
m = ds.molecule(2)
assert m["sigma_shake_id"] == 999808
np.testing.assert_array_equal(m["atomic_numbers"], [6, 7, 1]) # heteroatom passthrough
np.testing.assert_allclose(m["coordinates"], COORDS[3:6], atol=5e-5)
np.testing.assert_allclose(m["shielding_b3lyp_gas"], GAS[3:6], atol=5e-5)
np.testing.assert_allclose(m["shielding_b3lyp_pcm"], PCM[3:6], atol=5e-5)
assert not m["failed"]
def test_both_shieldings_distinct(ds):
m = ds.molecule(0)
# gas and pcm differ atom-by-atom (the solvent correction is the signal)
assert not np.allclose(m["shielding_b3lyp_gas"], m["shielding_b3lyp_pcm"])
def test_sigma_shake_id_accessor(ds):
assert ds.sigma_shake_id(0) == 1000265
assert ds.sigma_shake_id(1) == 42
def test_failed_flag(ds):
assert ds.molecule(1)["failed"] is True
assert ds.molecule(0)["failed"] is False
f = ds.molecule(1)
# failed molecule's shieldings are sentinel zeros; geometry/id still valid
assert np.all(f["shielding_b3lyp_gas"] == 0) and np.all(f["shielding_b3lyp_pcm"] == 0)
assert f["sigma_shake_id"] == 42
np.testing.assert_allclose(f["coordinates"], COORDS[2: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(f, n_atoms=[2], znums=[6, 1], sids=[7],
coords=np.array([[0, 0, 0], [1.0, 0, 0]], float),
gas=np.array([150.0, 30.0]), pcm=np.array([149.5, 29.8]), failed=[])
with SigmaConcentrate(p) as d:
assert d.molecule(0)["failed"] is False
def test_index_out_of_range_raises(ds):
with pytest.raises(IndexError):
ds.molecule(-1)
with pytest.raises(IndexError):
ds.molecule(3)
def test_context_manager_closes(mini_path):
d = SigmaConcentrate(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-concentrate.hdf5")
@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-concentrate.hdf5 not present")
def test_real_file():
with SigmaConcentrate(REAL) as ds:
assert ds.n_molecules == 50_000 and ds.n_atoms_total == 1_413_134
assert len(ds.f["failed_indices"]) == 0
# gas and pcm are distinct columns; PCM differs from gas on carbon
Z = ds.f["atomic_numbers"][:500_000]
gas = ds.f["shielding_b3lyp_gas"][:500_000].astype(float) * SCALE
pcm = ds.f["shielding_b3lyp_pcm"][:500_000].astype(float) * SCALE
C = Z == 6
assert gas[C].mean() != pcm[C].mean()
# sigma_shake_id is populated and plausibly a sigma-shake id (positive)
assert ds.sigma_shake_id(0) > 0
@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-concentrate.hdf5 not present")
def test_real_file_reproduces_table_s8_site_counts():
"""SI Table S8's 'sigma-concentrate' row is this file exactly, and (together with
sigma-pepper's wp04_wb97xd_pcSseg2 group -- see test_sigma_pepper.py's matching test) these
50,000 reused molecules make up the other half of the 'MagNET-Zero (fine-tuning)' row:
3,287,698 + 754,437 = 4,042,135 H sites; 2,028,642 + 528,569 = 2,557,211 C sites."""
with SigmaConcentrate(REAL) as ds:
z = ds.f["atomic_numbers"][:]
assert int((z == 1).sum()) == 754_437
assert int((z == 6).sum()) == 528_569
|