""" Tests for the sigma-shake reader (decode_sigma_shake.SigmaShake). 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 multi-GB real file. One opt-in test runs against the real `sigma-shake.hdf5` if it is present next to this file. Run: pytest test_sigma_shake.py -q Requires: pytest, numpy, h5py. """ import os import numpy as np import h5py import pytest from decode_sigma_shake import SigmaShake SCALE = 1e-4 # ---- known synthetic content: 3 molecules, 9 atoms; molecule 1 is incomplete ---- MOL_ID = np.array([100, 200, 300], np.int32) ATOM_START = np.array([0, 3, 5], np.int32) ATOM_END = np.array([3, 5, 9], np.int32) Z = np.array([6, 1, 1, 8, 1, 6, 6, 1, 1], np.int8) STAT_XYZ = np.array([ [0.0000, 0.0000, 0.0000], [1.0000, 0.0000, 0.0000], [0.0000, 1.0000, 0.0000], # mol 100 [2.0000, 2.0000, 2.0000], [2.5000, 2.0000, 2.0000], # mol 200 (incomplete) [3.0000, 0.0000, 0.0000], [4.0000, 0.0000, 0.0000], [3.0000, 1.0000, 0.0000], [4.0000, 1.0000, 0.0000], ], np.float64) STAT_SIG = np.array([150.1234, 30.5678, 31.0000, 200.0000, 25.0000, 140.0000, 141.0000, 29.0000, 29.5000]) PERT_XYZ = np.array([ [0.0100, 0.0000, 0.0000], [-0.0050, 0.0000, 0.0000], [0.0000, 0.0200, 0.0000], # mol 100 [0.0000, 0.0000, 0.0000], [0.0000, 0.0000, 0.0000], # mol 200 (zeroed) [0.0030, 0.0000, 0.0000], [-0.0030, 0.0000, 0.0000], [0.0000, 0.0050, 0.0000], [0.0000, -0.0050, 0.0000], ], np.float64) PERT_SIG = np.array([151.0000, 30.9999, 31.5000, 0.0, 0.0, 140.5000, 141.5000, 29.2000, 29.7000]) INCOMPLETE_MASK = np.array([0, 0, 0, 1, 1, 0, 0, 0, 0], bool) SPLIT = {"train": np.array([100, 200], np.int32), "val": np.array([300], np.int32), "test": np.array([], np.int32)} REMOVE = np.array([200], np.int32) INCOMPLETE_IDS = np.array([200], np.int32) def _i32(x): return np.round(np.asarray(x) / SCALE).astype(np.int32) @pytest.fixture(scope="module") def mini_path(tmp_path_factory): p = str(tmp_path_factory.mktemp("ss") / "mini.hdf5") with h5py.File(p, "w") as f: f.attrs["n_atoms"] = len(Z) f.attrs["n_molecules"] = len(MOL_ID) for name, arr in [("stationary_coords", STAT_XYZ), ("perturbation", PERT_XYZ), ("shielding_stationary", STAT_SIG), ("shielding_perturbed", PERT_SIG)]: d = f.create_dataset(name, data=_i32(arr)) d.attrs["scale"] = SCALE f.create_dataset("atomic_numbers", data=Z) f.create_dataset("incomplete_mask", data=INCOMPLETE_MASK) f.create_dataset("molecule_id", data=MOL_ID) f.create_dataset("atom_start", data=ATOM_START) f.create_dataset("atom_end", data=ATOM_END) f.create_dataset("remove_molecule_ids", data=REMOVE) f.create_dataset("incomplete_molecule_ids", data=INCOMPLETE_IDS) for k, v in SPLIT.items(): f.create_dataset(f"split_{k}_ids", data=v) return p @pytest.fixture def ds(mini_path): with SigmaShake(mini_path) as d: yield d def test_dimensions(ds): assert ds.n_molecules == 3 assert ds.n_atoms == 9 def test_complete_molecule_values(ds): m = ds.molecule(0) assert m["molecule_id"] == 100 assert not m["incomplete"] np.testing.assert_array_equal(m["atomic_numbers"], [6, 1, 1]) np.testing.assert_allclose(m["stationary_coords"], STAT_XYZ[0:3], atol=5e-5) np.testing.assert_allclose(m["shielding_stationary"], STAT_SIG[0:3], atol=5e-5) np.testing.assert_allclose(m["shielding_perturbed"], PERT_SIG[0:3], atol=5e-5) def test_perturbed_geometry_reconstruction(ds): m = ds.molecule(0) np.testing.assert_allclose(m["perturbed_coords"], STAT_XYZ[0:3] + PERT_XYZ[0:3], atol=1e-4) def test_incomplete_molecule_omits_perturbed_keys(ds): m = ds.molecule(1) assert m["molecule_id"] == 200 assert m["incomplete"] is True for k in ("perturbation", "perturbed_coords", "shielding_perturbed"): assert k not in m # stationary data still valid np.testing.assert_allclose(m["stationary_coords"], STAT_XYZ[3:5], atol=5e-5) np.testing.assert_allclose(m["shielding_stationary"], STAT_SIG[3:5], atol=5e-5) def test_scale_roundtrip_is_within_tolerance(ds): # every stored value should reconstruct to <= half the LSB (5e-5) m = ds.molecule(2) assert np.abs(m["stationary_coords"] - STAT_XYZ[5:9]).max() <= 5e-5 assert np.abs(m["shielding_stationary"] - STAT_SIG[5:9]).max() <= 5e-5 def test_by_id(ds): assert ds.by_id(100)["molecule_id"] == 100 assert ds.by_id(300)["molecule_id"] == 300 def test_by_id_unknown_raises(ds): with pytest.raises(KeyError): ds.by_id(999999) def test_index_out_of_range_raises(ds): with pytest.raises(IndexError): ds.molecule(-1) with pytest.raises(IndexError): ds.molecule(3) def test_split_effective_filters_remove_ids(ds): # raw train has 100 and 200; effective drops 200 (in remove_molecule_ids) np.testing.assert_array_equal(np.sort(ds.split("train", effective=False)), [100, 200]) np.testing.assert_array_equal(ds.split("train"), [100]) np.testing.assert_array_equal(ds.split("val"), [300]) assert len(ds.split("test")) == 0 def test_context_manager_closes(mini_path): d = SigmaShake(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-shake.hdf5") @pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-shake.hdf5 not present") def test_real_file_matches_si(): with SigmaShake(REAL) as ds: assert ds.n_molecules == 4_787_816 assert ds.n_atoms == 135_219_803 assert len(ds.split("train")) == 3_730_572 assert len(ds.split("val")) == 329_154 assert len(ds.split("test")) == 329_163 m = ds.molecule(0) assert m["molecule_id"] == 24553 # spot round-trip: perturbed = stationary + perturbation, bond-length-scale accuracy assert m["perturbed_coords"].shape == m["stationary_coords"].shape