analysis/code/test_delta22_analysis.py

"""Tests for analysis/code/delta22.py (the no-plotting analysis module).

The analysis functions operate on flat query DataFrames, so these tests build small synthetic
tables with a known linear relationship and check the fitting harness, the seeded splits, the
Pareto frontier, and the correlation matrix. One test proves the fast numpy `fit` gives the same
coefficients and RMSE as the original statsmodels fit.
"""
import os
import random
import sys

import numpy as np
import pandas as pd
import pytest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "..", "data", "delta22"))  # delta22_reader
sys.path.insert(0, HERE)
import paths

import delta22 as D  # noqa: E402


def make_query_df(n_solutes=12, solvents=("chloroform", "benzene"), noise=0.0, seed=0):
    """A flat query table whose experimental value is exactly 2*stationary + 0.5*pcm + 1 (plus
    optional noise), so a model with both terms fits near-perfectly and 'stationary' alone leaves
    the 0.5*pcm residual."""
    rng = np.random.default_rng(seed)
    rows = []
    for si in range(n_solutes):
        solute = f"m{si:02d}"
        for solvent in solvents:
            for k in range(3):
                stationary = rng.normal(100, 20)
                pcm = rng.normal(0, 2)
                experimental = 2.0 * stationary + 0.5 * pcm + 1.0 + rng.normal(0, noise)
                rows.append({
                    "solute": solute, "solvent": solvent, "nucleus": "H",
                    "site": f"{solute}_{solvent}_{k}", "experimental": experimental,
                    "stationary": stationary, "pcm": pcm,
                    "desmond": rng.normal(0, 2), "openMM": rng.normal(0, 2),
                    "qcd": rng.normal(0, 1), "desmond_vib": rng.normal(0, 1),
                    "openMM_vib": rng.normal(0, 1),
                })
    return pd.DataFrame(rows)


def test_fit_recovers_known_line():
    df = make_query_df(noise=0.0)
    train_rmse, test_rmse, params = D.fit(df, df, "stationary + pcm")
    assert params["Intercept"] == pytest.approx(1.0, abs=1e-6)
    assert params["stationary"] == pytest.approx(2.0, abs=1e-6)
    assert params["pcm"] == pytest.approx(0.5, abs=1e-6)
    assert train_rmse == pytest.approx(0.0, abs=1e-6)
    assert test_rmse == pytest.approx(0.0, abs=1e-6)


def test_fit_matches_statsmodels():
    df = make_query_df(noise=1.0, seed=3)
    train = df[df["solute"] <= "m07"]
    test = df[df["solute"] > "m07"]
    for formula in ["stationary", "stationary + pcm", "stationary + pcm + qcd"]:
        tr, te, params = D.fit(train, test, formula)
        s_tr, s_te, s_params = D._fit_statsmodels(train, test, formula)
        assert tr == pytest.approx(s_tr, abs=1e-8)
        assert te == pytest.approx(s_te, abs=1e-8)
        for name in s_params.index:
            assert params[name] == pytest.approx(s_params[name], abs=1e-6)


def test_fit_drops_missing_like_statsmodels():
    # fitting must drop rows with a missing response or predictor, exactly as statsmodels does,
    # and recover the same coefficients. (Scoring is checked on a clean test set; statsmodels would
    # otherwise NaN-poison an RMSE computed over rows that have missing predictors.)
    clean = make_query_df(noise=0.5, seed=5).reset_index(drop=True)
    train = clean.copy()
    train.loc[0, "experimental"] = np.nan      # missing response
    train.loc[1, "pcm"] = np.nan               # missing predictor
    _train_err, te, params = D.fit(train, clean, "stationary + pcm")
    s_tr, s_te, s_params = D._fit_statsmodels(train, clean, "stationary + pcm")
    for name in s_params.index:
        assert params[name] == pytest.approx(s_params[name], abs=1e-6)
    assert te == pytest.approx(s_te, abs=1e-8)


def test_generate_solute_splits_seeded_and_disjoint():
    df = make_query_df()
    solutes = [f"m{i:02d}" for i in range(12)]
    splits = D.generate_solute_splits(5, df, "chloroform", solutes, n_test=4)
    assert len(splits) == 5
    # reproduce the exact seed logic for split 0
    random.seed(0 + 100)
    expected = list(solutes)
    random.shuffle(expected)
    train_df, test_df = splits[0]
    assert set(test_df["solute"].unique()) == set(expected[:4])
    assert set(train_df["solute"].unique()) == set(expected[4:])
    # a solute is never in both halves, and only the requested solvent is present
    assert set(test_df["solute"]).isdisjoint(set(train_df["solute"]))
    assert set(test_df["solvent"].unique()) == {"chloroform"}
    # deterministic: same call gives the same partition
    again = D.generate_solute_splits(5, df, "chloroform", solutes, n_test=4)
    assert set(again[0][1]["solute"].unique()) == set(test_df["solute"].unique())


def test_run_fits_more_terms_lower_error():
    df = make_query_df(noise=0.3, seed=1)
    solutes = sorted(df["solute"].unique())
    results = D.run_fits(df, ["chloroform", "benzene"], D.STATIONARY_VS_PCM,
                         n_splits=5, solutes=solutes, n_test=4)
    assert set(results.columns) >= {"solvent", "formula", "seed", "train_RMSE", "test_RMSE"}
    med = D.median_test_rmse(results, ["formula"]).set_index("formula")["test_RMSE"]
    assert med["stationary + pcm"] < med["stationary"]   # the pcm term genuinely helps


def test_run_fits_group_cols():
    df = make_query_df(noise=0.3, seed=2)
    df = pd.concat([df.assign(sap_nmr_method="wp04"),
                    df.assign(sap_nmr_method="dsd_pbep86")], ignore_index=True)
    solutes = sorted(df["solute"].unique())
    results = D.run_fits(df, ["chloroform"], ["stationary + pcm"], n_splits=3,
                         solutes=solutes, group_cols=["sap_nmr_method"], n_test=4)
    assert "sap_nmr_method" in results.columns
    assert set(results["sap_nmr_method"].unique()) == {"wp04", "dsd_pbep86"}


def test_fit_coefficients_per_seed():
    df = make_query_df(noise=0.4, seed=7)
    solutes = sorted(df["solute"].unique())
    coeffs = D.fit_coefficients(df, "chloroform", "stationary + pcm", solutes,
                                n_splits=6, n_test=4)
    assert len(coeffs) == 6
    assert {"solvent", "seed", "Intercept", "stationary", "pcm"}.issubset(coeffs.columns)
    # the recovered slope should sit near the true value of 2
    assert coeffs["stationary"].mean() == pytest.approx(2.0, abs=0.1)


def test_full_fit_coefficients_shape_and_values():
    df = make_query_df(noise=0.0, solvents=("chloroform", "benzene"))
    coeffs = D.full_fit_coefficients(df, ["chloroform", "benzene"], "stationary + pcm")
    assert list(coeffs.columns) == ["chloroform", "benzene"]
    assert set(coeffs.index) == {"Intercept", "stationary", "pcm"}
    # the data is exactly 2*stationary + 0.5*pcm + 1, recovered per solvent
    for solvent in ["chloroform", "benzene"]:
        assert coeffs.loc["Intercept", solvent] == pytest.approx(1.0, abs=1e-6)
        assert coeffs.loc["stationary", solvent] == pytest.approx(2.0, abs=1e-6)
        assert coeffs.loc["pcm", solvent] == pytest.approx(0.5, abs=1e-6)


# ---- opt-in: checked against the STORED composite-model coefficients ----
# These coefficients live in applications.hdf5's composite_model/ols_coefficients group. Re-fitting
# the released delta-22 data with the fast numpy harness (nitromethane dropped) must reproduce them
# to within the data's integer-encoding precision (~5e-4 after error propagation through the small
# corrections).
REAL_H5 = paths.dataset_file("delta22", file=__file__)
REAL_XLSX = os.path.join(HERE, "..", "..", "data", "delta22", "delta22_experimental.xlsx")
APPLICATIONS_H5 = paths.dataset_file("applications", file=__file__)


@pytest.mark.skipif(
    not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX) and os.path.exists(APPLICATIONS_H5)),
    reason="real delta22.hdf5 / experimental xlsx / applications.hdf5 not present")
def test_full_fit_reproduces_stored_composite_coefficients():
    import io
    import h5py

    query = D.load_query_df_nn(REAL_H5, REAL_XLSX, exclude_solutes=["nitromethane"], verbose=False)
    worst = 0.0
    with h5py.File(APPLICATIONS_H5, "r") as f:
        ols = f["composite_model"]["ols_coefficients"]
        for nucleus in ["H", "C"]:
            sub = query[query["nucleus"] == nucleus]
            for formula in ["stationary", "stationary + openMM", "stationary + openMM + qcd"]:
                tag = "_AND_".join(term.strip() for term in formula.split("+"))
                raw = np.array(ols[tag][nucleus]).item()
                raw = raw.decode() if hasattr(raw, "decode") else raw
                stored = pd.read_csv(io.StringIO(raw)).set_index("parameter")
                mine = D.full_fit_coefficients(sub, list(stored.columns), formula)
                mine = mine.rename(index={"Intercept": "intercept"})
                for solvent in stored.columns:
                    for param in stored.index:
                        worst = max(worst, abs(mine.loc[param, solvent] - stored.loc[param, solvent]))
    assert worst < 5e-4, f"worst coefficient diff {worst:.2e} exceeds the int32 precision floor"


def make_dft_query_df(methods=("wp04", "dsd_pbep86"), solvents=("chloroform", "methanol", "benzene"),
                      n_solutes=12, seed=0):
    """A synthetic DFT query table with method/basis/geometry columns. Experimental is
    31 - (stationary + pcm) + tiny noise, i.e. a single shared slope on the gas shielding plus its
    implicit-solvent correction (the real structure), so the single-column 'stationary_plus_pcm'
    model fits well and 'stationary' alone leaves the pcm residual. desmond/openMM/qcd are
    independent noise that does not help."""
    rng = np.random.default_rng(seed)
    rows = []
    for method in methods:
        for si in range(n_solutes):
            solute = f"m{si:02d}"
            for solvent in solvents:
                for k in range(2):
                    stationary = rng.normal(100, 20)
                    pcm = rng.normal(0, 2)
                    rows.append({
                        "solute": solute, "solvent": solvent, "nucleus": "H",
                        "site": f"{solute}_{k}", "sap_nmr_method": method,
                        "sap_basis": "pcSseg2", "sap_geometry_type": "aimnet2",
                        "experimental": 31.0 - (stationary + pcm) + rng.normal(0, 0.05),
                        "stationary": stationary, "pcm": pcm,
                        "desmond": rng.normal(0, 1), "openMM": rng.normal(0, 1),
                        "qcd": rng.normal(0, 1), "desmond_vib": rng.normal(0, 1),
                        "openMM_vib": rng.normal(0, 1),
                    })
    return pd.DataFrame(rows)


def test_add_composite_columns():
    df = make_query_df()
    out = D.add_composite_columns(df)
    assert np.allclose(out["stationary_plus_pcm"], out["stationary"] + out["pcm"])
    assert np.allclose(out["stationary_plus_pcm_plus_qcd"],
                       out["stationary"] + out["pcm"] + out["qcd"])
    # an existing NN-style stationary_plus_pcm (with scaling factors) is preserved
    nn = df.copy()
    nn["stationary_plus_pcm"] = 999.0
    kept = D.add_composite_columns(nn)
    assert np.allclose(kept["stationary_plus_pcm"], 999.0)


def test_fig3b_shift_differences_exact():
    # one method, three solvents; experimental/pcm chosen so the differences are known
    rows = []
    vals = {"chloroform": (1.0, 0.1), "methanol": (1.5, 0.4), "benzene": (0.7, -0.2)}
    for solvent, (exp, pcm) in vals.items():
        rows.append({"solute": "m0", "site": "s0", "nucleus": "H", "solvent": solvent,
                     "sap_nmr_method": "wp04", "sap_basis": "pcSseg2",
                     "sap_geometry_type": "aimnet2", "experimental": exp, "pcm": pcm})
    df = pd.DataFrame(rows)
    out = D.fig3b_shift_differences(df, "wp04", "pcSseg2", "aimnet2", nucleus="H",
                                    x_solvent="methanol", y_solvent="benzene", reference="chloroform")
    assert len(out) == 1
    assert out["exp_diff_x"].iloc[0] == pytest.approx(1.5 - 1.0)
    assert out["exp_diff_y"].iloc[0] == pytest.approx(0.7 - 1.0)
    assert out["pcm_diff_x"].iloc[0] == pytest.approx(0.4 - 0.1)
    assert out["pcm_diff_y"].iloc[0] == pytest.approx(-0.2 - 0.1)


def test_fig3a_and_fig3d_pcm_helps():
    df = D.add_composite_columns(make_dft_query_df(seed=4))
    solutes = sorted(df["solute"].unique())
    solvents = ["chloroform", "methanol", "benzene"]
    a = D.fig3a_pcm_benefit(df, solvents, n_splits=4, solutes=solutes, n_test=4)
    assert {"sap_nmr_method", "solvent", "formula", "test_RMSE"}.issubset(a.columns)
    med = a.groupby("formula")["test_RMSE"].median()
    assert med["stationary_plus_pcm"] < med["stationary"]   # PCM lowers error
    d = D.fig3d_formula_regressions(df, "wp04", "pcSseg2", "aimnet2",
                                    ["stationary + pcm", "stationary + desmond"], solvents,
                                    n_splits=4, solutes=solutes, n_test=4)
    assert set(d["formula"].unique()) == {"stationary + pcm", "stationary + desmond"}


def test_fig3a_pcm_benefit_by_solvent_keeps_solvents_separate():
    # chloroform: pcm carries a real (though not perfectly scaled) part of the signal, so adding it
    # to the fit lowers test error. benzene: pcm is unrelated noise on top of an already-explained
    # response, so adding it as a predictor only overfits and raises test error. A solvent-pooled
    # average would blend these two directions away; fig3a_pcm_benefit_by_solvent must keep one row
    # per solvent so both signs are visible.
    rng = np.random.default_rng(11)
    rows = []
    n_solutes = 16
    for si in range(n_solutes):
        solute = f"m{si:02d}"
        for solvent in ["chloroform", "benzene"]:
            for k in range(3):
                stationary = rng.normal(100, 20)
                if solvent == "chloroform":
                    pcm = rng.normal(0, 2)
                    experimental = 31.0 - stationary - 0.8 * pcm + rng.normal(0, 0.05)
                else:
                    pcm = rng.normal(0, 8)   # large noise unrelated to experimental
                    experimental = 31.0 - stationary + rng.normal(0, 0.05)
                rows.append({
                    "solute": solute, "solvent": solvent, "nucleus": "H",
                    "site": f"{solute}_{solvent}_{k}", "sap_nmr_method": "wp04",
                    "sap_basis": "pcSseg2", "sap_geometry_type": "aimnet2",
                    "experimental": experimental, "stationary": stationary, "pcm": pcm,
                    "desmond": 0.0, "openMM": 0.0, "qcd": 0.0, "desmond_vib": 0.0, "openMM_vib": 0.0,
                })
    df = D.add_composite_columns(pd.DataFrame(rows))
    solutes = sorted(df["solute"].unique())
    out = D.fig3a_pcm_benefit_by_solvent(df, ["chloroform", "benzene"], n_splits=10,
                                         solutes=solutes, n_test=6)
    assert {"sap_nmr_method", "solvent", "seed", "percent_benefit"}.issubset(out.columns)
    chloro_benefit = out[out["solvent"] == "chloroform"]["percent_benefit"].median()
    benz_benefit = out[out["solvent"] == "benzene"]["percent_benefit"].median()
    assert chloro_benefit > 0    # PCM helps in chloroform
    assert benz_benefit < 0      # PCM hurts in benzene


def test_fig3a_does_not_pool_nuclei():
    # H sites (~30 ppm) and C sites (~150 ppm) with different relationships. If fig3a pooled both
    # nuclei into one fit the test RMSE would blow up; filtering to one nucleus keeps it small.
    rng = np.random.default_rng(0)
    rows = []
    for si in range(12):
        for solvent in ["chloroform", "methanol", "benzene"]:
            for nucleus, base, slope, offset in [("H", 30, 1.0, 31.0), ("C", 150, 1.05, 200.0)]:
                stationary = rng.normal(base, base * 0.1)
                pcm = rng.normal(0, 1)
                rows.append({
                    "solute": f"m{si:02d}", "site": f"m{si:02d}_{nucleus}", "nucleus": nucleus,
                    "solvent": solvent, "sap_nmr_method": "wp04", "sap_basis": "pcSseg2",
                    "sap_geometry_type": "aimnet2", "stationary": stationary, "pcm": pcm,
                    "experimental": offset - slope * (stationary + pcm) + rng.normal(0, 0.05),
                    "desmond": 0.0, "openMM": 0.0, "qcd": 0.0, "desmond_vib": 0.0, "openMM_vib": 0.0,
                })
    df = D.add_composite_columns(pd.DataFrame(rows))
    solutes = sorted(df["solute"].unique())
    res = D.fig3a_pcm_benefit(df, ["chloroform", "methanol", "benzene"], n_splits=4,
                              solutes=solutes, nucleus="H", n_test=4)
    assert res["test_RMSE"].median() < 1.0   # would be tens of ppm if C were pooled in


def test_add_solvent_mean():
    df = pd.DataFrame({
        "solute": ["m0", "m0"], "site": ["s0", "s0"], "nucleus": ["H", "H"],
        "solvent": ["chloroform", "benzene"], "experimental": [2.0, 4.0], "pcm": [0.2, 0.6],
    })
    out = D.add_solvent_mean(df)
    mean_row = out[out["solvent"] == "solvent_mean"]
    assert len(mean_row) == 1
    assert mean_row["experimental"].iloc[0] == pytest.approx(3.0)   # mean(2, 4)
    assert mean_row["pcm"].iloc[0] == pytest.approx(0.4)            # mean(0.2, 0.6)


def test_solvent_pair_differences_exact():
    def row(solvent, exp, pcm, des, dvib):
        return {"solute": "m0", "site": "s0", "nucleus": "H", "solvent": solvent,
                "experimental": exp, "pcm": pcm, "desmond": des, "desmond_vib": dvib}
    df = pd.DataFrame([row("methanol", 2.0, 0.4, 0.5, 0.1),    # check solvent
                       row("chloroform", 1.0, 0.1, 0.2, 0.05)])  # reference
    out = D.solvent_pair_differences(df, "methanol", "chloroform", nucleus="H", explicit="desmond")
    assert len(out) == 1
    r = out.iloc[0]
    assert r["exp_diff"] == pytest.approx(1.0 - 2.0)          # reference - check
    assert r["implicit_diff"] == pytest.approx(0.4 - 0.1)     # check - reference
    assert r["explicit_diff"] == pytest.approx(0.5 - 0.2)
    assert r["explicit_vib_diff"] == pytest.approx((0.5 + 0.1) - (0.2 + 0.05))


def test_solvent_induced_shifts_stacks_solvents():
    rng = np.random.default_rng(0)
    rows = []
    for solvent in ["chloroform", "methanol", "benzene"]:
        for si in range(3):
            rows.append({"solute": f"m{si}", "site": f"m{si}_s", "nucleus": "H", "solvent": solvent,
                         "experimental": rng.normal(), "pcm": rng.normal(), "desmond": rng.normal(),
                         "desmond_vib": rng.normal()})
    df = pd.DataFrame(rows)
    out = D.solvent_induced_shifts(df, "chloroform", ["chloroform", "methanol", "benzene"],
                                   nucleus="H", explicit="desmond")
    # two non-reference solvents x 3 sites = 6 rows, reference column constant
    assert len(out) == 6
    assert set(out["solvent"].unique()) == {"methanol", "benzene"}
    assert (out["reference"] == "chloroform").all()


def test_fit_differences_to_experimental():
    # exp_diff = -2 * predicted + 0.3 exactly
    pred = np.linspace(-1, 1, 20)
    df = pd.DataFrame({"exp_diff": -2.0 * pred + 0.3, "explicit_diff": pred})
    fit = D.fit_differences_to_experimental(df, "explicit_diff")
    assert fit["slope"] == pytest.approx(-2.0, abs=1e-9)
    assert fit["intercept"] == pytest.approx(0.3, abs=1e-9)
    assert fit["rmse"] == pytest.approx(0.0, abs=1e-9)
    assert fit["n"] == 20


def test_explicit_correction_pairs_dedups_and_drops_nan():
    rows = []
    for method in ["wp04", "dsd_pbep86"]:    # same site tiled across methods
        rows.append({"solute": "m0", "site": "s0", "nucleus": "H", "solvent": "chloroform",
                     "sap_nmr_method": method, "desmond": 0.5, "openMM": 0.4})
    rows.append({"solute": "m1", "site": "s1", "nucleus": "H", "solvent": "chloroform",
                 "sap_nmr_method": "wp04", "desmond": np.nan, "openMM": 0.2})   # dropped (NaN)
    out = D.explicit_correction_pairs(pd.DataFrame(rows), "chloroform", "H")
    assert len(out) == 1
    assert out.iloc[0]["desmond"] == pytest.approx(0.5)
    assert out.iloc[0]["openMM"] == pytest.approx(0.4)


def test_compare_dft_nn():
    # method-tiled correction; dedup by keys, then error = NN - DFT
    dft = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "sap_nmr_method": "wp04", "qcd": 0.10},
        {"solute": "m0", "site": "s0", "nucleus": "H", "sap_nmr_method": "mp2", "qcd": 0.10},
        {"solute": "m1", "site": "s1", "nucleus": "C", "sap_nmr_method": "wp04", "qcd": 0.50},
    ])
    nn = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "sap_nmr_method": "MagNET", "qcd": 0.13},
        {"solute": "m1", "site": "s1", "nucleus": "C", "sap_nmr_method": "MagNET", "qcd": 0.40},
    ])
    out = D.compare_dft_nn(dft, nn, "qcd", keys=("solute", "site", "nucleus"))
    assert len(out) == 2
    h = out[out["nucleus"] == "H"].iloc[0]
    assert h["qcd_dft"] == pytest.approx(0.10)
    assert h["qcd_nn"] == pytest.approx(0.13)
    assert h["error"] == pytest.approx(0.13 - 0.10)


def test_qcd_correction_by_site_filters_nucleus_and_sorts_by_solute():
    # three H sites across two solutes (out of alphabetical order) plus one C site that must be
    # dropped; check the nucleus filter, the solute sort, and that the DFT/NN values pass through.
    dft = pd.DataFrame([
        {"solute": "zzz", "site": "s0", "nucleus": "H", "sap_nmr_method": "wp04", "qcd": -0.70},
        {"solute": "AcOH", "site": "s1", "nucleus": "H", "sap_nmr_method": "wp04", "qcd": -0.65},
        {"solute": "AcOH", "site": "s2", "nucleus": "H", "sap_nmr_method": "wp04", "qcd": -0.80},
        {"solute": "AcOH", "site": "s3", "nucleus": "C", "sap_nmr_method": "wp04", "qcd": -4.0},
    ])
    nn = pd.DataFrame([
        {"solute": "zzz", "site": "s0", "nucleus": "H", "sap_nmr_method": "MagNET", "qcd": -0.72},
        {"solute": "AcOH", "site": "s1", "nucleus": "H", "sap_nmr_method": "MagNET", "qcd": -0.64},
        {"solute": "AcOH", "site": "s2", "nucleus": "H", "sap_nmr_method": "MagNET", "qcd": -0.79},
        {"solute": "AcOH", "site": "s3", "nucleus": "C", "sap_nmr_method": "MagNET", "qcd": -4.1},
    ])
    out = D.qcd_correction_by_site(dft, nn, nucleus="H")
    assert len(out) == 3                              # the C row is dropped
    assert (out["nucleus"] == "H").all()
    assert list(out["solute"]) == ["AcOH", "AcOH", "zzz"]   # sorted, "zzz" sorts after "AcOH"
    row = out[out["site"] == "s1"].iloc[0]
    assert row["qcd_dft"] == pytest.approx(-0.65)
    assert row["qcd_nn"] == pytest.approx(-0.64)
    assert row["error"] == pytest.approx(-0.64 - (-0.65))


def test_compare_dft_nn_by_engine_stacks_both_engines():
    # one shared site, a different desmond/openMM value each, so the two engines land in separate
    # rows with a common dft_value/nn_value pair of columns instead of engine-specific column names
    dft = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "solvent": "chloroform",
         "sap_nmr_method": "wp04", "desmond": 0.10, "openMM": 0.20},
    ])
    nn = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "solvent": "chloroform",
         "sap_nmr_method": "MagNET", "desmond": 0.13, "openMM": 0.17},
    ])
    out = D.compare_dft_nn_by_engine(dft, nn)
    assert len(out) == 2
    assert set(out["engine"]) == {"desmond", "openMM"}
    assert list(out.columns) == ["solute", "site", "nucleus", "solvent", "dft_value", "nn_value",
                                 "error", "engine"]
    desmond_row = out[out["engine"] == "desmond"].iloc[0]
    assert desmond_row["dft_value"] == pytest.approx(0.10)
    assert desmond_row["nn_value"] == pytest.approx(0.13)
    assert desmond_row["error"] == pytest.approx(0.13 - 0.10)
    openmm_row = out[out["engine"] == "openMM"].iloc[0]
    assert openmm_row["dft_value"] == pytest.approx(0.20)
    assert openmm_row["nn_value"] == pytest.approx(0.17)
    assert openmm_row["error"] == pytest.approx(0.17 - 0.20)


def test_explicit_correction_dft_nn_pairs_merges_sources():
    dft = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "solvent": "chloroform",
         "sap_nmr_method": "wp04", "desmond": 0.5, "openMM": 0.4},
        {"solute": "m1", "site": "s1", "nucleus": "H", "solvent": "chloroform",
         "sap_nmr_method": "wp04", "desmond": -0.2, "openMM": -0.3},
    ])
    nn = pd.DataFrame([
        {"solute": "m0", "site": "s0", "nucleus": "H", "solvent": "chloroform",
         "sap_nmr_method": "MagNET", "desmond": 0.55, "openMM": 0.42},
        # site s1 has no NN row, so it should not appear in the merged (inner-join) result
    ])
    out = D.explicit_correction_dft_nn_pairs(dft, nn, "chloroform", "H")
    assert len(out) == 1
    row = out.iloc[0]
    assert row["solute"] == "m0"
    assert row["desmond_dft"] == pytest.approx(0.5)
    assert row["openMM_dft"] == pytest.approx(0.4)
    assert row["desmond_nn"] == pytest.approx(0.55)
    assert row["openMM_nn"] == pytest.approx(0.42)


def test_si_s13d_fitting_accuracy_selects_formula_by_nucleus_and_engine():
    # stationary + qcd + desmond exactly reproduces experimental for H; stationary + desmond_vib +
    # openMM exactly reproduces it for C. Any other combination of (nucleus, engine) leaves a
    # nonzero residual from the term the formula does NOT include, so a near-zero test RMSE proves
    # si_s13d_fitting_accuracy picked the QCD-based formula for H and the vib-based one for C, and
    # picked the desmond/openMM explicit term matching the "engine" column.
    rng = np.random.default_rng(5)
    solvents = ["chloroform", "methanol", "TIP4P", "benzene"]
    rows = []
    for si in range(14):
        solute = f"m{si:02d}"
        for solvent in solvents:
            for k in range(2):
                stationary = rng.normal(50, 5)
                qcd = rng.normal(0, 1)
                desmond = rng.normal(0, 1)
                desmond_vib = rng.normal(0, 1)
                openMM = rng.normal(0, 1)
                openMM_vib = rng.normal(0, 1)
                rows.append({
                    "solute": solute, "solvent": solvent, "site": f"{solute}_{k}",
                    "sap_nmr_method": "dsd_pbep86", "sap_basis": "pcSseg3",
                    "sap_geometry_type": "pbe0_tz",
                    "stationary": stationary, "qcd": qcd, "pcm": 0.0,
                    "desmond": desmond, "openMM": openMM,
                    "desmond_vib": desmond_vib, "openMM_vib": openMM_vib,
                    "nucleus": "H",
                    "experimental": stationary + qcd + desmond,
                })
                rows.append({
                    "solute": solute, "solvent": solvent, "site": f"{solute}_{k}",
                    "sap_nmr_method": "dsd_pbep86", "sap_basis": "pcSseg3",
                    "sap_geometry_type": "pbe0_tz",
                    "stationary": stationary, "qcd": qcd, "pcm": 0.0,
                    "desmond": desmond, "openMM": openMM,
                    "desmond_vib": desmond_vib, "openMM_vib": openMM_vib,
                    "nucleus": "C",
                    "experimental": stationary + openMM_vib + openMM,
                })
    dft = pd.DataFrame(rows)
    nn = dft.copy()   # same synthetic values for both sources; the exact-fit check is source-agnostic
    solutes = sorted(dft["solute"].unique())

    h_results = D.si_s13d_fitting_accuracy(dft, nn, solvents, n_splits=3, solutes=solutes,
                                           nucleus="H")
    h_desmond = h_results[h_results["engine"] == "desmond"]["test_RMSE"]
    assert (h_desmond < 1e-8).all()
    h_openmm = h_results[h_results["engine"] == "openMM"]["test_RMSE"]
    assert (h_openmm > 0.1).all()   # missing the desmond term used to build "experimental"

    c_results = D.si_s13d_fitting_accuracy(dft, nn, solvents, n_splits=3, solutes=solutes,
                                           nucleus="C")
    c_openmm = c_results[c_results["engine"] == "openMM"]["test_RMSE"]
    assert (c_openmm < 1e-8).all()
    c_desmond = c_results[c_results["engine"] == "desmond"]["test_RMSE"]
    assert (c_desmond > 0.1).all()   # missing the openMM term used to build "experimental"


def test_pcm_benefit_per_solvent_positive():
    df = make_dft_query_df(seed=3)   # experimental = 31 - (stationary + pcm), so PCM genuinely helps
    solutes = sorted(df["solute"].unique())
    benefit = D.pcm_benefit_per_solvent(df, "wp04", "pcSseg2", "aimnet2",
                                        ["chloroform", "methanol", "benzene"], n_splits=4,
                                        solutes=solutes, nucleus="H", n_test=4)
    assert set(benefit.index) == {"chloroform", "methanol", "benzene"}
    assert (benefit > 0).all()   # adding PCM lowers the error in every solvent


def test_frame_convergence_math():
    # 3 atoms, 5 frames; site = atoms 0 and 2; solvated - isolated = a known per-frame value
    n_frames = 5
    isolated = np.zeros((n_frames, 3))
    solvated = np.zeros((n_frames, 3))
    per_frame_truth = np.array([0.2, 0.4, np.nan, 0.6, 0.8])
    for f in range(n_frames):
        solvated[f, 0] = per_frame_truth[f]
        solvated[f, 2] = per_frame_truth[f]
    solvated[2, :] = np.nan        # frame 2 invalid
    perturbed = np.stack([isolated, solvated], axis=-1)   # (5, 3, 2)
    fc = D.frame_corrections(perturbed, [0, 2])
    assert fc[0] == pytest.approx(0.2)
    assert np.isnan(fc[2])
    ra = D.running_average(fc)
    # running average over the valid frames 0.2, 0.4, (skip), 0.6, 0.8
    assert ra[0] == pytest.approx(0.2)
    assert ra[1] == pytest.approx(0.3)
    assert ra[2] == pytest.approx(0.3)             # invalid frame does not move the average
    assert ra[4] == pytest.approx((0.2 + 0.4 + 0.6 + 0.8) / 4)


def test_autocorrelation():
    ac = D.autocorrelation(np.array([1.0, -1.0, 1.0, -1.0, 1.0, -1.0]), max_lag=2)
    assert ac[0] == pytest.approx(1.0)             # lag 0 is always 1
    assert ac[1] < 0                               # alternating -> negative lag-1 autocorrelation


def test_frame_validity():
    # 4 frames, 2 atoms: frame 0 fully computed, frame 1 only the isolated value, frame 2 all NaN
    # (skipped entirely), frame 3 fully computed. A frame counts as valid if anything was computed.
    perturbed = np.full((4, 2, 2), np.nan)
    perturbed[0] = [[1.0, 2.0], [3.0, 4.0]]
    perturbed[1, 0, 0] = 5.0
    perturbed[3] = [[1.0, 2.0], [3.0, 4.0]]
    valid = D.frame_validity(perturbed)
    assert list(valid) == [True, True, False, True]


def test_frame_validity_grid_pads_to_the_longest_solvent(monkeypatch):
    # two "solvents" with different trajectory lengths; the grid must pad the shorter one with
    # False past its own frame count, and frame_counts must report each solvent's true length.
    fake_data = {
        "acetone": np.array([[[1.0, 2.0]], [[np.nan, np.nan]], [[1.0, 2.0]]]),   # 3 frames, 1 invalid
        "benzene": np.array([[[1.0, 2.0]], [[1.0, 2.0]]]),                        # 2 frames, all valid
    }

    def fake_load_perturbed_shieldings(hdf5_filepath, solute, solvent, explicit_solvation_model,
                                       shield_type):
        return fake_data[solvent]

    monkeypatch.setattr(D, "load_perturbed_shieldings", fake_load_perturbed_shieldings)
    grid, frame_counts = D.frame_validity_grid("unused.hdf5", "solute", ["acetone", "benzene"],
                                               "openMM")
    assert grid.shape == (2, 3)
    assert list(frame_counts) == [3, 2]
    assert list(grid[0]) == [True, False, True]
    assert list(grid[1]) == [True, True, False]        # padded with False past benzene's 2 frames


def test_pareto_frontier():
    points = pd.DataFrame({
        "label": ["a", "b", "c", "d", "e"],
        "total_time": [1.0, 2.0, 3.0, 4.0, 5.0],
        "test_RMSE": [5.0, 3.0, 4.0, 1.0, 2.0],
    })
    front = D.pareto_frontier(points, x="total_time", y="test_RMSE")
    assert list(front["label"]) == ["a", "b", "d"]   # c and e are dominated


def test_correlation_matrix():
    # two solvents whose per-site pcm corrections are perfectly correlated
    base = np.array([0.1, 0.5, -0.3, 0.7, -0.2, 0.4])
    rows = []
    for i, value in enumerate(base):
        rows.append({"solute": f"m{i}", "site": f"s{i}", "solvent": "chloroform", "pcm": value})
        rows.append({"solute": f"m{i}", "site": f"s{i}", "solvent": "benzene", "pcm": 2 * value + 1})
    df = pd.DataFrame(rows)
    corr = D.correlation_matrix(df, "pcm", ["solute", "site"], "solvent")
    assert corr.loc["chloroform", "benzene"] == pytest.approx(1.0, abs=1e-9)


# --- Figure 2A / S1 (Pareto) ----------------------------------------------------------------

def test_pcm_conversion_factors_recovers_known_slope():
    # build DFT rows where the output method's PCM in each solvent is a known multiple of the input
    # method's chloroform PCM, so the recomputed factor must equal that multiple.
    rng = np.random.default_rng(0)
    multiples = {"chloroform": 1.0, "benzene": 0.65, "methanol": 1.4}
    rows = []
    for si in range(8):
        ref = rng.normal(0, 2)   # the input method's chloroform PCM for this site
        for solvent, mult in multiples.items():
            common = {"solute": f"m{si}", "site": f"m{si}_s", "nucleus": "H",
                      "sap_geometry_type": "aimnet2", "sap_basis": "pcSseg2", "solvent": solvent}
            # input (reference) method: chloroform carries `ref`, other solvents irrelevant
            rows.append({**common, "sap_nmr_method": D.MAGNET_PCM_INPUT_METHOD,
                         "pcm": ref if solvent == "chloroform" else rng.normal()})
            # output method: this solvent's PCM is mult * the chloroform reference
            rows.append({**common, "sap_nmr_method": "wp04", "pcm": mult * ref})
    df = pd.DataFrame(rows)
    factors = D.pcm_conversion_factors(df, "H", output_method="wp04")
    for solvent, mult in multiples.items():
        assert factors[solvent] == pytest.approx(mult, abs=1e-9)


def test_add_nn_stationary_plus_pcm():
    df = pd.DataFrame({
        "solute": ["m0", "m0"], "site": ["a", "b"], "nucleus": ["H", "C"],
        "solvent": ["benzene", "benzene"], "stationary": [100.0, 50.0], "pcm": [2.0, 4.0],
    })
    factors = {"H": pd.Series({"benzene": 0.5}), "C": pd.Series({"benzene": 3.0})}
    out = D.add_nn_stationary_plus_pcm(df, factors)
    assert out["stationary_plus_pcm"].iloc[0] == pytest.approx(100.0 + 2.0 * 0.5)
    assert out["stationary_plus_pcm"].iloc[1] == pytest.approx(50.0 + 4.0 * 3.0)


def _pareto_synthetic():
    """A tiny DFT+NN pair where a 'fast' method is also more accurate, so it must sit on the
    frontier and a 'slow, worse' method must not."""
    rng = np.random.default_rng(1)
    solvents = ["chloroform", "benzene", "TIP4P"]
    dft_rows, nn_rows = [], []
    for si in range(12):
        solute = f"m{si:02d}"
        for solvent in solvents:
            stat = rng.normal(100, 20)
            pcm = rng.normal(0, 2)
            exp = 31.0 - (stat + pcm)
            for method, noise in [("good_fast", 0.02), ("bad_slow", 1.5)]:
                dft_rows.append({
                    "solute": solute, "site": f"{solute}_s", "nucleus": "H", "solvent": solvent,
                    "sap_nmr_method": method, "sap_basis": "pcSseg2", "sap_geometry_type": "aimnet2",
                    "experimental": exp + rng.normal(0, noise), "stationary": stat, "pcm": pcm})
            nn_rows.append({
                "solute": solute, "site": f"{solute}_s", "nucleus": "H", "solvent": solvent,
                "sap_nmr_method": "MagNET", "sap_basis": "N/A", "sap_geometry_type": "aimnet2",
                "experimental": exp + rng.normal(0, 0.05), "stationary": stat, "pcm": pcm})
    return pd.DataFrame(dft_rows), pd.DataFrame(nn_rows)


def test_assemble_pareto_fitting_df_and_rmse():
    dft, nn = _pareto_synthetic()
    factors = {"H": pd.Series({s: 1.0 for s in ["chloroform", "benzene", "TIP4P"]}),
               "C": pd.Series(dtype=float)}
    fitting = D.assemble_pareto_fitting_df(dft, nn, conversion_factors=factors, exclude_solutes=())
    # TIP4P renamed to water, all three methods present, single predictor column built
    assert "water" in set(fitting["solvent"]) and "TIP4P" not in set(fitting["solvent"])
    assert set(fitting["nmr_method"]) == {"good_fast", "bad_slow", "MagNET"}
    assert "stationary_plus_pcm" in fitting.columns
    rmse = D.pareto_solvent_averaged_rmse(fitting, n_splits=6, n_test=6)
    assert "solvent-averaged" in set(rmse["solvent"])
    sav = rmse[rmse["solvent"] == "solvent-averaged"].set_index("nmr_method")["fitting_RMSE"]
    assert sav["good_fast"] < sav["bad_slow"]     # the accurate method has the lower RMSE


def test_attach_pareto_timings_substitutes_magnet():
    rmse = pd.DataFrame({
        "geometry_type": ["aimnet2", "aimnet2"], "nmr_method": ["bad_slow", "MagNET"],
        "basis": ["pcSseg2", "N/A"], "nucleus": ["H", "H"], "solvent": ["solvent-averaged"] * 2,
        "fitting_RMSE": [1.0, 0.2],
    })
    dft_gas = pd.DataFrame({"geometry_time": [10.0], "nmr_time": [990.0], "total_time": [1000.0]},
                           index=pd.MultiIndex.from_tuples([("aimnet2", "bad_slow", "pcSseg2")],
                                 names=["geometry_type", "nmr_method", "basis"]))
    nn = pd.DataFrame({"geometry_time": [5.0], "nmr_time": [15.0], "total_time": [20.0]},
                      index=pd.MultiIndex.from_tuples([("aimnet2", "MagNET", "N/A")],
                            names=["geometry_type", "nmr_method", "basis"]))
    out = D.attach_pareto_timings(rmse, dft_gas, nn)
    slow = out[out["nmr_method"] == "bad_slow"].iloc[0]
    magnet = out[out["nmr_method"] == "MagNET"].iloc[0]
    assert slow["total_time"] == pytest.approx(1000.0)
    assert magnet["total_time"] == pytest.approx(20.0)     # substituted from the NN timings
    # MagNET dominates: faster and more accurate, so it is on the frontier and slow is not
    front = D.pareto_frontier(out, x="total_time", y="fitting_RMSE")
    assert list(front["nmr_method"]) == ["MagNET"]


def test_pareto_table_curated_keeps_only_the_documented_rows():
    points = pd.DataFrame([
        # kept: MagNET itself
        {"nucleus": "H", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "MagNET", "basis": "N/A", "total_time": 1.0},
        # kept: the one representative aimnet2 row (H's reference method wp04, pcSseg2)
        {"nucleus": "H", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "wp04", "basis": "pcSseg2", "total_time": 2.0},
        # dropped: aimnet2 + wp04 but the WRONG basis
        {"nucleus": "H", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "wp04", "basis": "pcSseg1", "total_time": 3.0},
        # dropped: aimnet2 + pcSseg2 but the WRONG method (not H's reference)
        {"nucleus": "H", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "hf", "basis": "pcSseg2", "total_time": 4.0},
        # kept: part of the full pbe0_tz grid, any method/basis
        {"nucleus": "H", "solvent": "chloroform", "geometry_type": "pbe0_tz", "nmr_method": "hf", "basis": "pcSseg1", "total_time": 5.0},
        # dropped: right nucleus/geometry, WRONG solvent
        {"nucleus": "H", "solvent": "benzene", "geometry_type": "pbe0_tz", "nmr_method": "hf", "basis": "pcSseg1", "total_time": 6.0},
        # dropped: right solvent/geometry, WRONG nucleus
        {"nucleus": "C", "solvent": "chloroform", "geometry_type": "pbe0_tz", "nmr_method": "hf", "basis": "pcSseg1", "total_time": 7.0},
    ])
    out = D.pareto_table_curated(points, "H", solvent="chloroform")
    assert len(out) == 3
    assert set(out["total_time"]) == {1.0, 2.0, 5.0}
    assert list(out["total_time"]) == sorted(out["total_time"])   # sorted by total_time ascending


def test_pareto_table_curated_uses_the_carbon_reference_method():
    # C's reference method is wb97xd (MAGNET_PCM_OUTPUT_METHODS), not H's wp04
    points = pd.DataFrame([
        {"nucleus": "C", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "wp04", "basis": "pcSseg2", "total_time": 1.0},
        {"nucleus": "C", "solvent": "chloroform", "geometry_type": "aimnet2", "nmr_method": "wb97xd", "basis": "pcSseg2", "total_time": 2.0},
    ])
    out = D.pareto_table_curated(points, "C", solvent="chloroform")
    assert list(out["nmr_method"]) == ["wb97xd"]


# Published Figure 2A values to reproduce from the released data (from the original
# fig2a_pareto_results.csv). MagNET sits at the cheap end of a flat frontier.
FIG2A_MAGNET_SOLVENT_AVERAGED = {"H": 0.1647203288369508, "C": 1.7268703504301826}
FIG2A_CONVERSION_FACTORS_H_BENZENE = 0.6513139306748533   # from magnet_pcm_conversion_factors_H.csv


@pytest.mark.skipif(
    not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX)),
    reason="real delta22.hdf5 / experimental xlsx not present")
def test_fig2a_reproduces_published_magnet_point():
    """Recompute the MagNET Pareto point and conversion factors from the released data and check
    them against the published Figure 2A numbers (the int32 encoding sets the precision floor)."""
    dft = D.load_query_df_dft(REAL_H5, REAL_XLSX, verbose=False)
    nn = D.load_query_df_nn(REAL_H5, REAL_XLSX, verbose=False)
    # conversion factors recomputed from the DFT PCM data match the stored CSV value
    cf_h = D.pcm_conversion_factors(dft, "H")
    assert cf_h["benzene"] == pytest.approx(FIG2A_CONVERSION_FACTORS_H_BENZENE, abs=1e-6)
    # build only the MagNET fitting rows and average over solvents (fast; avoids the full DFT sweep)
    factors = {nuc: D.pcm_conversion_factors(dft, nuc) for nuc in ("H", "C")}
    nn_fit = D.add_nn_stationary_plus_pcm(nn, factors).rename(columns=lambda c: c.replace("sap_", ""))
    nn_fit["solvent"] = nn_fit["solvent"].replace("TIP4P", "water")
    nn_fit["experimental"] = pd.to_numeric(nn_fit["experimental"], errors="coerce")
    nn_fit = nn_fit[nn_fit["solute"] != "nitromethane"]
    # the split solute order must match the full combined table (DFT solutes, in file order)
    solutes = D.assemble_pareto_fitting_df(dft, nn, conversion_factors=factors)["solute"].unique().tolist()
    rmse = D.pareto_solvent_averaged_rmse(nn_fit, solutes=solutes)
    sav = rmse[(rmse["nmr_method"] == "MagNET") & (rmse["solvent"] == "solvent-averaged")]
    for nucleus, expected in FIG2A_MAGNET_SOLVENT_AVERAGED.items():
        got = sav[sav["nucleus"] == nucleus]["fitting_RMSE"].iloc[0]
        assert got == pytest.approx(expected, abs=1e-4), f"{nucleus}: {got} vs {expected}"


# --- Composite-formula ablation workbook's "Correlations Between Features" -----------------------

def test_feature_correlation_matrix_synthetic():
    # a and b perfectly correlated, c independent noise uncorrelated with either
    n = 200
    rng = np.random.default_rng(0)
    a = rng.normal(size=n)
    df = pd.DataFrame({
        "nucleus": ["H"] * n, "solvent": ["chloroform"] * n,
        "stationary": a, "pcm": a * 2 + 1,  # perfectly correlated with stationary
        "desmond": rng.normal(size=n), "desmond_vib": rng.normal(size=n), "qcd": rng.normal(size=n),
    })
    corr = D.feature_correlation_matrix(df, "H", "chloroform")
    assert corr.loc["stationary", "pcm"] == pytest.approx(1.0, abs=1e-9)
    corr2 = D.feature_correlation_matrix(df, "H", "chloroform", squared=True)
    assert corr2.loc["stationary", "pcm"] == pytest.approx(1.0, abs=1e-9)


@pytest.mark.filterwarnings("ignore:Mean of empty slice:RuntimeWarning")
def test_average_feature_correlation_matrix_averages_across_solvents():
    df = pd.DataFrame({
        "nucleus": ["H"] * 4, "solvent": ["chloroform", "chloroform", "benzene", "benzene"],
        "stationary": [1.0, 2.0, 1.0, 3.0], "pcm": [1.0, 2.0, 3.0, 1.0],
        "desmond": [0.5, 0.5, 0.5, 0.5], "desmond_vib": [0.1, 0.2, 0.1, 0.2], "qcd": [0, 0, 0, 0],
    })
    avg = D.average_feature_correlation_matrix(df, "H", ["chloroform", "benzene"])
    chloroform_only = D.feature_correlation_matrix(df, "H", "chloroform")
    benzene_only = D.feature_correlation_matrix(df, "H", "benzene")
    expected = (chloroform_only.loc["stationary", "pcm"] + benzene_only.loc["stationary", "pcm"]) / 2
    assert avg.loc["stationary", "pcm"] == pytest.approx(expected)


@pytest.mark.filterwarnings("ignore:Mean of empty slice:RuntimeWarning")
def test_average_feature_correlation_matrix_ignores_a_solvent_with_too_few_sites():
    # benzene has only 1 site -> feature_correlation_matrix returns all-NaN for benzene alone
    # (< 2 sites). A plain sum()/len() average would make the WHOLE averaged cell NaN even though
    # chloroform (4 sites) has a perfectly good value; nanmean should use chloroform alone instead.
    df = pd.DataFrame({
        "nucleus": ["H"] * 5,
        "solvent": ["chloroform", "chloroform", "chloroform", "chloroform", "benzene"],
        "stationary": [1.0, 2.0, 3.0, 4.0, 1.0], "pcm": [2.0, 4.0, 6.0, 8.0, 5.0],
        "desmond": [0.5] * 5, "desmond_vib": [0.1] * 5, "qcd": [0.0] * 5,
    })
    avg = D.average_feature_correlation_matrix(df, "H", ["chloroform", "benzene"])
    assert not np.isnan(avg.loc["stationary", "pcm"])
    assert avg.loc["stationary", "pcm"] == pytest.approx(1.0)   # chloroform alone: perfectly correlated


@pytest.mark.skipif(
    not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX)),
    reason="real delta22.hdf5 / experimental xlsx not present")
def test_ablations_feature_correlations_reproduces_published_matrices():
    """Every entry of both nuclei's published Pearson R and R^2 "Correlations Between Features"
    matrices (the composite-formula ablation workbook), reproduced from the released data to 3
    decimal places. Guards against
    pooling all DFT methods/bases/geometries together instead of filtering to the MagNET-Zero
    reference level, which throws entries off by up to 0.03."""
    dft = D.load_query_df_dft(REAL_H5, REAL_XLSX, verbose=False)
    # R^2 here is the average of each solvent's r^2 (Jensen's inequality: != the square of the
    # averaged r), so it is hardcoded separately from published_r, not derived from it.
    published_r = {
        "H": {("pcm", "stationary"): 0.298, ("desmond", "stationary"): 0.262, ("desmond", "pcm"): 0.399,
              ("desmond_vib", "stationary"): -0.170, ("desmond_vib", "pcm"): -0.219, ("desmond_vib", "desmond"): -0.141,
              ("qcd", "stationary"): -0.851, ("qcd", "pcm"): -0.362, ("qcd", "desmond"): -0.394, ("qcd", "desmond_vib"): 0.405},
        "C": {("pcm", "stationary"): 0.584, ("desmond", "stationary"): 0.106, ("desmond", "pcm"): 0.667,
              ("desmond_vib", "stationary"): 0.531, ("desmond_vib", "pcm"): 0.369, ("desmond_vib", "desmond"): 0.075,
              ("qcd", "stationary"): -0.616, ("qcd", "pcm"): -0.635, ("qcd", "desmond"): -0.418, ("qcd", "desmond_vib"): -0.138},
    }
    published_r2 = {
        "H": {("pcm", "stationary"): 0.089, ("desmond", "stationary"): 0.093, ("desmond", "pcm"): 0.386,
              ("desmond_vib", "stationary"): 0.032, ("desmond_vib", "pcm"): 0.049, ("desmond_vib", "desmond"): 0.025,
              ("qcd", "stationary"): 0.724, ("qcd", "pcm"): 0.131, ("qcd", "desmond"): 0.184, ("qcd", "desmond_vib"): 0.167},
        "C": {("pcm", "stationary"): 0.341, ("desmond", "stationary"): 0.047, ("desmond", "pcm"): 0.463,
              ("desmond_vib", "stationary"): 0.284, ("desmond_vib", "pcm"): 0.139, ("desmond_vib", "desmond"): 0.019,
              ("qcd", "stationary"): 0.380, ("qcd", "pcm"): 0.404, ("qcd", "desmond"): 0.189, ("qcd", "desmond_vib"): 0.021},
    }
    for nucleus in ("H", "C"):
        corr = D.ablations_feature_correlations(dft, nucleus)
        for (a, b), expected in published_r[nucleus].items():
            assert corr["r"].loc[a, b] == pytest.approx(expected, abs=5e-4), f"{nucleus} r[{a},{b}]"
        for (a, b), expected in published_r2[nucleus].items():
            assert corr["r2"].loc[a, b] == pytest.approx(expected, abs=5e-4), f"{nucleus} r2[{a},{b}]"


@pytest.mark.skipif(
    not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX)),
    reason="real delta22.hdf5 / experimental xlsx not present")
def test_pcm_desmond_correlation_by_solvent_reproduces_published_table():
    dft = D.load_query_df_dft(REAL_H5, REAL_XLSX, verbose=False)
    table = D.pcm_desmond_correlation_by_solvent(dft)
    published = {("H", "chloroform"): 0.739, ("H", "TIP4P"): 0.728, ("H", "benzene"): -0.454,
                ("C", "chloroform"): 0.408, ("C", "TIP4P"): 0.840, ("C", "benzene"): 0.594}
    for (nucleus, solvent), expected in published.items():
        assert table.loc[nucleus, solvent] == pytest.approx(expected, abs=5e-4), f"{nucleus}/{solvent}"


@pytest.mark.skipif(
    not (os.path.exists(REAL_H5) and os.path.exists(REAL_XLSX)),
    reason="real delta22.hdf5 / experimental xlsx not present")
def test_pcm_benefit_reproduces_published_figure_s5_correlations():
    # Figure S5 aggregates the per-split PCM benefit with the mean across splits, not the median;
    # only the mean reproduces the published Pearson R values.
    dft = D.load_query_df_dft(REAL_H5, REAL_XLSX, verbose=False)
    solutes = sorted(dft["solute"].unique())

    def pearson_r(benefit, prop, exclude=()):
        s = benefit.drop(index=[x for x in exclude if x in benefit.index])
        x = np.array([prop[k] for k in s.index])
        return float(np.corrcoef(x, s.values)[0, 1])

    aromatics_tfe = ["benzene", "toluene", "chlorobenzene", "trifluoroethanol"]
    benefit_h = D.pcm_benefit_per_solvent(dft, D.MAGNET_PCM_OUTPUT_METHODS["H"], "pcSseg2", "aimnet2",
                                          D.DESMOND_SOLVENTS, n_splits=250, solutes=solutes, nucleus="H")
    benefit_c = D.pcm_benefit_per_solvent(dft, D.MAGNET_PCM_OUTPUT_METHODS["C"], "pcSseg2", "aimnet2",
                                          D.DESMOND_SOLVENTS, n_splits=250, solutes=solutes, nucleus="C")
    published = {  # (panel, axis): published Pearson R
        "S5A 1H dielectric": (pearson_r(benefit_h, D.SOLVENT_DIELECTRIC), 0.617),
        "S5A 1H polarizability": (pearson_r(benefit_h, D.SOLVENT_POLARIZABILITY), -0.801),
        "S5B 13C dielectric": (pearson_r(benefit_c, D.SOLVENT_DIELECTRIC), 0.722),
        "S5B 13C polarizability": (pearson_r(benefit_c, D.SOLVENT_POLARIZABILITY), -0.743),
        "S5C 1H dielectric": (pearson_r(benefit_h, D.SOLVENT_DIELECTRIC, aromatics_tfe), 0.884),
        "S5C 1H polarizability": (pearson_r(benefit_h, D.SOLVENT_POLARIZABILITY, aromatics_tfe), -0.587),
    }
    for name, (got, expected) in published.items():
        assert got == pytest.approx(expected, abs=2e-3), f"{name}: {got} != {expected}"