""" Tests for the dft8k reader (dft8k_reader.py). Most tests build a tiny synthetic dft8k.hdf5 in the exact on-disk format (so the fixture doubles as a spec of what the reader expects) and need no large data. One opt-in test runs against the real dft8k.hdf5 if it is present (skipped otherwise, e.g. in CI). Run: pytest test_dft8k.py -q Requires: pytest, numpy, h5py. """ import os import numpy as np import h5py import pytest from dft8k_reader import DFT8k, _decode, _MARKER SCALE = 1e4 def _enc(arr): """Encode floats the way the build script does: int32 = round(value * 1e4), NaN -> marker.""" a = np.asarray(arr, dtype=np.float64) nan = np.isnan(a) out = np.rint(np.where(nan, 0.0, a) * SCALE).astype(np.int32) out[nan] = _MARKER return out # --------------------------------------------------------------------------- # fixed-point decoding # --------------------------------------------------------------------------- def test_decode_round_trip_and_marker(): values = np.array([0.0, 207.07, -355.5, np.nan, 875.0], dtype=np.float64) decoded = _decode(_enc(values)) assert np.isnan(decoded[3]) np.testing.assert_allclose(decoded[~np.isnan(decoded)], values[~np.isnan(values)], atol=5e-5) def test_decode_passes_through_float(): # a plain-decimal array must read back unchanged (decode is a no-op on floats) values = np.array([1.0, 2.5, -7.25], dtype=np.float64) np.testing.assert_array_equal(_decode(values), values) def test_decode_handles_int64(): # any integer dtype decodes the same way as int32 vals = np.array([10000, -3550000, _MARKER], dtype=np.int64) out = _decode(vals) np.testing.assert_allclose(out[:2], [1.0, -355.0]) assert np.isnan(out[2]) # --------------------------------------------------------------------------- # synthetic dft8k.hdf5 in the real two-group layout # --------------------------------------------------------------------------- # two molecules: ids 100 and 200, with 2 and 3 atoms (atoms are H, C, ...) A_IDS = [100, 200] A_COUNTS = [2, 3] A_Z = [1, 6, 6, 1, 7] # mol100: H C ; mol200: C H N # per-atom DFT shieldings (concatenated): pbe0, wp04, wb97xd A_PBE0 = [30.0, 150.0, 151.0, 31.0, 200.0] A_WP04 = [30.5, 150.5, 151.5, 31.5, 200.5] # the 1H reference A_WB97 = [30.7, 150.7, 151.7, 31.7, 200.7] # the 13C reference # MagNET predicts H with the WP04 model and C with the wB97X-D model; other atoms blank A_NN_WP04 = [30.6, np.nan, np.nan, 31.6, np.nan] # only at H atoms (index 0, 3) A_NN_WB97 = [np.nan, 150.6, 151.6, np.nan, np.nan] # only at C atoms (index 1, 2) def _make_synthetic_dft8k(path): with h5py.File(path, "w") as f: g = f.create_group("aimnet2_wp04_wb97xd_pcsseg2") g.attrs["geometry"] = "AIMNet2" g.attrs["level_of_theory"] = "PBE0/pcSseg-1; WP04/pcSseg-2; wB97X-D/pcSseg-2" g.attrs["n_molecules"] = len(A_IDS) g.attrs["n_atoms"] = sum(A_COUNTS) g.attrs["scale"] = 1e-4 g.create_dataset("molecule_ids", data=np.array(A_IDS, np.uint32)) g.create_dataset("n_atoms", data=np.array(A_COUNTS, np.uint8)) # one SMILES per molecule; mol100 has a real one, mol200 carries the literal "none" g.create_dataset("smiles", data=np.array(["CC", "none"], dtype=object), dtype=h5py.string_dtype(encoding="utf-8")) g.create_dataset("atomic_numbers", data=np.array(A_Z, np.int8)) coords = np.arange(sum(A_COUNTS) * 3, dtype=np.float64).reshape(-1, 3) * 0.1 g.create_dataset("coordinates", data=_enc(coords)) g.create_dataset("shielding_pbe0_pcSseg1", data=_enc(A_PBE0)) g.create_dataset("shielding_wp04_pcSseg2", data=_enc(A_WP04)) g.create_dataset("shielding_wb97xd_pcSseg2", data=_enc(A_WB97)) g.create_dataset("nn_shielding_wp04_pcSseg2", data=_enc(A_NN_WP04)) g.create_dataset("nn_shielding_wb97xd_pcSseg2", data=_enc(A_NN_WB97)) # second group, SAME ids but reversed row order, three geometries per molecule b = f.create_group("b3lyp_pbe0_pcsseg1") b.attrs["geometry"] = "B3LYP/pcSseg-1" b.attrs["level_of_theory"] = "PBE0/pcSseg-1 on B3LYP/pcSseg-1 geometries" b.attrs["n_molecules"] = 2 b.attrs["n_atoms"] = 5 b.attrs["n_geometries"] = 3 b.attrs["scale"] = 1e-4 b.create_dataset("molecule_ids", data=np.array([200, 100], np.uint32)) # reversed b.create_dataset("n_atoms", data=np.array([3, 2], np.uint8)) b.create_dataset("atomic_numbers", data=np.array([6, 1, 7, 1, 6], np.int8)) bcoords = np.arange(3 * 5 * 3, dtype=np.float64).reshape(3, 5, 3) * 0.01 b.create_dataset("coordinates", data=_enc(bcoords)) bshield = np.arange(3 * 5, dtype=np.float64).reshape(3, 5) + 100.0 b.create_dataset("shielding", data=_enc(bshield)) b.create_dataset("weights", data=np.array([1, 1, 1, 0.5, 0.5], np.float32)) # symmetry groups are 0-based atom indices, local to each molecule. # mol200 (first row, 3 atoms): two groups [[1, 2]] and [[0]]; # mol100 (second row, 2 atoms): one group [[0, 1]]. b.create_dataset("symmetric_atoms_n_groups", data=np.array([2, 1], np.int32)) b.create_dataset("symmetric_atoms_group_lengths", data=np.array([2, 1, 2], np.int32)) b.create_dataset("symmetric_atoms_members", data=np.array([1, 2, 0, 0, 1], np.int32)) def test_aimnet2_group(tmp_path): p = tmp_path / "dft8k.hdf5" _make_synthetic_dft8k(str(p)) with DFT8k(str(p)) as ds: assert len(ds.aimnet2) == 2 m = ds.aimnet2.molecule(0) assert m["id"] == 100 assert m["smiles"] == "CC" assert ds.aimnet2.smiles(1) == "none" # mol200 has no SMILES np.testing.assert_array_equal(m["atomic_numbers"], [1, 6]) np.testing.assert_allclose(m["wp04_pcSseg2"], [30.5, 150.5], atol=5e-5) # H atom (index 0) has a WP04 prediction; C atom (index 1) does not assert not np.isnan(m["nn_wp04_pcSseg2"][0]) assert np.isnan(m["nn_wp04_pcSseg2"][1]) # C atom has a wB97X-D prediction; H atom does not assert np.isnan(m["nn_wb97xd_pcSseg2"][0]) assert not np.isnan(m["nn_wb97xd_pcSseg2"][1]) np.testing.assert_allclose(m["nn_wb97xd_pcSseg2"][1], 150.6, atol=5e-5) def test_all_shieldings_flat(tmp_path): p = tmp_path / "dft8k.hdf5" _make_synthetic_dft8k(str(p)) with DFT8k(str(p)) as ds: flat = ds.aimnet2.all_shieldings() # one entry per atom across both molecules, in stored order np.testing.assert_array_equal(flat["atomic_numbers"], [1, 6, 6, 1, 7]) np.testing.assert_allclose(flat["wp04_pcSseg2"], A_WP04, atol=5e-5) # MagNET prediction is present only at the atoms its model covers assert np.isnan(flat["nn_wp04_pcSseg2"]).tolist() == [False, True, True, False, True] # the 1H residual (DFT minus MagNET) at the two H atoms is -0.1 Z = flat["atomic_numbers"] h = (Z == 1) & np.isfinite(flat["nn_wp04_pcSseg2"]) np.testing.assert_allclose(flat["wp04_pcSseg2"][h] - flat["nn_wp04_pcSseg2"][h], [-0.1, -0.1], atol=5e-5) def test_b3lyp_group_and_symmetry(tmp_path): p = tmp_path / "dft8k.hdf5" _make_synthetic_dft8k(str(p)) with DFT8k(str(p)) as ds: b = ds.b3lyp assert b.n_geometries == 3 m = b.molecule(0) # first row is id 200 assert m["id"] == 200 assert m["coordinates"].shape == (3, 3, 3) # (geometries, atoms, xyz) assert m["shielding"].shape == (3, 3) assert m["symmetric_atoms"] == [[1, 2], [0]] m2 = b.molecule(1) # id 100 assert m2["symmetric_atoms"] == [[0, 1]] np.testing.assert_allclose(m2["weights"], [0.5, 0.5]) # the symmetry indices are valid 0-based positions into this molecule's atoms for group in m["symmetric_atoms"] + m2["symmetric_atoms"]: for atom_index in group: assert 0 <= atom_index < len(m["atomic_numbers"]) or \ 0 <= atom_index < len(m2["atomic_numbers"]) def test_match_across_groups_by_id(tmp_path): p = tmp_path / "dft8k.hdf5" _make_synthetic_dft8k(str(p)) with DFT8k(str(p)) as ds: # rows are ordered differently in the two groups; matching by id must still work a = ds.aimnet2.molecule_by_id(200) b = ds.b3lyp.molecule_by_id(200) assert a["id"] == b["id"] == 200 np.testing.assert_array_equal(a["atomic_numbers"], b["atomic_numbers"]) # --------------------------------------------------------------------------- # opt-in test against the real file (skipped when absent, e.g. in CI) # --------------------------------------------------------------------------- _HERE = os.path.dirname(os.path.abspath(__file__)) _REAL = os.path.join(_HERE, "dft8k.hdf5") @pytest.mark.skipif(not os.path.exists(_REAL), reason="real dft8k.hdf5 not present") def test_real_dft8k(): with DFT8k(_REAL) as ds: assert len(ds.aimnet2) == 7111 assert len(ds.b3lyp) == 7111 assert ds.b3lyp.n_geometries == 3 # the two groups carry the same id set assert set(int(x) for x in ds.aimnet2.molecule_ids) == \ set(int(x) for x in ds.b3lyp.molecule_ids) m = ds.aimnet2.molecule(0) # SMILES: 6780 of the 7111 molecules have one; the rest carry the literal "none" n_real_smiles = sum(1 for i in range(len(ds.aimnet2)) if ds.aimnet2.smiles(i) != "none") assert n_real_smiles == 6780 assert isinstance(m["smiles"], str) and m["smiles"] != "none" z = m["atomic_numbers"] h = z == 1 c = z == 6 other = ~(h | c) # the WP04 model predicts every hydrogen and nothing else; the wB97X-D model # predicts every carbon and nothing else; non-target atoms read back blank assert not np.isnan(m["nn_wp04_pcSseg2"][h]).any() assert np.isnan(m["nn_wp04_pcSseg2"][~h]).all() assert not np.isnan(m["nn_wb97xd_pcSseg2"][c]).any() assert np.isnan(m["nn_wb97xd_pcSseg2"][~c]).all() assert other.sum() == 0 or np.isnan(m["nn_wp04_pcSseg2"][other]).all() # the same molecule resolves in the other group by id b = ds.b3lyp.molecule_by_id(m["id"]) np.testing.assert_array_equal(m["atomic_numbers"], b["atomic_numbers"]) # every symmetry index is a valid 0-based position into this molecule's atoms for group in b["symmetric_atoms"]: for atom_index in group: assert 0 <= atom_index < len(b["atomic_numbers"])