""" 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