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