data/sigma-fresh/test_sigma_fresh.py

"""
Tests for the sigma-fresh reader (decode_sigma_fresh.SigmaFresh).

A tiny synthetic dataset is built in the *exact* on-disk format (so the fixture also
serves as an executable spec of the layout) and the decoder is exercised against it.
One opt-in test runs against the real `sigma-fresh.hdf5` if present next to this file.

Run:  pytest test_sigma_fresh.py -q
Requires: pytest, numpy, h5py.
"""
import os
import numpy as np
import h5py
import pytest

from decode_sigma_fresh import SigmaFresh, SENTINEL

SCALE = 1e-4


def _qc(x):
    return np.round(np.asarray(x, float) / SCALE).astype(np.int32)


def _qs(x):
    x = np.asarray(x, float)
    out = np.where(np.isnan(x), SENTINEL, np.round(x / SCALE)).astype(np.int32)
    return out


def _solute(g, name, *, idx, n_solute, znums, coords, iso, solv, status, nsolv, prad, frad, charge=0, typ="train"):
    s = g.create_group(f"solute_{idx:05d}")
    s.attrs["name"] = name
    s.attrs["type"] = typ
    s.attrs["n_solute_atoms"] = n_solute
    s.attrs["solute_charge"] = charge
    s.create_dataset("atomic_numbers", data=np.asarray(znums, np.int8))
    s.create_dataset("coordinates", data=_qc(coords))
    s.create_dataset("shielding_isolated", data=_qs(iso))
    s.create_dataset("shielding_solvated", data=_qs(solv))
    s.create_dataset("status", data=np.asarray(status, np.uint8))
    s.create_dataset("n_solvents_partial", data=np.asarray(nsolv, np.int32))
    s.create_dataset("partial_radius", data=_qc(prad))
    s.create_dataset("full_radius", data=_qc(frad))


@pytest.fixture(scope="module")
def mini_path(tmp_path_factory):
    p = str(tmp_path_factory.mktemp("sf") / "mini.hdf5")
    with h5py.File(p, "w") as f:
        f.attrs["dataset"] = "sigma-fresh"
        f.attrs["solvents"] = "chloroform,benzene"
        f.attrs["sentinel"] = int(SENTINEL)
        f.attrs["scale"] = SCALE
        f.attrs["level_of_theory"] = "PBE0/pcSseg-1 (solute), PBE0/MIDI! (solvent)"
        # chloroform: solute_00001 has 3 frames; solute (2 atoms) + 1 chloroform (5 atoms) = 7
        gc = f.create_group("chloroform"); gc.attrs["n_solvent_atoms"] = 5
        Z = [6, 1, 6, 1, 17, 17, 17]                 # 2 solute + 5 solvent
        n_solute = 2
        # 4 frames: complete, geometry-only (all NaN), complete, partial (isolated finite, solvated NaN)
        coords = np.zeros((4, 7, 3), float)
        for fr in range(4):
            coords[fr] = np.arange(7 * 3).reshape(7, 3) * 0.1 + fr   # raw, frame-shifted (uncentered)
        iso = np.array([[150.0, 30.0], [np.nan, np.nan], [151.0, 31.0], [152.0, 32.0]])
        solv = np.array([[149.5, 29.8], [np.nan, np.nan], [150.6, 30.7], [np.nan, np.nan]])
        _solute(gc, "GLY000001", idx=1, n_solute=n_solute, znums=Z, coords=coords,
                iso=iso, solv=solv, status=[3, 1, 3, 2], nsolv=[1, 1, 1, 1],
                prad=[3.1, 3.1, 3.2, 3.2], frad=[9.9, 9.9, 9.9, 9.9], typ="train")
        # benzene: a test solute, 1 frame complete
        gb = f.create_group("benzene"); gb.attrs["n_solvent_atoms"] = 12
        Zb = [8, 1]                                   # 2-atom solute, 0 solvent (degenerate but valid)
        _solute(gb, "MOL000009", idx=9, n_solute=2, znums=Zb,
                coords=np.array([[[0, 0, 0], [1.0, 0, 0]]], float),
                iso=np.array([[280.0, 32.0]]), solv=np.array([[279.0, 31.5]]),
                status=[3], nsolv=[0], prad=[0.0], frad=[0.0], typ="test")
    return p


@pytest.fixture
def ds(mini_path):
    with SigmaFresh(mini_path) as d:
        yield d


def test_catalogue(ds):
    assert ds.solvents == ["chloroform", "benzene"]
    assert ds.n_solutes("chloroform") == 1 and ds.n_solutes("benzene") == 1
    assert ds.solutes("chloroform") == ["solute_00001"]


def test_pose_values_roundtrip(ds):
    p = ds.pose("chloroform", 1, 1)        # 1-based frame
    assert p["name"] == "GLY000001" and p["type"] == "train" and p["status"] == "complete"
    np.testing.assert_array_equal(p["atomic_numbers"], [6, 1, 6, 1, 17, 17, 17])
    np.testing.assert_allclose(p["shielding_isolated"], [150.0, 30.0], atol=5e-5)
    np.testing.assert_allclose(p["shielding_solvated"], [149.5, 29.8], atol=5e-5)
    assert p["n_solvents_partial"] == 1
    np.testing.assert_allclose(p["partial_radius"], 3.1, atol=5e-5)


def test_shieldings_solute_only(ds):
    p = ds.pose("chloroform", 1, 1)
    # shieldings have one entry per solute atom, not per cluster atom
    assert p["shielding_isolated"].shape == (2,)
    assert p["coordinates"].shape == (7, 3)


def test_geometry_only_frame_is_nan(ds):
    p = ds.pose("chloroform", 1, 2)        # frame 2 = geometry-only
    assert p["status"] == "geometries_only"
    assert np.isnan(p["shielding_isolated"]).all()
    assert np.isnan(p["shielding_solvated"]).all()
    # geometry is still valid on a geometry-only frame
    assert np.isfinite(p["coordinates"]).all()


def test_raw_coordinates_not_centered(ds):
    # frames are stored raw: distinct frames have distinct centroids (no per-frame centering)
    c1 = ds.pose("chloroform", 1, 1)["coordinates"].mean(0)
    c3 = ds.pose("chloroform", 1, 3)["coordinates"].mean(0)
    assert not np.allclose(c1, c3)


def test_solute_bundle_and_mask(ds):
    s = ds.solute("chloroform", 1)
    assert s["coordinates"].shape == (4, 7, 3)
    assert s["shielding_isolated"].shape == (4, 2)
    assert s["solute_mask"].tolist() == [True, True, False, False, False, False, False]
    assert list(s["status"]) == ["complete", "geometries_only", "complete", "partial_shieldings"]
    # correction is finite on complete frames, NaN on the geometry-only and partial frames
    corr = s["shielding_solvated"] - s["shielding_isolated"]
    assert np.isfinite(corr[0]).all() and np.isnan(corr[1]).all() and np.isnan(corr[3]).all()


def test_partial_status_frame(ds):
    # status==2: isolated shielding present, solvated not computed (sentinel -> NaN)
    p = ds.pose("chloroform", 1, 4)
    assert p["status"] == "partial_shieldings"
    assert np.isfinite(p["shielding_isolated"]).all()
    assert np.isnan(p["shielding_solvated"]).all()


def test_test_solute_type(ds):
    p = ds.pose("benzene", 9, 1)
    assert p["type"] == "test" and p["name"] == "MOL000009"


def test_bad_lookups(ds):
    with pytest.raises(IndexError):
        ds.pose("chloroform", 1, 0)        # frames are 1-based
    with pytest.raises(IndexError):
        ds.pose("chloroform", 1, 99)
    with pytest.raises(IndexError):
        ds.pose("chloroform", 999, 1)
    with pytest.raises(KeyError):
        ds.pose("acetone", 1, 1)


def test_context_manager_closes(mini_path):
    d = SigmaFresh(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-fresh.hdf5")


@pytest.mark.skipif(not os.path.exists(REAL), reason="real sigma-fresh.hdf5 not present")
def test_real_file():
    with SigmaFresh(REAL) as ds:
        assert set(ds.solvents) == {"chloroform", "benzene", "methanol", "TIP4P"}
        for s in ds.solvents:
            assert ds.n_solutes(s) > 9000
        p = ds.pose("chloroform", 3, 1)
        assert p["coordinates"].shape[0] == p["atomic_numbers"].shape[0]
        assert p["shielding_isolated"].shape == (p["n_solute_atoms"],)
        # at least one complete frame exists for this solute, with finite shieldings
        s = ds.solute("chloroform", 3)
        complete = s["status_code"] == 3
        assert complete.any()
        assert np.isfinite(s["shielding_solvated"][complete]).any()