"""Tests for the applications (natural-products) loader. These build a tiny synthetic applications.hdf5 + experimental spreadsheet in a temp dir that mimic the real schema, so the loader and its decoding can be checked in CI without the multi- hundred-megabyte release files. """ import numpy as np import pandas as pd import h5py import pytest from applications_reader import Applications, _decode, SHELL_SIZES, SOLVENTS SCALE = 1e4 MISS = np.int32(-2147483648) def _enc(a): a = np.asarray(a, dtype=np.float64) nan = np.isnan(a) out = np.where(nan, 0.0, a) out = np.round(out * SCALE).astype(np.int32) out[nan] = MISS return out @pytest.fixture def fake_release(tmp_path): """A 2-solute synthetic release: vomicine (with a missing shell) + isomer_1E.""" h5 = tmp_path / "applications.hdf5" xlsx = tmp_path / "applications_experimental.xlsx" rng = np.random.default_rng(0) def mk_solute(g, n_atoms, atomic_numbers, shells): g.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32)) mz = g.create_group("magnet_zero") # MagNET predicts H/C only -> heteroatoms read 0 (mirror the real data) hetero = ~np.isin(np.asarray(atomic_numbers), [1, 6]) stat = rng.normal(100, 30, n_atoms); stat[hetero] = 0.0 pcm = rng.normal(0, 2, n_atoms); pcm[hetero] = 0.0 mz.create_dataset("stationary", data=_enc(stat)) mz.create_dataset("pcm_correction", data=_enc(pcm)) cf = mz.create_group("conformers").create_group("conformer_0") cf.create_dataset("energy", data=np.float64(-1234.5)) # float, not encoded cf.create_dataset("geometry", data=_enc(rng.normal(0, 5, (n_atoms, 3)))) cf.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32)) cf.create_dataset("magnet_zero", data=_enc(rng.normal(100, 30, n_atoms))) cf.create_dataset("pcm", data=_enc(rng.normal(0, 2, n_atoms))) q = g.create_group("qcd") q.create_dataset("stationary", data=_enc(rng.normal(0, 50, (n_atoms, 4)))) q.create_dataset("trajectories", data=_enc(rng.normal(0, 50, (3, 4, n_atoms, 4)))) mx = g.create_group("magnet_x") for sv in SOLVENTS: svg = mx.create_group(sv) for sh in shells: svg.create_dataset(f"shell_{sh}", data=_enc(rng.normal(100, 30, (5, n_atoms, 2)))) svg.create_dataset("stationary", data=_enc(rng.normal(100, 30, n_atoms))) dr = g.create_group("dft_reference") # dft_reference uses gas + CPCM-continuum 'water' (not the explicit 'TIP4P'), like the real data for sv in ["gas", "water", "benzene", "chloroform", "methanol"]: dr.create_group(sv).create_dataset("shieldings", data=_enc(rng.normal(100, 40, n_atoms))) gg = dr.create_group("geometries").create_group("conformer_001") gg.create_dataset("coordinates", data=_enc(rng.normal(0, 5, (n_atoms, 3)))) gg.create_dataset("atomic_numbers", data=np.asarray(atomic_numbers, dtype=np.int32)) gg.create_dataset("energy", data=np.float64(-100.0)) with h5py.File(h5, "w") as f: sg = f.create_group("solvents") for sv in SOLVENTS: sg.create_group(sv).create_dataset("atomic_numbers", data=np.array([6, 1, 1], dtype=np.int32)) # vomicine: drop one shell (mimic the isomer_1Z/chloroform/250 corruption) mk_solute(f.create_group("vomicine"), 4, [8, 7, 6, 1], [s for s in SHELL_SIZES if s != 250]) mk_solute(f.create_group("isomer_1E"), 3, [6, 1, 1], SHELL_SIZES) # composite_model group: fit tables stored verbatim as CSV text cm = f.create_group("composite_model") cm.create_group("pcm_conversion_factors").create_dataset( "H", data="solvent,pcm_conversion_factor\nchloroform,1.0\nbenzene,0.6\n") cm.create_group("ols_coefficients").create_group("stationary_plus_pcm").create_dataset( "H", data="parameter,chloroform,benzene\nintercept,31.2,31.3\nstationary,-0.98,-0.99\n") cm.create_group("bootstrap_coefficients").create_group("stationary_plus_pcm").create_dataset( "H", data="solvent,seed,Intercept,stationary,pcm\nchloroform,0,31.2,-0.98,-1.0\nbenzene,0,31.3,-0.99,-0.6\n") cm.create_group("rmse_distributions").create_dataset( "H", data="solvent,nucleus,formula,seed,Bootstrap_RMSE\nchloroform,H,pcm2,0,0.07\nbenzene,H,pcm2,0,0.21\n") # experiment xlsx: one sheet per solute with pd.ExcelWriter(xlsx) as w: for sol in ["vomicine", "isomer_1E"]: pd.DataFrame({ "site": ["a_C", "b_H"], "nucleus": ["C", "H"], "atom_numbers": ["3", "4"], "chloroform": [1.0, 2.0], "benzene": [1.1, 2.1], "methanol": [1.2, 2.2], "water": [1.3, 2.3], }).to_excel(w, sheet_name=sol, index=False) return Applications(h5, xlsx) def test_decode_roundtrip_and_missing(): vals = np.array([1.2345, -50.0, 0.0], dtype=np.float64) enc = _enc(vals) out = _decode(enc) assert np.allclose(out, vals, atol=5e-5) enc_with_miss = np.append(enc, MISS).astype(np.int32) assert np.isnan(_decode(enc_with_miss)[-1]) def test_decode_int64_is_still_scaled(): # a rebuild that stored the columns as int64 (any integer width) must decode the same way as # int32; the old exact `dtype == np.int32` check fell through and returned raw scaled integers. vals = np.array([1.2345, -50.0, 0.0], dtype=np.float64) enc = _enc(vals).astype(np.int64) out = _decode(enc) assert np.allclose(out, vals, atol=5e-5) enc_with_miss = np.append(enc, np.int64(-2147483648)) assert np.isnan(_decode(enc_with_miss)[-1]) def test_solutes_and_solvents_group_excluded(fake_release): assert set(fake_release.solutes()) == {"vomicine", "isomer_1E"} def test_magnet_zero_heteroatoms_read_zero(fake_release): # vomicine atomic_numbers [8, 7, 6, 1]: the O and N positions must decode to 0 (MagNET does H/C only) mz = fake_release.magnet_zero("vomicine") hetero = ~np.isin(fake_release.atomic_numbers("vomicine"), [1, 6]) assert np.allclose(mz["stationary"][hetero], 0.0) assert np.allclose(mz["pcm_correction"][hetero], 0.0) assert np.any(mz["stationary"][~hetero] != 0.0) # H/C carry real values def test_qcd_shielding_only_drops_xyz(fake_release): full = fake_release.qcd("isomer_1E") just = fake_release.qcd("isomer_1E", shielding_only=True) assert full["stationary"].shape == (3, 4) assert just["stationary"].shape == (3,) assert just["trajectories"].shape == (3, 4, 3) assert np.allclose(just["stationary"], full["stationary"][..., 3], atol=5e-5) def test_magnet_x_and_shell_guard(fake_release): mx = fake_release.magnet_x("isomer_1E", "benzene", 650) assert mx.shape == (5, 3, 2) # vomicine is missing shell 250 -> available_shells must exclude it, not crash assert 250 not in fake_release.available_shells("vomicine", "chloroform") assert fake_release.available_shells("isomer_1E", "chloroform") == SHELL_SIZES def test_dft_reference_and_conformers(fake_release): assert fake_release.dft_reference("vomicine", "gas").shape == (4,) conf = fake_release.magnet_zero_conformers("vomicine") assert conf["conformer_0"]["energy"] == pytest.approx(-1234.5) assert conf["conformer_0"]["geometry"].shape == (4, 3) def test_experiment_water_rename(fake_release): df = fake_release.experiment() assert "TIP4P" in df.columns and "water" not in df.columns assert set(df["solute"].unique()) == {"vomicine", "isomer_1E"} def test_site_atom_table(fake_release): tbl = fake_release.site_atom_table() assert list(tbl.columns) == ["solute", "site", "nucleus", "atom_numbers"] assert set(tbl["solute"].unique()) == {"vomicine", "isomer_1E"} assert set(str(v) for v in tbl["atom_numbers"]) == {"3", "4"} def test_composite_model_accessors(fake_release): conv = fake_release.pcm_conversion_factors("H") assert list(conv.columns) == ["solvent", "pcm_conversion_factor"] assert conv.set_index("solvent").loc["chloroform", "pcm_conversion_factor"] == 1.0 ols = fake_release.ols_coefficients("stationary_plus_pcm", "H") assert "parameter" in ols.columns and "chloroform" in ols.columns assert fake_release.composite_formulas("ols_coefficients") == ["stationary_plus_pcm"] boot = fake_release.bootstrap_coefficients("stationary_plus_pcm", "H") assert {"solvent", "seed", "Intercept", "stationary", "pcm"}.issubset(boot.columns) rd = fake_release.rmse_distribution("H") assert {"solvent", "formula", "Bootstrap_RMSE"}.issubset(rd.columns) def test_dft_reference_uses_water_not_tip4p(fake_release): # the DFT reference is CPCM-continuum 'water'; the explicit 'TIP4P' key does not exist there assert fake_release.dft_reference("vomicine", "water").shape == (4,) with pytest.raises(KeyError): fake_release.dft_reference("vomicine", "TIP4P") def test_dft_reference_geometries(fake_release): geo = fake_release.dft_reference_geometries("vomicine") assert "conformer_001" in geo c = geo["conformer_001"] assert c["coordinates"].shape == (4, 3) assert c["energy"] == pytest.approx(-100.0) assert list(c["atomic_numbers"]) == [8, 7, 6, 1] def test_md_geometry_companion(tmp_path): md = tmp_path / "applications_md_geometries.hdf5" with h5py.File(md, "w") as f: g = f.create_group("vomicine") g.create_dataset("openMM_stationary_geometry", data=_enc(np.zeros((4, 3)))) sv = g.create_group("chloroform") sv.create_dataset("geometries_50", data=_enc(np.ones((5, 7, 3)))) sv.create_dataset("atomic_numbers_50", data=np.arange(7, dtype=np.int32)) app = Applications(tmp_path / "applications.hdf5", md_geometries_path=md) assert app.available_md_shells("vomicine", "chloroform") == [50] assert app.md_geometry("vomicine", "chloroform", 50).shape == (5, 7, 3) assert app.md_geometry_atomic_numbers("vomicine", "chloroform", 50).shape == (7,) assert app.md_stationary_geometry("vomicine").shape == (4, 3) def test_md_geometry_absent_raises_clearly(fake_release): # the big companion is optional; accessing it without it gives a clear, actionable error with pytest.raises(ValueError, match="companion"): fake_release.available_md_shells("vomicine", "chloroform") def test_uuid_prefix_h_sites(fake_release): # each solute's single H site (fixture site "b_H") becomes "01 b_H"; C sites untouched tbl = fake_release.site_atom_table(uuid_prefix_h_sites=True) h_sites = tbl[tbl["nucleus"] == "H"]["site"].tolist() assert all(s.startswith("01 ") for s in h_sites) c_sites = tbl[tbl["nucleus"] == "C"]["site"].tolist() assert all(not s[:2].isdigit() for s in c_sites)