Add files using upload-large-folder tool

analysis/code/applications.py

@@ -0,0 +1,549 @@
+"""Core analysis for the natural-products ("applications") figures (Figures 5C/5D, SI S8, S15). No plotting.
+
+Takes an Applications loader (data/applications/applications_reader.py) and computes the numbers the
+figure notebooks and tests share. A "composite model" predicts a shift as an OLS-weighted sum of
+physics terms (gas-phase shielding + implicit correction + rovibrational correction); these figures
+test how well weights fit on delta-22 carry to the natural-products test set.
+"""
+import numpy as np
+import pandas as pd
+import statsmodels.formula.api as smf
+
+from stats import rmse as _rmse
+from spreadsheet import site_atom_indices
+from applications_reader import SOLVENTS, SHELL_SIZES
+
+# composite-model formula sets (column names refer to build_query_df_nn output)
+FITTING_FORMULAS_H = [
+    "stationary", "stationary_plus_pcm", "stationary + pcm", "stationary + pcm + qcd",
+    "stationary_plus_qcd + pcm", "stationary_plus_pcm_plus_qcd", "stationary + openMM",
+    "stationary + openMM + qcd", "stationary_plus_qcd + openMM", "stationary_plus_openMM_plus_qcd",
+]
+FITTING_FORMULAS_C = [
+    "stationary", "stationary_plus_pcm", "stationary + pcm", "stationary + pcm + openMM_vib",
+    "stationary_plus_op_vib + pcm", "stationary_plus_pcm_plus_op_vib", "stationary + openMM",
+    "stationary + openMM + openMM_vib", "stationary_plus_op_vib + openMM", "stationary_plus_openMM_plus_op_vib",
+]
+FORMULA_REMAP = {
+    "stationary_plus_pcm": "pcm2", "stationary + pcm": "pcm3", "stationary + pcm + qcd": "pcm4",
+    "stationary + pcm + openMM_vib": "pcm4", "stationary_plus_qcd + pcm": "semi_parsimonious_pcm",
+    "stationary_plus_op_vib + pcm": "semi_parsimonious_pcm", "stationary_plus_pcm_plus_qcd": "parsimonious_pcm",
+    "stationary_plus_pcm_plus_op_vib": "parsimonious_pcm", "stationary + openMM": "expl3",
+    "stationary + openMM + qcd": "expl4", "stationary + openMM + openMM_vib": "expl4",
+    "stationary_plus_qcd + openMM": "semi_parsimonious_expl", "stationary_plus_op_vib + openMM": "semi_parsimonious_expl",
+    "stationary_plus_openMM_plus_qcd": "parsimonious_expl", "stationary_plus_openMM_plus_op_vib": "parsimonious_expl",
+}
+# manuscript test-set membership, keyed by the canonical solute names build_query_df_nn emits
+ALL_IN_ONE_BIN = {"Test Set": [
+    "isomer_1E", "isomer_1Z", "isomer_2E", "isomer_2Z", "isomer_3E", "isomer_3Z", "isomer_4N", "isomer_4O",
+    "vomicine", "prednisone", "peptide", "flavone", "dihydrotanshinone_I",
+]}
+# canonical solute -> manuscript display label (used as figure x-tick labels)
+SOLUTE_DISPLAY = {
+    "isomer_1E": "Isomer 1 (E)", "isomer_1Z": "Isomer 1 (Z)", "isomer_2E": "Isomer 2 (E)",
+    "isomer_2Z": "Isomer 2 (Z)", "isomer_3E": "Isomer 3 (E)", "isomer_3Z": "Isomer 3 (Z)",
+    "isomer_4N": "1,4-dimethylpyridin-\n2(1H)one", "isomer_4O": "2-methoxy-\n4-methylpyridine",
+    "vomicine": "Vomicine", "prednisone": "Prednisone", "peptide": "Acetyl-L-alanyl-L-\nglutamine",
+    "flavone": "Flavone", "dihydrotanshinone_I": "Dihydrotanshinone I",
+}
+# the peptide (and, for carbon, the pyridone isomer_4N) is highlighted green as the high-error outlier
+PEPTIDE_HIGHLIGHT_H = {"peptide": "#B0FF40"}
+PEPTIDE_HIGHLIGHT_C = {"isomer_4N": "#B0FF40", "peptide": "#B0FF40"}
+
+# bootstrap-coefficient file token -> remapped formula name, per nucleus
+_BOOTSTRAP_TOKENS = {
+    "H": {"stationary_plus_pcm": "pcm2", "stationary_AND_pcm_AND_qcd": "pcm4",
+          "stationary_plus_qcd_AND_pcm": "semi_parsimonious_pcm", "stationary_plus_pcm_plus_qcd": "parsimonious_pcm",
+          "stationary_AND_openMM_AND_qcd": "expl4", "stationary_plus_qcd_AND_openMM": "semi_parsimonious_expl",
+          "stationary_plus_openMM_plus_qcd": "parsimonious_expl"},
+    "C": {"stationary_plus_pcm": "pcm2", "stationary_AND_pcm_AND_openMM_vib": "pcm4",
+          "stationary_plus_op_vib_AND_pcm": "semi_parsimonious_pcm", "stationary_plus_pcm_plus_op_vib": "parsimonious_pcm",
+          "stationary_AND_openMM_AND_openMM_vib": "expl4", "stationary_plus_op_vib_AND_openMM": "semi_parsimonious_expl",
+          "stationary_plus_openMM_plus_op_vib": "parsimonious_expl"},
+}
+
+
+def shell_convergence_corrections(loader, solute, nucleus, solvents=SOLVENTS):
+    """MagNET-X explicit-solvent correction (solvated minus isolated shielding) for each NMR
+    site, averaged over the site's atoms and over the MD frames, as a function of solvent shell
+    size. Backs SI Figure S8 panel A (shown for vomicine, proton sites).
+
+    Returns a DataFrame indexed by (site, solvent) with one column per shell ("shell_50" ...
+    "shell_650"). Shells absent for a given solute/solvent (e.g. the isomer_1Z/chloroform/250
+    gap) come back as NaN.
+    """
+    sites = loader.site_atom_table(uuid_prefix_h_sites=True)
+    sites = sites[(sites["solute"] == solute) & (sites["nucleus"] == nucleus)]
+    rows = []
+    for _, r in sites.iterrows():
+        idx = site_atom_indices(r["atom_numbers"])
+        for solvent in solvents:
+            present = set(loader.available_shells(solute, solvent))
+            vals = {"site": r["site"], "solvent": solvent}
+            for shell in SHELL_SIZES:
+                if shell not in present:
+                    vals[f"shell_{shell}"] = np.nan
+                    continue
+                mx = loader.magnet_x(solute, solvent, shell)          # (frames, atoms, 2)
+                corr = mx[:, idx, 1] - mx[:, idx, 0]                  # solvated - isolated
+                corr = np.mean(corr, axis=1)                          # average over the site's atoms
+                corr = corr[~np.isnan(corr)]                          # drop frames with missing data
+                vals[f"shell_{shell}"] = float(np.mean(corr)) if len(corr) else np.nan
+            rows.append(vals)
+    return pd.DataFrame(rows).set_index(["site", "solvent"])
+
+
+def experimental_stack(loader):
+    """Experimental shifts as a long DataFrame indexed (solute, nucleus, site, solvent)."""
+    exp = loader.experiment(uuid_prefix_h_sites=True).drop(columns=["atom_numbers"])
+    exp = exp.set_index(["solute", "nucleus", "site"])
+    stack = exp.stack(future_stack=True)
+    stack = pd.DataFrame(stack, columns=["experimental"])
+    stack.index.names = ["solute", "nucleus", "site", "solvent"]
+    return stack.sort_index()
+
+
+def _site_mean(arr, idx):
+    return float(np.asarray(arr)[idx].mean())
+
+
+def build_query_df_nn(loader):
+    """Assemble the per-(solute, nucleus, site, solvent) NN feature table that the composite-model
+    figures fit and plot (Figure 5C/5D, SI S15). Columns: experimental, stationary, pcm, qcd,
+    openMM, openMM_vib, plus the composite columns added by add_composite_columns
+    (stationary_plus_pcm/_qcd/_op_vib and their pcm+/openMM+ combinations). All inputs (MagNET
+    shieldings and the PCM conversion factors) are read through the loader.
+    """
+    sites = loader.site_atom_table(uuid_prefix_h_sites=True)
+    solvents = [c for c in loader.experiment().columns
+                if c not in ("solute", "site", "nucleus", "atom_numbers")]
+
+    # Cache the per-solute / per-(solute, solvent) loader reads. Each accessor opens the HDF5 file
+    # fresh, and the loops below touch only ~13 solutes (or ~52 solute-solvent pairs) across
+    # hundreds of sites, so without caching this re-opens the file thousands of times (~12x slower).
+    _mz_cache, _qcd_cache, _shell_cache, _mx_cache, _mxs_cache = {}, {}, {}, {}, {}
+    def get_mz(solute):
+        if solute not in _mz_cache: _mz_cache[solute] = loader.magnet_zero(solute)
+        return _mz_cache[solute]
+    def get_qcd(solute):
+        if solute not in _qcd_cache: _qcd_cache[solute] = loader.qcd(solute)
+        return _qcd_cache[solute]
+    def get_shells(solute, solvent):
+        k = (solute, solvent)
+        if k not in _shell_cache: _shell_cache[k] = loader.available_shells(solute, solvent)
+        return _shell_cache[k]
+    def get_mx(solute, solvent):
+        k = (solute, solvent)
+        if k not in _mx_cache: _mx_cache[k] = loader.magnet_x(solute, solvent, 650)
+        return _mx_cache[k]
+    def get_mxs(solute, solvent):
+        k = (solute, solvent)
+        if k not in _mxs_cache: _mxs_cache[k] = loader.magnet_x_stationary(solute, solvent)
+        return _mxs_cache[k]
+
+    # --- MagNET-Zero stationary + PCM (PCM predicted for chloroform, broadcast to all solvents) ---
+    mz_rows = []
+    for _, r in sites.iterrows():
+        idx = site_atom_indices(r["atom_numbers"])
+        mz = get_mz(r["solute"])
+        mz_rows.append([r["solute"], r["nucleus"], r["site"],
+                        _site_mean(mz["stationary"], idx), _site_mean(mz["pcm_correction"], idx)])
+    mz = pd.DataFrame(mz_rows, columns=["solute", "nucleus", "site", "stationary", "pcm"])
+    mz = pd.concat([mz.assign(solvent=s) for s in solvents], ignore_index=True)
+    mz = mz.set_index(["solute", "nucleus", "site", "solvent"]).sort_index()
+
+    # --- QCD correction (solvent-independent, tiled across solvents) ---
+    qcd_rows = []
+    for _, r in sites.iterrows():
+        idx = site_atom_indices(r["atom_numbers"])
+        q = get_qcd(r["solute"])
+        corr = q["trajectories"].mean(0).mean(0)[idx, 3] - q["stationary"][idx, 3]
+        qcd_rows.append([r["solute"], r["nucleus"], r["site"],
+                         float(corr.mean()) if len(corr) else np.nan])
+    qcd = pd.DataFrame(qcd_rows, columns=["solute", "nucleus", "site", "qcd"])
+    qcd = pd.concat([qcd.assign(solvent=s) for s in solvents], ignore_index=True)
+    qcd = qcd.set_index(["solute", "nucleus", "site", "solvent"]).sort_index()
+
+    # --- MagNET-X explicit-solvent correction + vibrational correction (per solvent, 650-atom shell) ---
+    ex_rows = []
+    for _, r in sites.iterrows():
+        idx = site_atom_indices(r["atom_numbers"])
+        for solvent in solvents:
+            shells = get_shells(r["solute"], solvent)
+            if 650 not in shells:
+                ex_rows.append([r["solute"], r["nucleus"], r["site"], solvent, np.nan, np.nan]); continue
+            mx = get_mx(r["solute"], solvent)        # (frames, atoms, [isolated, solvated])
+            corr = mx[:, idx, 1] - mx[:, idx, 0]
+            corr = np.mean(corr, axis=1)
+            corr = corr[~np.isnan(corr)]
+            openmm = float(np.mean(corr)) if len(corr) else np.nan
+            stat = get_mxs(r["solute"], solvent)
+            vib = float(np.mean(mx.mean(0)[idx, 0] - np.asarray(stat)[idx]))
+            ex_rows.append([r["solute"], r["nucleus"], r["site"], solvent, openmm, vib])
+    ex = pd.DataFrame(ex_rows, columns=["solute", "nucleus", "site", "solvent", "openMM", "openMM_vib"])
+    ex = ex.set_index(["solute", "nucleus", "site", "solvent"]).sort_index()
+
+    # --- combine + composite columns ---
+    comb = pd.concat([experimental_stack(loader), mz, ex, qcd], axis=1).reset_index()
+    comb = comb.set_index(["solute", "nucleus", "site", "solvent"])
+    # stationary_plus_pcm: scale the chloroform PCM correction to each solvent (factors from the loader,
+    # which renames water -> TIP4P to match the explicit-solvent naming)
+    conv = {}
+    for nuc in ["H", "C"]:
+        s = loader.pcm_conversion_factors(nuc).set_index("solvent")["pcm_conversion_factor"]
+        conv[nuc] = s.rename({"water": "TIP4P"})
+    nn = comb.reset_index()
+    factors = nn.apply(lambda r: conv[r["nucleus"]].get(r["solvent"], np.nan), axis=1)
+    nn["stationary_plus_pcm"] = nn["stationary"] + nn["pcm"] * factors
+    nn = nn.set_index(comb.index.names)
+    comb["stationary_plus_pcm"] = nn["stationary_plus_pcm"]
+    comb["stationary_plus_qcd"] = comb["stationary"] + comb["qcd"]
+    comb["stationary_plus_op_vib"] = comb["stationary"] + comb["openMM_vib"]
+    comb["stationary_plus_pcm_plus_qcd"] = comb["stationary_plus_pcm"] + comb["qcd"]
+    comb["stationary_plus_pcm_plus_op_vib"] = comb["stationary_plus_pcm"] + comb["openMM_vib"]
+    comb["stationary_plus_openMM_plus_qcd"] = comb["stationary"] + comb["openMM"] + comb["qcd"]
+    comb["stationary_plus_openMM_plus_op_vib"] = comb["stationary"] + comb["openMM"] + comb["openMM_vib"]
+    return comb.reset_index()
+
+
+# ----------------------------------------------------------------------------
+# Composite-model fitting (Figure 5C/5D, SI S15). The coefficients applied at scale come from the
+# loader's composite_model group (fit on delta-22).
+# ----------------------------------------------------------------------------
+
+def fit(fit_df, formula):
+    """OLS of experimental shift on a composite formula; returns (RMSE, params).
+
+    params is a pandas Series indexed by "Intercept" and the predictor names (matching statsmodels'
+    naming, which the composite-model coefficient code downstream relies on). This does the ordinary
+    least squares directly with numpy (np.linalg.lstsq), which is identical to the notebooks'
+    statsmodels OLS but far faster in this fitting loop: statsmodels re-parses the formula with patsy
+    on every call (~1000x slower per call). The `fit`-matches-`_fit_statsmodels` test guards the
+    equivalence.
+
+    Rows with a missing experimental value or a missing predictor are dropped before fitting and are
+    not scored: statsmodels drops those rows when fitting too, and its `predict` returns NaN for a row
+    with a missing predictor, so a single missing value never turns the whole RMSE into NaN.
+    """
+    terms = [term.strip() for term in formula.split("+") if term.strip()]
+    response = fit_df["experimental"].to_numpy(dtype=float)
+    columns = [fit_df[term].to_numpy(dtype=float) for term in terms]
+    design = np.column_stack([np.ones(len(fit_df))] + columns)
+    keep = np.isfinite(response) & np.all(np.isfinite(design), axis=1)
+    design, response = design[keep], response[keep]
+    beta, *_ = np.linalg.lstsq(design, response, rcond=None)
+    params = pd.Series(beta, index=["Intercept"] + terms)
+    return _rmse(design @ beta, response), params
+
+
+def _fit_statsmodels(fit_df, formula):
+    """The original statsmodels OLS, kept only as a test oracle for `fit` (see the equivalence test).
+    Slow because statsmodels re-parses the formula with patsy on every call; do not use in the harness."""
+    result = smf.ols(formula=f"experimental ~ {formula}", data=fit_df).fit()
+    fit_df = fit_df.copy()
+    fit_df["predicted"] = result.predict(fit_df)
+    scored = fit_df[fit_df["experimental"].notna() & fit_df["predicted"].notna()]
+    return _rmse(scored["predicted"], scored["experimental"]), result.params
+
+
+def fit_formulas_per_solvent(df, formulas, solvents, formula_remap):
+    """One OLS fit per (formula, solvent), pooled across all solutes; returns rmse and params rows."""
+    rows = []
+    for formula in formulas:
+        for solvent in solvents:
+            d = df[df["solvent"] == solvent].copy()
+            rmse, params = fit(d, formula)
+            rows.append({"formula": formula_remap.get(formula, formula), "solvent": solvent,
+                         "rmse": rmse, "params": params})
+    return pd.DataFrame(rows)
+
+
+def fit_formulas_per_solvent_and_solute(df, formulas, solvents, solutes, formula_remap):
+    """One OLS fit per (formula, solvent, solute); skips a group with under 2 rows or no
+    experimental data, and records the exception instead of raising if a fit fails."""
+    rows = []
+    for formula in formulas:
+        for solvent in solvents:
+            for solute in solutes:
+                d = df[(df["solvent"] == solvent) & (df["solute"] == solute)].copy()
+                if len(d) < 2 or d["experimental"].isnull().all():
+                    continue
+                try:
+                    rmse, params = fit(d, formula)
+                    rows.append({"formula": formula_remap.get(formula, formula), "solvent": solvent,
+                                 "solute": solute, "rmse": rmse, "params": params})
+                except Exception as e:  # noqa: BLE001 - mirror source behaviour
+                    rows.append({"formula": formula_remap.get(formula, formula), "solvent": solvent,
+                                 "solute": solute, "rmse": np.nan, "params": str(e)})
+    return pd.DataFrame(rows)
+
+
+def per_solvent_fits(query_df_nn):
+    """all_solute_fitting_results: full per-solvent OLS fits for both nuclei ('scaled to test set')."""
+    out = {}
+    formulas = {"H": FITTING_FORMULAS_H, "C": FITTING_FORMULAS_C}
+    solvents = list(query_df_nn["solvent"].unique())
+    for nuc in ["H", "C"]:
+        d = query_df_nn[query_df_nn["nucleus"] == nuc].copy()
+        out[nuc] = fit_formulas_per_solvent(d, formulas[nuc], solvents, FORMULA_REMAP)
+    return out
+
+
+def per_solute_fits(query_df_nn):
+    """per_solute_fitting_results: per-solvent-and-solute OLS fits ('scaled to solute' baseline)."""
+    out = {}
+    formulas = {"H": FITTING_FORMULAS_H, "C": FITTING_FORMULAS_C}
+    solvents = list(query_df_nn["solvent"].unique())
+    solutes = list(query_df_nn["solute"].unique())
+    for nuc in ["H", "C"]:
+        d = query_df_nn[query_df_nn["nucleus"] == nuc].copy()
+        out[nuc] = fit_formulas_per_solvent_and_solute(d, formulas[nuc], solvents, solutes, FORMULA_REMAP)
+    return out
+
+
+def build_bootstrap_seed_coeffs(loader):
+    """Per-nucleus DataFrame of bootstrap coefficients with a remapped 'formula' column, read from
+    the loader's composite_model group (the delta-22 bootstrap fits applied to the natural products)."""
+    out = {}
+    for nuc in ["H", "C"]:
+        frames = []
+        for token, name in _BOOTSTRAP_TOKENS[nuc].items():
+            df = loader.bootstrap_coefficients(token, nuc).copy()
+            df["formula"] = name
+            frames.append(df)
+        combined = pd.concat(frames, ignore_index=True, sort=False)
+        front = [c for c in ["solvent", "formula", "seed"] if c in combined.columns]
+        out[nuc] = combined[front + [c for c in combined.columns if c not in front]]
+    return out
+
+
+def apply_bootstrap_params_to_full_dataset(data_df, params_df, nucleus=None):
+    """Apply each bootstrap coefficient set (per formula/seed/solvent) to the full NP dataset,
+    returning predictions: solute, nucleus, site, solvent, formula, seed, experimental, predicted."""
+    coeffs_df = params_df.copy()
+    if nucleus is not None:
+        data_df = data_df[data_df["nucleus"] == nucleus]
+    keep = ["solute", "nucleus", "site", "solvent", "experimental"]
+    keep += [c for c in data_df.columns if c in coeffs_df.columns and c not in keep]
+    data_df = data_df[keep]
+    merged = coeffs_df.merge(data_df, how="outer", on="solvent", suffixes=("_coeff", "_data"))
+    idx = ["solute", "nucleus", "site", "solvent", "formula", "seed"]
+    merged = merged[idx + [c for c in merged.columns if c not in idx]]
+    # vectorized prediction: intercept (0 if missing) plus, for each parameter, coeff*data where the
+    # coeff is present (a missing coeff contributes nothing; a present coeff with missing data
+    # propagates NaN).
+    pred = merged["Intercept"].astype(float).fillna(0.0) if "Intercept" in merged.columns else 0.0
+    for cc in [c for c in merged.columns if c.endswith("_coeff")]:
+        dc = cc.replace("_coeff", "_data")
+        if dc not in merged.columns:
+            continue
+        term = merged[cc].astype(float) * merged[dc].astype(float)
+        pred = pred + term.where(merged[cc].notna(), 0.0)
+    merged["predicted"] = pred
+    return merged[idx + ["experimental", "predicted"]]
+
+
+def compute_solute_rmses(df, solute_col="solute"):
+    """Per (solvent, nucleus, formula, seed, solute) RMSE over a solute's sites."""
+    valid = df[df["experimental"].notnull() & df["predicted"].notnull()].copy()
+    # RMSE is sqrt(mean(squared error)), so a single vectorized groupby-mean over a precomputed
+    # squared-error column is identical to grouped.apply(rmse) but ~19x faster on the ~200k groups
+    # of a full bootstrap table (stats.rmse is exactly np.sqrt(np.mean(np.square(...)))).
+    valid["_squared_error"] = np.square(valid["predicted"] - valid["experimental"])
+    result = (valid.groupby(["solvent", "nucleus", "formula", "seed", solute_col])["_squared_error"]
+              .mean().pow(0.5).reset_index(name="Bootstrap_RMSE"))
+    return result.rename(columns={solute_col: "solute"}) if solute_col != "solute" else result
+
+
+def compute_grouped_rmse(df, solute_groups, solute_col="solute"):
+    """Per (solvent, nucleus, formula, seed, solute_group) RMSE over all sites in the group."""
+    lookup = {s: g for g, ss in solute_groups.items() for s in ss}
+    df = df.copy()
+    df["solute_group"] = df[solute_col].map(lookup).fillna("Unknown")
+    valid = df[df["experimental"].notnull() & df["predicted"].notnull()].copy()
+    # vectorized RMSE, same identity and speedup as compute_solute_rmses above
+    valid["_squared_error"] = np.square(valid["predicted"] - valid["experimental"])
+    return (valid.groupby(["solvent", "nucleus", "formula", "seed", "solute_group"])["_squared_error"]
+            .mean().pow(0.5).reset_index(name="Bootstrap_RMSE"))
+
+
+# ---------------------------------------------------------------------------------------------
+# SI Figure S15's "Fitting RMSE Comparisons (chloroform, H/C)" bar chart: per test-set solute,
+# three ways of choosing the composite-model coefficients -- "Scaled to Solute" (fit each solute
+# on its own), "Scaled to Test Set" (one fit pooled across all 12 test-set solutes), "Extrapolated
+# from delta22" (delta-22's bootstrap fits applied cold, never seeing the test set). Built from
+# per_solute_fits (Scaled to Solute), per_solvent_fits (Scaled to Test Set, split back out per
+# solute by scaled_to_test_set_per_solute_rmse below), and apply_bootstrap_params_to_full_dataset +
+# compute_solute_rmses (Extrapolated from delta22, the same data Figure 5D's boxplot uses).
+
+def scaled_to_test_set_per_solute_rmse(query_df_nn, all_solute_fits, nucleus, solvent, formula):
+    """per_solvent_fits' output has one pooled rmse per (formula, solvent): the model is fit once
+    across every test-set solute together. To get a comparable per-solute rmse (needed to sit
+    "Scaled to Test Set" next to "Scaled to Solute" and "Extrapolated from delta22" in the same bar
+    chart), this applies that one pooled fit's coefficients to each solute's own sites and scores
+    them individually, reusing apply_bootstrap_params_to_full_dataset's prediction machinery by
+    reshaping the fitted statsmodels Params Series into that function's expected wide
+    "<term>_coeff" row format."""
+    remapped = FORMULA_REMAP.get(formula, formula)
+    fits = all_solute_fits[nucleus]
+    match = fits[(fits["formula"] == remapped) & (fits["solvent"] == solvent)]
+    if match.empty:
+        raise KeyError(f"no per_solvent_fits row for formula={remapped!r} solvent={solvent!r}")
+    params = match.iloc[0]["params"]
+    # bare term names, matching data_df's own raw column names -- apply_bootstrap_params_to_full_
+    # dataset's merge(..., suffixes=("_coeff", "_data")) is what appends "_coeff"/"_data" (only
+    # triggered by the name COLLIDING with data_df's column of the same name); pre-suffixing here
+    # would just create an unmatched "<term>_coeff" column no data_df column collides with.
+    wide = {"Intercept": float(params.get("Intercept", 0.0)), "solvent": solvent,
+           "formula": remapped, "seed": 0}
+    for term in params.index:
+        if term != "Intercept":
+            wide[term] = float(params[term])
+    params_df = pd.DataFrame([wide])
+    # Restrict to this one solvent before the merge: apply_bootstrap_params_to_full_dataset's outer
+    # join zero-fills a missing coefficient instead of raising, so an unfiltered call would
+    # silently score every OTHER solvent's rows against this solvent's plane as if the missing
+    # coefficients were all zero.
+    solvent_only = query_df_nn[query_df_nn["solvent"] == solvent]
+    preds = apply_bootstrap_params_to_full_dataset(solvent_only, params_df, nucleus=nucleus)
+    return compute_solute_rmses(preds)
+
+
+def fitting_rmse_comparison_table(query_df_nn, per_solute_fits_result, per_solvent_fits_result,
+                                  bootstrap_rmses, nucleus, solvent, formula):
+    """Assembles SI Figure S15's "Fitting RMSE Comparisons" table: one row per test-set solute,
+    columns "Scaled to Solute", "Scaled to Test Set", "Extrapolated from delta22" (ppm RMSE), all
+    at one nucleus/solvent/formula. bootstrap_rmses is compute_solute_rmses' output for `nucleus`
+    (e.g. Figure 5D's bootstrap_rmses_h), averaged here over its bootstrap seeds."""
+    remapped = FORMULA_REMAP.get(formula, formula)
+
+    solute_fit = per_solute_fits_result[nucleus]
+    scaled_to_solute = (solute_fit[(solute_fit["formula"] == remapped) & (solute_fit["solvent"] == solvent)]
+                        .set_index("solute")["rmse"])
+
+    scaled_to_test_set = (scaled_to_test_set_per_solute_rmse(query_df_nn, per_solvent_fits_result, nucleus, solvent, formula)
+                          .query("solvent == @solvent and formula == @remapped")
+                          .set_index("solute")["Bootstrap_RMSE"])
+
+    boot = bootstrap_rmses[(bootstrap_rmses["nucleus"] == nucleus) & (bootstrap_rmses["solvent"] == solvent)
+                           & (bootstrap_rmses["formula"] == remapped)]
+    extrapolated = boot.groupby("solute")["Bootstrap_RMSE"].mean()
+
+    table = pd.DataFrame({"Scaled to Solute": scaled_to_solute, "Scaled to Test Set": scaled_to_test_set,
+                          "Extrapolated from delta22": extrapolated})
+    return table.dropna(how="all")
+
+
+def distribution_shift_by_solvent_table(per_solute_fits_result, per_solvent_fits_result,
+                                        bootstrap_rmses, nucleus, formula, solvents=None):
+    """SI Figure S15's per-solvent "distribution shift is minor" cross-check -- the solvent-averaged
+    companion to fitting_rmse_comparison_table (which is per-solute for a single solvent). For one
+    nucleus and formula, the mean test-set RMSE under the three ways of choosing the composite-model
+    coefficients, one row per solvent: "Extrapolated from delta22" (delta-22's bootstrap
+    coefficients applied cold), "Scaled to Test Set" (one fit pooled across the whole test set), and
+    "Scaled to Solute" (a separate fit per solute). Where the three sit close, re-optimizing the
+    coefficients barely helps, so they transfer and the residual error is physics; the one solvent
+    where "Extrapolated" rises well above the refits is TIP4P (water), dominated by the peptide's
+    conformer-population sensitivity, which is why S15 reports solvent-averaged performance.
+
+    bootstrap_rmses is compute_solute_rmses' output for `nucleus` (e.g. the si_figure_s15 notebook's
+    bootstrap_rmses["H"]); per_solute_fits_result / per_solvent_fits_result are the per_solute_fits /
+    per_solvent_fits dicts keyed by nucleus. Averages over bootstrap seeds and over each solvent's
+    solutes."""
+    remapped = FORMULA_REMAP.get(formula, formula)
+    boot = bootstrap_rmses[(bootstrap_rmses["nucleus"] == nucleus)
+                           & (bootstrap_rmses["formula"] == remapped)]
+    test_set = per_solvent_fits_result[nucleus]
+    test_set = test_set[test_set["formula"] == remapped]
+    solute = per_solute_fits_result[nucleus]
+    solute = solute[solute["formula"] == remapped]
+    if solvents is None:
+        solvents = sorted(boot["solvent"].unique())
+    rows = []
+    for solvent in solvents:
+        rows.append({
+            "solvent": solvent,
+            "Extrapolated from delta22": boot[boot["solvent"] == solvent]["Bootstrap_RMSE"].mean(),
+            "Scaled to Test Set": test_set[test_set["solvent"] == solvent]["rmse"].mean(),
+            "Scaled to Solute": solute[solute["solvent"] == solvent]["rmse"].mean(),
+        })
+    return pd.DataFrame(rows).set_index("solvent")
+
+
+# ---------------------------------------------------------------------------------------------
+# SI Figure S15's "Feature Space Coverage by Solvent" and "Residuals for Delta22 Fitting
+# Coefficients" panels: do the test set and delta-22 cover the same region of feature space (so
+# the test set's higher error can't be blamed on extrapolating outside what delta-22's fitting
+# coefficients ever saw)? Both panels combine the test set (this module's query_df_nn) with
+# delta-22's own per-site data at the same two features (delta22.py's query_df_nn, passed in by
+# the caller so this module doesn't import delta22.py directly).
+
+# Feature Space Coverage only compares the 4 solvents both datasets have explicit-solvent (OpenMM)
+# coverage for -- delta-22 has Desmond+OpenMM MD for all 12 solvents, but the natural-products test
+# set (sigma-fresh-derived) only has 4.
+FEATURE_SPACE_SOLVENTS = ("chloroform", "benzene", "methanol", "TIP4P")
+
+
+def _feature_space_config(nucleus):
+    if nucleus == "H":
+        return {"x_feature": "stationary_plus_qcd", "y_feature": "openMM"}
+    if nucleus == "C":
+        return {"x_feature": "stationary_plus_op_vib", "y_feature": "openMM"}
+    raise ValueError(f"unsupported nucleus {nucleus!r}, expected 'H' or 'C'")
+
+
+def _test_set_and_delta22(query_df_nn, delta22_query_df_nn, nucleus, cols):
+    """The shared setup both panels below need: the test set's own rows (excludes any row
+    literally named "delta22", a defensive filter) tagged "Test Set", concatenated with
+    delta22_query_df_nn's rows (already delta-22-only) tagged "Delta22"."""
+    test = query_df_nn[query_df_nn["nucleus"] == nucleus].copy()
+    test = test[~test["solute"].astype(str).str.strip().str.lower().eq("delta22")]
+    test["dataset"] = "Test Set"
+    d22 = delta22_query_df_nn[delta22_query_df_nn["nucleus"] == nucleus].copy()
+    d22["dataset"] = "Delta22"
+    return pd.concat([test[cols + ["dataset"]], d22[cols + ["dataset"]]], ignore_index=True)
+
+
+def feature_space_coverage_table(query_df_nn, delta22_query_df_nn, nucleus, solvents=FEATURE_SPACE_SOLVENTS):
+    """SI Figure S15's "Feature Space Coverage by Solvent": test-set and delta-22 per-site feature
+    values, mean-centered globally -- one scalar mean per feature, pooled over every solvent and
+    both datasets (the notebook's separate fitting cell instead centers per solvent). Returns a
+    tidy DataFrame: solvent, dataset ("Test Set"/"Delta22"), x (centered), y (centered)."""
+    cfg = _feature_space_config(nucleus)
+    x_feature, y_feature = cfg["x_feature"], cfg["y_feature"]
+    combined = _test_set_and_delta22(query_df_nn, delta22_query_df_nn, nucleus, ["solvent", x_feature, y_feature])
+    combined = combined[combined["solvent"].isin(solvents)].dropna(subset=[x_feature, y_feature])
+    combined["x"] = combined[x_feature] - combined[x_feature].mean()
+    combined["y"] = combined[y_feature] - combined[y_feature].mean()
+    return combined[["solvent", "dataset", "x", "y"]].reset_index(drop=True)
+
+
+def delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, nucleus, solvents=FEATURE_SPACE_SOLVENTS):
+    """SI Figure S15's "Residuals for Delta22 Fitting Coefficients": fits a 2-feature OLS plane
+    (experimental ~ x_feature + y_feature, the same two features feature_space_coverage_table
+    plots) to delta-22's data alone, separately per solvent, then applies that one plane to both
+    delta-22 and the test set, returning residual = experimental - predicted. For a linear model,
+    minimizing RMSE directly (Nelder-Mead) is the same convex problem as ordinary least squares, so
+    np.linalg.lstsq gives the identical plane. Returns a tidy DataFrame: solvent, dataset,
+    experimental, residual."""
+    cfg = _feature_space_config(nucleus)
+    x_feature, y_feature = cfg["x_feature"], cfg["y_feature"]
+    combined = _test_set_and_delta22(query_df_nn, delta22_query_df_nn, nucleus,
+                                     ["solvent", x_feature, y_feature, "experimental"])
+    combined = combined[combined["solvent"].isin(solvents)].dropna(subset=[x_feature, y_feature, "experimental"])
+
+    rows = []
+    for solvent in solvents:
+        sdf = combined[combined["solvent"] == solvent]
+        d22_sdf = sdf[sdf["dataset"] == "Delta22"]
+        if len(d22_sdf) < 3:
+            continue
+        design = np.column_stack([np.ones(len(d22_sdf)), d22_sdf[x_feature], d22_sdf[y_feature]])
+        c0, c1, c2 = np.linalg.lstsq(design, d22_sdf["experimental"].to_numpy(), rcond=None)[0]
+        predicted = c0 + c1 * sdf[x_feature] + c2 * sdf[y_feature]
+        out = sdf[["solvent", "dataset", "experimental"]].copy()
+        out["residual"] = sdf["experimental"] - predicted
+        rows.append(out)
+    if not rows:
+        return pd.DataFrame(columns=["solvent", "dataset", "experimental", "residual"])
+    return pd.concat(rows, ignore_index=True).reset_index(drop=True)

analysis/code/applications_plots.py

@@ -0,0 +1,752 @@
+"""Plotting for the natural-products figures (main-text Figures 5C/5D, SI Figure S15).
+
+Holds every plotting engine for the applications (natural-products) notebooks: the delta-22-vs-test-set
+benefit bar chart (Figure 5C), the shared per-solute box-plot engine (Figure 5D and SI Figure S15's
+bootstrap-RMSE panels), and SI Figure S15's fitting-RMSE comparison bars, feature-space/residual scatter
+grids, and distribution-shift bars. Engines are carried over from the source notebooks with notebook
+globals turned into explicit arguments. The numbers come from applications.py; this module is kept
+separate so applications.py stays plotting-free.
+"""
+import os
+
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+
+def plot_nps_on_boxplot_delta22_simplified(
+    regression_ss_df,
+    nps_rmse_df,
+    best_possible_df,  # expects a 'rmse' column
+    nucleus,
+    formulas,
+    colors,
+    site_counts,
+    solvents_filter=None,
+    box_width=0.20,
+    box_gap=0.05,
+    formula_gap=0.15,
+    figsize=(14, 8),
+    formula_remap=None,
+    solute_remap=None,
+    solute_order=None,
+    title=None,
+    save_path=None,
+    solute_color_remap=None,
+    max_bar_height=0.05,
+
+    # --- SELF-SCALING BASELINE CONTROLS ---
+    show_baseline=True,
+    baseline_annotation_text='Self-scaling baseline (dotted)',
+    baseline_annotation_x=0.985,
+    baseline_annotation_y=0.93,
+    baseline_annotation_fontsize=9,
+    
+    # --- FULL FIT LINE CONTROLS ---
+    show_full_fit_line=False,
+    all_solute_fitting_results=None,  # {nucleus: per-solvent fit df}; required when show_full_fit_line
+    full_fit_line_label='Full Fit RMSE',
+    full_fit_line_label_x=0.99,
+
+    # --- INSET AXIS CONTROLS ---
+    site_count_axis_mode='attached',
+    site_count_inset_axis_offset=0.0,
+    site_count_inset_area_frac=0.1,
+    site_count_inset_axis_label_pad=0.04,
+    site_count_alpha=0.25,
+    site_count_label='Site count',
+    site_count_preserve_visual_height=True,
+
+    # ---- font sizes ----
+    xtick_fontsize=10,
+    delta22_fontsize=10,
+    ylabel_fontsize=13,
+    title_fontsize=16,
+
+    # ---- SAVE SVG ----
+    save_svg_path=None,
+    default_svg_name="Figure_5C_extras.svg",
+    svg_transparent=True,
+):
+    """
+    Simplified box and whisker plot:
+    - Plots all solutes separately (no grouping)
+    - Averages RMSE across solvents for each bootstrap seed
+    - Shows "best possible" RMSE as dashed baseline for each solute
+
+    Parameters:
+    - regression_ss_df: Delta-22 regression samples (Bootstrap_RMSE)
+    - nps_rmse_df: Should have columns: solvent, nucleus, formula, seed, solute, Bootstrap_RMSE
+    - best_possible_df: DataFrame with best possible RMSEs per solute/solvent/formula
+                       Should have columns: formula, solvent, solute, rmse, params
+    - nucleus: Which nucleus to filter for
+    - formulas: List of formulas to compare
+    - colors: Colors for each formula
+    - site_counts: DataFrame indexed by solute with nucleus count columns (e.g., 'H', 'C')
+    - solvents_filter: Optional list of solvents to include in averaging
+    - show_baseline: Whether to show the 'best possible df' baseline
+    - show_full_fit_line: Whether to draw a line representing the full-dataset fit RMSE across all solutes
+    - full_fit_line_label: Label for the full fit line
+    - full_fit_line_label_x: X-coordinate for the full fit line label
+    - box_width: Width of each box
+    - box_gap: Gap between boxes
+    - formula_gap: Gap between formulas
+    - figsize: Figure size
+    - formula_remap: Optional dict to remap formula names
+    - max_bar_height: Maximum visual height target for site-count bars relative to RMSE axis
+    - site_count_axis_mode: One of {'attached', 'inset', 'off'}
+    - site_count_inset_axis_offset: Horizontal inset-axis adjustment in axes coordinates (clipped to stay inside in inset mode)
+    - site_count_inset_area_frac: Fraction of total figure area to allocate to inset axis (inset mode only)
+    - site_count_alpha: Alpha for site-count bars
+    - site_count_label: Label for site-count axis
+    - site_count_preserve_visual_height: Keep the old visual bar footprint when using a right axis
+    """
+    if formula_remap is None:
+        formula_remap = {}
+
+    if solute_remap is None:
+        solute_remap = {}
+
+    if solute_color_remap is None:
+        solute_color_remap = {}
+
+    valid_axis_modes = {'attached', 'inset', 'off'}
+    if site_count_axis_mode not in valid_axis_modes:
+        raise ValueError(f"site_count_axis_mode must be one of {sorted(valid_axis_modes)}")
+
+    if site_counts is not None and not isinstance(site_counts, pd.DataFrame):
+        raise TypeError("site_counts must be a pandas DataFrame or None")
+
+    if site_count_inset_area_frac < 0 or site_count_inset_area_frac > 1:
+        raise ValueError("site_count_inset_area_frac must be between 0 and 1")
+
+    def get_rounded_count_ticks(max_count, target_ticks=7):
+        if max_count <= 0:
+            return np.array([0.0, 1.0]), 1.0
+
+        rough_step = max_count / max(target_ticks - 1, 1)
+        magnitude = 10 ** np.floor(np.log10(rough_step)) if rough_step > 0 else 1.0
+        normalized = rough_step / magnitude
+        if normalized <= 1:
+            nice_base = 1
+        elif normalized <= 2:
+            nice_base = 2
+        elif normalized <= 5:
+            nice_base = 5
+        else:
+            nice_base = 10
+
+        step = nice_base * magnitude
+        rounded_top = np.ceil(max_count / step) * step
+        ticks = np.arange(0, rounded_top + 0.5 * step, step)
+        return ticks, rounded_top
+
+    n_formulas = len(formulas)
+
+    if solvents_filter is not None:
+        regression_ss_df = regression_ss_df[regression_ss_df['solvent'].isin(solvents_filter)].copy()
+        nps_rmse_df = nps_rmse_df[nps_rmse_df['solvent'].isin(solvents_filter)].copy()
+        best_possible_df = best_possible_df[best_possible_df['solvent'].isin(solvents_filter)].copy()
+
+        if regression_ss_df.empty:
+            raise ValueError(f"No regression data found for solvents: {solvents_filter}")
+        if nps_rmse_df.empty:
+            raise ValueError(f"No NPS RMSE data found for solvents: {solvents_filter}")
+
+    def get_delta22_averaged_rmse(regression_df, formula):
+        df = regression_df[regression_df['formula'] == formula]
+        return df.groupby('seed')['Bootstrap_RMSE'].mean().values
+
+    def get_solute_bootstrap_averaged(nps_df, formula, nucleus):
+        np_formula = formula_remap.get(formula, formula)
+        df = nps_df[
+            (nps_df['formula'] == np_formula) &
+            (nps_df['nucleus'] == nucleus)
+        ]
+        if df.empty:
+            return {}
+
+        # average over solvent per solute/seed, then collect each solute's per-seed array
+        averaged_df = df.groupby(['solute', 'seed'])['Bootstrap_RMSE'].mean().reset_index()
+        solute_dict = {}
+        for solute in averaged_df['solute'].unique():
+            arr = averaged_df[averaged_df['solute'] == solute]['Bootstrap_RMSE'].values
+            if arr.size > 0:
+                solute_dict[solute] = arr
+        return solute_dict
+
+    def get_best_possible_averaged(best_df, solute, formula):
+        np_formula = formula_remap.get(formula, formula)
+        df = best_df[
+            (best_df['formula'] == np_formula) &
+            (best_df['solute'] == solute)
+        ]
+
+        if df.empty:
+            # fall back to the raw formula name in case best_possible_df wasn't remapped
+            df = best_df[
+                (best_df['formula'] == formula) &
+                (best_df['solute'] == solute)
+            ]
+
+        if df.empty:
+            return None
+        return df['rmse'].mean()
+
+    all_solutes = set()
+    for formula in formulas:
+        np_formula = formula_remap.get(formula, formula)
+        df_filtered = nps_rmse_df[
+            (nps_rmse_df['formula'] == np_formula) &
+            (nps_rmse_df['nucleus'] == nucleus)
+        ]
+        all_solutes.update(df_filtered['solute'].unique())
+    all_solutes = sorted(list(all_solutes))
+    if solute_order is not None:                 # explicit ordering (e.g. main-text Figure 5D)
+        ordered = [s for s in solute_order if s in all_solutes]
+        all_solutes = ordered + [s for s in all_solutes if s not in ordered]
+
+    delta22_by_formula = [get_delta22_averaged_rmse(regression_ss_df, f) for f in formulas]
+
+    # each formula block is 1 delta-22 box plus one box per solute
+    total_boxes_per_formula = 1 + len(all_solutes)
+
+    all_data = []
+    all_positions = []
+    all_colors = []
+    all_alphas = []
+    tick_positions = []
+    tick_labels = []
+    baseline_segments = []  # Store (x_start, x_end, y_value, solute_label) for baseline segments
+
+    count_bar_positions = []
+    count_bar_values = []
+    delta22_count_by_nucleus = {'H': 43.0, 'C': 50.0}
+
+    current_position = 0.0
+
+    for j, formula in enumerate(formulas):
+
+        # 1) Delta-22 box (solvent-averaged)
+        box_pos = current_position + box_width / 2
+        all_data.append(delta22_by_formula[j])
+        all_positions.append(box_pos)
+        all_colors.append(colors[j])
+        all_alphas.append(1.0)
+        tick_positions.append(box_pos)
+        tick_labels.append('delta-22')
+
+        delta22_count = delta22_count_by_nucleus.get(nucleus)
+        if delta22_count is not None:
+            count_bar_positions.append(box_pos)
+            count_bar_values.append(float(delta22_count))
+
+        current_position += box_width + box_gap
+
+        # 2) Individual solute boxes (solvent-averaged, lighter alpha)
+        solute_data = get_solute_bootstrap_averaged(nps_rmse_df, formula, nucleus)
+
+        count_col = None
+        if site_counts is not None:
+            if nucleus in site_counts.columns:
+                count_col = nucleus
+            elif 'H' in site_counts.columns:
+                count_col = 'H'
+
+        for solute in all_solutes:
+            box_pos = current_position + box_width / 2
+            if solute in solute_data and len(solute_data[solute]) > 0:
+                all_data.append(solute_data[solute])
+                all_positions.append(box_pos)
+                all_colors.append(solute_color_remap.get(solute, colors[j]))
+                all_alphas.append(0.5)
+                tick_positions.append(box_pos)
+                tick_labels.append(solute_remap.get(solute, solute))
+
+                if (
+                    count_col is not None
+                    and solute in site_counts.index
+                    and pd.notna(site_counts.loc[solute, count_col])
+                ):
+                    count_bar_positions.append(box_pos)
+                    count_bar_values.append(float(site_counts.loc[solute, count_col]))
+                else:
+                    count_bar_positions.append(box_pos)
+                    count_bar_values.append(0.0)
+
+                # best_possible_df has no nucleus column, so this doesn't filter by nucleus
+                if show_baseline:
+                    best_rmse = get_best_possible_averaged(best_possible_df, solute, formula)
+                    if best_rmse is not None:
+                        # baseline segment spans the box's full width including its gap
+                        x_start = current_position - box_gap / 2
+                        x_end = current_position + box_width + box_gap / 2
+                        baseline_segments.append((x_start, x_end, best_rmse, solute))
+
+            current_position += box_width + box_gap
+
+        # Add gap after formula (except last)
+        if j < n_formulas - 1:
+            current_position += formula_gap
+
+    plot_figsize = figsize
+    fig, ax = plt.subplots(figsize=plot_figsize)
+
+    # skip boxes with no data
+    valid_indices = [i for i, data in enumerate(all_data) if len(data) > 0]
+    valid_data = [all_data[i] for i in valid_indices]
+    valid_positions = [all_positions[i] for i in valid_indices]
+    valid_colors = [all_colors[i] for i in valid_indices]
+    valid_alphas = [all_alphas[i] for i in valid_indices]
+
+    if valid_data:
+        box = ax.boxplot(
+            valid_data,
+            positions=valid_positions,
+            widths=box_width,
+            showfliers=False,
+            patch_artist=True,
+            whiskerprops=dict(color='#404040'),
+            capprops=dict(color='#404040')
+        )
+
+        for patch, color, alpha in zip(box['boxes'], valid_colors, valid_alphas):
+            patch.set_facecolor(color)
+            patch.set_alpha(alpha)
+        for median in box['medians']:
+            median.set_color('#404040')
+
+    if show_baseline:
+        for i, (x_start, x_end, y_value, _) in enumerate(baseline_segments):
+            ax.plot([x_start, x_end], [y_value, y_value],
+                    color='gray', linestyle='--', linewidth=1.5, alpha=0.7)
+
+            # connect to the previous segment unless this is the first solute of a formula block
+            # (connects across formula gaps too, not just within one)
+            if len(formulas) > 0 and len(baseline_segments) > 0 and i % (len(baseline_segments) / len(formulas)) != 0:
+                _, prev_x_end, prev_y, _ = baseline_segments[i - 1]
+                ax.plot([prev_x_end, x_start], [prev_y, y_value],
+                        color='gray', linestyle='--', linewidth=1.5, alpha=0.7)
+
+    if show_baseline and baseline_segments and baseline_annotation_text:
+        ax.text(
+            baseline_annotation_x,
+            baseline_annotation_y,
+            baseline_annotation_text,
+            transform=ax.transAxes,
+            ha='right',
+            va='top',
+            fontsize=baseline_annotation_fontsize,
+            color='gray',
+            style='italic',
+            alpha=0.85,
+        )
+
+    # the full-fit-line RMSE uses the first formula's delta-22 fit as its reference
+    if show_full_fit_line:
+        formula_remapped = formula_remap.get(formulas[0], formulas[0])
+        query = f'formula == "{formula_remapped}"' + (f' and solvent in {solvents_filter}' if solvents_filter is not None else '')
+        full_fit_rmse = all_solute_fitting_results[nucleus].query(query)['rmse'].mean()
+        ax.axhline(y=full_fit_rmse, color='red', linestyle='-.', linewidth=1.5, alpha=0.7)
+        ax.text(full_fit_line_label_x,
+                full_fit_rmse, f'{full_fit_line_label}: {full_fit_rmse:.3f} ppm',
+                transform=ax.get_yaxis_transform(),
+                ha='right', va='bottom', fontsize=9, color='red', style='italic', alpha=0.7)
+
+    ax.set_xticks(tick_positions)
+    ax.set_xticklabels(tick_labels, rotation=75, fontsize=xtick_fontsize, ha='center')
+    for i, label_obj in enumerate(ax.get_xticklabels()):
+        if tick_labels[i] == 'delta-22':
+            label_obj.set_weight('bold')
+            label_obj.set_fontsize(delta22_fontsize)
+            label_obj.set_style('italic')
+
+    # vertical separators between formula blocks
+    for j in range(1, n_formulas):
+        x_pos = j * (total_boxes_per_formula * box_width + formula_gap) - formula_gap / 2
+        ax.axvline(x=x_pos, color='gray', linestyle='--', linewidth=1, alpha=0.5)
+
+    if solvents_filter is not None:
+        solvents_str = ', '.join(solvents_filter)
+        ylabel = f"1H RMSE (ppm)\n(Averaged across: {solvents_str})"
+    else:
+        ylabel = "1H RMSE (ppm)"
+    ax.set_ylabel(ylabel, fontsize=ylabel_fontsize, fontweight='bold')
+    ax.set_ylim(bottom=0)
+
+    # caller may override via title=; the default text reflects the solvent filter
+    if title is None:
+        if solvents_filter is not None:
+            title = f"Delta22 vs Individual Solute Bootstrap Sample RMSEs ({nucleus} nucleus; {solvents_filter})"
+        else:
+            title = f"Delta22 vs Individual Solute Bootstrap Sample RMSEs ({nucleus} nucleus, solvent-averaged)"
+    ax.set_title(title, fontsize=title_fontsize, fontweight='bold')
+
+    ax.set_xlim(-0.2, current_position + 0.2)
+
+    # site-count bars use a dedicated right (twin) axis
+    added_site_count_bars = False
+    site_count_ax = None
+    max_count = max(count_bar_values) if count_bar_values else 0.0
+    if site_count_axis_mode != 'off' and max_bar_height > 0 and len(count_bar_positions) > 0 and max_count > 0:
+        site_count_ax = ax.twinx()
+
+        tick_counts, rounded_count_top = get_rounded_count_ticks(max_count)
+        if site_count_preserve_visual_height:
+            rmse_top = ax.get_ylim()[1]
+            rmse_top = max(rmse_top, 1e-6)
+            bar_target = max(max_bar_height, 1e-6)
+            # Keep bars in the lower band while right-axis tick labels remain true counts.
+            axis_top = max(rounded_count_top * (rmse_top / bar_target), rounded_count_top * 1.05)
+        else:
+            axis_top = max(rounded_count_top, 1.0)
+
+        site_count_ax.bar(
+            count_bar_positions,
+            count_bar_values,
+            width=box_width + box_gap,
+            bottom=0,
+            color='#404040',
+            alpha=site_count_alpha,
+            edgecolor='none',
+            zorder=0,
+        )
+
+        site_count_ax.set_ylim(0, axis_top)
+
+        if site_count_axis_mode == 'inset':
+            # Hide inset twin-axis renderer and draw a short attached axis manually.
+            for spine_name in ('left', 'right', 'top', 'bottom'):
+                site_count_ax.spines[spine_name].set_visible(False)
+            site_count_ax.patch.set_visible(False)
+            site_count_ax.yaxis.set_visible(False)
+
+            inset_axis_top_frac = min(rounded_count_top / axis_top, 1.0)
+            # clip to keep the inset axis inside the plotting area near the right edge
+            inset_x = float(np.clip(1.0 + site_count_inset_axis_offset, 0.0, 1.0))
+
+            ax.plot(
+                [inset_x, inset_x],
+                [0.0, inset_axis_top_frac],
+                transform=ax.transAxes,
+                color='#404040',
+                linewidth=1.0,
+                clip_on=False,
+                zorder=3,
+            )
+
+            # Manual ticks and tick labels, confined to the short segment.
+            tick_len = 0.008
+            tick_label_pad = 0.012
+            for t in tick_counts:
+                if t == 0:
+                    continue
+                y_frac = t / axis_top
+                if y_frac <= inset_axis_top_frac + 1e-9:
+                    ax.plot(
+                        [inset_x, inset_x + tick_len],
+                        [y_frac, y_frac],
+                        transform=ax.transAxes,
+                        color='#404040',
+                        linewidth=0.9,
+                        clip_on=False,
+                        zorder=3,
+                    )
+                    ax.text(
+                        inset_x + tick_len + tick_label_pad,
+                        y_frac,
+                        f"{int(t)}",
+                        transform=ax.transAxes,
+                        va='center',
+                        ha='left',
+                        fontsize=10,
+                        color='#404040',
+                        clip_on=False,
+                    )
+
+            label_x = inset_x + site_count_inset_axis_label_pad
+            ax.text(
+                label_x,
+                inset_axis_top_frac / 2,
+                site_count_label,
+                transform=ax.transAxes,
+                rotation=90,
+                va='center',
+                ha='left',
+                fontsize=11,
+                fontweight='bold',
+                color='#404040',
+                clip_on=False,
+            )
+        else:
+            site_count_ax.set_yticks(tick_counts)
+            site_count_ax.set_yticklabels([f"{int(t)}" for t in tick_counts])
+            site_count_ax.set_ylabel(site_count_label, fontsize=13, fontweight='bold', color='#404040')
+            site_count_ax.tick_params(axis='y', labelcolor='#404040')
+            site_count_ax.spines['right'].set_color('#404040')
+
+        site_count_ax.grid(False)
+
+        # Keep RMSE artists visually above site-count bars.
+        ax.set_zorder(2)
+        ax.patch.set_alpha(0)
+        site_count_ax.set_zorder(1)
+        added_site_count_bars = True
+
+
+    if site_count_axis_mode == 'inset' and added_site_count_bars:
+        ax.set_xlim(left=-box_width/2-box_gap/2, right=4.55)
+    plt.tight_layout()
+    # PNG is the release format; SVG is also supported for standalone vector export
+    if save_path is not None:
+        outdir = os.path.dirname(save_path)
+        if outdir:
+            os.makedirs(outdir, exist_ok=True)
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+
+    if save_svg_path is not None:
+        # If user passed a directory, save into it with default name
+        if os.path.isdir(save_svg_path) or save_svg_path.endswith(os.sep):
+            save_svg_path = os.path.join(save_svg_path, default_svg_name)
+
+        if not save_svg_path.lower().endswith(".svg"):
+            save_svg_path = save_svg_path + ".svg"
+
+        outdir = os.path.dirname(save_svg_path)
+        if outdir:
+            os.makedirs(outdir, exist_ok=True)
+
+        fig.savefig(
+            save_svg_path,
+            format="svg",
+            bbox_inches="tight",
+            pad_inches=0.02,
+            transparent=svg_transparent,
+        )
+        print(f"Saved SVG to: {os.path.abspath(save_svg_path)}")
+
+    plt.show()
+
+
+def plot_nps_benefit_barplot(regression_ss_df, nps_rmse_df, nucleus, solvents, np_solute_groups,
+                             formulas, labels, colors, formula_remap=None, figsize=(6, 5),
+                             box_width=0.18, formula_gap=0.04, solvent_gap=0.35, spacer_width=0.0,
+                             y_min=0.0, y_max=None, save_path=None):
+    """Grouped bar chart of mean bootstrap RMSE per (solvent, dataset, formula), with 2.5-97.5
+    percentile bootstrap error bars. Delta-22 is drawn first within each solvent, then the
+    NP/complex bin. Used by main-text Figure 5C (implicit vs explicit correction, delta-22 vs the
+    test set)."""
+    sns.set_theme(style="ticks", context="paper")
+    formula_remap = formula_remap or {}
+    n_formulas = len(formulas)
+    label_remap = {labels[0]: "Implicit", labels[1]: "Explicit"}
+    label_to_color = {label_remap.get(labels[i], labels[i]): colors[i] for i in range(len(labels))}
+    flat_solvents = [s for g in solvents for s in g] if isinstance(solvents[0], list) else list(solvents)
+    solvent_groups_iter = solvents if isinstance(solvents[0], list) else [solvents]
+
+    def d22_by_solvent(formula):
+        df = regression_ss_df[regression_ss_df["formula"] == formula]
+        return {s: df[df["solvent"] == s]["Bootstrap_RMSE"].dropna().values for s in flat_solvents}
+
+    def np_grouped(formula, solvent):
+        df = nps_rmse_df[(nps_rmse_df["formula"] == formula) & (nps_rmse_df["solvent"] == solvent)
+                         & (nps_rmse_df["nucleus"] == nucleus)]
+        if df.empty:
+            return {}
+        out = {}
+        for g in np_solute_groups.keys():
+            arr = df[df["solute_group"] == g]["Bootstrap_RMSE"].dropna().values
+            if arr.size:
+                out[g] = arr
+        return out
+
+    def summ(vals):
+        vals = np.asarray(vals, float); vals = vals[~np.isnan(vals)]
+        if vals.size == 0:
+            return None
+        m = np.mean(vals); lo, hi = np.percentile(vals, [2.5, 97.5])
+        return m, m - lo, hi - m
+
+    d22 = [d22_by_solvent(f) for f in formulas]
+    fig, ax = plt.subplots(figsize=figsize)
+    category_gap = 0.55
+    x = 0.0
+    ticks, ticklabels, solvent_centers, seen = [], [], {}, {}
+
+    def draw_pair(center, getter):
+        drew = False
+        for j, formula in enumerate(formulas):
+            stats = getter(j, formula)
+            if stats is None:
+                continue
+            mean, elo, ehi = stats
+            disp = label_remap.get(labels[j], labels[j])
+            offset = (j - (n_formulas - 1) / 2) * (box_width + formula_gap)
+            bar = ax.bar(center + offset, mean, width=box_width,
+                         yerr=np.array([[elo], [ehi]]), capsize=3, color=label_to_color[disp],
+                         edgecolor="black", linewidth=0.9, zorder=3,
+                         label=disp if disp not in seen else None)
+            seen.setdefault(disp, bar)
+            drew = True
+        return drew
+
+    for solvent_group in solvent_groups_iter:
+        for solvent in solvent_group:
+            centers = [x]
+            draw_pair(x, lambda j, f: summ(d22[j].get(solvent, np.array([]))))
+            ticks.append(x); ticklabels.append("delta-22"); x += category_gap
+            for bin_label in np_solute_groups.keys():
+                if draw_pair(x, lambda j, f, s=solvent, b=bin_label:
+                             summ(np_grouped(formula_remap.get(f, f), s).get(b, [])) if
+                             np_grouped(formula_remap.get(f, f), s).get(b) is not None else None):
+                    centers.append(x); ticks.append(x)
+                    ticklabels.append("complex" if bin_label == "Test Set" else bin_label)
+                    x += category_gap
+            solvent_centers[solvent] = np.mean(centers)
+            x += solvent_gap
+        x += spacer_width
+
+    ax.set_xticks(ticks)
+    ax.set_xticklabels(ticklabels, rotation=45, fontsize=9, style="italic", ha="center")
+    for lab in ax.get_xticklabels():
+        if lab.get_text() == "delta-22":
+            lab.set_fontweight("bold"); lab.set_fontsize(10); lab.set_style("normal")
+    for solvent, center in solvent_centers.items():
+        ax.text(center, -0.17, solvent, transform=ax.get_xaxis_transform(),
+                ha="center", va="top", fontsize=12, fontweight="bold")
+    nuc = "^{1}\\mathrm{H}" if nucleus == "H" else "^{13}\\mathrm{C}"
+    ax.set_ylabel(r"RMSE $(" + nuc + r"$ ppm)", fontsize=13)
+    ax.set_ylim(bottom=y_min, top=y_max)
+    disp_nuc = "1H" if nucleus == "H" else "13C"
+    ax.set_title(f"Delta-22 vs NP/Isomer Bootstrap RMSE by Solvent ({disp_nuc} nucleus)",
+                 fontsize=14, fontweight="bold")
+    ax.tick_params(axis="both", which="major", labelsize=10, length=4, width=0.8)
+    ax.yaxis.grid(True, linestyle="-", linewidth=0.4, alpha=0.35, zorder=0)
+    ax.xaxis.grid(False)
+    sns.despine(ax=ax, top=True, right=True)
+    for side in ["left", "bottom", "right", "top"]:
+        ax.spines[side].set_linewidth(0.9)
+    leg = ax.legend(frameon=False, fontsize=10)
+    if leg is not None:
+        leg.set_title("")
+    fig.tight_layout()
+    if save_path:
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def plot_fitting_rmse_comparison_bars(table, nucleus, solvent, figsize=(11, 4), save_path=None):
+    """One group of three bars per test-set solute: fit on that solute alone ("Scaled to Solute"),
+    fit once across the whole test set ("Scaled to Test Set"), and delta-22's bootstrap
+    coefficients applied cold ("Extrapolated from delta22"). Used by SI Figure S15's "Fitting RMSE
+    Comparisons" panel."""
+    colors = {"Scaled to Solute": "#4C72B0", "Scaled to Test Set": "#DD8452",
+              "Extrapolated from delta22": "#55A868"}
+    columns = list(colors.keys())
+    x = np.arange(len(table))
+    width = 0.27
+    fig, ax = plt.subplots(figsize=figsize)
+    for i, col in enumerate(columns):
+        ax.bar(x + (i - 1) * width, table[col].to_numpy(), width,
+               label=col.replace("delta22", "Δ22"), color=colors[col])
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(s).replace("\n", " ") for s in table.index], rotation=30, ha="right")
+    ax.set_ylabel(f"{'¹H' if nucleus == 'H' else '¹³C'} RMSE (ppm)")
+    ax.set_title(f"Fitting RMSE Comparisons ({solvent}, {nucleus})")
+    ax.legend()
+    fig.tight_layout()
+    if save_path:
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    return fig, ax
+
+
+# --- SI Figure S15 feature-space / residual 2x2 grids + distribution-shift bars ---
+# style constants private to the two _dataset_scatter_grid callers below.
+_DATASET_COLORS = {"Test Set": "#7e57c2", "Delta22": "#2ca02c"}   # purple circles / green squares
+_DATASET_MARKERS = {"Test Set": "o", "Delta22": "s"}
+_SOLVENT_DISPLAY = {"chloroform": "Chloroform", "benzene": "Benzene",
+                    "methanol": "Methanol", "TIP4P": "Water (TIP4P)"}
+
+
+def _dataset_scatter_grid(table, x_col, y_col, solvents, title, subtitle, x_label, y_label,
+                          figsize, zero_lines):
+    """2x2 solvent grid overlaying the Test Set and delta-22 point clouds, shared x/y across panels."""
+    fig, axes = plt.subplots(2, 2, figsize=figsize, sharex=True, sharey=True)
+    axes = axes.ravel()
+    for ax, solvent in zip(axes, solvents):
+        sdf = table[table["solvent"] == solvent]
+        for dataset in ["Test Set", "Delta22"]:
+            sub = sdf[sdf["dataset"] == dataset]
+            if not sub.empty:
+                ax.scatter(sub[x_col], sub[y_col], s=32, alpha=0.75, marker=_DATASET_MARKERS[dataset],
+                           c=_DATASET_COLORS[dataset], edgecolors="white", linewidths=0.5, label=dataset)
+        if zero_lines:
+            ax.axhline(0, color="0.35", linestyle="--", linewidth=1.0, zorder=0)
+            ax.axvline(0, color="0.35", linestyle="--", linewidth=1.0, zorder=0)
+        else:
+            ax.axhline(0, color="0.5", linestyle="--", linewidth=1.0, zorder=0)
+        ax.set_title(_SOLVENT_DISPLAY.get(solvent, solvent), fontsize=11, fontweight="semibold")
+        ax.grid(alpha=0.2)
+    handles, labels = axes[0].get_legend_handles_labels()
+    fig.legend(handles, labels, loc="center left", bbox_to_anchor=(0.88, 0.5), frameon=False, title="Dataset")
+    fig.supxlabel(x_label, fontsize=11)
+    fig.supylabel(y_label, fontsize=11)
+    fig.suptitle(title, fontsize=13, fontweight="bold")
+    fig.text(0.5, 0.94, subtitle, ha="center", fontsize=10, color="0.35")
+    fig.tight_layout(rect=[0.02, 0.02, 0.86, 0.92])
+    return fig, axes
+
+
+def plot_feature_space_coverage_grid(table, nucleus, x_label,
+                                     solvents=("chloroform", "benzene", "methanol", "TIP4P"),
+                                     figsize=(11, 9), save_path=None):
+    """SI Figure S15's "Feature Space Coverage by Solvent" panel."""
+    subtitle = f"{'Hydrogen' if nucleus == 'H' else 'Carbon'} sites: Test Set vs Delta22"
+    fig, axes = _dataset_scatter_grid(table, "x", "y", solvents, "Feature Space Coverage by Solvent",
+                                      subtitle, x_label, "OpenMM Correction (centered, ppm)",
+                                      figsize, zero_lines=True)
+    if save_path:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return fig, axes
+
+
+def plot_delta22_plane_residuals_grid(table, nucleus,
+                                      solvents=("chloroform", "benzene", "methanol", "TIP4P"),
+                                      figsize=(11, 9), save_path=None):
+    """SI Figure S15's "Residuals for Delta22 Fitting Coefficients" panel."""
+    subtitle = f"{'Hydrogen' if nucleus == 'H' else 'Carbon'} sites: Test Set vs Delta22"
+    fig, axes = _dataset_scatter_grid(table, "experimental", "residual", solvents,
+                                      "Residuals for Delta22 Fitting Coefficients", subtitle,
+                                      "Experimental Shielding (ppm)",
+                                      "Residual vs Delta22 Plane (experimental - predicted)",
+                                      figsize, zero_lines=False)
+    if save_path:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return fig, axes
+
+
+def plot_distribution_shift_by_solvent_bars(table, nucleus, figsize=(7, 5), save_path=None):
+    """One group of three bars per solvent: extrapolated-from-delta22 / scaled-to-test-set /
+    scaled-to-solute RMSE. The three sitting close together (except water) is the transferability
+    story. Used by SI Figure S15's "Distribution Shift by Solvent" panel."""
+    colors = {"Extrapolated from delta22": "#A72608", "Scaled to Test Set": "#C93240",
+              "Scaled to Solute": "#61a89a"}
+    columns = list(colors.keys())
+    x = np.arange(len(table))
+    width = 0.26
+    fig, ax = plt.subplots(figsize=figsize)
+    for i, col in enumerate(columns):
+        ax.bar(x + (i - 1) * width, table[col].to_numpy(), width,
+               label=col.replace("delta22", "Δ22"), color=colors[col])
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(s) for s in table.index])
+    ax.set_ylabel(f"{'¹H' if nucleus == 'H' else '¹³C'} RMSE (ppm)")
+    ax.set_title(f"Distribution Shift by Solvent ({nucleus})")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+    fig.tight_layout()
+    if save_path:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return fig, ax
+

analysis/code/build_composite_model.py

@@ -0,0 +1,290 @@
+"""Regenerate the composite_model/ group in applications.hdf5 from the shipped delta-22 data.
+
+Replaces the old build_coefficients.py, which copied pre-fit CSVs out of a private tree via
+environment variables. Everything here is fit in-repo from delta22.hdf5, so the natural-products
+composite-model coefficients (Figure 5C/5D, SI S15) are reproducible from released inputs alone.
+
+The stored group's structure (formula keys, parameter order, solvent order, seed count) is read
+from the current applications.hdf5 and used as a template; only the numbers are regenerated.
+"""
+import os
+import random
+import sys
+
+import h5py
+import numpy as np
+import pandas as pd
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+REPO = os.path.abspath(os.path.join(HERE, "..", ".."))
+for _p in ("analysis/code", "analysis/code/shared", "data/delta22", "data/applications"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+
+import delta22 as D           # noqa: E402
+import paths                  # noqa: E402
+
+# The 22 delta-22 solutes in first-appearance order: the universe random.choices() draws from.
+# nitromethane stays in the draw pool (index 14) but its rows are dropped from the fit data, exactly
+# as the original pipeline did.
+SOLUTE_ORDER = ["AcOH", "EtOAc", "EtOH", "THF", "acetone", "acetonitrile", "benzaldehyde", "dmf",
+                "dmso", "ethane", "ethylene", "glycol", "imidazole", "isopropanol", "nitromethane",
+                "propylene", "pyridine", "pyrrole", "pyrrolidine", "tbme", "toluene", "triethylamine"]
+SEED_OFFSET = 100
+BASE_FEATURES = ["experimental", "stationary", "qcd", "pcm", "desmond", "desmond_vib",
+                 "openMM", "openMM_vib"]
+ALIASES = {"des_vib": "desmond_vib", "op_vib": "openMM_vib"}
+
+
+def _load_tables():
+    """Three delta-22 NN tables (nitromethane rows removed) plus the pcm conversion factors:
+
+      full_tab  - every row; the OLS point fits use it and drop only each formula's own missing
+                  features, so pcm/qcd formulas fit all 12 solvents.
+      fit_tab   - strict: all base features present (only the 4 explicit-solvent solvents survive);
+                  the bootstrap resamples fit here, matching the stored bootstrap's 4 solvents.
+      apply_tab - desmond or openMM present; the RMSE predictions are evaluated over it.
+    """
+    h5 = paths.dataset_file("delta22", root=REPO)
+    xlsx = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+    nn = D.load_query_df_nn(h5, xlsx, verbose=False)
+    dft = D.load_query_df_dft(h5, xlsx, verbose=False)
+    factors = {"H": D.pcm_conversion_factors(dft, "H"), "C": D.pcm_conversion_factors(dft, "C")}
+    nn = nn[nn["solute"] != "nitromethane"].copy()
+    # stationary_plus_pcm folds in the per-solvent conversion factor; every other _plus_ column is a
+    # plain sum, built on demand by _ensure_column.
+    fac = nn.apply(lambda r: factors[r["nucleus"]].get(r["solvent"], np.nan), axis=1)
+    nn["stationary_plus_pcm"] = nn["stationary"] + nn["pcm"] * fac
+    full_tab = nn.reset_index(drop=True)
+    fit_tab = nn.dropna(subset=BASE_FEATURES).reset_index(drop=True)
+    apply_tab = nn.dropna(subset=["desmond", "openMM"], how="all").reset_index(drop=True)
+    return full_tab, fit_tab, apply_tab, factors
+
+
+def _canonical(name):
+    return ALIASES.get(name, name)
+
+
+def _regressor_columns(formula_key):
+    """The design-matrix regressors for a formula key. '_AND_' separates independent regressors;
+    a '_plus_' name is one precomputed sum column. Returns the list of column names to fit on."""
+    return [g for g in formula_key.split("_AND_")]
+
+
+def _plus_parts(regressor):
+    """Base features a '_plus_' regressor decomposes into (canonical names); [] if it is itself a
+    base feature."""
+    if "_plus_" not in regressor:
+        return []
+    return [_canonical(p) for p in regressor.split("_plus_") if p]
+
+
+def _ensure_column(df, regressor):
+    """Make sure a precomputed _plus_ regressor column exists on df (sum of its parts; the pcm part
+    already carries the conversion factor via stationary_plus_pcm)."""
+    if regressor in df.columns:
+        return
+    parts = regressor.split("_plus_")
+    assert _canonical(parts[0]) == "stationary", f"unexpected base term in {regressor}"
+    # build cumulatively so stationary_plus_pcm_plus_op_vib = stationary_plus_pcm + openMM_vib. The
+    # pcm chains start from the scaled stationary_plus_pcm column; everything else from raw stationary.
+    if parts[1:2] == ["pcm"]:
+        col = df["stationary_plus_pcm"].to_numpy(float).copy()
+        rest = parts[2:]
+    else:
+        col = df["stationary"].to_numpy(float).copy()
+        rest = parts[1:]
+    for p in rest:
+        col = col + df[_canonical(p)].to_numpy(float)
+    df[regressor] = col
+
+
+def _fit(df, regressors):
+    """OLS of experimental on the given regressor columns (with intercept); returns {name: coeff}
+    keyed by 'Intercept' and each regressor. NaN rows are dropped, matching statsmodels."""
+    y = pd.to_numeric(df["experimental"], errors="coerce").to_numpy(float)
+    cols = [df[r].to_numpy(float) for r in regressors]
+    X = np.column_stack([np.ones(len(df))] + cols)
+    mask = np.isfinite(y) & np.all(np.isfinite(X), axis=1)
+    beta, *_ = np.linalg.lstsq(X[mask], y[mask], rcond=None)
+    return dict(zip(["Intercept"] + list(regressors), beta))
+
+
+def _decompose(coeffs, regressors, solvent, nucleus, factors):
+    """Expand each _plus_ regressor's coefficient onto its base features (same value copied). The pcm
+    row is scaled by the per-solvent conversion factor ONLY when the formula used a
+    stationary_plus_pcm term (a free _AND_pcm regressor stays raw). Returns {base_feature: coeff}."""
+    out = {"intercept": coeffs["Intercept"]}
+    for reg, val in coeffs.items():
+        if reg == "Intercept":
+            continue
+        parts = _plus_parts(reg)
+        for base in (parts if parts else [_canonical(reg)]):
+            out[base] = val
+    uses_spp = any(_plus_parts(r)[1:2] == ["pcm"] for r in regressors)
+    if uses_spp and "pcm" in out:
+        out["pcm"] = out["pcm"] * float(factors[nucleus].get(solvent, np.nan))
+    return out
+
+
+def _bootstrap_sample(fit_tab, nucleus, solvent, seed):
+    """The rows fit for one bootstrap draw: seed the RNG with seed+100, draw 22 solutes with
+    replacement from SOLUTE_ORDER, keep that solvent/nucleus's rows for the drawn solutes."""
+    random.seed(int(seed) + SEED_OFFSET)
+    drawn = random.choices(SOLUTE_ORDER, k=22)
+    return fit_tab[(fit_tab["nucleus"] == nucleus) & (fit_tab["solvent"] == solvent)
+                   & (fit_tab["solute"].isin(drawn))]
+
+
+def _rmse(fit_tab, apply_tab, display_formula, nucleus, solvent, seed, factors):
+    """Bootstrap RMSE: fit the formula on one bootstrap draw, apply those raw regressor coefficients
+    to every apply-table row for the solvent, and take the RMSE over labeled rows. The RMSE display
+    names join regressors with ' + ' (each is a column, possibly a precomputed _plus_ one)."""
+    regs = [t.strip() for t in display_formula.split(" + ")]
+    for r in (fit_tab, apply_tab):
+        for reg in regs:
+            _ensure_column(r, reg)
+    coeffs = _fit(_bootstrap_sample(fit_tab, nucleus, solvent, seed), regs)
+    a = apply_tab[(apply_tab["nucleus"] == nucleus) & (apply_tab["solvent"] == solvent)]
+    pred = coeffs["Intercept"] + sum(coeffs[r] * a[r].to_numpy(float) for r in regs)
+    exp = pd.to_numeric(a["experimental"], errors="coerce").to_numpy(float)
+    mask = np.isfinite(pred) & np.isfinite(exp)
+    return float(np.sqrt(np.mean((pred[mask] - exp[mask]) ** 2)))
+
+
+def _pcm_table(factors, nucleus):
+    """pcm_conversion_factors subgroup for one nucleus: solvent, pcm_conversion_factor. Stored under
+    the 'water' name rather than the loader's 'TIP4P'."""
+    s = factors[nucleus].rename(index={"TIP4P": "water"})
+    return pd.DataFrame({"solvent": s.index, "pcm_conversion_factor": s.to_numpy()})
+
+
+def regenerate(reader):
+    """Regenerate every composite_model subgroup from delta-22, keyed and ordered to match the stored
+    group. Returns {"ols_coefficients"|"bootstrap_coefficients"|"rmse_distributions"|
+    "pcm_conversion_factors": {key: DataFrame}}."""
+    full_tab, fit_tab, apply_tab, factors = _load_tables()
+    out = {k: {} for k in ("ols_coefficients", "bootstrap_coefficients",
+                           "rmse_distributions", "pcm_conversion_factors")}
+
+    for f in reader.composite_formulas("ols_coefficients"):
+        regs = _regressor_columns(f)
+        for reg in regs:
+            _ensure_column(full_tab, reg)
+        for nuc in ("H", "C"):
+            try:
+                template = reader.ols_coefficients(f, nuc)
+            except KeyError:
+                continue
+            params = template["parameter"].tolist()
+            solvents = [c for c in template.columns if c != "parameter"]
+            data = {"parameter": params}
+            for solvent in solvents:
+                sub = full_tab[(full_tab["nucleus"] == nuc) & (full_tab["solvent"] == solvent)]
+                dec = _decompose(_fit(sub, regs), regs, solvent, nuc, factors)
+                data[solvent] = [dec[p] for p in params]
+            out["ols_coefficients"][(f, nuc)] = pd.DataFrame(data)
+
+    for f in reader.composite_formulas("bootstrap_coefficients"):
+        regs = _regressor_columns(f)
+        for reg in regs:
+            _ensure_column(fit_tab, reg)
+        for nuc in ("H", "C"):
+            try:
+                template = reader.bootstrap_coefficients(f, nuc)
+            except KeyError:
+                continue
+            params = [c for c in template.columns if c not in ("solvent", "seed")]
+            rows = []
+            for solvent in dict.fromkeys(template["solvent"]):
+                seeds = sorted(template[template["solvent"] == solvent]["seed"].unique())
+                for seed in seeds:
+                    dec = _decompose(_fit(_bootstrap_sample(fit_tab, nuc, solvent, seed), regs),
+                                     regs, solvent, nuc, factors)
+                    row = {"solvent": solvent, "seed": seed}
+                    for p in params:
+                        row[p] = dec["intercept"] if p == "Intercept" else dec[p]
+                    rows.append(row)
+            out["bootstrap_coefficients"][(f, nuc)] = pd.DataFrame(rows)[template.columns.tolist()]
+
+    for nuc in ("H", "C"):
+        template = reader.rmse_distribution(nuc)
+        rows = []
+        for f in dict.fromkeys(template["formula"]):
+            sub = template[template["formula"] == f]
+            for solvent in dict.fromkeys(sub["solvent"]):
+                for seed in sorted(sub[sub["solvent"] == solvent]["seed"].unique()):
+                    rows.append({"solvent": solvent, "nucleus": nuc, "formula": f, "seed": seed,
+                                 "Bootstrap_RMSE": _rmse(fit_tab, apply_tab, f, nuc, solvent, seed, factors)})
+        out["rmse_distributions"][nuc] = pd.DataFrame(rows)[template.columns.tolist()]
+        out["pcm_conversion_factors"][nuc] = _pcm_table(factors, nuc)
+
+    return out
+
+
+def verify(reader, regenerated, atol=5e-3):
+    """Compare regenerated numeric values against the stored group. Returns the global max abs diff.
+    The floor is the fixed-point encoding granularity (~1e-3), far below any recipe error."""
+    worst = 0.0
+    for (f, nuc), df in regenerated["ols_coefficients"].items():
+        st = reader.ols_coefficients(f, nuc)
+        worst = max(worst, np.abs(df.drop(columns="parameter").to_numpy(float)
+                                  - st.drop(columns="parameter").to_numpy(float)).max())
+    for (f, nuc), df in regenerated["bootstrap_coefficients"].items():
+        st = reader.bootstrap_coefficients(f, nuc).sort_values(["solvent", "seed"]).reset_index(drop=True)
+        df = df.sort_values(["solvent", "seed"]).reset_index(drop=True)
+        num = [c for c in df.columns if c not in ("solvent",)]
+        worst = max(worst, np.abs(df[num].to_numpy(float) - st[num].to_numpy(float)).max())
+    for nuc, df in regenerated["rmse_distributions"].items():
+        st = reader.rmse_distribution(nuc).sort_values(["formula", "solvent", "seed"]).reset_index(drop=True)
+        df = df.sort_values(["formula", "solvent", "seed"]).reset_index(drop=True)
+        worst = max(worst, np.abs(df["Bootstrap_RMSE"].to_numpy(float)
+                                  - st["Bootstrap_RMSE"].to_numpy(float)).max())
+    for nuc, df in regenerated["pcm_conversion_factors"].items():
+        st = reader.pcm_conversion_factors(nuc).set_index("solvent")["pcm_conversion_factor"]
+        got = df.set_index("solvent")["pcm_conversion_factor"]
+        worst = max(worst, np.abs(got - st.reindex(got.index)).max())
+    return float(worst)
+
+
+_NOTE = ("Linear composite-model coefficients fit on the delta-22 dataset and applied to the "
+         "natural products (Figure 5C/5D, SI S15). Stored as CSV text; read via applications_reader. "
+         "Regenerated in-repo from delta22.hdf5 by build_composite_model.py.")
+
+
+def write(hdf5_path, regenerated):
+    """Overwrite the composite_model/ group in applications.hdf5 with the regenerated CSV text."""
+    with h5py.File(hdf5_path, "a") as h:
+        if "composite_model" in h:
+            del h["composite_model"]
+        cm = h.create_group("composite_model")
+        cm.attrs["note"] = _NOTE
+        for (formula, nucleus), df in regenerated["ols_coefficients"].items():
+            cm.require_group(f"ols_coefficients/{formula}").create_dataset(nucleus, data=df.to_csv(index=False))
+        for (formula, nucleus), df in regenerated["bootstrap_coefficients"].items():
+            cm.require_group(f"bootstrap_coefficients/{formula}").create_dataset(nucleus, data=df.to_csv(index=False))
+        for nucleus, df in regenerated["rmse_distributions"].items():
+            cm.require_group("rmse_distributions").create_dataset(nucleus, data=df.to_csv(index=False))
+        for nucleus, df in regenerated["pcm_conversion_factors"].items():
+            cm.require_group("pcm_conversion_factors").create_dataset(nucleus, data=df.to_csv(index=False))
+
+
+def main():
+    import argparse
+    ap = argparse.ArgumentParser(description="Regenerate composite_model/ in applications.hdf5 from delta-22.")
+    ap.add_argument("--write", action="store_true",
+                    help="overwrite the composite_model group in applications.hdf5 (default: verify only)")
+    args = ap.parse_args()
+    sys.path.insert(0, os.path.join(REPO, "data", "applications"))
+    from applications_reader import Applications
+    hdf5_path = paths.dataset_file("applications", root=REPO)
+    reader = Applications(hdf5_path)
+    regenerated = regenerate(reader)
+    max_diff = verify(reader, regenerated)
+    print(f"composite_model regenerated from delta-22; max abs diff vs stored = {max_diff:.2e}")
+    if args.write:
+        write(hdf5_path, regenerated)
+        print(f"wrote composite_model/ into {hdf5_path}")
+
+
+if __name__ == "__main__":
+    main()

analysis/code/build_nb_applications.py

@@ -0,0 +1,282 @@
+# build_nb_applications.py -- generator for the natural-products figure notebooks.
+#
+# Source of truth for the notebooks: edit the cell sources below, then regenerate the notebook(s) you
+# touched (regenerating overwrites the .ipynb, so edit this script, not the notebook):
+#
+#     python3 build_nb_applications.py si_figure_s08_vomicine
+#
+# Run with no arguments to list the notebook names; the NOTEBOOKS dict maps each to its destination.
+# The numbers come from applications.py; the notebooks arrange and plot them.
+from nb_build import md, code, save_notebook as _save
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S8, panel A: explicit-solvent shell-size convergence for vomicine
+# ----------------------------------------------------------------------------
+si_figure_s08 = [
+    md(r"""
+# SI Figure S8 (panel A): explicit-solvent correction vs solvent-shell size (vomicine)
+
+The MagNET-x explicit-solvent correction (solvated minus isolated, averaged over a proton site's
+atoms and MD frames) vs solvent-shell size (heavy-atom count of the solute-plus-solvent cluster),
+for vomicine, one panel per solvent.
+"""),
+    code(r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/applications", "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""),
+    code(r"""
+import matplotlib.pyplot as plt
+
+from applications_reader import Applications, SHELL_SIZES
+import applications
+import paths
+"""),
+    code(r"""
+DATA = os.path.join(REPO, "data", "applications")
+
+def figure_path(name):
+    os.makedirs("figures", exist_ok=True)
+    return os.path.join("figures", name)
+"""),
+    code(r"""
+loader = Applications(paths.dataset_file("applications", root=REPO),
+                      os.path.join(DATA, "applications_experimental.xlsx"))
+
+# compute the per-site, per-solvent shell-convergence corrections for vomicine (proton sites)
+SOLUTE, NUCLEUS = "vomicine", "H"
+corrections = applications.shell_convergence_corrections(loader, SOLUTE, NUCLEUS)
+corrections.head()
+"""),
+    code(r"""
+# panel layout and font styling for the 2x2 grid (one panel per solvent)
+solvent_order = ["chloroform", "methanol", "TIP4P", "benzene"]
+shell_cols = [f"shell_{s}" for s in SHELL_SIZES]
+plt.rcParams.update({"font.family": "serif", "font.serif": ["Arial", "DejaVu Serif"],
+                     "axes.labelsize": 12, "axes.titlesize": 14})
+"""),
+    code(r"""
+# 2x2 panel, one solvent each; one line per proton site
+fig, axes = plt.subplots(2, 2, figsize=(11, 8), sharex=True, sharey=True)
+for ax, solvent in zip(axes.ravel(), solvent_order):
+    sub = corrections.xs(solvent, level="solvent")
+    sub = sub.dropna(how="all")
+    ax.set_title("TIP4P" if solvent == "TIP4P" else solvent.capitalize(), fontweight="bold")
+    for site, row in sub.iterrows():
+        ax.plot(SHELL_SIZES, row[shell_cols].values.astype(float),
+                marker="o", linewidth=1.5, alpha=0.75, label=site)
+    ax.grid(True, alpha=0.3)
+    ax.set_xlabel("Shell Size (Heavy Atom Count)", fontweight="bold")
+    ax.set_ylabel("OpenMM correction (ppm)", fontweight="bold")
+fig.suptitle("Shell-size Convergence for Vomicine", fontweight="bold", fontsize=18)
+fig.tight_layout()
+fig.savefig(figure_path("si_figure_s08_vomicine.png"), dpi=300, bbox_inches="tight")
+plt.show()
+"""),
+]
+
+# ----------------------------------------------------------------------------
+# SI Figure S15: transferability of composite-model coefficients to the test set (1H + 13C)
+# ----------------------------------------------------------------------------
+si_figure_s15 = [
+    md(r"""
+# SI Figure S15: per-solute RMSE with delta-22 composite coefficients applied
+
+Per-solute ¹H and ¹³C RMSE for the natural-products / olefin-isomer test set with delta-22
+coefficients applied (reproduces Figure 5C in ¹H, plus the ¹³C analogue), plus fitting-RMSE
+comparisons across coefficient choices, feature-space coverage by solvent, delta-22-plane residuals,
+and the RMSE distribution shift by solvent.
+"""),
+    code(r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/applications", "data/delta22", "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""),
+    code(r"""
+import matplotlib.pyplot as plt
+
+from applications_reader import Applications
+import applications
+import applications_plots
+import delta22
+import paths
+"""),
+    code(r"""
+DATA = os.path.join(REPO, "data", "applications")
+DELTA22_HDF5 = paths.dataset_file("delta22", root=REPO)
+DELTA22_XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+def figure_path(name):
+    os.makedirs("figures", exist_ok=True)
+    return os.path.join("figures", name)
+"""),
+    code(r"""
+loader = Applications(paths.dataset_file("applications", root=REPO),
+                      os.path.join(DATA, "applications_experimental.xlsx"))
+
+# assemble the feature table and run the composite-model fits + bootstrap
+query_df_nn = applications.build_query_df_nn(loader)
+site_counts = (query_df_nn.drop_duplicates(subset=["solute", "nucleus", "site"])
+               .groupby(["solute", "nucleus"]).size().unstack(fill_value=0))
+
+per_solute = applications.per_solute_fits(query_df_nn)      # "scaled to solute" baseline
+all_solute = applications.per_solvent_fits(query_df_nn)     # "scaled to test set" full fit
+
+seed = applications.build_bootstrap_seed_coeffs(loader)
+bootstrap_rmses = {}
+for nuc in ["H", "C"]:
+    preds = applications.apply_bootstrap_params_to_full_dataset(query_df_nn, seed[nuc], nucleus=nuc)
+    bootstrap_rmses[nuc] = applications.compute_solute_rmses(preds)
+"""),
+    code(r"""
+# published solute order (isomers, then the two pyridines, then the natural products); matches
+# main-text Figure 5D. delta-22 is drawn first by the box-plot engine, so it is not listed here.
+S15_SOLUTE_ORDER = ["isomer_1E", "isomer_1Z", "isomer_2E", "isomer_2Z", "isomer_3E", "isomer_3Z",
+                    "isomer_4N", "isomer_4O", "vomicine", "prednisone", "peptide", "flavone",
+                    "dihydrotanshinone_I"]
+"""),
+    code(r"""
+# 1H panel (reproduces Figure 5C)
+applications_plots.plot_nps_on_boxplot_delta22_simplified(
+    loader.rmse_distribution("H"), bootstrap_rmses["H"], per_solute["H"],
+    nucleus="H", formulas=["stationary_plus_qcd + openMM"], colors=["#61a89a"],
+    site_counts=site_counts, formula_remap=applications.FORMULA_REMAP,
+    solute_remap=applications.SOLUTE_DISPLAY, solute_order=S15_SOLUTE_ORDER,
+    solute_color_remap=applications.PEPTIDE_HIGHLIGHT_H,
+    figsize=(14, 8), box_width=0.20, box_gap=0.1,
+    show_baseline=True, baseline_annotation_text="Scaled to solute",
+    baseline_annotation_x=0.215, baseline_annotation_y=0.39,
+    show_full_fit_line=True, full_fit_line_label="Scaled to Test Set", full_fit_line_label_x=0.835,
+    all_solute_fitting_results=all_solute,
+    max_bar_height=0.06, site_count_axis_mode="inset", site_count_inset_area_frac=0.1,
+    site_count_inset_axis_offset=-0.06, site_count_inset_axis_label_pad=0.045,
+    save_path=figure_path("si_figure_s15_1H.png"),
+)
+"""),
+    code(r"""
+# 13C panel
+applications_plots.plot_nps_on_boxplot_delta22_simplified(
+    loader.rmse_distribution("C"), bootstrap_rmses["C"], per_solute["C"],
+    nucleus="C", formulas=["stationary_plus_op_vib + openMM"], colors=["#61a89a"],
+    site_counts=site_counts, formula_remap=applications.FORMULA_REMAP,
+    solute_remap=applications.SOLUTE_DISPLAY, solute_order=S15_SOLUTE_ORDER,
+    solute_color_remap=applications.PEPTIDE_HIGHLIGHT_C,
+    figsize=(14, 8), box_width=0.20, box_gap=0.1,
+    show_baseline=True, baseline_annotation_text="Scaled to solute",
+    baseline_annotation_x=0.28, baseline_annotation_y=0.355,
+    show_full_fit_line=True, full_fit_line_label="Scaled to Test Set", full_fit_line_label_x=0.98,
+    all_solute_fitting_results=all_solute,
+    max_bar_height=0.8, site_count_axis_mode="inset", site_count_inset_area_frac=0.1,
+    site_count_inset_axis_offset=-0.06, site_count_inset_axis_label_pad=0.045,
+    save_path=figure_path("si_figure_s15_13C.png"),
+)
+"""),
+    md(r"""
+## Fitting RMSE Comparisons (chloroform)
+
+Per-solute RMSE under three coefficient choices (Scaled to Solute / Scaled to Test Set / Extrapolated
+from delta22).
+"""),
+    code(r"""
+for nucleus, formula in [("H", "stationary_plus_qcd + openMM"), ("C", "stationary_plus_op_vib + openMM")]:
+    table = applications.fitting_rmse_comparison_table(query_df_nn, per_solute, all_solute, bootstrap_rmses[nucleus],
+                                            nucleus, "chloroform", formula)
+    table = table.rename(index=applications.SOLUTE_DISPLAY)
+    order = [applications.SOLUTE_DISPLAY.get(k, k) for k in S15_SOLUTE_ORDER]   # published order; peptide has no chloroform data
+    table = table.reindex([n for n in order if n in table.index])
+    print(f"--- {nucleus} ---")
+    display(table.round(3))
+    applications_plots.plot_fitting_rmse_comparison_bars(
+        table, nucleus, "chloroform",
+        save_path=figure_path(f"si_figure_s15_fitting_rmse_comparisons_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+    md(r"""
+## Feature Space Coverage by Solvent
+
+Test-set vs delta-22 mean-centered feature values across the four explicit-solvent solvents.
+"""),
+    code(r"""
+delta22_query_df_nn = delta22.add_composite_columns(delta22.load_query_df_nn(
+    DELTA22_HDF5, DELTA22_XLSX, verbose=False))
+"""),
+    code(r"""
+_FEATURE_X_LABELS = {
+    "H": "Gas-Phase Shielding + QCD Correction (centered, ppm)",
+    "C": "Gas-Phase Shielding + OpenMM Vibrational Correction (centered, ppm)",
+}
+for nucleus in ["H", "C"]:
+    coverage = applications.feature_space_coverage_table(query_df_nn, delta22_query_df_nn, nucleus)
+    applications_plots.plot_feature_space_coverage_grid(
+        coverage, nucleus, _FEATURE_X_LABELS[nucleus],
+        save_path=figure_path(f"si_figure_s15_feature_space_coverage_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+    md(r"""
+## Residuals for Delta22 Fitting Coefficients
+
+Residuals of a delta-22-only 2-feature plane applied to both delta-22 and the test set, vs
+experimental shielding.
+"""),
+    code(r"""
+for nucleus in ["H", "C"]:
+    residuals = applications.delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, nucleus)
+    applications_plots.plot_delta22_plane_residuals_grid(
+        residuals, nucleus,
+        save_path=figure_path(f"si_figure_s15_delta22_plane_residuals_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+    md(r"""
+## Distribution Shift by Solvent
+
+Per-solvent mean test-set RMSE under the three coefficient choices.
+"""),
+    code(r"""
+for nucleus, formula in [("H", "stationary_plus_qcd + openMM"), ("C", "stationary_plus_op_vib + openMM")]:
+    shift = applications.distribution_shift_by_solvent_table(per_solute, all_solute, bootstrap_rmses[nucleus],
+                                                  nucleus, formula)
+    print(f"--- {nucleus} ---")
+    display(shift.round(3))
+    applications_plots.plot_distribution_shift_by_solvent_bars(
+        shift, nucleus,
+        save_path=figure_path(f"si_figure_s15_distribution_shift_by_solvent_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+]
+
+# NOTE: main-text Figure 5 panels C and D (analysis/manuscript_figures/fig5c.ipynb, fig5d.ipynb) are
+# HAND-AUTHORED, not generated here. Their plotting engines (plot_nps_benefit_barplot for 5C,
+# plot_nps_on_boxplot_delta22_simplified for 5D, also used by SI S15) live in applications_plots.py;
+# the notebooks only load data and call them. The SI notebooks above are still generated here.
+
+
+NOTEBOOKS = {
+    "si_figure_s08_vomicine": (si_figure_s08, "analysis/si_figures/si_figure_s08_panelA_vomicine.ipynb"),
+    "si_figure_s15": (si_figure_s15, "analysis/si_figures/si_figure_s15.ipynb"),
+}
+
+if __name__ == "__main__":
+    import os
+    import sys
+    here = os.path.dirname(os.path.abspath(__file__))
+    repo = os.path.abspath(os.path.join(here, "..", ".."))
+    names = sys.argv[1:]
+    if not names:
+        print("usage: python3 build_nb_applications.py <name> [<name> ...]")
+        print("regenerates ONLY the named notebook(s) -- pick just the one(s) you edited,")
+        print("since regenerating clears a notebook's execution outputs.")
+        print("available names:", ", ".join(NOTEBOOKS))
+        raise SystemExit(1)
+    unknown = [n for n in names if n not in NOTEBOOKS]
+    if unknown:
+        raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}")
+    for name in names:
+        cells, relpath = NOTEBOOKS[name]
+        _save(cells, os.path.join(repo, relpath))

analysis/code/build_nb_composite_models.py

@@ -0,0 +1,124 @@
+"""Source of truth for the composite-model ablation-study notebook. Edit the cell sources here,
+then regenerate:
+
+    python3 build_nb_composite_models.py --write
+    jupyter nbconvert --to notebook --execute --inplace ../si_figures/si_composite_models_ablations.ipynb
+
+Regenerating overwrites the .ipynb, so do not hand-edit the notebook. All real code lives in
+composite_models.py (numbers and the openpyxl table-writing layer, which itself calls delta22.py's
+existing fitting harness); the notebook only runs it. The notebook itself lives in
+analysis/si_figures/ (grouped by role), not here.
+
+Run with no arguments to see this usage message; pass --write to regenerate. This generator produces
+a single notebook and takes no name argument, so a bare run must not silently overwrite it.
+"""
+import os
+import sys
+
+from nb_build import md, code, save_notebook as _save
+
+
+_BOOTSTRAP = r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/delta22", "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""
+
+_IMPORTS = r"""
+import matplotlib.pyplot as plt
+
+import delta22
+import composite_models
+import composite_plots
+import paths
+"""
+
+_SETUP = r"""
+DELTA22_HDF5 = paths.dataset_file("delta22", root=REPO)
+XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+def figure_path(name):
+    os.makedirs("figures", exist_ok=True)
+    return os.path.join("figures", name)
+
+def document_path(name):
+    os.makedirs("documents", exist_ok=True)
+    return os.path.join("documents", name)
+
+# the shipped ablations.xlsx used 250 seeded train/test splits per (formula, solvent)
+N_SPLITS = 250
+"""
+
+si_composite_models_ablations = [
+    md(r"""
+# SI Figure S15: composite-formula ablation workbook (`ablations.xlsx`)
+
+Rebuilds the `ablations.xlsx` workbook behind ~32 of SI Figure S15's images: ~25 composite-formula
+variants (stationary geometry plus some combination of implicit PCM, explicit Desmond, and
+rovibrational QCD corrections) across all 12 delta-22 solvents, at two reference levels:
+DSD-PBEP86/pcSseg-3 (geometry PBE0/tz) and the MagNET-Zero training reference (WP04/pcSseg-2 for ¹H,
+wB97X-D/pcSseg-2 for ¹³C).
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(r"""
+query_df_dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))
+solutes = delta22.delta22_solutes(DELTA22_HDF5)
+print(len(query_df_dft), "rows;", len(solutes), "solutes")
+"""),
+    md("## Preview: one formula's mean test RMSE at the DSD-PBEP86 reference level"),
+    code(r"""
+# preview before the full build (all ~25 formulas x 12 solvents x 250 splits x 2 nuclei x 2
+# reference levels)
+level = composite_models.REFERENCE_LEVELS["dsd"]
+preview = composite_models.ablation_rmse_table(query_df_dft, "H", level["method_h"], level["basis_h"],
+                                 level["geometry_h"], solutes, n_splits=20)
+preview[["chloroform", "benzene", "Mean Test RMSE"]].round(3)
+"""),
+    md("## Build the full ablations workbook (both reference levels, both nuclei)"),
+    code(r"""
+output_path = document_path("ablations.xlsx")
+composite_models.build_ablations_workbook(query_df_dft, solutes, output_path, n_splits=N_SPLITS)
+"""),
+    md(r"""
+## Correlations Between Features
+
+Solvent-averaged Pearson r correlation matrix between the five composite-model features (stationary
+shielding, PCM, Desmond, its vibrational analogue, QCD), both nuclei, plus a per-solvent
+PCM-vs-Desmond table.
+"""),
+    code(r"""
+for nucleus, label in [("H", "Proton"), ("C", "Carbon")]:
+    corr = delta22.si_s15_feature_correlations(query_df_dft, nucleus)
+    nuc_label = "1H" if nucleus == "H" else "13C"
+    nuc_title = "$^{1}$H" if nucleus == "H" else "$^{13}$C"
+    composite_plots.plot_feature_correlation_heatmap(
+        corr["r"], vmin=-1, vmax=1, cmap="RdBu",
+        title=f"Solvent-Averaged Pearson $r$ Correlation Matrix for {nuc_title}",
+        save_path=figure_path(f"si_figure_s15_feature_corr_r_{nuc_label}.png"))
+plt.show()
+"""),
+    code(r"""
+pcm_desmond_corr = delta22.pcm_desmond_correlation_by_solvent(query_df_dft)
+print("PCM vs. Desmond Pearson R by nucleus and solvent:")
+display(pcm_desmond_corr.round(3))
+composite_plots.plot_pcm_desmond_correlation_table(pcm_desmond_corr,
+                                   save_path=figure_path("si_figure_s15_pcm_desmond_corr_table.png"))
+plt.show()
+"""),
+]
+
+
+if __name__ == "__main__":
+    if "--write" not in sys.argv:
+        print("usage: python3 build_nb_composite_models.py --write")
+        print("regenerates analysis/si_figures/si_composite_models_ablations.ipynb -- pass --write to actually do it,")
+        print("since regenerating clears the notebook's execution outputs.")
+        sys.exit(0)
+    here = os.path.dirname(os.path.abspath(__file__))
+    repo = os.path.abspath(os.path.join(here, "..", ".."))
+    _save(si_composite_models_ablations, os.path.join(repo, "analysis", "si_figures", "si_composite_models_ablations.ipynb"))

analysis/code/build_nb_delta22.py

@@ -0,0 +1,798 @@
+# build_nb_delta22.py -- generator for the delta-22 figure notebooks.
+#
+# This script is the source of truth for the notebooks. Edit the cell sources here, then
+# regenerate and re-execute ONLY the notebook(s) you touched, e.g.:
+#
+#     python3 build_nb_delta22.py si_figure_s04
+#     jupyter nbconvert --to notebook --execute --inplace ../si_figures/si_figure_s04.ipynb
+#
+# Regenerating overwrites the .ipynb; do not hand-edit it. All real code lives in delta22.py
+# (numbers); plotting is inline in each notebook.
+#
+# Run with no arguments to list the available notebook names. Regenerating a notebook clears its
+# execution outputs, so there is no "build everything" shortcut -- pick just the notebook(s) touched.
+from nb_build import md, code, save_notebook as _save
+
+
+# under nbconvert the CWD is the notebook's own directory, so sys.path is set relative to it
+_BOOTSTRAP = r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/delta22", "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""
+
+# pure imports for the figure notebooks
+_IMPORTS = r"""
+import matplotlib.pyplot as plt
+
+import delta22
+import delta22_plots
+import paths
+"""
+
+# data-file path constants + the figure_path helper, for the figure notebooks
+_SETUP = r"""
+DELTA22_HDF5 = paths.dataset_file("delta22", root=REPO)
+XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+def figure_path(name):
+    os.makedirs("figures", exist_ok=True)
+    return os.path.join("figures", name)
+"""
+
+# pure imports for the table notebooks (no matplotlib -- they only print and write spreadsheets)
+_IMPORTS_TABLE = r"""
+import delta22
+import paths
+"""
+
+# figure imports plus numpy, for the few notebooks that use numpy directly in a notebook cell
+_IMPORTS_NP = _IMPORTS.replace("import matplotlib.pyplot as plt",
+                               "import numpy as np\nimport matplotlib.pyplot as plt", 1)
+# table imports plus numpy/pandas (or just pandas), used directly when building the table dataframes
+_IMPORTS_TABLE_NPPD = _IMPORTS_TABLE.replace("import delta22",
+                                             "import numpy as np\nimport pandas as pd\n\nimport delta22", 1)
+_IMPORTS_TABLE_PD = _IMPORTS_TABLE.replace("import delta22",
+                                           "import pandas as pd\n\nimport delta22", 1)
+
+# data-file path constants + the document_path helper, for the table notebooks
+_SETUP_TABLE = r"""
+DELTA22_HDF5 = paths.dataset_file("delta22", root=REPO)
+XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+def document_path(name):
+    os.makedirs("documents", exist_ok=True)
+    return os.path.join("documents", name)
+"""
+
+# only the notebooks that use N_SPLITS include this cell (fig2a_pareto uses its own, smaller
+# PARETO_N_SPLITS instead).
+_N_SPLITS_CELL = r"""
+# the published panels use 250 seeded train/test splits
+N_SPLITS = 250
+"""
+
+
+# Figure 3 (analysis/manuscript_figures/fig3.ipynb) is hand-authored with its plotting inline; its
+# data comes from delta22.py (fig3a_pcm_benefit_by_solvent, fig3b_shift_differences,
+# fig3d_formula_regressions).
+
+
+# ----------------------------------------------------------------------------
+# Figure 4 + SI S16/S17/S18: Explicit Solvent Corrections are Physically Meaningful
+# ----------------------------------------------------------------------------
+_FIG4_LOAD = r"""
+# Figure 4 uses one DFT method for the proton sites; differences are taken against a reference
+# solvent, and the solvent_mean pseudo-solvent is added for the solvent-averaged reference.
+METHOD, BASIS, GEOM = "b3lyp_d3bj", "pcSseg2", "aimnet2"
+q = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+one = q[(q["sap_nmr_method"] == METHOD) & (q["sap_basis"] == BASIS) & (q["sap_geometry_type"] == GEOM)]
+one = delta22.add_solvent_mean(one)
+print(len(one), "rows for", METHOD, BASIS, GEOM)
+"""
+
+# shared display labels for the 12 delta-22 solvents, used by Figure 4C and SI Figures S16-S18
+_SOLVENT_LABELS = r"""
+FIG4C_LABELS = {
+    "chloroform": "CDCl3", "dichloromethane": "DCM", "tetrahydrofuran": "THF",
+    "acetonitrile": "MeCN", "dimethylsulfoxide": "DMSO", "acetone": "acetone",
+    "methanol": "MeOH", "TIP4P": "TIP4P", "trifluoroethanol": "TFE",
+    "benzene": "benzene", "toluene": "toluene", "chlorobenzene": "chlorobenzene",
+}
+"""
+
+# Figure 4 (analysis/manuscript_figures/fig4.ipynb) is hand-authored with its plotting inline; its
+# data comes from delta22.py (solvent_induced_shifts, solvent_pair_differences,
+# fit_differences_to_experimental). The _FIG4_LOAD / _SOLVENT_LABELS cells above are shared with the
+# SI S16-S18 notebook below.
+
+
+si_figure_s16_s18 = [
+    md(r"""
+# SI Figures S16-S18: predicted vs experimental solvent-induced shifts by reference solvent
+
+Per solvent, implicit (PCM) and explicit (Desmond) predicted solvent-induced shift differences vs
+measured (y=x ideal), for three reference solvents: **S16** chloroform, **S17** benzene, **S18**
+solvent-averaged. One panel per solvent, each point a proton site.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(_FIG4_LOAD),
+    code(_SOLVENT_LABELS),
+    code(r"""
+# panel order matches the canonical figures; a different order from delta22.DESMOND_SOLVENTS
+SI_S16_18_SOLVENT_ORDER = [
+    "chloroform", "tetrahydrofuran", "dichloromethane", "acetone", "acetonitrile",
+    "dimethylsulfoxide", "trifluoroethanol", "methanol", "TIP4P", "benzene", "toluene", "chlorobenzene",
+]
+
+# panel titles use the full solvent name; axis labels use the compact NMR abbreviations (THF, DCM,
+# DMSO, TFE, CDCl3); the reference is CDCl3 / Benzene / "Avg Corr." for S16 / S17 / S18.
+TITLE_LABEL = {"chloroform": "Chloroform", "tetrahydrofuran": "Tetrahydrofuran",
+               "dichloromethane": "Dichloromethane", "acetone": "Acetone", "acetonitrile": "Acetonitrile",
+               "dimethylsulfoxide": "Dimethylsulfoxide", "trifluoroethanol": "Trifluoroethanol",
+               "methanol": "Methanol", "TIP4P": "TIP4P", "benzene": "Benzene", "toluene": "Toluene",
+               "chlorobenzene": "Chlorobenzene"}
+AXIS_LABEL = {**TITLE_LABEL, "chloroform": "CDCl3", "tetrahydrofuran": "THF",
+              "dichloromethane": "DCM", "dimethylsulfoxide": "DMSO", "trifluoroethanol": "TFE"}
+REF_LABEL = {"chloroform": "CDCl3", "benzene": "Benzene", "solvent_mean": "Avg Corr."}
+"""),
+    code(r"""
+for figure, reference in [("S16", "chloroform"), ("S17", "benzene"), ("S18", "solvent_mean")]:
+    sh = delta22.solvent_induced_shifts(one, reference, SI_S16_18_SOLVENT_ORDER, nucleus="H", explicit="desmond")
+    implicit_fit = delta22.fit_differences_to_experimental(sh, "implicit_diff")
+    explicit_fit = delta22.fit_differences_to_experimental(sh, "explicit_diff")
+    print(f"{figure}: reference={reference:12s} n={explicit_fit['n']:4d}  "
+          f"implicit fit RMSE={implicit_fit['rmse']:.3f}  explicit fit RMSE={explicit_fit['rmse']:.3f}")
+    panel_solvents = [s for s in SI_S16_18_SOLVENT_ORDER if s != reference]
+    delta22_plots.plot_shift_prediction_scatter_grid(
+        sh, panel_solvents, REF_LABEL[reference], AXIS_LABEL, TITLE_LABEL,
+        save_path=figure_path(f"si_figure_{figure.lower()}.png"))
+    plt.show()
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S7: Explicit Solvent Corrections are Method-Independent
+# ----------------------------------------------------------------------------
+si_figure_s07 = [
+    md(r"""
+# SI Figure S7: DFT ¹³C explicit-solvent correction, Desmond vs OpenMM
+
+For the four OpenMM solvents, the DFT ¹³C explicit-solvent correction from Desmond MD vs from OpenMM
+(y=x reference line).
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS_NP),
+    code(_SETUP),
+    code(r"""
+dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+"""),
+    md("## DFT explicit corrections: Desmond vs OpenMM"),
+    code(r"""
+SOLVENT_ORDER = ["chloroform", "methanol", "TIP4P", "benzene"]
+dft_pairs = {sv: delta22.explicit_correction_pairs(dft, sv, "C") for sv in SOLVENT_ORDER}
+for sv, p in dft_pairs.items():
+    print(f"DFT {sv:12s} n={len(p):3d}  r={np.corrcoef(p['desmond'], p['openMM'])[0,1]:.3f}")
+delta22_plots.plot_desmond_vs_openmm_grid(dft_pairs, save_path=figure_path("si_figure_s07.png"))
+plt.show()
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S4: Implicit Solvent Corrections are Highly Correlated
+# ----------------------------------------------------------------------------
+si_figure_s04 = [
+    md(r"""
+# SI Figure S4: correlation of implicit (PCM) corrections across solvents and methods
+
+**S4A** correlation across solvents (both nuclei), **S4B** one solvent vs another (¹H), **S4C**
+correlation across methods (both nuclei). Cells show -log10(1-r), so 3 means r=0.999.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS_NP),
+    code(_SETUP),
+    code(r"""
+METHOD, BASIS, GEOM = "b3lyp_d3bj", "pcSseg2", "pbe0_tz"
+dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+base = dft[(dft["sap_nmr_method"] == METHOD) & (dft["sap_basis"] == BASIS)
+           & (dft["sap_geometry_type"] == GEOM)]
+one = base[base["nucleus"] == "H"]
+
+# solvent order that groups polar-aprotic -> polar-protic -> aromatic (matches the published panel)
+ORDERED_SOLVENTS = ["chloroform", "tetrahydrofuran", "dichloromethane", "acetone", "acetonitrile",
+                    "dimethylsulfoxide", "trifluoroethanol", "methanol", "TIP4P",
+                    "benzene", "toluene", "chlorobenzene"]
+"""),
+    md("## S4A: solvent-vs-solvent PCM correlation (¹H and ¹³C)"),
+    code(r"""
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    solvent_corr = delta22.correlation_matrix(base[base["nucleus"] == nucleus], "pcm", ["solute", "site"], "solvent")
+    solvent_corr = solvent_corr.reindex(index=ORDERED_SOLVENTS, columns=ORDERED_SOLVENTS)
+    vals = solvent_corr.values[np.triu_indices_from(solvent_corr.values, k=1)]
+    print(f"solvent PCM correlation ({label}): mean r={np.nanmean(vals):.4f}  min r={np.nanmin(vals):.4f}")
+    caption = ("Correlation coefficients between solvents across all 22 solutes.\n"
+               f"PCM corrections computed with {METHOD} with the {BASIS} basis.\n"
+               "A value of 3 means the coefficient is 0.999.")
+    delta22_plots.plot_correlation_matrix(solvent_corr, f"Solvents are Highly Correlated ({nucleus})", caption,
+                                          colormap="Reds", show_values=True,
+                                          save_path=figure_path(f"si_figure_s04a_{label}.png"))
+"""),
+    md("## S4B: PCM corrections, chloroform vs acetonitrile (1H) -- one square of S4A"),
+    code(r"""
+pair = one.pivot_table(index=["solute", "site"], columns="solvent", values="pcm")[
+    ["chloroform", "acetonitrile"]].dropna()
+delta22_plots.plot_pcm_scatter(pair["chloroform"], pair["acetonitrile"], "chloroform", "acetonitrile", "H",
+                               save_path=figure_path("si_figure_s04b_1H.png"))
+"""),
+    md("## S4C: method-vs-method correlation (chloroform, double hybrids excluded, ¹H and ¹³C)"),
+    code(r"""
+# one solvent, exclude the double hybrids whose PCM is the substituted reference value
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    chcl3 = dft[(dft["sap_basis"] == BASIS) & (dft["sap_geometry_type"] == GEOM)
+                & (dft["nucleus"] == nucleus) & (dft["solvent"] == "chloroform")
+                & (~dft["sap_nmr_method"].isin(delta22.DOUBLE_HYBRID_METHODS))]
+    method_corr = delta22.correlation_matrix(chcl3, "pcm", ["solute", "site"], "sap_nmr_method")
+    method_order = sorted(method_corr.index)          # alphabetical, matches the published panel
+    method_corr = method_corr.reindex(index=method_order, columns=method_order)
+    mvals = method_corr.values[np.triu_indices_from(method_corr.values, k=1)]
+    print(f"method PCM correlation ({label}): mean r={np.nanmean(mvals):.4f}  min r={np.nanmin(mvals):.4f}")
+    caption = ("Correlation coefficients between NMR methods across all 22 solutes.\n"
+               f"PCM corrections for chloroform with the {BASIS} basis.")
+    delta22_plots.plot_correlation_matrix(method_corr, f"NMR Methods are Highly Correlated ({nucleus})", caption,
+                                          colormap="Reds", show_values=True,
+                                          save_path=figure_path(f"si_figure_s04c_{label}.png"))
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figures S11 and S13: MagNET reproduces the DFT corrections
+# ----------------------------------------------------------------------------
+_DFT_NN_LOAD = r"""
+dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)
+"""
+
+si_figure_s11 = [
+    md(r"""
+# SI Figure S11: MagNET vs DFT rovibrational (QCD) corrections
+
+Four panels: **A** MagNET (NN) vs B3LYP/cc-pVDZ rovibrational (QCD) correction per site, both nuclei;
+**B** the NN-minus-DFT error distribution; **C/D** every proton/carbon site's DFT and NN QCD
+correction, one column per solute. (QCD is gas-phase, so no solvent axis.)
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(_DFT_NN_LOAD),
+    md("## Panels A and B: MagNET vs DFT QCD correction, both nuclei"),
+    code(r"""
+qcd = delta22.compare_dft_nn(dft, nn, "qcd", keys=("solute", "site", "nucleus"))
+delta22_plots.plot_qcd_scatter(qcd, save_path=figure_path("si_figure_s11a_scatter.png"))
+delta22_plots.plot_qcd_error_histogram(qcd, save_path=figure_path("si_figure_s11b_hist.png"))
+plt.show()
+"""),
+    md("## Panel C: every proton site's QCD correction, one column per solute"),
+    code(r"""
+delta22_plots.plot_qcd_correction_by_site(delta22.qcd_correction_by_site(dft, nn, nucleus="H"), "H", "upper left",
+                                          save_path=figure_path("si_figure_s11c_sites_1H.png"))
+plt.show()
+"""),
+    md("## Panel D: every carbon site's QCD correction, one column per solute"),
+    code(r"""
+delta22_plots.plot_qcd_correction_by_site(delta22.qcd_correction_by_site(dft, nn, nucleus="C"), "C", "upper right",
+                                          save_path=figure_path("si_figure_s11d_sites_13C.png"))
+plt.show()
+"""),
+]
+
+si_figure_s13 = [
+    md(r"""
+# SI Figure S13: MagNET-x vs DFT explicit-solvent corrections
+
+Four panels, both nuclei: **A** MagNET-x (NN) vs DFT explicit-solvent correction per site/solvent,
+Desmond and OpenMM overlaid; **B** the NN-minus-DFT error distribution; **C** every site's DFT and NN
+correction in chloroform, one column per solute; **D** the semi-parsimonious composite model's fit accuracy vs experiment across the four OpenMM solvents, split by engine and
+DFT-vs-NN source.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(_DFT_NN_LOAD),
+    code(r"""
+# openMM/Desmond engine hues (used by the scatter/histogram/site/fitting-accuracy panels below)
+ENGINE_COLORS = {"openMM": "#2E86AB", "desmond": "#A23B72"}
+ENGINE_LABELS = {"desmond": "Desmond", "openMM": "OpenMM"}
+"""),
+    md("## Panels A and B: MagNET-x vs DFT explicit corrections, Desmond and OpenMM overlaid"),
+    code(r"""
+explicit_by_engine = delta22.compare_dft_nn_by_engine(dft, nn, keys=("solute", "site", "nucleus", "solvent"))
+delta22_plots.plot_dft_nn_scatter_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,
+                                            save_path=figure_path("si_figure_s13a_scatter.png"))
+delta22_plots.plot_dft_nn_error_histogram_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,
+                                                    save_path=figure_path("si_figure_s13b_hist.png"))
+plt.show()
+"""),
+    md("## Panel C: every proton/carbon site's explicit correction in chloroform"),
+    code(r"""
+for nucleus, label in [("H", "Proton"), ("C", "Carbon")]:
+    pairs = delta22.explicit_correction_dft_nn_pairs(dft, nn, "chloroform", nucleus)
+    delta22_plots.plot_explicit_correction_by_site(
+        pairs, f"All {label} Sites: DFT vs NN Explicit Corrections (chloroform)", ENGINE_COLORS,
+        save_path=figure_path(f"si_figure_s13c_sites_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+    md("## Panel D: fitting accuracy of the semi-parsimonious composite model"),
+    code(_N_SPLITS_CELL),
+    code(r"""
+SOLVENT_ORDER = ["chloroform", "methanol", "TIP4P", "benzene"]
+solutes = sorted(dft["solute"].unique())
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    fitting = delta22.si_s13d_fitting_accuracy(dft, nn, SOLVENT_ORDER, N_SPLITS, solutes, nucleus=nucleus)
+    delta22_plots.plot_fitting_accuracy_boxplot(
+        fitting, SOLVENT_ORDER, label, ENGINE_COLORS,
+        save_path=figure_path(f"si_figure_s13d_fitting_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S5: PCM Captures Bulk Solvent Effects
+# ----------------------------------------------------------------------------
+si_figure_s05 = [
+    md(r"""
+# SI Figure S5: per-solvent PCM benefit vs bulk dielectric constant and polarizability
+
+Per-solvent PCM benefit (% test-RMSE reduction) vs bulk dielectric constant and polarizability, for
+the MagNET-Zero reference method per nucleus (WP04 ¹H, wB97X-D ¹³C). Panels A (¹H) and B (¹³C)
+show all solvents; C repeats ¹H excluding aromatics and trifluoroethanol.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(_N_SPLITS_CELL),
+    code(r"""
+dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+solutes = sorted(dft["solute"].unique())
+benefit = {}
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    method = delta22.MAGNET_PCM_OUTPUT_METHODS[nucleus]
+    benefit[nucleus] = delta22.pcm_benefit_per_solvent(dft, method, "pcSseg2", "aimnet2",
+                                                       delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS,
+                                                       solutes=solutes, nucleus=nucleus)
+    print(f"{label} ({method}):"); print(benefit[nucleus].round(1).to_string())
+"""),
+    md("## S5A (1H, all solvents), S5B (13C, all solvents), and S5C (1H, excluding aromatics and trifluoroethanol)"),
+    code(r"""
+# panels A (1H) and B (13C): all solvents
+for nucleus, letter, nuc_label in [("H", "A", "H"), ("C", "B", "C")]:
+    delta22_plots.plot_pcm_benefit_vs_properties(
+        benefit[nucleus], delta22.SOLVENT_DIELECTRIC, delta22.SOLVENT_POLARIZABILITY, nuc_label,
+        save_path=figure_path(f"si_figure_s05{letter}_all.png"))
+# panel C is 1H only, excluding the aromatic solvents and trifluoroethanol (matches the published SI)
+delta22_plots.plot_pcm_benefit_vs_properties(
+    benefit["H"], delta22.SOLVENT_DIELECTRIC, delta22.SOLVENT_POLARIZABILITY, "H",
+    exclude=["benzene", "toluene", "chlorobenzene", "trifluoroethanol"],
+    title_extra="Excluding Aromatic Solvents and Trifluoroethanol",
+    save_path=figure_path("si_figure_s05A_no_aromatics_tfe.png"))
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S8 (panels B-D): Convergence of Explicit Solvent Corrections
+# ----------------------------------------------------------------------------
+si_figure_s08 = [
+    md(r"""
+# SI Figure S8 (panels B-D): explicit-solvent correction vs MD-frame count
+
+For an example site (AcOH proton in chloroform), the explicit-solvent correction vs number of MD frames included (running average): OpenMM vs Desmond (B1) and DFT vs MagNET-x/"NN" at fixed OpenMM (B3), with per-frame distributions (B2/B4) and autocorrelation (C). Panel D shows which frames have computed DFT data per solvent (also the source of B's histogram normalization).
+"""),
+    code(_BOOTSTRAP),
+    code(r"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+import delta22
+import delta22_reader
+import delta22_plots
+import paths
+"""),
+    code(_SETUP),
+    code(r"""
+site_atoms = delta22_reader.load_site_atom_data(XLSX, verbose=False)
+idx = delta22.site_atom_indices(site_atoms.loc[("AcOH", "H", "H"), "atom_numbers"])
+SOLUTE, SOLVENT = "AcOH", "chloroform"
+"""),
+    code(r"""
+# OpenMM/Desmond engine hues; DFT vs NN reuse the OpenMM hue (DFT lightened, NN full strength)
+ENGINE_COLORS = {"openMM": "#2E86AB", "desmond": "#A23B72"}
+ENGINE_LABELS = {"openMM": "OpenMM", "desmond": "Desmond"}
+
+SOURCE_COLORS = {"dft": delta22_plots.lighten_color(ENGINE_COLORS["openMM"]), "nn": ENGINE_COLORS["openMM"]}
+SOURCE_LABELS = {"dft": "DFT", "nn": "NN"}
+LABEL_COLORS = {**ENGINE_COLORS, **SOURCE_COLORS}
+LABEL_NAMES = {**ENGINE_LABELS, **SOURCE_LABELS}
+"""),
+    md(r"""
+## Panel B1/B2: OpenMM vs Desmond
+
+Running-average convergence (B1) and per-frame distribution (B2), comparing the two MD engines.
+"""),
+    code(r"""
+running_by_engine, finals_by_engine, per_frame_by_engine, n_frames_by_engine = {}, {}, {}, {}
+for engine in ["openMM", "desmond"]:
+    perturbed = delta22_reader.load_perturbed_shieldings(DELTA22_HDF5, SOLUTE, SOLVENT, engine, "dft")
+    per_frame = delta22.frame_corrections(perturbed, idx)
+    running = delta22.running_average(per_frame)
+    running_by_engine[engine] = running
+    finals_by_engine[engine] = running[~np.isnan(running)][-1]
+    per_frame_by_engine[engine] = per_frame
+    n_frames_by_engine[engine] = len(per_frame)          # the total trajectory length, valid or not
+    print(f"{engine}: {int(np.sum(~np.isnan(per_frame)))} / {len(per_frame)} valid frames, "
+          f"converged correction {finals_by_engine[engine]:.4f} ppm")
+
+delta22_plots.plot_frame_convergence(
+    running_by_engine, finals_by_engine, LABEL_COLORS, LABEL_NAMES,
+    title=f"Convergence of Explicit Solvent Corrections\n({SOLUTE} in {SOLVENT})",
+    save_path=figure_path("si_figure_s08_b1_convergence_engine.png"))
+plt.show()
+
+delta22_plots.plot_frame_correction_histogram(
+    per_frame_by_engine, n_frames_by_engine, LABEL_COLORS, LABEL_NAMES,
+    title="Distribution of Frame-wise Corrections",
+    save_path=figure_path("si_figure_s08_b2_histogram_engine.png"))
+plt.show()
+"""),
+    md(r"""
+## Panel B3/B4: DFT vs NN
+
+Same running average (B3) and per-frame distribution (B4), at fixed OpenMM engine, comparing DFT vs
+MagNET-x ("NN") shieldings.
+"""),
+    code(r"""
+running_by_source, finals_by_source, per_frame_by_source, n_frames_by_source = {}, {}, {}, {}
+for source in ["dft", "nn"]:
+    perturbed = delta22_reader.load_perturbed_shieldings(DELTA22_HDF5, SOLUTE, SOLVENT, "openMM", source)
+    per_frame = delta22.frame_corrections(perturbed, idx)
+    running = delta22.running_average(per_frame)
+    running_by_source[source] = running
+    finals_by_source[source] = running[~np.isnan(running)][-1]
+    per_frame_by_source[source] = per_frame
+    n_frames_by_source[source] = len(per_frame)
+    print(f"{source}: {int(np.sum(~np.isnan(per_frame)))} / {len(per_frame)} valid frames, "
+          f"converged correction {finals_by_source[source]:.4f} ppm")
+
+delta22_plots.plot_frame_convergence(
+    running_by_source, finals_by_source, LABEL_COLORS, LABEL_NAMES, xlabel="Number of MD Frames",
+    title=f"OpenMM DFT vs. NN Convergence Comparison\n({SOLUTE} in {SOLVENT})",
+    save_path=figure_path("si_figure_s08_b3_convergence_source.png"))
+plt.show()
+
+delta22_plots.plot_frame_correction_histogram(
+    per_frame_by_source, n_frames_by_source, LABEL_COLORS, LABEL_NAMES,
+    title="Distribution of Frame-wise Corrections",
+    save_path=figure_path("si_figure_s08_b4_histogram_source.png"))
+plt.show()
+"""),
+    md(r"""
+## Panel C: autocorrelation, DFT vs NN
+
+Autocorrelation of the per-frame correction to lag 200 (~100 frames per trajectory repeat), DFT vs NN,
+same OpenMM trajectory as B3/B4.
+"""),
+    code(r"""
+autocorr_by_source = {source: delta22.autocorrelation(per_frame_by_source[source], max_lag=200)
+                      for source in ["dft", "nn"]}
+for source, autocorr in autocorr_by_source.items():
+    print(f"{source}: lag-1 autocorrelation {autocorr[1]:.3f}")
+
+delta22_plots.plot_frame_autocorrelation(
+    autocorr_by_source, LABEL_COLORS, LABEL_NAMES,
+    title=f"Autocorrelation of Frame-wise Corrections\n({SOLUTE} in {SOLVENT})",
+    save_path=figure_path("si_figure_s08_c_autocorrelation.png"))
+plt.show()
+"""),
+    md(r"""
+## Panel D: frame data availability
+
+Which trajectory frames have computed DFT shielding for AcOH, across all 12 Desmond and 4 OpenMM
+solvents. DFT was computed for a non-contiguous subset (compute-budget limits; jobs queued randomly).
+"""),
+    code(r"""
+# row order matching the published panel: chloroform first, then by solvent class (aprotic ->
+# protic -> aromatic); OpenMM only has the four explicit-solvent solvents
+DESMOND_ORDER = ["chloroform", "tetrahydrofuran", "dichloromethane", "acetone", "acetonitrile",
+                 "dimethylsulfoxide", "trifluoroethanol", "methanol", "TIP4P",
+                 "benzene", "toluene", "chlorobenzene"]
+OPENMM_ORDER = ["chloroform", "methanol", "TIP4P", "benzene"]
+
+grids_by_engine, solvents_by_engine = {}, {}
+for engine, solvents in [("desmond", DESMOND_ORDER), ("openMM", OPENMM_ORDER)]:
+    grid, frame_counts = delta22.frame_validity_grid(DELTA22_HDF5, SOLUTE, solvents, engine, shield_type="dft")
+    grids_by_engine[engine] = grid
+    solvents_by_engine[engine] = [delta22_plots.display_solvent_name(s) for s in solvents]
+    print(f"{engine}: {len(solvents)} solvents, frame counts {dict(zip(solvents, frame_counts))}")
+
+delta22_plots.plot_frame_validity_heatmaps(grids_by_engine, solvents_by_engine, ENGINE_LABELS, solute=SOLUTE,
+                                           save_path=figure_path("si_figure_s08_d_frame_validity.png"))
+plt.show()
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# Figure 2A + SI Figure S1: the accuracy-vs-cost Pareto frontier
+# ----------------------------------------------------------------------------
+_PARETO_LOAD = r"""
+# Load the flat DFT and MagNET tables and the timing tables.
+dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)
+nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)
+dft_gas_timings, nn_timings = delta22.load_pareto_timings(DELTA22_HDF5)
+
+# One point per method/basis/geometry/nucleus/solvent, plus the solvent-averaged rows, each with its
+# mean test RMSE and total compute time. MagNET appears as a single method (aimnet2, basis "N/A").
+# Figure 2A uses 100 seeded splits, not the 250 the other delta-22 panels use, and trains on the
+# first 10 shuffled solutes / tests on the rest; fig2a_pareto_points handles that split convention.
+PARETO_N_SPLITS = 100
+points = delta22.fig2a_pareto_points(dft, nn, dft_gas_timings, nn_timings, n_splits=PARETO_N_SPLITS)
+print(points["nmr_method"].nunique(), "methods;",
+      "MagNET total time", float(points.query("nmr_method=='MagNET'")["total_time"].iloc[0]), "s")
+"""
+
+# Figure 2A (analysis/manuscript_figures/fig2a_pareto.ipynb) is hand-authored with its plotting
+# inline; its data comes from delta22.py.
+
+# SI Figure S1 (analysis/si_figures/si_figure_s01.ipynb) is hand-authored too, with the same Pareto
+# panel engine as Figure 2A run for both nuclei; its data comes from delta22.py.
+
+# ----------------------------------------------------------------------------
+# SI Tables S1/S2: the Pareto plot data (same chloroform rows as SI Figure S1)
+# ----------------------------------------------------------------------------
+si_table_s01_s02_pareto = [
+    md(r"""
+# SI Tables S1 and S2: the Pareto plot data
+
+Each method's total compute time and CDCl3 test RMSE, for proton (S1) and carbon (S2): a curated
+55-row subset (MagNET, one AIMNet2-geometry row, and the DFT grid on PBE0/cc-pVTZ geometries).
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS_TABLE_NPPD),
+    code(_SETUP_TABLE),
+    code(_PARETO_LOAD),
+    md("## Tables S1 and S2"),
+    code(r"""
+# full column set matching the SI tables: identity + fitting RMSE + the three time components and
+# their log10 (all solvents here are chloroform). geometry_time and nmr_time sum to total_time.
+COLUMNS = ["geometry_type", "nmr_method", "basis", "solvent", "fitting_RMSE",
+           "geometry_time", "nmr_time", "total_time", "log10_total_time"]
+
+curated = {}
+for nucleus, label in [("H", "S1"), ("C", "S2")]:
+    sub = delta22.pareto_table_curated(points, nucleus).copy()
+    sub["log10_total_time"] = np.log10(sub["total_time"])
+    sub = sub[COLUMNS]
+    curated[label] = sub
+    print(f"Table {label} ({nucleus}): {len(sub)} rows x {sub.shape[1]} columns")
+    print(sub.to_string(index=False), "\n")
+
+# write the two tables to this notebook's documents/ folder, one sheet per SI table
+out = document_path("si_table_s01_s02_pareto.xlsx")
+with pd.ExcelWriter(out) as writer:
+    curated["S1"].to_excel(writer, sheet_name="Table S1 (1H)", index=False)
+    curated["S2"].to_excel(writer, sheet_name="Table S2 (13C)", index=False)
+print("wrote", os.path.relpath(out, REPO))
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S6: Explicit Solvation is Broadly Beneficial
+# ----------------------------------------------------------------------------
+si_figure_s06 = [
+    md(r"""
+# SI Figure S6: per-solvent test RMSE of six solvent-correction models
+
+Per nucleus, the per-solvent test RMSE of six solvent-correction models: SotA implicit (single-slope
+PCM), the paper's implicit fit (stationary + PCM as separate terms), and explicit (Desmond), each
+with and without a vibrational correction (QCD for ¹H, Desmond vibration for ¹³C). Level of theory:
+dsd_pbep86/pcSseg3 on pbe0_tz geometries, PCM substituted from b3lyp_d3bj/pcSseg3.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(r"""
+query = delta22.add_composite_columns(delta22.load_query_df_dft(
+    DELTA22_HDF5, XLSX, pcm_reference_method="b3lyp_d3bj", pcm_reference_basis="pcSseg3", verbose=False))
+solutes = sorted(query["solute"].unique())
+
+METHOD, BASIS, GEOM = "dsd_pbep86", "pcSseg3", "pbe0_tz"
+
+# same six categories for both nuclei; the "+ Vibrations" term differs (qcd for H, desmond_vib for C)
+LADDER = {
+    "H": {
+        "stationary_plus_pcm": "SotA Implicit Solvent",
+        "stationary_plus_pcm_plus_qcd": "SotA Implicit Solvent + Vibrations",
+        "stationary + pcm": "Implicit Solvent",
+        "stationary_plus_qcd + pcm": "Implicit Solvent + Vibrations",
+        "stationary + desmond": "Explicit Solvent",
+        "stationary_plus_qcd + desmond": "Explicit Solvent + Vibrations"},
+    "C": {
+        "stationary_plus_pcm": "SotA Implicit Solvent",
+        "stationary_plus_pcm_plus_des_vib": "SotA Implicit Solvent + Vibrations",
+        "stationary + pcm": "Implicit Solvent",
+        "stationary_plus_des_vib + pcm": "Implicit Solvent + Vibrations",
+        "stationary + desmond": "Explicit Solvent",
+        "stationary_plus_des_vib + desmond": "Explicit Solvent + Vibrations"},
+}
+"""),
+    code(r"""
+# one dark/light pair per solvent-treatment family (SotA implicit / implicit / explicit), light = +Vibrations
+S6_COLORS = ["#C4B037", "#F5EDA0", "#A72608", "#E4A0A0", "#5F7C8A", "#B9E7DF"]
+"""),
+    md("## Formula ladder per nucleus"),
+    code(_N_SPLITS_CELL),
+    code(r"""
+for nucleus, labels in LADDER.items():
+    results = delta22.fig3d_formula_regressions(query, METHOD, BASIS, GEOM, list(labels),
+                                                delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS,
+                                                solutes=solutes, nucleus=nucleus)
+    med = results.groupby("formula")["test_RMSE"].median()
+    print(f"--- {nucleus} ({METHOD}) median test RMSE ---")
+    for formula in labels:
+        print(f"  {labels[formula]:38s} {med[formula]:.4f}")
+    delta22_plots.plot_formula_ladder_boxplot(
+        results, labels, delta22.SOLVENT_GROUPS, nucleus=nucleus, colors=S6_COLORS,
+        title=f"Explicit Solvation + Vibrational Effects Improve Predictions Across Solvents ({nucleus} Nucleus)",
+        save_path=figure_path(f"si_figure_s06_ladder_{'1H' if nucleus == 'H' else '13C'}.png"))
+plt.show()
+"""),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Table S6: Explicit solvent corrections are superior for predicting solvent-induced shifts
+# ----------------------------------------------------------------------------
+# SI table row order (chloroform is the reference and so is excluded)
+_S6_SOLVENT_ORDER = [
+    "tetrahydrofuran", "dichloromethane", "acetone", "acetonitrile", "dimethylsulfoxide",
+    "trifluoroethanol", "methanol", "TIP4P", "benzene", "toluene", "chlorobenzene",
+]
+si_table_s06 = [
+    md(r"""
+# SI Table S6: per-solvent implicit vs explicit fit RMSE (chloroform reference, ¹H)
+
+Per solvent (chloroform reference), fit RMSE of experimental solvent-induced ¹H shifts using implicit
+(PCM) vs explicit (Desmond), and the percent improvement.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS_TABLE_PD),
+    code(_SETUP_TABLE),
+    code(_FIG4_LOAD),
+    md("## Table S6: per-solvent implicit vs explicit fit RMSE (chloroform reference, 1H)"),
+    code(r"""
+rows = []
+for solvent in _S6_SOLVENT_ORDER:
+    sp = delta22.solvent_pair_differences(one, solvent, "chloroform", nucleus="H", explicit="desmond")
+    implicit = delta22.fit_differences_to_experimental(sp, "implicit_diff")["rmse"]
+    explicit = delta22.fit_differences_to_experimental(sp, "explicit_diff")["rmse"]
+    rows.append({"solvent": solvent,
+                 "implicit_fitting_rmse": round(implicit, 4),
+                 "explicit_fitting_rmse": round(explicit, 4),
+                 "explicit_improvement_pct": round(100.0 * (implicit - explicit) / implicit, 2)})
+table_s6 = pd.DataFrame(rows)
+print(table_s6.to_string(index=False))
+
+# write the table to this notebook's documents/ folder
+out = document_path("si_table_s06_solvent_corrections.xlsx")
+with pd.ExcelWriter(out) as writer:
+    table_s6.to_excel(writer, sheet_name="Table S6", index=False)
+print("wrote", os.path.relpath(out, REPO))
+""".replace("_S6_SOLVENT_ORDER", repr(_S6_SOLVENT_ORDER))),
+]
+
+
+# ----------------------------------------------------------------------------
+# SI Figure S14: End-to-End Neural Network Acceleration Delivers Equivalent Performance
+# ----------------------------------------------------------------------------
+si_figure_s14 = [
+    md(r"""
+# SI Figure S14: test RMSE with DFT vs MagNET features, by solvent
+
+For the four explicit-solvent solvents (chloroform, methanol, TIP4P water, benzene), per-solvent test
+RMSE of predicting experimental ¹H/¹³C shifts with **DFT** features (solid) vs end-to-end **MagNET/NN**
+features (lightened), each under three conditions (implicit SOTA, explicit, explicit + vibrations).
+Two pages, one per MD engine (Desmond, OpenMM); nitromethane dropped.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(_N_SPLITS_CELL),
+    code(r"""
+dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))
+nn = delta22.add_composite_columns(delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False))
+print(len(dft), "DFT rows;", len(nn), "NN rows")
+
+LABELS = ["Implicit Solvent (SOTA)", "Explicit Solvent", "Explicit Solvent + Vibrations"]
+SS_LABELS = {"TIP4P": "Water (TIP4P)"}
+S14_SOLVENTS = ["chloroform", "methanol", "TIP4P", "benzene"]   # the 4 explicit solvents, SI order
+"""),
+    md("## Desmond page (DFT vs NN)"),
+    code(r"""
+# implicit = 2-term (stationary + pcm); explicit = stationary + desmond;
+# explicit + vibrations = stationary_plus_qcd + desmond (1H) / stationary_plus_des_vib + desmond (13C)
+DESMOND_FORMULAS = {
+    "H":  ["stationary + pcm", "stationary + desmond", "stationary_plus_qcd + desmond"],
+    "C":  ["stationary + pcm", "stationary + desmond", "stationary_plus_des_vib + desmond"],
+}
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    formulas = DESMOND_FORMULAS[nucleus]
+    delta22_plots.plot_ss_boxplot_dft_vs_nn(
+        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),
+        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),
+        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,
+        save_path=figure_path(f"si_figure_s14_desmond_{label}.png"))
+plt.show()
+"""),
+    md("## OpenMM page (DFT vs NN)"),
+    code(r"""
+# implicit = 1-term composite (stationary_plus_pcm); explicit = stationary + openMM;
+# explicit + vibrations = stationary_plus_qcd + openMM (1H) / stationary_plus_op_vib + openMM (13C)
+OPENMM_FORMULAS = {
+    "H":  ["stationary_plus_pcm", "stationary + openMM", "stationary_plus_qcd + openMM"],
+    "C":  ["stationary_plus_pcm", "stationary + openMM", "stationary_plus_op_vib + openMM"],
+}
+for nucleus, label in [("H", "1H"), ("C", "13C")]:
+    formulas = OPENMM_FORMULAS[nucleus]
+    delta22_plots.plot_ss_boxplot_dft_vs_nn(
+        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),
+        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),
+        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,
+        save_path=figure_path(f"si_figure_s14_openmm_{label}.png"))
+plt.show()
+"""),
+]
+
+
+# name -> (cells, path relative to repo root). Notebooks are grouped by role (manuscript/si_figures/
+# si_tables), not by dataset.
+NOTEBOOKS = {
+    "si_figure_s06": (si_figure_s06, "analysis/si_figures/si_figure_s06.ipynb"),
+    "si_figure_s14": (si_figure_s14, "analysis/si_figures/si_figure_s14.ipynb"),
+    "si_table_s06": (si_table_s06, "analysis/si_tables/si_table_s06_solvent_corrections.ipynb"),
+    "si_table_s01_s02_pareto": (si_table_s01_s02_pareto, "analysis/si_tables/si_table_s01_s02_pareto.ipynb"),
+    "si_figure_s16_s18": (si_figure_s16_s18, "analysis/si_figures/si_figure_s16_s18.ipynb"),
+    "si_figure_s07": (si_figure_s07, "analysis/si_figures/si_figure_s07.ipynb"),
+    "si_figure_s04": (si_figure_s04, "analysis/si_figures/si_figure_s04.ipynb"),
+    "si_figure_s11": (si_figure_s11, "analysis/si_figures/si_figure_s11.ipynb"),
+    "si_figure_s13": (si_figure_s13, "analysis/si_figures/si_figure_s13.ipynb"),
+    "si_figure_s05": (si_figure_s05, "analysis/si_figures/si_figure_s05.ipynb"),
+    "si_figure_s08": (si_figure_s08, "analysis/si_figures/si_figure_s08_panelsBCD.ipynb"),
+}
+
+if __name__ == "__main__":
+    import os
+    import sys
+    here = os.path.dirname(os.path.abspath(__file__))
+    repo = os.path.abspath(os.path.join(here, "..", ".."))
+    names = sys.argv[1:]
+    if not names:
+        print("usage: python3 build_nb_delta22.py <name> [<name> ...]")
+        print("regenerates ONLY the named notebook(s) -- pick just the one(s) you edited,")
+        print("since regenerating clears a notebook's execution outputs.")
+        print("available names:", ", ".join(NOTEBOOKS))
+        raise SystemExit(1)
+    unknown = [n for n in names if n not in NOTEBOOKS]
+    if unknown:
+        raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}")
+    for name in names:
+        cells, relpath = NOTEBOOKS[name]
+        _save(cells, os.path.join(repo, relpath))

analysis/code/build_nb_magnet_benchmark.py

@@ -0,0 +1,159 @@
+"""Source of truth for the MagNET-benchmark table notebooks (Tables S3/S4 and Table S5). Edit the
+cell sources here, then regenerate and re-execute ONLY the notebook you touched, e.g.:
+
+    python3 build_nb_magnet_benchmark.py si_table_s03_s04_performance
+    jupyter nbconvert --to notebook --execute --inplace ../si_tables/si_table_s03_s04_performance.ipynb
+
+    python3 build_nb_magnet_benchmark.py si_table_s05_qcd
+    jupyter nbconvert --to notebook --execute --inplace ../si_tables/si_table_s05_qcd.ipynb
+
+Regenerating overwrites the .ipynb (clearing its execution outputs), so do not hand-edit the
+notebooks and regenerate only what you changed. All real code lives in magnet_benchmark.py (the
+PUBLISHED reference values and the exact-reproduction group mapping); the notebooks only read
+data/magnet_test_predictions/ through it. The notebooks live in analysis/si_tables/ (grouped by
+role), not here.
+
+Run with no arguments to list the available notebook names.
+"""
+import os
+import sys
+
+from nb_build import md, code, save_notebook as _save
+
+
+_BOOTSTRAP = r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/magnet_test_predictions", "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""
+
+_IMPORTS = r"""
+import pandas as pd
+import magnet_test_predictions_reader
+import magnet_benchmark
+import paths
+"""
+
+_SETUP = r"""
+PREDICTIONS = paths.dataset_file("magnet_test_predictions", root=REPO)
+
+def document_path(name):
+    # table/spreadsheet outputs go under this notebook's documents/ folder, created on first save
+    os.makedirs("documents", exist_ok=True)
+    return os.path.join("documents", name)
+"""
+
+si_table_s03_s04_performance = [
+    md(r"""
+# Tables S3 and S4: MagNET performance across geometries
+
+How accurately MagNET reproduces its DFT training reference (PBE0/pcSseg-1) as geometries move from
+stationary to vibrated to solvated, for the foundation model and the chloroform MagNET-x.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(r"""
+rows = magnet_benchmark.exact_stats_table(PREDICTIONS, magnet_test_predictions_reader)
+res = pd.DataFrame(rows).set_index(["nucleus", "model", "test_set"])
+
+table_rows = []
+for nucleus in ("1H", "13C"):
+    for model in magnet_benchmark.MODELS:
+        for ts in magnet_benchmark.TEST_SETS:
+            r = res.loc[(nucleus, model, ts)]
+            pmed, pmae, prmse = magnet_benchmark.PUBLISHED[nucleus][(model, ts)]
+            table_rows.append(dict(nucleus=nucleus, model=model, test_set=ts, n=int(r.n),
+                                   median_repro=round(r.median_ae, 6), median_SI=round(pmed, 6),
+                                   mae_repro=round(r.mae, 6), mae_SI=round(pmae, 6),
+                                   rmse_repro=round(r.rmse, 5), rmse_SI=round(prmse, 5)))
+table = pd.DataFrame(table_rows)
+table_s3 = table[table.nucleus == "1H"].drop(columns="nucleus")
+table_s4 = table[table.nucleus == "13C"].drop(columns="nucleus")
+print("Table S3 (1H):"); display(table_s3)
+print("Table S4 (13C):"); display(table_s4)
+
+# write the reproduced tables to this notebook's documents/ folder, one sheet per SI table
+out = document_path("si_table_s03_s04_performance.xlsx")
+with pd.ExcelWriter(out) as writer:
+    table_s3.to_excel(writer, sheet_name="Table S3 (1H)", index=False)
+    table_s4.to_excel(writer, sheet_name="Table S4 (13C)", index=False)
+print("wrote", os.path.relpath(out, REPO))
+"""),
+    md("## Exact-reproduction check"),
+    code(r"""
+# every row's reproduced median/MAE/RMSE should match the published SI value to a few parts per million
+max_dev = (table[["median_repro", "median_SI"]].diff(axis=1).iloc[:, -1].abs().max(),
+          table[["mae_repro", "mae_SI"]].diff(axis=1).iloc[:, -1].abs().max(),
+          table[["rmse_repro", "rmse_SI"]].diff(axis=1).iloc[:, -1].abs().max())
+print("largest reproduced-vs-published deviation (median, MAE, RMSE):", max_dev)
+assert max(max_dev) < 1e-3, "a row diverged from the SI by more than float rounding"
+"""),
+]
+
+
+si_table_s05_qcd = [
+    md(r"""
+# Table S5: Performance statistics for predicting QCD corrections
+
+Foundation MagNET's accuracy at predicting the rovibrational (QCD) correction (stationary vs
+trajectory-averaged shielding) over qcdtraj2500 (2500 molecules), ¹H and ¹³C, all shieldings at
+PBE0/pcSseg-1.
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(r"""
+rows = magnet_benchmark.qcd_stats_table(PREDICTIONS, magnet_test_predictions_reader)
+table_rows = []
+for r in rows:
+    nucleus = r["model"].split("(")[1].rstrip(")")   # "1H" / "13C"
+    pmed, pmae, prmse = magnet_benchmark.PUBLISHED_S5[nucleus]
+    table_rows.append(dict(model=r["model"], n=int(r["n"]),
+                           median_repro=round(r["median_ae"], 8), median_SI=round(pmed, 8),
+                           mae_repro=round(r["mae"], 8), mae_SI=round(pmae, 8),
+                           rmse_repro=round(r["rmse"], 8), rmse_SI=round(prmse, 8)))
+table_s5 = pd.DataFrame(table_rows)
+print("Table S5 (QCD corrections):"); display(table_s5)
+
+# write the reproduced table to this notebook's documents/ folder
+out = document_path("si_table_s05_qcd.xlsx")
+with pd.ExcelWriter(out) as writer:
+    table_s5.to_excel(writer, sheet_name="Table S5", index=False)
+print("wrote", os.path.relpath(out, REPO))
+"""),
+    md("## Exact-reproduction check"),
+    code(r"""
+# every reproduced median/MAE/RMSE should match the published SI value to a few parts per million
+dev = max((table_s5[["median_repro", "median_SI"]].diff(axis=1).iloc[:, -1].abs().max(),
+           table_s5[["mae_repro", "mae_SI"]].diff(axis=1).iloc[:, -1].abs().max(),
+           table_s5[["rmse_repro", "rmse_SI"]].diff(axis=1).iloc[:, -1].abs().max()))
+print("largest reproduced-vs-published deviation:", dev)
+assert dev < 1e-3, "a row diverged from the SI by more than float rounding"
+"""),
+]
+
+
+# name -> (cells, path relative to repo root)
+NOTEBOOKS = {
+    "si_table_s03_s04_performance": (si_table_s03_s04_performance,
+                                     "analysis/si_tables/si_table_s03_s04_performance.ipynb"),
+    "si_table_s05_qcd": (si_table_s05_qcd, "analysis/si_tables/si_table_s05_qcd.ipynb"),
+}
+
+if __name__ == "__main__":
+    here = os.path.dirname(os.path.abspath(__file__))
+    repo = os.path.abspath(os.path.join(here, "..", ".."))
+    names = sys.argv[1:]
+    if not names:
+        print("available names:", ", ".join(NOTEBOOKS))
+        sys.exit(0)
+    unknown = [n for n in names if n not in NOTEBOOKS]
+    if unknown:
+        raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}")
+    for name in names:
+        cells, relpath = NOTEBOOKS[name]
+        _save(cells, os.path.join(repo, relpath))

analysis/code/build_nb_scaling.py

@@ -0,0 +1,126 @@
+"""Source of truth for the Tables S10/S11 notebook (MagNET-Zero/MagNET-PCM scaling factors). Edit
+the cell sources here, then regenerate:
+
+    python3 build_nb_scaling.py si_table_s10_s11_scaling
+    jupyter nbconvert --to notebook --execute --inplace ../si_tables/si_table_s10_s11_scaling.ipynb
+
+Regenerating overwrites the .ipynb (clearing its execution outputs). All real code lives in
+scaling_factors.py (the published tables and the re-derivation from delta-22); the notebook only
+runs it. The notebook lives in analysis/si_tables/ (grouped by role), not here.
+
+Run with no arguments to list the available notebook names.
+"""
+import os
+import sys
+
+from nb_build import md, code, save_notebook as _save
+
+
+_BOOTSTRAP = r"""
+import os, sys
+
+# make the in-repo modules importable (not pip-installed)
+REPO = os.path.abspath("../..")
+for _p in ("data/delta22", "data/scaling_factors", "data/applications",
+           "analysis/code", "analysis/code/shared"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+"""
+
+_IMPORTS = r"""
+import numpy as np
+import pandas as pd
+import scaling_factors
+import scaling_factors_reader
+import build_composite_model
+import paths
+from applications_reader import Applications
+"""
+
+_SETUP = r"""
+DELTA22 = paths.dataset_file("delta22", root=REPO)
+XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+def document_path(name):
+    os.makedirs("documents", exist_ok=True)
+    return os.path.join("documents", name)
+"""
+
+si_table_s10_s11_scaling = [
+    md(r"""
+# Tables S10 and S11: MagNET-Zero/MagNET-PCM scaling factors
+
+Per-solvent linear coefficients (intercept, stationary, pcm) mapping MagNET-Zero shieldings plus a
+MagNET-PCM correction to predicted shifts: ¹H (S10) and ¹³C (S11).
+"""),
+    code(_BOOTSTRAP),
+    code(_IMPORTS),
+    code(_SETUP),
+    code(r"""
+tables = scaling_factors.published_scaling_tables()   # {"H": Table S10, "C": Table S11}
+print("Table S10 (1H):"); display(tables["H"])
+print("Table S11 (13C):"); display(tables["C"])
+
+out = document_path("si_table_s10_s11_scaling.xlsx")
+with pd.ExcelWriter(out) as writer:
+    tables["H"].reset_index().to_excel(writer, sheet_name="Table S10 (1H)", index=False)
+    tables["C"].reset_index().to_excel(writer, sheet_name="Table S11 (13C)", index=False)
+print("wrote", os.path.relpath(out, REPO))
+"""),
+    md("## Reproducibility check: re-derive from delta-22"),
+    code(r"""
+derived = scaling_factors.build_scaling_tables(DELTA22, XLSX)
+for nucleus in ("H", "C"):
+    p = tables[nucleus]
+    d = derived[nucleus].reindex(p.index)[p.columns]
+    max_dev = float(np.abs(p.values - d.values).max())
+    print(f"{nucleus}: largest published-vs-rederived deviation = {max_dev:.2e}")
+    assert max_dev < 1e-5, f"{nucleus} scaling table diverged from the published values"
+print("both tables reproduce from delta-22")
+"""),
+    md(r"""
+## Deployment tables with reflection symmetrization
+
+The tables above match the published SI exactly. For serving new molecules, the recommended tables
+average each prediction with its mirror image, correcting a reflection-parity error. Reproducing
+them needs the model checkpoints; the shipped tables are shown below.
+"""),
+    code(r"""
+symmetrized = scaling_factors_reader.load_symmetrized_tables()
+print("Symmetrized deployment Table S10 (1H):"); display(symmetrized["H"])
+print("Symmetrized deployment Table S11 (13C):"); display(symmetrized["C"])
+"""),
+    md(r"""
+## Composite-model coefficients (Figure 5C/5D, SI S15)
+
+The natural-products figures use a larger family of per-solvent coefficients fit on delta-22 (the
+`composite_model` group in `applications.hdf5`): OLS fits, 1000-seed bootstrap resamples, their RMSE
+distributions, and the PCM conversion factors. The group regenerates from released inputs alone.
+"""),
+    code(r"""
+reader = Applications(paths.dataset_file("applications", root=REPO))
+regenerated = build_composite_model.regenerate(reader)
+max_dev = build_composite_model.verify(reader, regenerated)
+print(f"composite_model largest regenerated-vs-stored deviation = {max_dev:.2e}")
+"""),
+]
+
+
+# name -> (cells, path relative to repo root)
+NOTEBOOKS = {
+    "si_table_s10_s11_scaling": (si_table_s10_s11_scaling,
+                                 "analysis/si_tables/si_table_s10_s11_scaling.ipynb"),
+}
+
+if __name__ == "__main__":
+    here = os.path.dirname(os.path.abspath(__file__))
+    repo = os.path.abspath(os.path.join(here, "..", ".."))
+    names = sys.argv[1:]
+    if not names:
+        print("available names:", ", ".join(NOTEBOOKS))
+        sys.exit(0)
+    unknown = [n for n in names if n not in NOTEBOOKS]
+    if unknown:
+        raise SystemExit(f"unknown notebook name(s): {unknown}; available: {', '.join(NOTEBOOKS)}")
+    for name in names:
+        cells, relpath = NOTEBOOKS[name]
+        _save(cells, os.path.join(repo, relpath))

analysis/code/composite_models.py

@@ -0,0 +1,395 @@
+"""Reproduces the delta-22 composite-formula ablation study behind SI Figure S15's Excel screenshots.
+
+Rebuilds `ablations.xlsx`: ~25 composite-formula variants (stationary geometry plus some combination
+of implicit PCM, explicit Desmond, and rovibrational QCD corrections) across all 12 delta-22 solvents,
+at two reference levels: DSD-PBEP86/pcSseg-3 (geometry PBE0/tz) and the MagNET-Zero training reference
+(WP04/pcSseg-2 for 1H, wB97X-D/pcSseg-2 for 13C, AIMNet2 geometries). The fitting reuses delta22.py's
+harness (same as si_figure_s06 at the DSD level, widened to the full formula set and both reference
+levels); the openpyxl table-writing layer (color scales, %benefit formulas, best-model block)
+matches the published screenshots.
+
+Nitromethane: kept in (unlike scaling_factors.py, which drops it as an ML outlier). Per the SI, this
+is a plain OLS fit, not ML, so all 22 solutes are used (10 test / 12 train per split).
+"""
+from openpyxl import Workbook
+from openpyxl.styles import Font, PatternFill
+from openpyxl.utils import get_column_letter
+from openpyxl.formatting.rule import ColorScaleRule, FormulaRule
+
+import delta22 as D
+
+# ---------------------------------------------------------------------------
+# formulas and reference-level configuration
+
+# the full ablation formula set (one-term through three-term composites), Desmond explicit-solvent
+# variant. This order is also the display order in "Multi-Term Models".
+FORMULAS = [
+    "stationary",
+    "stationary + pcm",
+    "stationary + desmond",
+    "stationary + qcd",
+    "stationary + desmond_vib",
+    "stationary + pcm + qcd",
+    "stationary + pcm + desmond_vib",
+    "stationary + desmond + qcd",
+    "stationary + desmond + desmond_vib",
+    "stationary_plus_pcm",
+    "stationary_plus_desmond",
+    "stationary_plus_qcd",
+    "stationary_plus_des_vib",
+    "stationary_plus_pcm + qcd",
+    "stationary_plus_pcm + desmond_vib",
+    "stationary_plus_desmond + qcd",
+    "stationary_plus_desmond + desmond_vib",
+    "stationary_plus_qcd + pcm",
+    "stationary_plus_qcd + desmond",
+    "stationary_plus_des_vib + pcm",
+    "stationary_plus_des_vib + desmond",
+    "stationary_plus_pcm_plus_qcd",
+    "stationary_plus_pcm_plus_des_vib",
+    "stationary_plus_desmond_plus_qcd",
+    "stationary_plus_desmond_plus_des_vib",
+]
+
+# display column order for the spreadsheet (distinct from D.DESMOND_SOLVENTS' iteration order)
+ORDERED_SOLVENTS = [
+    "chloroform", "tetrahydrofuran", "dichloromethane", "acetone", "acetonitrile",
+    "dimethylsulfoxide", "trifluoroethanol", "methanol", "TIP4P", "benzene", "toluene",
+    "chlorobenzene",
+]
+
+# the two reference levels the shipped workbook compares. "geometry" is the sap_geometry_type the
+# stationary/pcm/etc columns were computed on; PCM's own reference level (used for the "PCM ="
+# header line) is fixed at B3LYP-D3(BJ)/pcSseg-3 for both, matching load_delta22_dft_data's default.
+# sheet_suffix_h/sheet_suffix_c differ within a level (the shipped workbook names the carbon sheet
+# after wB97X-D, not WP04, since each nucleus uses its own MagNET-Zero training reference method).
+REFERENCE_LEVELS = {
+    "dsd": {
+        "sheet_suffix_h": "DSD", "sheet_suffix_c": "DSD",
+        "method_h": "dsd_pbep86", "basis_h": "pcSseg3", "geometry_h": "pbe0_tz",
+        "method_c": "dsd_pbep86", "basis_c": "pcSseg3", "geometry_c": "pbe0_tz",
+    },
+    "wp04_wb97xd": {
+        "sheet_suffix_h": "WP04", "sheet_suffix_c": "wB97XD",
+        "method_h": "wp04", "basis_h": "pcSseg2", "geometry_h": "aimnet2",
+        "method_c": "wb97xd", "basis_c": "pcSseg2", "geometry_c": "aimnet2",
+    },
+}
+PCM_REFERENCE_METHOD = "b3lyp_d3bj"
+PCM_REFERENCE_BASIS = "pcSseg3"
+
+# vmin/vmax for the RMSE color scale, per nucleus (fixed, not data-dependent -- matches the shipped
+# workbook so repeated runs render identically)
+_VMIN_VMAX = {"H": (0.06, 0.28), "C": (1.2, 2.8)}
+
+# per-nucleus labeled formula pairs the "story" sections of the sheet compare
+PROTON_FORMULA_CONFIG = {
+    "pair1_a": "stationary + pcm + qcd",
+    "pair1_b": "stationary + desmond + qcd",
+    "pair2_title": "benefit of qcd over desmond_vib",
+    "pair2_a": "stationary + desmond_vib",
+    "pair2_b": "stationary + qcd",
+    "pair3_a": "stationary + desmond + desmond_vib",
+    "pair3_b": "stationary + desmond + qcd",
+    "best_multi": "stationary + desmond + qcd",
+    "best_semi": "stationary_plus_qcd + desmond",
+    "best_pars": "stationary_plus_desmond_plus_qcd",
+}
+CARBON_FORMULA_CONFIG = {
+    "pair1_a": "stationary + pcm + desmond_vib",
+    "pair1_b": "stationary + desmond + desmond_vib",
+    "pair2_title": "benefit of desmond_vib over qcd",
+    "pair2_a": "stationary + qcd",
+    "pair2_b": "stationary + desmond_vib",
+    "pair3_a": "stationary + desmond + qcd",
+    "pair3_b": "stationary + desmond + desmond_vib",
+    "best_multi": "stationary + desmond + desmond_vib",
+    "best_semi": "stationary_plus_des_vib + desmond",
+    "best_pars": "stationary_plus_desmond_plus_des_vib",
+}
+_FORMULA_CONFIG = {"H": PROTON_FORMULA_CONFIG, "C": CARBON_FORMULA_CONFIG}
+
+# This ablation study fits plain OLS composite formulas, not an ML model. The canonical SI text
+# for Figure S15 states that "since no ML was used in this particular analysis, all 22 molecules
+# were considered, including nitromethane" -- other, ML-related analyses (e.g. scaling_factors.py's
+# EXCLUDE_SOLUTES) drop nitromethane as an ML-training outlier; this module's default excludes nothing.
+EXCLUDE_SOLUTES = ()
+
+
+# ---------------------------------------------------------------------------
+# fitting: reuses delta22.py's harness entirely
+
+def _filtered_query_df(query_df_dft, method, basis, geometry, exclude_solutes=EXCLUDE_SOLUTES):
+    """One (method, basis, geometry) slice of the DFT query table, composite columns added,
+    any solutes in `exclude_solutes` dropped (default: none -- see EXCLUDE_SOLUTES above; this
+    analysis is documented in the SI as using all 22 delta-22 solutes, since it's a plain OLS fit,
+    not an ML model)."""
+    sub = query_df_dft[(query_df_dft["sap_nmr_method"] == method)
+                       & (query_df_dft["sap_basis"] == basis)
+                       & (query_df_dft["sap_geometry_type"] == geometry)]
+    sub = sub[~sub["solute"].isin(set(exclude_solutes))]
+    return D.add_composite_columns(sub)
+
+
+def ablation_rmse_table(query_df_dft, nucleus, method, basis, geometry, solutes,
+                        formulas=FORMULAS, solvents=None, n_splits=250,
+                        exclude_solutes=EXCLUDE_SOLUTES):
+    """The formula x solvent test-RMSE table for one nucleus at one reference level: mean test RMSE
+    over `n_splits` seeded splits (delta22.generate_solute_splits/run_fits), pivoted to formula rows
+    x solvent columns, with an appended "Mean Test RMSE" column (the row mean across solvents).
+    Row/column order matches FORMULAS / ORDERED_SOLVENTS. This is the DataFrame write_nucleus_table
+    reads from."""
+    solvents = list(solvents) if solvents is not None else ORDERED_SOLVENTS
+    nuc_df = query_df_dft[query_df_dft["nucleus"] == nucleus]
+    sub = _filtered_query_df(nuc_df, method, basis, geometry, exclude_solutes)
+    results = D.run_fits(sub, solvents, formulas, n_splits, solutes)
+    mean_rmse = results.groupby(["solvent", "formula"])["test_RMSE"].mean().reset_index()
+    table = mean_rmse.pivot(index="formula", columns="solvent", values="test_RMSE")
+    table = table.reindex(index=formulas, columns=solvents)
+    table["Mean Test RMSE"] = table[solvents].mean(axis=1)
+    return table
+
+
+# ---------------------------------------------------------------------------
+# openpyxl table-writing layer
+
+_TABLE_COLUMNS = ORDERED_SOLVENTS + ["Mean Test RMSE"]
+_WHITE_TEXT_THRESHOLD_FRACTION = 0.50
+_PERCENT_COLOR_LIMIT = 0.40
+_PERCENT_WHITE_TEXT_MIN_ABS = 0.25
+_PERCENT_WHITE_TEXT_MAX_ABS = 0.40
+
+
+def _canonical_formula_name(name):
+    return "".join(ch for ch in name.lower() if ch.isalnum())
+
+
+def _build_formula_lookup(df):
+    return {_canonical_formula_name(idx): idx for idx in df.index}
+
+
+def _get_rmse(df, lookup, formula_name, col):
+    key = _canonical_formula_name(formula_name)
+    if key not in lookup:
+        raise KeyError(f"Formula {formula_name!r} was not found in test RMSE dataframe index.")
+    return float(df.loc[lookup[key], col])
+
+
+def _write_benefit_formula_row(ws, row, ref_row, target_row, start_col, n_cols):
+    # percent benefit = (reference - target) / reference, written as a live Excel formula so the
+    # spreadsheet stays self-consistent if a cell above it is ever hand-edited
+    for c in range(start_col, start_col + n_cols):
+        ref = f"{get_column_letter(c)}{ref_row}"
+        target = f"{get_column_letter(c)}{target_row}"
+        ws.cell(row=row, column=c, value=f"=IFERROR(({ref}-{target})/{ref},0)")
+        ws.cell(row=row, column=c).number_format = "0%"
+
+
+def _format_method_label(method_name, basis_name):
+    return f"{method_name.upper().replace('_', '-')}/{basis_name}"
+
+
+def _write_color_scale_note(ws, vmin, vmax, threshold_fraction, pct_limit, pct_white_min, pct_white_max):
+    note_col = 16  # column P
+    threshold_value = vmin + threshold_fraction * (vmax - vmin)
+
+    ws.cell(row=1, column=note_col, value="Color Scale Note").font = Font(bold=True)
+    ws.cell(row=2, column=note_col, value=f"RMSE gradient uses fixed range [{vmin:.2f}, {vmax:.2f}] from notebook settings.")
+    ws.cell(row=3, column=note_col, value=f"White text for RMSE cells below {threshold_fraction:.0%} of range (<= {threshold_value:.3f}).")
+    ws.cell(row=4, column=note_col, value=f"%benefit uses fixed RdBu range [-{pct_limit:.0%}, +{pct_limit:.0%}], centered at 0%.")
+    ws.cell(row=5, column=note_col, value=f"White text for %benefit where abs(value) is at least {pct_white_min:.0%}.")
+
+    ws.cell(row=7, column=note_col, value="RMSE Legend").font = Font(bold=True)
+    ws.cell(row=8, column=note_col, value=f"Low ({vmin:.2f})")
+    ws.cell(row=9, column=note_col, value=f"Mid ({(vmin + vmax) / 2:.2f})")
+    ws.cell(row=10, column=note_col, value=f"High ({vmax:.2f})")
+    ws.cell(row=8, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="2D004B")
+    ws.cell(row=9, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="CC4778")
+    ws.cell(row=10, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="F0F921")
+
+    ws.cell(row=12, column=note_col, value="%Benefit Legend").font = Font(bold=True)
+    ws.cell(row=13, column=note_col, value=f"Negative (-{pct_limit:.0%})")
+    ws.cell(row=14, column=note_col, value="Zero")
+    ws.cell(row=15, column=note_col, value=f"Positive (+{pct_limit:.0%})")
+    ws.cell(row=13, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="B2182B")
+    ws.cell(row=14, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="F7F7F7")
+    ws.cell(row=15, column=note_col + 1).fill = PatternFill(fill_type="solid", fgColor="2166AC")
+
+    ws.column_dimensions[get_column_letter(note_col)].width = 82
+    ws.column_dimensions[get_column_letter(note_col + 1)].width = 14
+
+
+def _apply_rmse_conditional_format(ws, rmse_rows, start_col, end_col, vmin, vmax):
+    rule = ColorScaleRule(start_type="num", start_value=vmin, start_color="2D004B",
+                          mid_type="num", mid_value=(vmin + vmax) / 2, mid_color="CC4778",
+                          end_type="num", end_value=vmax, end_color="F0F921")
+    for r in rmse_rows:
+        ws.conditional_formatting.add(f"{get_column_letter(start_col)}{r}:{get_column_letter(end_col)}{r}", rule)
+
+
+def _apply_percent_conditional_format(ws, pct_rows, start_col, end_col, pct_limit):
+    rule = ColorScaleRule(start_type="num", start_value=-pct_limit, start_color="B2182B",
+                          mid_type="num", mid_value=0, mid_color="F7F7F7",
+                          end_type="num", end_value=pct_limit, end_color="2166AC")
+    for r in pct_rows:
+        ws.conditional_formatting.add(f"{get_column_letter(start_col)}{r}:{get_column_letter(end_col)}{r}", rule)
+
+
+def _apply_percent_font_contrast(ws, pct_rows, start_col, end_col, pct_white_min, pct_white_max):
+    min_s = f"{pct_white_min:.6f}"
+    for r in pct_rows:
+        for c in range(start_col, end_col + 1):
+            addr = f"{get_column_letter(c)}{r}"
+            rule = FormulaRule(formula=[f"ABS({addr})>={min_s}"], stopIfTrue=False, font=Font(color="FFFFFFFF"))
+            ws.conditional_formatting.add(addr, rule)
+
+
+def _apply_rmse_font_contrast(ws, rmse_rows, start_col, end_col, vmin, vmax, threshold_fraction):
+    threshold_value = vmin + threshold_fraction * (vmax - vmin)
+    for r in rmse_rows:
+        for c in range(start_col, end_col + 1):
+            cell = ws.cell(row=r, column=c)
+            if isinstance(cell.value, (int, float)):
+                cell.font = Font(color="FFFFFFFF" if float(cell.value) <= threshold_value else "FF000000")
+
+
+def write_nucleus_table(ws, start_row, nucleus_name, stationary_desc, pcm_desc, test_rmse_df,
+                        vmin, vmax, formula_config,
+                        threshold_fraction=_WHITE_TEXT_THRESHOLD_FRACTION,
+                        pct_limit=_PERCENT_COLOR_LIMIT,
+                        pct_white_min=_PERCENT_WHITE_TEXT_MIN_ABS,
+                        pct_white_max=_PERCENT_WHITE_TEXT_MAX_ABS):
+    """Write one nucleus's ablation table into `ws` starting at `start_row`; returns the next free
+    row. Layout: a header block, a "Multi-Term Models" section (reference stationary row; benefit of
+    desmond over pcm; three labeled formula-pair comparisons from formula_config), then a "Best
+    Model: Approach Comparison" section (SOTA/multi-term/semi-parsimonious/parsimonious, each vs the
+    others). Every RMSE row gets a fixed-range color scale; every %benefit row is a live Excel
+    formula with a diverging color scale. The layout matches the published workbook screenshots
+    section-for-section."""
+    lookup = _build_formula_lookup(test_rmse_df)
+    c0, c1 = 1, 2
+    c_end = c1 + len(_TABLE_COLUMNS) - 1
+    rmse_rows, pct_rows = [], []
+
+    def _rmse_row(row, label, formula):
+        ws.cell(row=row, column=c0, value=label)
+        rmse_rows.append(row)
+        for i, col_name in enumerate(_TABLE_COLUMNS, start=c1):
+            ws.cell(row=row, column=i, value=_get_rmse(test_rmse_df, lookup, formula, col_name))
+            ws.cell(row=row, column=i).number_format = "0.000"
+
+    def _benefit_row(row, ref_row, target_row):
+        ws.cell(row=row, column=c0, value="%benefit")
+        pct_rows.append(row)
+        _write_benefit_formula_row(ws, row, ref_row, target_row, c1, len(_TABLE_COLUMNS))
+
+    ws.cell(row=start_row, column=c0, value=nucleus_name).font = Font(bold=True)
+    ws.cell(row=start_row, column=c1, value=f"stationary = {stationary_desc}")
+    ws.cell(row=start_row + 1, column=c1, value=f"PCM = {pcm_desc}")
+
+    ws.cell(row=start_row + 3, column=c0, value="Multi-Term Models").font = Font(bold=True)
+    ws.cell(row=start_row + 4, column=c0, value="Test RMSEs (Solvent-Specific)").font = Font(bold=True)
+    for i, col_name in enumerate(_TABLE_COLUMNS, start=c1):
+        ws.cell(row=start_row + 4, column=i, value=col_name).font = Font(bold=True)
+    ws.cell(row=start_row + 5, column=c0, value="formula").font = Font(bold=True)
+
+    ws.cell(row=start_row + 7, column=c0, value="reference").font = Font(bold=True)
+    _rmse_row(start_row + 8, "stationary", "stationary")
+
+    ws.cell(row=start_row + 10, column=c0, value="benefit of desmond over pcm").font = Font(bold=True)
+    _rmse_row(start_row + 11, "stationary + pcm", "stationary + pcm")
+    _rmse_row(start_row + 12, "stationary + desmond", "stationary + desmond")
+    _benefit_row(start_row + 13, start_row + 11, start_row + 12)
+
+    _rmse_row(start_row + 15, formula_config["pair1_a"], formula_config["pair1_a"])
+    _rmse_row(start_row + 16, formula_config["pair1_b"], formula_config["pair1_b"])
+    _benefit_row(start_row + 17, start_row + 15, start_row + 16)
+
+    ws.cell(row=start_row + 19, column=c0, value=formula_config["pair2_title"]).font = Font(bold=True)
+    _rmse_row(start_row + 20, formula_config["pair2_a"], formula_config["pair2_a"])
+    _rmse_row(start_row + 21, formula_config["pair2_b"], formula_config["pair2_b"])
+    _benefit_row(start_row + 22, start_row + 20, start_row + 21)
+
+    _rmse_row(start_row + 24, formula_config["pair3_a"], formula_config["pair3_a"])
+    _rmse_row(start_row + 25, formula_config["pair3_b"], formula_config["pair3_b"])
+    _benefit_row(start_row + 26, start_row + 24, start_row + 25)
+
+    ws.cell(row=start_row + 29, column=c0, value="Best Model: Approach Comparison").font = Font(bold=True)
+    _rmse_row(start_row + 31, "SOTA: stationary_plus_pcm", "stationary_plus_pcm")
+    _rmse_row(start_row + 32, f"Multi-term: {formula_config['best_multi']}", formula_config["best_multi"])
+    _rmse_row(start_row + 33, f"Semi-parsimonious: {formula_config['best_semi']}", formula_config["best_semi"])
+    _rmse_row(start_row + 34, f"Parsimonious: {formula_config['best_pars']}", formula_config["best_pars"])
+
+    ws.cell(row=start_row + 36, column=c0, value="%benefit Multi-term vs. SOTA")
+    pct_rows.append(start_row + 36)
+    _write_benefit_formula_row(ws, start_row + 36, start_row + 31, start_row + 32, c1, len(_TABLE_COLUMNS))
+    ws.cell(row=start_row + 37, column=c0, value="%benefit Semi-parsimonious vs. SOTA")
+    pct_rows.append(start_row + 37)
+    _write_benefit_formula_row(ws, start_row + 37, start_row + 31, start_row + 33, c1, len(_TABLE_COLUMNS))
+    ws.cell(row=start_row + 38, column=c0, value="%benefit Parsimonious vs. SOTA")
+    pct_rows.append(start_row + 38)
+    _write_benefit_formula_row(ws, start_row + 38, start_row + 31, start_row + 34, c1, len(_TABLE_COLUMNS))
+    ws.cell(row=start_row + 40, column=c0, value="%benefit Semi-parsimonious vs. Multi-term")
+    pct_rows.append(start_row + 40)
+    _write_benefit_formula_row(ws, start_row + 40, start_row + 32, start_row + 33, c1, len(_TABLE_COLUMNS))
+    ws.cell(row=start_row + 41, column=c0, value="%benefit Parsimonious vs. Multi-term")
+    pct_rows.append(start_row + 41)
+    _write_benefit_formula_row(ws, start_row + 41, start_row + 32, start_row + 34, c1, len(_TABLE_COLUMNS))
+
+    ws.column_dimensions["A"].width = 52
+    for col_idx in range(2, 2 + len(_TABLE_COLUMNS)):
+        ws.column_dimensions[get_column_letter(col_idx)].width = 14
+
+    _apply_rmse_conditional_format(ws, rmse_rows, c1, c_end, vmin, vmax)
+    _apply_percent_conditional_format(ws, pct_rows, c1, c_end, pct_limit)
+    _apply_rmse_font_contrast(ws, rmse_rows, c1, c_end, vmin, vmax, threshold_fraction)
+    _apply_percent_font_contrast(ws, pct_rows, c1, c_end, pct_white_min, pct_white_max)
+    _write_color_scale_note(ws, vmin, vmax, threshold_fraction, pct_limit, pct_white_min, pct_white_max)
+
+    return start_row + 45
+
+
+# ---------------------------------------------------------------------------
+# top-level pipeline
+
+def build_ablations_workbook(query_df_dft, solutes, output_path,
+                             reference_levels=("dsd", "wp04_wb97xd"), n_splits=250,
+                             exclude_solutes=EXCLUDE_SOLUTES, verbose=True):
+    """Computes the ablation RMSE tables for both nuclei at each requested reference level and
+    writes them into one workbook, one sheet per (nucleus, reference level) -- 4 sheets for the
+    default two reference levels, matching the shipped ablations_updated.xlsx: "Proton (DSD)",
+    "Carbon (DSD)", "Proton (WP04)", "Carbon (wB97XD)". Pass a single-element reference_levels to
+    produce the plain "Proton"/"Carbon" two-sheet form for one reference level. query_df_dft is
+    delta22.load_query_df_dft(...)'s output (unfiltered by method/basis); this function does the
+    per-reference-level filtering itself."""
+    wb = Workbook()
+    first = True
+    for level_key in reference_levels:
+        level = REFERENCE_LEVELS[level_key]
+        for nucleus, method_key, basis_key, geom_key, suffix_key in (
+            ("H", "method_h", "basis_h", "geometry_h", "sheet_suffix_h"),
+            ("C", "method_c", "basis_c", "geometry_c", "sheet_suffix_c"),
+        ):
+            method, basis, geometry = level[method_key], level[basis_key], level[geom_key]
+            if verbose:
+                print(f"fitting {nucleus} at {method}/{basis} ({level_key}), {n_splits} splits/solvent...",
+                     flush=True)
+            table = ablation_rmse_table(query_df_dft, nucleus, method, basis, geometry, solutes,
+                                        n_splits=n_splits, exclude_solutes=exclude_solutes)
+            nucleus_name = "Proton" if nucleus == "H" else "Carbon"
+            sheet_title = (f"{nucleus_name} ({level[suffix_key]})" if len(reference_levels) > 1
+                          else nucleus_name)
+            ws = wb.active if first else wb.create_sheet()
+            ws.title = sheet_title
+            first = False
+            write_nucleus_table(
+                ws=ws, start_row=1, nucleus_name=nucleus_name,
+                stationary_desc=_format_method_label(method, basis),
+                pcm_desc=_format_method_label(PCM_REFERENCE_METHOD, PCM_REFERENCE_BASIS),
+                test_rmse_df=table, vmin=_VMIN_VMAX[nucleus][0], vmax=_VMIN_VMAX[nucleus][1],
+                formula_config=_FORMULA_CONFIG[nucleus],
+            )
+    wb.save(output_path)
+    if verbose:
+        print(f"wrote {output_path}", flush=True)
+    return output_path

analysis/code/composite_plots.py

@@ -0,0 +1,41 @@
+"""Correlation-heatmap panels for the composite-model ablation SI figure (Figure S15).
+
+Two seaborn heatmaps behind Figure S15's feature-correlation panels: a per-nucleus lower-triangle
+correlation matrix among the five composite-model features (stationary shielding, PCM, Desmond, its
+vibrational analogue, QCD), and a PCM-vs-Desmond Pearson r table by nucleus and solvent. The notebook
+supplies the correlation table(s) computed by delta22.py and calls these functions directly.
+"""
+import numpy as np
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+
+def plot_feature_correlation_heatmap(corr_matrix, vmin, vmax, cmap, title="", figsize=(6, 5), save_path=None):
+    """Lower-triangle seaborn heatmap of a correlation matrix (Pearson r)."""
+    # In a strict lower triangle the first row and last column hold no cells, so drop them to avoid
+    # a blank leading/trailing strip; the diagonal of the trimmed matrix is a real below-diagonal value.
+    trimmed = corr_matrix.iloc[1:, :-1]
+    mask = np.triu(np.ones(trimmed.shape, dtype=bool), k=1)   # keep the diagonal of the trimmed matrix
+    fig, ax = plt.subplots(figsize=figsize)
+    sns.heatmap(trimmed, mask=mask, annot=True, fmt=".3f", cmap=cmap, vmin=vmin, vmax=vmax,
+                square=True, cbar=False, ax=ax)
+    ax.set_title(title, fontsize=11)
+    fig.tight_layout()
+    if save_path:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return fig
+
+
+def plot_pcm_desmond_correlation_table(corr_table, figsize=(13, 4.5), save_path=None):
+    """Annotated seaborn heatmap of PCM-vs-Desmond Pearson r: rows = nucleus, columns = solvent."""
+    # tall figsize keeps the two rows from being squashed by the rotated solvent labels
+    fig, ax = plt.subplots(figsize=figsize)
+    sns.heatmap(corr_table, annot=True, fmt=".3f", cmap="RdBu", vmin=-1, vmax=1, square=True,
+                linewidths=0.5, cbar=False, ax=ax)
+    ax.set_title("PCM vs. Desmond Pearson $r$ by Solvent")
+    _sup = {"H": "$^{1}$H", "C": "$^{13}$C"}
+    ax.set_yticklabels([_sup.get(t.get_text(), t.get_text()) for t in ax.get_yticklabels()], rotation=0)
+    fig.tight_layout()
+    if save_path:
+        fig.savefig(save_path, dpi=300, bbox_inches="tight")
+    return fig

analysis/code/convergence_plot.py

@@ -0,0 +1,621 @@
+"""3D accuracy-vs-slope-vs-year convergence panel for Figure 2D.
+
+The notebook supplies the per-method statistics table, the method->rung map, and the rung->color map as
+arguments (`results_df`, `methods_categorization`, `method_color_map`) and calls
+`plot_statistics_3d_emphasize_rungs(...)`.
+"""
+import os
+
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+from matplotlib.ticker import MaxNLocator
+from matplotlib import colors as mcolors
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  (registers the 3d projection)
+
+
+def plot_statistics_3d_emphasize_rungs(
+    x_param="Slope",
+    y_param="RMSE (scaled)",
+    z_param=None,
+    results_df=None,
+    methods_categorization=None,
+    method_color_map=None,
+    figsize=(9, 7),
+    elev=12,
+    azim=20,
+
+    # =========================
+    # SAVING
+    # =========================
+    save_path=None,
+    save_formats=("png",),
+    dpi=300,
+    svg_transparent=True,
+    png_transparent=False,
+
+    # =========================
+    # DROPLINES (ALL POINTS)
+    # =========================
+    droplines_for_all=True,
+    dropline_alpha=0.1,
+    dropline_lw=0.8,
+    dropline_color="0.4",
+    dropline_shadow_alpha=0.25,
+    dropline_shadow_size=18,
+
+    xlim=None,
+    ylim=None,
+    zlim=None,
+
+    # =========================
+    # YEAR CONTROL
+    # =========================
+    year_map=None,
+    year_default=np.nan,
+    year_tick_step=2,
+    drop_missing_years=False,
+
+    # =========================
+    # Z LANES
+    # =========================
+    z_lane_by_category=True,
+    z_lane_width=0.35,
+    z_jitter_by_category=0.0,
+
+    # =========================
+    # XY LANES
+    # =========================
+    xy_lane_by_category=True,
+    xy_lane_strength=0.25,
+    xy_lane_style="circle",
+    lane_exclude_categories=("ab initio",),
+
+    auto_expand_xy_limits=True,
+    xy_limit_pad_frac=0.06,
+
+    # =========================
+    # BROKEN Z
+    # =========================
+    z_break=None,
+    z_break_gap=2.0,
+    drop_years_in_break=False,
+
+    # =========================
+    # DECADE GROUPING
+    # =========================
+    group_by_decade=True,
+    decade_bin=10,
+    decade_alpha=0.85,
+    decade_marker_size=80,
+    decade_sort_ascending=True,
+
+    # =========================
+    # XY GRID PLANE
+    # =========================
+    show_xy_grid_plane=True,
+    xy_grid_z=0.0,
+    grid_alpha=0.15,
+    grid_lw=0.4,
+    show_plot=True,
+    auto_xy_grid_z_if_years=True,
+
+    # =========================
+    # BASELINE SHADOW COLOR
+    # =========================
+    baseline_tint_strength=0.35,
+    baseline_gray="#7A7A7A",
+
+    # =========================
+    # MAJOR TICKS (CUSTOMIZABLE)
+    # =========================
+    x_major_nbins=5,
+    y_major_nbins=5,
+    major_tick_length=0.1,
+    major_tick_width=0.1,
+    major_tick_labelsize=10,
+    major_tick_pad=2,
+    major_tick_color="0.2",
+    major_tick_labelcolor="0.1",
+    major_tick_rotation_x=0,
+    major_tick_rotation_y=0,
+    major_tick_rotation_z=0,
+
+    # =========================
+    # "SPINES"/BORDERS FOR 3D AXES (CUSTOMIZABLE)
+    # =========================
+    axisline_lw=0.3,
+    axisline_color="0.05",
+
+    pane_fill=False,
+    pane_alpha=0.0,
+    pane_color="white",
+    pane_edgecolor="0.15",
+    pane_edge_lw=0.3,
+
+    # alias (so passing axis_linewidth won't crash)
+    axis_linewidth=None,
+
+    # =========================
+    # HIGHLIGHT + LABEL METHODS
+    # =========================
+    highlight_methods=None,          # list/tuple of method names to highlight
+    highlight_marker_size=140,       # size for highlighted markers
+    highlight_edgecolor="k",         # outline color
+    highlight_edgewidth=1.2,         # outline thickness
+    highlight_alpha=1.0,             # alpha for highlighted markers
+    highlight_text=True,             # whether to draw labels
+    highlight_text_size=8,           # label fontsize
+    highlight_text_color="k",        # label color
+    highlight_text_dx_frac=0.008,    # x-offset for text (fraction of x-range)
+    highlight_text_dy_frac=0.008,    # y-offset for text (fraction of y-range)
+    highlight_text_dz=0.0,           # z-offset in z-units (mapped units)
+    highlight_text_bbox=True,        # draw a small white box behind text
+    highlight_text_bbox_alpha=0.65,  # bbox alpha
+
+    # legend
+    show_legend=False,
+    legend_title="Rung Type",
+    legend_fontsize=8,
+    legend_title_fontsize=9,
+    legend_loc="center left",
+    legend_bbox_to_anchor=(1.02, 0.5),
+):
+    """Draw the Figure 2D 3D scatter of RMSE vs. slope vs. year, grouped and colored by rung."""
+    sns.set_theme(context="paper", style="ticks", font_scale=1.0)
+
+    def _coerce_pair(val, name, allow_none=True):
+        if val is None:
+            if allow_none:
+                return None
+            raise ValueError(f"{name} cannot be None")
+        if np.isscalar(val):
+            raise ValueError(
+                f"{name} must be a 2-tuple like ({name}_min, {name}_max), not a scalar: {val}"
+            )
+        if len(val) != 2:
+            raise ValueError(f"{name} must be length-2 like (min, max). Got: {val}")
+        return (float(val[0]), float(val[1]))
+
+    # apply alias if user passed axis_linewidth
+    if axis_linewidth is not None:
+        axisline_lw = float(axis_linewidth)
+
+    # coerce user limits early
+    xlim2 = _coerce_pair(xlim, "xlim", allow_none=True)
+    ylim2 = _coerce_pair(ylim, "ylim", allow_none=True)
+    zlim2 = _coerce_pair(zlim, "zlim", allow_none=True)
+
+    def _draw_xy_grid_plane(ax, xlim, ylim, z=0.0, alpha=0.25, lw=0.8):
+        x_ticks = ax.get_xticks()
+        y_ticks = ax.get_yticks()
+        for x in x_ticks:
+            if xlim[0] <= x <= xlim[1]:
+                ax.plot([x, x], [ylim[0], ylim[1]], [z, z],
+                        color="0.0", alpha=alpha, linewidth=lw, zorder=1)
+        for y in y_ticks:
+            if ylim[0] <= y <= ylim[1]:
+                ax.plot([xlim[0], xlim[1]], [y, y], [z, z],
+                        color="0.0", alpha=alpha, linewidth=lw, zorder=1)
+
+    def _blend_hex(c1, c2, t):
+        c1 = np.array(mcolors.to_rgb(c1))
+        c2 = np.array(mcolors.to_rgb(c2))
+        return (1 - t) * c1 + t * c2
+
+    def _make_z_break_mapper(z_break, gap, drop_in_break):
+        z_break2 = _coerce_pair(z_break, "z_break", allow_none=True)
+        if z_break2 is None:
+            return (lambda y: float(y)), (lambda z: float(z))
+
+        lo, hi = z_break2
+        if not (hi > lo):
+            raise ValueError(f"z_break must be (lo, hi) with hi>lo; got {z_break2}")
+        span = hi - lo
+        gap = float(gap)
+
+        def fwd(y):
+            y = float(y)
+            if lo < y < hi:
+                if drop_in_break:
+                    return None
+                y = lo if (y - lo) <= (hi - y) else hi
+            if y <= lo:
+                return y
+            return y - span + gap
+
+        def inv(z):
+            z = float(z)
+            if z <= lo:
+                return z
+            return z + span - gap
+
+        return fwd, inv
+
+    def _norm_cat(s: str) -> str:
+        s = str(s).lower()
+        return "".join(ch for ch in s if ch.isalnum())
+
+    exclude_norm = {_norm_cat(c) for c in (lane_exclude_categories or ())}
+
+    z_map, _ = _make_z_break_mapper(z_break, z_break_gap, drop_years_in_break)
+
+    # --- Prepare dataframe ---
+    df = results_df.copy()
+    df["category"] = df["Method"].map(methods_categorization).fillna("Unknown")
+    df["_cat_norm"] = df["category"].apply(_norm_cat)
+
+    rung_order = [c for c in method_color_map.keys() if c in df["category"].unique()]
+    rung_order += [c for c in df["category"].unique() if c not in rung_order]
+
+    # --- Z assignment ---
+    if z_param is not None:
+        df["_z_raw"] = df[z_param].astype(float)
+        df["_year"] = df["_z_raw"].copy()
+    else:
+        if year_map is None:
+            raise ValueError("z_param is None, so you must provide year_map={Method: year}.")
+        df["_z_raw"] = df["Method"].map(year_map).fillna(year_default)
+
+        missing_methods = df[df["_z_raw"].isna()]["Method"].unique()
+        if len(missing_methods) > 0:
+            if drop_missing_years:
+                df = df.dropna(subset=["_z_raw"]).copy()
+            else:
+                examples = missing_methods[:12].tolist()
+                raise ValueError(
+                    f"Missing years for {len(missing_methods)} methods. "
+                    f"Examples: {examples}. "
+                    f"Either add them to year_map or set drop_missing_years=True."
+                )
+
+        df["_z_raw"] = df["_z_raw"].astype(float)
+        df["_year"] = df["_z_raw"].copy()
+
+        if z_lane_by_category:
+            cat_index = {cat: i for i, cat in enumerate(rung_order)}
+            offsets = df["category"].map(lambda c: cat_index.get(c, 0)).astype(float)
+            offsets = offsets - offsets.mean()
+            is_excluded = df["_cat_norm"].isin(exclude_norm)
+            offsets = offsets.where(~is_excluded, 0.0)
+            df["_z_raw"] = df["_z_raw"] + offsets * float(z_lane_width)
+
+        if z_jitter_by_category and z_jitter_by_category != 0.0:
+            rng = np.random.default_rng(0)
+            df["_z_raw"] = df["_z_raw"] + rng.uniform(-1, 1, size=len(df)) * float(z_jitter_by_category)
+
+    # --- apply z break mapping ---
+    if z_break is not None:
+        z_mapped = df["_z_raw"].apply(z_map)
+        if drop_years_in_break:
+            df = df.loc[~z_mapped.isna()].copy()
+            z_mapped = df["_z_raw"].apply(z_map)
+        df["_z"] = z_mapped.astype(float)
+    else:
+        df["_z"] = df["_z_raw"].astype(float)
+
+    # --- Base limits from ORIGINAL (unshifted) data ---
+    x_min, x_max = df[x_param].min(), df[x_param].max()
+    y_min, y_max = df[y_param].min(), df[y_param].max()
+    x_pad = 0.15 * (x_max - x_min) if x_max > x_min else 0.5
+    y_pad = 0.15 * (y_max - y_min) if y_max > y_min else 0.5
+    xm_default = (x_min - x_pad, x_max + x_pad)
+    ym_default = (y_min - y_pad, y_max + y_pad)
+
+    # Use provided limits (or defaults) to define lane sizes
+    x0l, x1l = (xm_default if xlim2 is None else xlim2)
+    y0l, y1l = (ym_default if ylim2 is None else ylim2)
+    xr = (x1l - x0l) if (x1l - x0l) != 0 else 1.0
+    yr = (y1l - y0l) if (y1l - y0l) != 0 else 1.0
+
+    # --- XY category lanes ---
+    if xy_lane_by_category:
+        n = max(1, len(rung_order))
+        cat_index = {cat: i for i, cat in enumerate(rung_order)}
+
+        if xy_lane_style.lower() == "grid":
+            cols = int(np.ceil(np.sqrt(n)))
+            rows = int(np.ceil(n / cols))
+
+            def grid_xy(i):
+                r = i // cols
+                c = i % cols
+                cx = c - (cols - 1) / 2
+                cy = r - (rows - 1) / 2
+                return cx, cy
+
+            coords = {cat: grid_xy(cat_index[cat]) for cat in rung_order}
+        else:
+            angles = np.linspace(0, 2*np.pi, n, endpoint=False)
+            coords = {cat: (np.cos(angles[cat_index[cat]]), np.sin(angles[cat_index[cat]])) for cat in rung_order}
+
+        dx = float(xy_lane_strength) * xr
+        dy = float(xy_lane_strength) * yr
+        is_excluded = df["_cat_norm"].isin(exclude_norm)
+
+        def _dx(c): return coords.get(c, (0.0, 0.0))[0] * dx
+        def _dy(c): return coords.get(c, (0.0, 0.0))[1] * dy
+
+        df["_x_plot"] = df[x_param].astype(float) + np.where(is_excluded, 0.0, df["category"].map(_dx).astype(float))
+        df["_y_plot"] = df[y_param].astype(float) + np.where(is_excluded, 0.0, df["category"].map(_dy).astype(float))
+    else:
+        df["_x_plot"] = df[x_param].astype(float)
+        df["_y_plot"] = df[y_param].astype(float)
+
+    # Expand axis limits to include shifted points
+    if auto_expand_xy_limits:
+        xmn, xmx = float(np.nanmin(df["_x_plot"])), float(np.nanmax(df["_x_plot"]))
+        ymn, ymx = float(np.nanmin(df["_y_plot"])), float(np.nanmax(df["_y_plot"]))
+        xr2 = (xmx - xmn) if (xmx - xmn) != 0 else 1.0
+        yr2 = (ymx - ymn) if (ymx - ymn) != 0 else 1.0
+        xpad2 = xy_limit_pad_frac * xr2
+        ypad2 = xy_limit_pad_frac * yr2
+        xm_shifted = (xmn - xpad2, xmx + xpad2)
+        ym_shifted = (ymn - ypad2, ymx + ypad2)
+    else:
+        xm_shifted = xm_default
+        ym_shifted = ym_default
+
+    # baseline plane fix
+    if zlim2 is not None:
+        z_base = z_map(zlim2[0]) if (z_break is not None) else float(zlim2[0])
+    else:
+        z_base = float(np.nanmin(df["_z"]))
+    Z_BASELINE = z_base
+    if auto_xy_grid_z_if_years and (z_param is None):
+        xy_grid_z = Z_BASELINE
+
+    # decade grouping
+    if group_by_decade and (z_param is None):
+        df["_decade"] = (np.floor(df["_year"] / decade_bin) * decade_bin).astype(int)
+        decade_order = sorted(df["_decade"].unique(), reverse=not decade_sort_ascending)
+    else:
+        df["_decade"] = np.nan
+        decade_order = []
+
+    # --- Figure ---
+    fig = plt.figure(figsize=figsize, constrained_layout=True)
+    ax = fig.add_subplot(111, projection="3d")
+    ax.view_init(elev=elev, azim=azim)
+    ax.set_box_aspect((1.5, 2, 1.5))
+
+    # axis lines
+    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
+        try:
+            axis.line.set_linewidth(axisline_lw)
+            axis.line.set_color(axisline_color)
+        except Exception:
+            pass
+
+    # panes
+    for axis in (ax.xaxis, ax.yaxis, ax.zaxis):
+        try:
+            axis.pane.set_edgecolor(pane_edgecolor)
+            axis.pane.set_linewidth(pane_edge_lw)
+            if pane_fill:
+                axis.pane.set_facecolor(mcolors.to_rgba(pane_color, pane_alpha))
+            else:
+                axis.pane.set_facecolor((1, 1, 1, 0))
+        except Exception:
+            pass
+
+    # --- SCATTER ---
+    if group_by_decade and (z_param is None):
+        for d in decade_order:
+            sub_d = df[df["_decade"] == d]
+            if sub_d.empty:
+                continue
+            for cat in rung_order:
+                sub = sub_d[sub_d["category"] == cat]
+                if sub.empty:
+                    continue
+                ax.scatter(
+                    sub["_x_plot"], sub["_y_plot"], sub["_z"],
+                    s=decade_marker_size,
+                    depthshade=False,
+                    color=method_color_map.get(cat, "#777777"),
+                    edgecolor="none",
+                    alpha=decade_alpha,
+                    zorder=5,
+                )
+    else:
+        for cat in rung_order:
+            sub = df[df["category"] == cat]
+            if sub.empty:
+                continue
+            ax.scatter(
+                sub["_x_plot"], sub["_y_plot"], sub["_z"],
+                s=80,
+                depthshade=False,
+                color=method_color_map.get(cat, "#777777"),
+                edgecolor="none",
+                alpha=0.85,
+                zorder=5,
+            )
+
+    # --- DROPLINES FOR ALL POINTS ---
+    if droplines_for_all:
+        cat_cols = df["category"].map(lambda c: method_color_map.get(c, "#777777"))
+        shadow_cols = [_blend_hex(baseline_gray, c, baseline_tint_strength) for c in cat_cols]
+        for (x0, y0, z0, shcol) in zip(df["_x_plot"].to_numpy(float),
+                                      df["_y_plot"].to_numpy(float),
+                                      df["_z"].to_numpy(float),
+                                      shadow_cols):
+            ax.plot([x0, x0], [y0, y0], [z0, Z_BASELINE],
+                    color=dropline_color, alpha=dropline_alpha,
+                    linewidth=dropline_lw, zorder=2)
+            ax.scatter([x0], [y0], [Z_BASELINE],
+                       s=dropline_shadow_size,
+                       color=shcol,
+                       alpha=dropline_shadow_alpha,
+                       depthshade=False,
+                       zorder=2)
+
+    # --- limits ---
+    ax.set_xlim(xm_shifted if xlim2 is None else xlim2)
+    ax.set_ylim(ym_shifted if ylim2 is None else ylim2)
+    if zlim2 is not None:
+        ax.set_zlim(z_map(zlim2[0]), z_map(zlim2[1])) if (z_break is not None) else ax.set_zlim(zlim2)
+    else:
+        zmin = float(np.nanmin(df["_z"]))
+        zmax = float(np.nanmax(df["_z"]))
+        ax.set_zlim(zmin - 0.5, zmax + 0.5)
+
+    ax.set_ylabel("Scaled RMSE (" + r"$^{1}\mathrm{H}$" + " ppm)", fontsize=13, labelpad=10)
+    ax.set_xlabel("Slope vs. CCSD(T)", fontsize=13, labelpad=10)
+    ax.set_zlabel("Year", fontsize=13, labelpad=8)
+
+    # --- major tick locations ---
+    ax.xaxis.set_major_locator(MaxNLocator(nbins=x_major_nbins))
+    ax.yaxis.set_major_locator(MaxNLocator(nbins=y_major_nbins))
+
+    # --- major tick appearance ---
+    for axis_name in ("x", "y", "z"):
+        ax.tick_params(
+            axis=axis_name,
+            which="major",
+            length=major_tick_length,
+            width=major_tick_width,
+            pad=major_tick_pad,
+            colors=major_tick_color,
+            labelsize=major_tick_labelsize,
+            labelcolor=major_tick_labelcolor,
+        )
+
+    for lbl in ax.get_xticklabels():
+        lbl.set_rotation(major_tick_rotation_x)
+    for lbl in ax.get_yticklabels():
+        lbl.set_rotation(major_tick_rotation_y)
+    for lbl in ax.get_zticklabels():
+        lbl.set_rotation(major_tick_rotation_z)
+
+    # --- Z ticks as years ---
+    if (z_param is None) or (z_break is not None):
+        if zlim2 is not None:
+            real_min, real_max = float(zlim2[0]), float(zlim2[1])
+        else:
+            real_min, real_max = float(np.nanmin(df["_year"])), float(np.nanmax(df["_year"]))
+
+        if z_break is not None:
+            lo, hi = map(float, _coerce_pair(z_break, "z_break", allow_none=False))
+            step = int(year_tick_step)
+            ticks_real = list(np.arange(np.floor(real_min), np.floor(min(real_max, lo)) + 1, step))
+            ticks_real += list(np.arange(np.ceil(max(real_min, hi)), np.ceil(real_max) + 1, step))
+            ticks_real = [int(t) for t in ticks_real]
+            ticks_pos = [z_map(t) for t in ticks_real]
+            ax.set_zticks(ticks_pos)
+            ax.set_zticklabels([str(t) for t in ticks_real], fontsize=major_tick_labelsize, color=major_tick_labelcolor)
+        else:
+            ticks = np.arange(int(np.floor(real_min)), int(np.ceil(real_max)) + 1, int(year_tick_step))
+            ax.set_zticks(ticks)
+            ax.set_zticklabels([str(t) for t in ticks], fontsize=major_tick_labelsize, color=major_tick_labelcolor)
+
+    ax.grid(False)
+
+    if show_xy_grid_plane:
+        _draw_xy_grid_plane(ax, xlim=ax.get_xlim(), ylim=ax.get_ylim(), z=xy_grid_z,
+                            alpha=grid_alpha, lw=grid_lw)
+
+    # =========================
+    # HIGHLIGHT + LABEL
+    # =========================
+    if highlight_methods:
+        hm = set(highlight_methods)
+        sub_h = df[df["Method"].isin(hm)].copy()
+        if not sub_h.empty:
+            ax.scatter(
+                sub_h["_x_plot"], sub_h["_y_plot"], sub_h["_z"],
+                s=highlight_marker_size,
+                depthshade=False,
+                color=sub_h["category"].map(lambda c: method_color_map.get(c, "#777777")),
+                edgecolor=highlight_edgecolor,
+                linewidths=highlight_edgewidth,
+                alpha=highlight_alpha,
+                zorder=20,
+            )
+
+            if highlight_text:
+                dx_text = float(highlight_text_dx_frac) * xr
+                dy_text = float(highlight_text_dy_frac) * yr
+                for _, r in sub_h.iterrows():
+                    px, py, pz = float(r["_x_plot"]), float(r["_y_plot"]), float(r["_z"])
+                    tx, ty, tz = px + dx_text, py + dy_text, pz + float(highlight_text_dz)
+                    # thin leader line so the label reads in clear space, not under the marker
+                    ax.plot([px, tx], [py, ty], [pz, tz], color="0.45", linewidth=0.6, zorder=25)
+                    txt_kwargs = dict(
+                        fontsize=highlight_text_size,
+                        color=highlight_text_color,
+                        zorder=30, ha="left", va="center",
+                    )
+                    if highlight_text_bbox:
+                        txt_kwargs["bbox"] = dict(
+                            boxstyle="round,pad=0.15",
+                            facecolor=(1, 1, 1, highlight_text_bbox_alpha),
+                            edgecolor="none",
+                        )
+                    ax.text(tx, ty, tz, str(r["Method"]), **txt_kwargs)
+
+    if show_legend:
+        import matplotlib.patches as mpatches
+        patches = [mpatches.Patch(color=method_color_map.get(cat, "#777777"), label=cat) for cat in rung_order]
+        ax.legend(handles=patches,
+                  title=legend_title,
+                  fontsize=legend_fontsize,
+                  title_fontsize=legend_title_fontsize,
+                  loc=legend_loc,
+                  bbox_to_anchor=legend_bbox_to_anchor,
+                  frameon=True)
+
+    # =========================
+    # SAVE
+    # =========================
+    def _save_figure(fig, save_path, formats=("svg", "png"), dpi=300,
+                     svg_transparent=True, png_transparent=False):
+        if not save_path:
+            return []
+        root, ext = os.path.splitext(save_path)
+        os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True)
+        saved = []
+
+        def _save_one(outpath, fmt):
+            fmt = fmt.lower().lstrip(".")
+            if fmt == "svg":
+                fig.savefig(outpath, format="svg", bbox_inches="tight", pad_inches=0.12,
+                            dpi=dpi, transparent=bool(svg_transparent))
+            elif fmt in {"png", "jpg", "jpeg", "tif", "tiff"}:
+                fig.savefig(outpath, format=fmt, bbox_inches="tight", pad_inches=0.12,
+                            dpi=dpi, transparent=bool(png_transparent))
+            else:
+                fig.savefig(outpath, format=fmt, bbox_inches="tight", pad_inches=0.12,
+                            dpi=dpi, transparent=False)
+
+        if ext.lower() in {".png", ".svg", ".pdf", ".jpg", ".jpeg", ".tif", ".tiff"}:
+            fmt = ext.lower().lstrip(".")
+            _save_one(save_path, fmt)
+            saved.append(save_path)
+            return saved
+
+        for f in formats:
+            f = f.lower().lstrip(".")
+            out = f"{save_path}.{f}"
+            _save_one(out, f)
+            saved.append(out)
+        return saved
+
+    saved_files = _save_figure(
+        fig,
+        save_path=save_path,
+        formats=save_formats,
+        dpi=dpi,
+        svg_transparent=svg_transparent,
+        png_transparent=png_transparent,
+    )
+    if saved_files:
+        print("Saved:", *saved_files, sep="\n  - ")
+
+    if show_plot:
+        plt.show()
+    else:
+        plt.close(fig)
+
+    return fig, ax, saved_files

analysis/code/dataset_summary.py

@@ -0,0 +1,78 @@
+"""SI Table S8 (Dataset Summary Statistics): molecule and 1H/13C site counts for the training datasets.
+
+Site counts from each file's flat ``atomic_numbers`` (1H = 1, 13C = 6); molecule counts from each
+file's per-molecule index. Vectorized reads, so seconds even on the multi-million-molecule files. The
+MagNET-Zero rows combine the two sigma-pepper groups (initial PBE0/pcSseg-1, WP04/wB97X-D/pcSseg-2
+fine-tuning) with sigma-concentrate.
+"""
+import os
+
+import h5py
+import pandas as pd
+
+import paths
+
+# Published SI Table S8 values, for the exact-reproduction check: label -> (molecules, 1H, 13C).
+PUBLISHED_S8 = {
+    "sigma-shake": (4_787_816, 72_175_813, 50_542_154),
+    "MagNET-Zero (initial round)": (2_167_210, 24_047_999, 15_330_361),
+    "MagNET-Zero (fine-tuning)": (370_105, 4_042_135, 2_557_211),
+    "MagNET-Zero (both rounds combined)": (2_537_315, 28_090_134, 17_887_572),
+    "sigma-concentrate": (50_000, 754_437, 528_569),
+}
+
+
+def _counts(node):
+    """(molecules, 1H sites, 13C sites) for an open h5py file or group."""
+    z = node["atomic_numbers"][:]
+    n_1h = int((z == 1).sum())
+    n_13c = int((z == 6).sum())
+    if "molecule_id" in node:
+        n_mol = node["molecule_id"].shape[0]
+    elif "smiles_offsets" in node:            # offsets are n_molecules + 1
+        n_mol = node["smiles_offsets"].shape[0] - 1
+    elif "sigma_shake_id" in node:
+        n_mol = node["sigma_shake_id"].shape[0]
+    else:
+        raise KeyError("node has no recognized per-molecule key "
+                       "(molecule_id / smiles_offsets / sigma_shake_id)")
+    return n_mol, n_1h, n_13c
+
+
+def counts(path, group=None):
+    """(molecules, 1H sites, 13C sites) for one HDF5 file, optionally within one group."""
+    with h5py.File(path, "r") as f:
+        return _counts(f[group] if group else f)
+
+
+def summary_table(data_dir):
+    """SI Table S8 as a DataFrame with columns dataset / molecules / n_1H_sites / n_13C_sites.
+
+    data_dir is the release repo's ``data/`` directory; the three source files are read from their
+    dataset subfolders. The two MagNET-Zero training rounds and their combination are assembled from
+    the sigma-pepper groups plus sigma-concentrate.
+    """
+    # data_dir is the repo's data/ folder; its parent is the repo root the deposit fallback needs.
+    root = os.path.dirname(os.path.normpath(data_dir))
+    shake = paths.dataset_file("sigma-shake", root=root)
+    pepper = paths.dataset_file("sigma-pepper", root=root)
+    concentrate = paths.dataset_file("sigma-concentrate", root=root)
+
+    sigma_shake = counts(shake)
+    initial = counts(pepper, "pbe0_pcSseg1")                 # MagNET-Zero initial round
+    pepper_ft = counts(pepper, "wp04_wb97xd_pcSseg2")        # fine-tuning: sigma-pepper part
+    sigma_concentrate = counts(concentrate)                  # fine-tuning: sigma-concentrate part
+    fine_tuning = tuple(a + b for a, b in zip(pepper_ft, sigma_concentrate))
+    combined = tuple(a + b for a, b in zip(initial, fine_tuning))
+
+    rows = [
+        ("sigma-shake", sigma_shake),
+        ("MagNET-Zero (initial round)", initial),
+        ("MagNET-Zero (fine-tuning)", fine_tuning),
+        ("MagNET-Zero (both rounds combined)", combined),
+        ("sigma-concentrate", sigma_concentrate),
+    ]
+    return pd.DataFrame(
+        [{"dataset": label, "molecules": m, "n_1H_sites": h, "n_13C_sites": c}
+         for label, (m, h, c) in rows]
+    )

analysis/code/delta22.py

@@ -0,0 +1,927 @@
+"""Core analysis for the delta-22 figures (no plotting).
+
+Shared machinery behind the delta-22 figures (Figures 2A, 3, 4; SI S1, S4-S8, S11, S13, S16-S18):
+a seeded per-solute train/test fitting harness, the Pareto computation, and correlation matrices.
+Reads the released data via `delta22_reader`; the figure notebooks and tests import from here.
+
+Splits are seeded with random.seed(seed + 100), shuffle, first n_test solutes are the test set, so
+results are deterministic. `fit` uses numpy lstsq directly instead of statsmodels smf.ols in the
+inner loop (patsy re-parsing makes that ~1000x slower); `_fit_statsmodels` is kept as a test oracle.
+"""
+import random
+
+import numpy as np
+import pandas as pd
+
+from stats import rmse, mae, ols_fit, linear_fit_1d
+from spreadsheet import site_atom_indices  # noqa: F401  re-exported for notebooks (D.site_atom_indices)
+from delta22_reader import (
+    load_delta22_dft_data,
+    load_delta22_nn_data,
+    load_solutes,
+    load_perturbed_shieldings,
+)
+
+# The 12 solvents with explicit (Desmond) data, in a stable order.
+DESMOND_SOLVENTS = [
+    "tetrahydrofuran", "dichloromethane", "chloroform", "toluene", "benzene", "chlorobenzene",
+    "acetone", "dimethylsulfoxide", "acetonitrile", "trifluoroethanol", "methanol", "TIP4P",
+]
+
+# the double hybrids (and DLPNO-MP2) have no implicit-solvent shielding of their own; the reader
+# substitutes a B3LYP-D3(BJ) PCM correction for them, so their pcm column is not method-specific.
+# Exclude them from any analysis that compares the PCM correction across methods.
+DOUBLE_HYBRID_METHODS = ["B2GP_PLYP", "B2PLYP", "dlpno_mp2", "dsd_pbep86",
+                         "mPW2PLYP", "revdsd_pbep86"]
+
+STATIONARY_VS_PCM = ["stationary", "stationary + pcm"]
+
+
+# ---------------------------------------------------------------------------
+# loading / assembling the flat query table
+
+def delta22_solutes(delta22_path):
+    """The list of solute names stored in the file."""
+    return load_solutes(delta22_path)
+
+
+def build_query_df(combined_df):
+    """Flatten a reader combined table (load_delta22_dft_data / load_delta22_nn_data output,
+    which is multi-indexed) into a plain DataFrame with the index levels as columns. Fitting
+    and split helpers operate on this flat table."""
+    return combined_df.reset_index()
+
+
+def load_query_df_dft(delta22_path, experimental_path, **kwargs):
+    """Convenience: load the DFT combined table and flatten it in one call."""
+    return build_query_df(load_delta22_dft_data(delta22_path, experimental_path, **kwargs))
+
+
+def load_query_df_nn(delta22_path, experimental_path, **kwargs):
+    """Convenience: load the MagNET (NN) combined table and flatten it in one call."""
+    return build_query_df(load_delta22_nn_data(delta22_path, experimental_path, **kwargs))
+
+
+# ---------------------------------------------------------------------------
+# composite predictor columns
+
+# the precomputed sum columns the formulas reference as single terms (e.g. "stationary_plus_pcm" is
+# one predictor = stationary + pcm, distinct from the two-term "stationary + pcm"). For MagNET (NN)
+# data, stationary_plus_pcm is precomputed by the loader using the solvent-specific PCM scaling
+# factors; add_composite_columns leaves an existing stationary_plus_pcm column untouched.
+# NOTE: the pcm three-term chains here (stationary_plus_pcm_plus_qcd etc.) are RAW sums with an
+# unscaled pcm term. build_composite_model.py needs the pcm term carried at its scaled value, so it
+# builds those columns itself on top of the scaled stationary_plus_pcm. Do not swap it to this
+# helper: the results differ by order 1 ppm on every pcm chain.
+_COMPOSITE_COLUMNS = {
+    "stationary_plus_pcm": ("stationary", "pcm"),
+    "stationary_plus_desmond": ("stationary", "desmond"),
+    "stationary_plus_openMM": ("stationary", "openMM"),
+    "stationary_plus_qcd": ("stationary", "qcd"),
+    "stationary_plus_des_vib": ("stationary", "desmond_vib"),
+    "stationary_plus_op_vib": ("stationary", "openMM_vib"),
+    "stationary_plus_pcm_plus_qcd": ("stationary", "pcm", "qcd"),
+    "stationary_plus_pcm_plus_des_vib": ("stationary", "pcm", "desmond_vib"),
+    "stationary_plus_pcm_plus_op_vib": ("stationary", "pcm", "openMM_vib"),
+    "stationary_plus_desmond_plus_qcd": ("stationary", "desmond", "qcd"),
+    "stationary_plus_desmond_plus_des_vib": ("stationary", "desmond", "desmond_vib"),
+    "stationary_plus_openMM_plus_qcd": ("stationary", "openMM", "qcd"),
+    "stationary_plus_openMM_plus_op_vib": ("stationary", "openMM", "openMM_vib"),
+}
+
+
+def add_composite_columns(query_df, keep_existing=("stationary_plus_pcm",)):
+    """Add the precomputed sum columns the composite formulas reference as single terms. Returns a
+    copy. Columns named in keep_existing that are already present (e.g. an NN stationary_plus_pcm
+    built with the solvent scaling factors) are not overwritten."""
+    out = query_df.copy()
+    for name, parts in _COMPOSITE_COLUMNS.items():
+        if name in keep_existing and name in out.columns:
+            continue
+        out[name] = sum(out[part] for part in parts)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# the linear fitting harness
+
+def _formula_terms(formula):
+    """The additive predictor column names in 'a + b + c'. These figures use only additive main
+    effects (no interactions), so a plain split on '+' is exact."""
+    return [term.strip() for term in formula.split("+") if term.strip()]
+
+
+def _design_matrix(df, terms):
+    """Design matrix [1, term_1, term_2, ...] as float, plus the response, with any row that has a
+    missing value in the response or a predictor dropped (statsmodels drops missing rows too)."""
+    response = df["experimental"].to_numpy(dtype=float)
+    columns = [df[term].to_numpy(dtype=float) for term in terms]
+    design = np.column_stack([np.ones(len(df))] + columns)
+    keep = np.isfinite(response) & np.all(np.isfinite(design), axis=1)
+    return design[keep], response[keep]
+
+
+def fit(train_df, test_df, formula, mode="rmse"):
+    """Fit experimental ~ formula by ordinary least squares on train_df and score on both sets.
+
+    Returns (train_error, test_error, params), where params is a pandas Series indexed by
+    "Intercept" and the predictor names (matching statsmodels' naming). mode is "rmse" (default)
+    or "mae". Uses numpy least squares directly (see the module docstring for why)."""
+    terms = _formula_terms(formula)
+    train_x, train_y = _design_matrix(train_df, terms)
+    beta = ols_fit(train_x, train_y)
+    params = pd.Series(beta, index=["Intercept"] + terms)
+
+    test_x, test_y = _design_matrix(test_df, terms)
+    train_pred = train_x @ beta
+    test_pred = test_x @ beta
+    score = rmse if mode == "rmse" else mae if mode == "mae" else None
+    if score is None:
+        raise ValueError(f"mode must be 'rmse' or 'mae', got {mode!r}")
+    return score(train_pred, train_y), score(test_pred, test_y), params
+
+
+def _fit_statsmodels(train_df, test_df, formula, mode="rmse"):
+    """Statsmodels OLS, kept only as a test oracle for `fit`. Slow because statsmodels re-parses
+    the formula with patsy on every call; do not use it in the harness."""
+    import statsmodels.formula.api as smf
+    result = smf.ols(formula=f"experimental ~ {formula}", data=train_df).fit()
+    score = rmse if mode == "rmse" else mae
+    train_err = score(result.predict(train_df).to_numpy(), train_df["experimental"].to_numpy())
+    test_err = score(result.predict(test_df).to_numpy(), test_df["experimental"].to_numpy())
+    return train_err, test_err, result.params
+
+
+def generate_solute_splits(n_splits, query_df, solvent, solutes,
+                           n_test=10, seed_offset=100):
+    """Build n_splits seeded train/test partitions for one solvent, splitting on SOLUTE so a
+    molecule is never in both sets. For split k: seed the RNG with k + seed_offset, shuffle the
+    solute list, take the first n_test as the test set and the rest as the training set, then select
+    that solvent's rows. Returns a list of (train_df, test_df) pairs."""
+    solute_list = list(solutes)
+    # filter this solvent once, then group its rows by solute, so each seed just concatenates the
+    # relevant groups instead of masking the whole table again.
+    solvent_df = query_df[query_df["solvent"] == solvent]
+    by_solute = {solute: rows for solute, rows in solvent_df.groupby("solute")}
+    empty = solvent_df.iloc[0:0]
+
+    def _gather(molecules):
+        frames = [by_solute[m] for m in molecules if m in by_solute]
+        return pd.concat(frames) if frames else empty
+
+    splits = []
+    for seed in range(n_splits):
+        random.seed(seed + seed_offset)
+        molecules = list(solute_list)
+        random.shuffle(molecules)
+        splits.append((_gather(molecules[n_test:]), _gather(molecules[:n_test])))
+    return splits
+
+
+def _solute_partitions(n_splits, solutes, n_test, seed_offset=100):
+    """The seeded (train_solutes, test_solutes) name partitions the harness fits over, matching
+    generate_solute_splits exactly: for split k, seed random with k + seed_offset, shuffle the full
+    solute list, test = the first n_test names, train = the rest. Row selection and solvent filtering
+    are left to the caller (see _split_fit_scores), so these partitions are computed once and reused
+    across formulas."""
+    partitions = []
+    base = list(solutes)
+    for seed in range(n_splits):
+        random.seed(seed + seed_offset)
+        molecules = list(base)
+        random.shuffle(molecules)
+        partitions.append((molecules[n_test:], molecules[:n_test]))
+    return partitions
+
+
+def _split_fit_scores(group_df, solvent, formulas, partitions, mode="rmse", min_train=2, swap=False):
+    """Fit each formula over the given seeded solute partitions for one group and solvent, returning
+    {(formula, seed): (train_error, test_error)} for every split that was fit.
+
+    The fast path behind run_fits / pareto_solvent_averaged_rmse: converts each column to numpy once
+    and selects each split's rows by integer index, avoiding the per-split pd.concat that dominates
+    runtime in generate_solute_splits + fit. Mathematically identical to fit() on the same rows.
+    Splits whose train set is smaller than min_train, whose test set is empty, or that fail to solve
+    are omitted from the result. swap=True trains on the first n_test solutes instead of the rest
+    (Figure 2A)."""
+    solvent_df = group_df[group_df["solvent"] == solvent]
+    n_rows = len(solvent_df)
+    score = rmse if mode == "rmse" else mae if mode == "mae" else None
+    if score is None:
+        raise ValueError(f"mode must be 'rmse' or 'mae', got {mode!r}")
+
+    # row positions (0..n_rows-1 within solvent_df) for each solute, and the numeric columns as numpy
+    solute_values = solvent_df["solute"].to_numpy()
+    by_solute = {}
+    for position, solute in enumerate(solute_values):
+        by_solute.setdefault(solute, []).append(position)
+    by_solute = {s: np.array(v, dtype=np.intp) for s, v in by_solute.items()}
+    empty = np.array([], dtype=np.intp)
+
+    def gather(names):
+        arrays = [by_solute[m] for m in names if m in by_solute]
+        return np.concatenate(arrays) if arrays else empty
+
+    # each split's train/test row indices, computed once and shared across formulas
+    split_indices = []
+    for train_names, test_names in partitions:
+        train_idx, test_idx = gather(train_names), gather(test_names)
+        if swap:
+            train_idx, test_idx = test_idx, train_idx
+        split_indices.append((train_idx, test_idx))
+
+    response = solvent_df["experimental"].to_numpy(dtype=float)
+    term_columns = {term: solvent_df[term].to_numpy(dtype=float)
+                    for formula in formulas for term in _formula_terms(formula)}
+
+    scores = {}
+    for formula in formulas:
+        terms = _formula_terms(formula)
+        design = np.column_stack([np.ones(n_rows)] + [term_columns[t] for t in terms])
+        finite = np.isfinite(response) & np.all(np.isfinite(design), axis=1)
+        for seed, (train_idx, test_idx) in enumerate(split_indices):
+            if len(train_idx) < min_train or len(test_idx) < 1:
+                continue
+            train_fit = train_idx[finite[train_idx]]
+            test_fit = test_idx[finite[test_idx]]
+            try:
+                beta = ols_fit(design[train_fit], response[train_fit])
+                train_err = score(design[train_fit] @ beta, response[train_fit])
+                test_err = score(design[test_fit] @ beta, response[test_fit])
+            except Exception:
+                continue
+            scores[(formula, seed)] = (train_err, test_err)
+    return scores
+
+
+def run_fits(query_df, solvents, formulas, n_splits, solutes,
+             group_cols=None, n_test=10, mode="rmse", min_train=2):
+    """Run the fitting harness over every (group, solvent, formula, split) combination and return
+    a tidy DataFrame with train_RMSE / test_RMSE per row.
+
+    group_cols: optional list of columns to fit within separately (e.g.
+    ["sap_nmr_method", "sap_basis", "sap_geometry_type"] for the per-method figures). When None,
+    the whole query_df is used as a single group (the caller has pre-filtered it). Splits that are
+    too small to fit, or that cannot be solved, are skipped.
+    """
+    rows = []
+    if group_cols:
+        group_iter = list(query_df.groupby(group_cols))
+    else:
+        group_iter = [((), query_df)]
+    partitions = _solute_partitions(n_splits, solutes, n_test)
+    for group_key, group_df in group_iter:
+        meta = {}
+        if group_cols:
+            key = group_key if isinstance(group_key, tuple) else (group_key,)
+            meta = dict(zip(group_cols, key))
+        for solvent in solvents:
+            scores = _split_fit_scores(group_df, solvent, formulas, partitions,
+                                       mode=mode, min_train=min_train)
+            for formula in formulas:
+                for seed in range(n_splits):
+                    got = scores.get((formula, seed))
+                    if got is None:
+                        continue
+                    rows.append({**meta, "solvent": solvent, "formula": formula,
+                                 "seed": seed, "train_RMSE": got[0], "test_RMSE": got[1]})
+    return pd.DataFrame(rows)
+
+
+def full_fit_coefficients(query_df, solvents, formula):
+    """Fit `formula` on the full data (no train/test split) for each solvent and return a DataFrame
+    indexed by parameter name ("Intercept" plus the predictors) with one column per solvent. This is
+    the deterministic full-fit OLS whose coefficients are the composite model's coefficients, stored
+    in applications.hdf5's composite_model/ols_coefficients group. Fitting the released delta-22 NN
+    data here (with nitromethane dropped) reproduces those coefficients to within the data's
+    integer-encoding precision."""
+    columns = {}
+    for solvent in solvents:
+        sub = query_df[query_df["solvent"] == solvent]
+        if len(sub) < 2:
+            continue
+        _, _, params = fit(sub, sub, formula)
+        columns[solvent] = params
+    return pd.DataFrame(columns)
+
+
+def fit_coefficients(query_df, solvent, formula, solutes, n_splits,
+                     n_test=10, seed_offset=100):
+    """Fit one formula for one solvent across all seeded splits and return the per-split fitted
+    coefficients as a DataFrame (one row per seed, columns = the statsmodels parameter names plus
+    the seed). This is the building block for the bootstrap coefficient distributions and for
+    verifying the composite-model coefficients fit on delta-22."""
+    splits = generate_solute_splits(n_splits, query_df, solvent, solutes,
+                                    n_test=n_test, seed_offset=seed_offset)
+    records = []
+    for seed, (train_df, test_df) in enumerate(splits):
+        if len(train_df) < 2:
+            continue
+        _, _, params = fit(train_df, test_df, formula)
+        record = {"solvent": solvent, "seed": seed}
+        record.update({name: float(value) for name, value in params.items()})
+        records.append(record)
+    return pd.DataFrame(records)
+
+
+# ---------------------------------------------------------------------------
+# Figure S8 B-D (Convergence of Explicit Solvent Corrections)
+
+def frame_corrections(perturbed, atom_indices):
+    """Per-frame solvent correction for one site: the mean over the site's atoms of
+    (solvated - isolated), from a raw (n_frames, n_atoms, 2) perturbed-shielding array
+    (load_perturbed_shieldings). Frames whose site atoms were not computed read as NaN."""
+    site = (perturbed[:, :, 1] - perturbed[:, :, 0])[:, list(atom_indices)]
+    valid = ~np.isnan(site)
+    counts = valid.sum(axis=1)
+    sums = np.where(valid, site, 0.0).sum(axis=1)
+    return np.where(counts > 0, sums / np.maximum(counts, 1), np.nan)
+
+
+def running_average(values):
+    """Cumulative mean over frames ignoring NaN frames - the running value a correction converges to
+    as more molecular-dynamics frames are included (Figure S8 panel B)."""
+    values = np.asarray(values, dtype=float)
+    valid = ~np.isnan(values)
+    counts = np.cumsum(valid)
+    sums = np.cumsum(np.where(valid, values, 0.0))
+    with np.errstate(invalid="ignore", divide="ignore"):
+        return np.where(counts > 0, sums / counts, np.nan)
+
+
+def autocorrelation(values, max_lag=None):
+    """Autocorrelation of the valid per-frame corrections at lags 0..max_lag (Figure S8 panel C).
+    A fast drop to zero means frames are effectively independent."""
+    x = np.asarray(values, dtype=float)
+    x = x[~np.isnan(x)]
+    x = x - x.mean()
+    n = len(x)
+    if n < 2:
+        return np.array([1.0])
+    if max_lag is None:
+        max_lag = n - 1
+    denom = np.sum(x * x)
+    return np.array([np.sum(x[:n - lag] * x[lag:]) / denom for lag in range(max_lag + 1)])
+
+
+def frame_validity(perturbed):
+    """Per frame, whether any shielding value was computed for that frame (a raw (n_frames,
+    n_atoms, 2) perturbed-shielding array from load_perturbed_shieldings has NaN entries for atoms
+    that were not computed). A frame with every entry NaN was skipped entirely by the DFT queue;
+    Figure S8 panel D visualizes this mask across solvents for a representative solute."""
+    return ~np.all(np.isnan(perturbed), axis=(1, 2))
+
+
+def frame_validity_grid(hdf5_filepath, solute, solvents, explicit_solvation_model, shield_type="dft"):
+    """Figure S8 panel D data: for one solute and one molecular-dynamics engine, the per-frame
+    validity mask (frame_validity) for each solvent, stacked into a single boolean grid of shape
+    (n_solvents, max_frames). Solvents with fewer trajectory frames than the longest one are padded
+    with False (no data) past their own frame count, so every row lines up on a shared frame-index
+    axis. Returns (grid, frame_counts), where frame_counts[i] is the true (unpadded) frame count for
+    solvents[i], needed to normalize the frame-wise correction histograms in panel B."""
+    masks = []
+    for solvent in solvents:
+        perturbed = load_perturbed_shieldings(hdf5_filepath, solute, solvent,
+                                              explicit_solvation_model, shield_type)
+        masks.append(frame_validity(perturbed))
+    frame_counts = np.array([len(m) for m in masks])
+    max_len = int(frame_counts.max()) if len(frame_counts) else 0
+    grid = np.zeros((len(solvents), max_len), dtype=bool)
+    for i, mask in enumerate(masks):
+        grid[i, :len(mask)] = mask
+    return grid, frame_counts
+
+
+# ---------------------------------------------------------------------------
+# Figures S11 / S13 (MagNET vs DFT corrections)
+
+def compare_dft_nn(dft_df, nn_df, value_col, keys=("solute", "site", "nucleus")):
+    """Merge the DFT and MagNET (NN) tables on the given keys and return, per site, the DFT and NN
+    value of one correction column plus their difference (error = NN - DFT). Used for the
+    rovibrational (qcd) comparison of Figure S11 and the explicit-solvent (desmond/openMM) comparison
+    of Figure S13. The correction is tiled across NMR methods, so duplicate keyed rows are dropped
+    before merging."""
+    keys = list(keys)
+    dft = dft_df[keys + [value_col]].drop_duplicates(subset=keys).dropna(subset=[value_col])
+    nn = nn_df[keys + [value_col]].drop_duplicates(subset=keys).dropna(subset=[value_col])
+    merged = dft.merge(nn, on=keys, suffixes=("_dft", "_nn"))
+    merged["error"] = merged[f"{value_col}_nn"] - merged[f"{value_col}_dft"]
+    return merged
+
+
+def qcd_correction_by_site(dft_df, nn_df, nucleus="H"):
+    """Per site, the DFT and MagNET (NN) rovibrational (QCD) correction for one nucleus, sorted by
+    solute (SI Figure S11 panel C, the per-solute view behind panel A's aggregate scatter). QCD is a
+    gas-phase rovibrational correction, so unlike the explicit-solvent per-site view (Figure S13 panel
+    C) there is no solvent to split on: this is a plain wrapper around compare_dft_nn that keeps one
+    nucleus and orders the rows for the per-solute strip plot."""
+    merged = compare_dft_nn(dft_df, nn_df, "qcd", keys=("solute", "site", "nucleus"))
+    sub = merged[merged["nucleus"] == nucleus]
+    return sub.sort_values("solute", kind="stable").reset_index(drop=True)
+
+
+def compare_dft_nn_by_engine(dft_df, nn_df, engines=("desmond", "openMM"),
+                             keys=("solute", "site", "nucleus", "solvent")):
+    """SI Figure S13 panels A and B data: compare_dft_nn run separately for each molecular-dynamics
+    engine (Desmond and OpenMM), stacked into one table with a shared 'dft_value' / 'nn_value' /
+    'error' column and an 'engine' label. compare_dft_nn alone cannot be stacked directly because it
+    names its value columns after value_col (e.g. 'desmond_dft' vs 'openMM_dft'); this renames them
+    to a common pair of columns first so both engines can share one axis."""
+    frames = []
+    for engine in engines:
+        one = compare_dft_nn(dft_df, nn_df, engine, keys=keys)
+        one = one.rename(columns={f"{engine}_dft": "dft_value", f"{engine}_nn": "nn_value"})
+        one["engine"] = engine
+        frames.append(one[list(keys) + ["dft_value", "nn_value", "error", "engine"]])
+    return pd.concat(frames, ignore_index=True)
+
+
+# ---------------------------------------------------------------------------
+# Figure S5 (PCM Captures Bulk Solvent Effects)
+
+# bulk solvent properties used to test what the PCM benefit tracks (TIP4P uses water's values)
+SOLVENT_DIELECTRIC = {
+    "chloroform": 4.81, "dichloromethane": 9.08, "methanol": 32.6, "TIP4P": 78.54,
+    "trifluoroethanol": 26.7, "acetone": 21.01, "acetonitrile": 36.64, "dimethylsulfoxide": 47.0,
+    "tetrahydrofuran": 7.52, "benzene": 2.28, "chlorobenzene": 5.69, "toluene": 2.38,
+}
+SOLVENT_POLARIZABILITY = {
+    "chloroform": 8.5, "dichloromethane": 6.5, "methanol": 3.3, "TIP4P": 1.5,
+    "trifluoroethanol": 4.5, "acetone": 6.4, "acetonitrile": 4.4, "dimethylsulfoxide": 7.9,
+    "tetrahydrofuran": 8.0, "benzene": 10.3, "chlorobenzene": 12.4, "toluene": 12.3,
+}
+
+
+def pcm_benefit_per_solvent(query_df_dft, method, basis, geometry, solvents, n_splits, solutes,
+                            nucleus="H", n_test=10):
+    """Panel S5 data: per solvent, the percent reduction in mean test RMSE from adding the
+    implicit-solvent (PCM) term, for one DFT method. Returns a Series indexed by solvent. A positive
+    value means PCM lowered the error. Plotted against the bulk solvent properties above.
+
+    The published Figure S5 aggregates the per-split test RMSEs with the mean (not the median); this
+    is what reproduces the paper's Pearson R values exactly."""
+    one = query_df_dft[(query_df_dft["sap_nmr_method"] == method)
+                       & (query_df_dft["sap_basis"] == basis)
+                       & (query_df_dft["sap_geometry_type"] == geometry)
+                       & (query_df_dft["nucleus"] == nucleus)]
+    one = add_composite_columns(one)
+    results = run_fits(one, solvents, ["stationary", "stationary_plus_pcm"],
+                       n_splits=n_splits, solutes=solutes, n_test=n_test)
+    mean = results.groupby(["solvent", "formula"])["test_RMSE"].mean().unstack("formula")
+    return -((mean["stationary_plus_pcm"] - mean["stationary"]) / mean["stationary"]) * 100.0
+
+
+# ---------------------------------------------------------------------------
+# Figure S7 (Explicit Solvent Corrections are Method-Independent)
+
+def explicit_correction_pairs(combined_df, solvent, nucleus="H"):
+    """Per site, the Desmond and OpenMM explicit-solvent corrections for one solvent (Figure S7).
+    The explicit correction does not depend on the NMR method, so the same value is tiled across the
+    method rows; duplicate site rows are dropped. Works on either the DFT or the NN combined table.
+    A tight Desmond-vs-OpenMM agreement is the method-independence claim."""
+    sub = combined_df[(combined_df["nucleus"] == nucleus) & (combined_df["solvent"] == solvent)]
+    out = sub[["solute", "site", "desmond", "openMM"]].drop_duplicates(subset=["solute", "site"])
+    return out.dropna(subset=["desmond", "openMM"]).reset_index(drop=True)
+
+
+def explicit_correction_dft_nn_pairs(dft_df, nn_df, solvent, nucleus="H"):
+    """Per site, the Desmond and OpenMM explicit-solvent corrections from both DFT and MagNET-x (NN)
+    for one solvent and nucleus (Figure S13 panel C, the per-site view behind panel A's aggregate
+    scatter). Built from two explicit_correction_pairs calls (one per source) merged on
+    (solute, site), giving columns desmond_dft, openMM_dft, desmond_nn, openMM_nn."""
+    dft_pairs = explicit_correction_pairs(dft_df, solvent, nucleus)
+    nn_pairs = explicit_correction_pairs(nn_df, solvent, nucleus)
+    return dft_pairs.merge(nn_pairs, on=["solute", "site"], suffixes=("_dft", "_nn"))
+
+
+def si_s13d_fitting_accuracy(query_df_dft, query_df_nn, solvents, n_splits, solutes, nucleus="H",
+                             method="dsd_pbep86", basis="pcSseg3", geometry="pbe0_tz", n_test=10):
+    """Figure S13 panel D data: the semi-parsimonious composite model's (composite_models.py's
+    best_semi formula) test RMSE against experiment, for every solvent, split by molecular-dynamics
+    engine (Desmond vs OpenMM) and by the source of the explicit-solvent/vibrational term (DFT vs
+    MagNET-x, "NN"). Protons use the QCD rovibrational correction (stationary_plus_qcd), which the
+    SI text says is the same B3LYP/cc-pVDZ value for both DFT and NN bars; carbons use the classical
+    vibrational correction computed from the same engine's trajectories (stationary_plus_desmond_vib
+    or stationary_plus_openMM_vib), which does differ between the DFT and NN combined tables. Returns
+    a tidy DataFrame with one row per (engine, source, solvent, seed) and a test_RMSE column."""
+    engine_vib_column = {"desmond": "stationary_plus_des_vib", "openMM": "stationary_plus_op_vib"}
+    rows = []
+    for source_name, query_df in [("DFT", query_df_dft), ("NN", query_df_nn)]:
+        sub = query_df[query_df["nucleus"] == nucleus]
+        if source_name == "DFT":
+            sub = sub[(sub["sap_nmr_method"] == method) & (sub["sap_basis"] == basis)
+                     & (sub["sap_geometry_type"] == geometry)]
+        sub = add_composite_columns(sub)
+        for engine in ("desmond", "openMM"):
+            base_term = "stationary_plus_qcd" if nucleus == "H" else engine_vib_column[engine]
+            formula = f"{base_term} + {engine}"
+            results = run_fits(sub, solvents, [formula], n_splits=n_splits, solutes=solutes,
+                               n_test=n_test)
+            results["engine"] = engine
+            results["source"] = source_name
+            rows.append(results)
+    return pd.concat(rows, ignore_index=True)
+
+
+# ---------------------------------------------------------------------------
+# Figure 2A / Figure S1 (the accuracy-vs-cost Pareto frontier)
+
+# MagNET predicts the implicit-solvent (PCM) correction for chloroform only. To get a value for the
+# other solvents, the chloroform correction is scaled by a per-solvent factor, fit on the delta-22
+# DFT PCM corrections themselves: the slope (through the origin) of one method's PCM correction in
+# that solvent against the reference method's PCM correction in chloroform. The reference (input)
+# method is B3LYP-D3(BJ); the output method is the one MagNET was trained to reproduce (wP04 for
+# proton, wB97X-D for carbon). These factors are also stored in applications.hdf5, but recomputing
+# them here keeps Figure 2A self-contained within the delta-22 data.
+MAGNET_PCM_INPUT_METHOD = "b3lyp_d3bj"
+MAGNET_PCM_OUTPUT_METHODS = {"H": "wp04", "C": "wb97xd"}
+
+
+def pcm_conversion_factors(query_df_dft, nucleus, output_method=None,
+                           input_method=MAGNET_PCM_INPUT_METHOD,
+                           geometry="aimnet2", basis="pcSseg2"):
+    """The per-solvent factor that scales MagNET's chloroform PCM correction to each solvent, for
+    one nucleus. Recomputed from the delta-22 DFT PCM corrections: for each solvent, the
+    through-the-origin slope of the output method's PCM correction in that solvent against the input
+    (reference) method's PCM correction in chloroform, over the shared sites. Returns a Series
+    indexed by solvent in the reader's native naming (water is "TIP4P"). Matches the factors stored
+    in applications.hdf5 to machine precision."""
+    if output_method is None:
+        output_method = MAGNET_PCM_OUTPUT_METHODS[nucleus]
+    base = query_df_dft[(query_df_dft["sap_geometry_type"] == geometry)
+                        & (query_df_dft["sap_basis"] == basis)
+                        & (query_df_dft["nucleus"] == nucleus)]
+    reference = (base[(base["sap_nmr_method"] == input_method) & (base["solvent"] == "chloroform")]
+                 .set_index(["solute", "site"])["pcm"])
+    output = base[base["sap_nmr_method"] == output_method]
+    factors = {}
+    for solvent in DESMOND_SOLVENTS:
+        y = output[output["solvent"] == solvent].set_index(["solute", "site"])["pcm"]
+        joined = pd.concat([reference.rename("x"), y.rename("y")], axis=1).dropna()
+        xv = joined["x"].to_numpy(dtype=float)
+        yv = joined["y"].to_numpy(dtype=float)
+        factors[solvent] = float((xv * yv).sum() / (xv * xv).sum()) if len(xv) else np.nan
+    return pd.Series(factors, name="pcm_conversion_factor")
+
+
+def add_nn_stationary_plus_pcm(query_df_nn, factors_by_nucleus):
+    """Add the MagNET stationary_plus_pcm column to a flat NN query table: the gas-phase shielding
+    plus the chloroform PCM correction scaled by the per-solvent, per-nucleus factor from
+    pcm_conversion_factors. factors_by_nucleus maps "H"/"C" to those Series. Returns a copy."""
+    out = query_df_nn.copy()
+    factor = out.apply(lambda r: factors_by_nucleus[r["nucleus"]].get(r["solvent"], np.nan), axis=1)
+    out["stationary_plus_pcm"] = out["stationary"] + out["pcm"] * factor
+    return out
+
+
+# the columns the Pareto fit needs from each flat table, before the "sap_" prefix is stripped
+_PARETO_COLUMNS = ["sap_geometry_type", "sap_nmr_method", "sap_basis", "nucleus", "site",
+                   "solute", "solvent", "experimental", "stationary_plus_pcm"]
+
+
+def assemble_pareto_fitting_df(query_df_dft, query_df_nn, conversion_factors=None,
+                               exclude_solutes=("nitromethane",)):
+    """Build the combined DFT + MagNET table the Pareto fit runs on: experimental shift against the
+    single stationary_plus_pcm predictor, for every method/basis/geometry/nucleus/solvent. The DFT
+    predictor is stationary + pcm; the MagNET predictor uses the scaled chloroform correction
+    (conversion_factors, recomputed from query_df_dft if not given). TIP4P is renamed to water and
+    the excluded solutes (nitromethane by default, an outlier) are removed. The DFT rows come first
+    so the solute order is stable for the seeded splits."""
+    if conversion_factors is None:
+        conversion_factors = {nuc: pcm_conversion_factors(query_df_dft, nuc) for nuc in ("H", "C")}
+    dft = query_df_dft.copy()
+    dft["stationary_plus_pcm"] = dft["stationary"] + dft["pcm"]
+    nn = add_nn_stationary_plus_pcm(query_df_nn, conversion_factors)
+    combined = pd.concat([dft[_PARETO_COLUMNS], nn[_PARETO_COLUMNS]], ignore_index=True)
+    combined = combined.rename(columns=lambda c: c.replace("sap_", ""))
+    combined["solvent"] = combined["solvent"].replace("TIP4P", "water")
+    combined["experimental"] = pd.to_numeric(combined["experimental"], errors="coerce")
+    if exclude_solutes:
+        combined = combined[~combined["solute"].isin(set(exclude_solutes))]
+    return combined.reset_index(drop=True)
+
+
+def pareto_solvent_averaged_rmse(fitting_df, solutes=None, n_splits=100, n_test=10,
+                                 random_state=1, formula="stationary_plus_pcm"):
+    """For every (geometry, method, basis, nucleus, solvent), the mean test RMSE of the
+    experimental ~ stationary_plus_pcm fit over n_splits seeded per-solute splits, then a
+    solvent-averaged row per (geometry, method, basis, nucleus). This is the accuracy axis of
+    Figure 2A / Figure S1.
+
+    Figure 2A splits differently from the other figures: it trains on the first n_test shuffled
+    solutes and tests on the rest, seeding with random_state + split index (the other figures train
+    on the remainder and seed with +100); swap=True on _split_fit_scores gets that train-on-first-
+    n_test behavior. The per-split errors are averaged with the mean, not the median used elsewhere."""
+    if solutes is None:
+        solutes = fitting_df["solute"].unique().tolist()
+    group_cols = ["geometry_type", "nmr_method", "basis", "nucleus", "solvent"]
+    rows = []
+    partitions = _solute_partitions(n_splits, solutes, n_test, seed_offset=random_state)
+    for key, group_df in fitting_df.groupby(group_cols):
+        solvent = key[group_cols.index("solvent")]
+        scores = _split_fit_scores(group_df, solvent, [formula], partitions, swap=True)
+        errors = [scores.get((formula, seed), (np.nan, np.nan))[1] for seed in range(n_splits)]
+        with np.errstate(invalid="ignore"):
+            mean_err = float(np.nanmean(errors)) if errors else np.nan
+        rows.append({**dict(zip(group_cols, key)), "fitting_RMSE": mean_err})
+    raw = pd.DataFrame(rows)
+    averaged = (raw.groupby(["geometry_type", "nmr_method", "basis", "nucleus"])["fitting_RMSE"]
+                .mean().reset_index())
+    averaged["solvent"] = "solvent-averaged"
+    return pd.concat([raw, averaged], ignore_index=True)
+
+
+def load_pareto_timings(delta22_path):
+    """The compute-time axis of Figure 2A: (dft_gas_timings, nn_timings). DFT timings are the
+    gas-phase ('none' solvent) total time per (geometry, method, basis); MagNET timings sum the NMR
+    time across its solvents and keep one geometry time. Reads the timing groups through the reader."""
+    from delta22_reader import load_delta22_dft_timings, load_delta22_nn_timings
+    dft = load_delta22_dft_timings(delta22_path)
+    dft.index = dft.index.set_names([n.replace("sap_", "") for n in dft.index.names])
+    dft = dft.droplevel("solvent_model")
+    dft_gas = dft[dft.index.get_level_values("solvent") == "none"].droplevel("solvent")
+    nn = load_delta22_nn_timings(delta22_path)
+    nn.index = nn.index.set_names([n.replace("sap_", "") for n in nn.index.names])
+    nn = nn.droplevel("solvent_model")
+    nn = nn.groupby(["geometry_type", "nmr_method", "basis"]).agg(
+        {"geometry_time": "first", "nmr_time": "sum"})
+    nn["total_time"] = nn["geometry_time"] + nn["nmr_time"]
+    return dft_gas, nn
+
+
+def attach_pareto_timings(rmse_df, dft_gas_timings, nn_timings):
+    """Join the gas-phase DFT timings onto the per-method RMSE rows and substitute the MagNET
+    timings (which are not in the DFT gas table) for the MagNET rows. Returns rmse_df with
+    geometry_time / nmr_time / total_time columns added."""
+    out = rmse_df.join(dft_gas_timings, on=["geometry_type", "nmr_method", "basis"], how="left")
+    magnet = out["nmr_method"] == "MagNET"
+    nn_row = nn_timings.loc[("aimnet2", "MagNET", "N/A")]
+    out.loc[magnet, ["geometry_time", "nmr_time", "total_time"]] = \
+        nn_row[["geometry_time", "nmr_time", "total_time"]].to_numpy()
+    return out
+
+
+def fig2a_pareto_points(query_df_dft, query_df_nn, dft_gas_timings, nn_timings,
+                        solutes=None, n_splits=100, n_test=10, random_state=1,
+                        conversion_factors=None):
+    """The full Figure 2A / Figure S1 table: one row per (geometry, method, basis, nucleus, solvent)
+    plus the solvent-averaged rows, each carrying its mean test RMSE and its total compute time.
+    Pass the flat DFT and NN query tables (build_query_df output) and the two timing tables from
+    load_pareto_timings. The frontier for a given nucleus is pareto_frontier on the
+    solvent-averaged rows with x="total_time", y="fitting_RMSE"."""
+    fitting = assemble_pareto_fitting_df(query_df_dft, query_df_nn, conversion_factors)
+    rmse = pareto_solvent_averaged_rmse(fitting, solutes, n_splits=n_splits, n_test=n_test,
+                                        random_state=random_state)
+    return attach_pareto_timings(rmse, dft_gas_timings, nn_timings)
+
+
+def pareto_table_curated(points_df, nucleus, solvent="chloroform"):
+    """SI Tables S1 (1H, nucleus="H") and S2 (13C, nucleus="C"): the curated 55-row subset of
+    fig2a_pareto_points' full cross product that the source tables print. fig2a_pareto_points (and
+    the SI Figure S1 scatter built from it) wants every aimnet2-geometry method/basis combination,
+    each an unlabeled point on the plot; the printed tables instead keep only MagNET plus a single
+    representative aimnet2-geometry DFT row (the same reference method each nucleus's MagNET-Zero
+    output is trained against, MAGNET_PCM_OUTPUT_METHODS, at pcSseg2), alongside the full
+    pbe0_tz-geometry grid (18 methods x 3 basis sets, minus mp2/pcSseg3 which was never computed, so
+    53 rows). Sorted by total_time ascending to match the source table's row order."""
+    sub = points_df[(points_df["nucleus"] == nucleus) & (points_df["solvent"] == solvent)]
+    aimnet2_method = MAGNET_PCM_OUTPUT_METHODS[nucleus]
+    keep = (
+        (sub["geometry_type"] == "pbe0_tz")
+        | (sub["nmr_method"] == "MagNET")
+        | ((sub["geometry_type"] == "aimnet2") & (sub["nmr_method"] == aimnet2_method)
+           & (sub["basis"] == "pcSseg2"))
+    )
+    return sub[keep].sort_values("total_time").reset_index(drop=True)
+
+
+def pareto_frontier(points_df, x="total_time", y="test_RMSE"):
+    """Return the rows of points_df on the lower-left Pareto frontier, i.e. the points for which
+    no other point has both a smaller x (compute time) and a smaller y (error). Sorted by x.
+    """
+    ordered = points_df.sort_values([x, y]).reset_index(drop=True)
+    frontier_rows = []
+    best_y = np.inf
+    for _, row in ordered.iterrows():
+        if row[y] < best_y:
+            frontier_rows.append(row)
+            best_y = row[y]
+    return pd.DataFrame(frontier_rows).reset_index(drop=True)
+
+
+# ---------------------------------------------------------------------------
+# correlation matrices (Figure 2B / Figure S4)
+
+def correlation_matrix(query_df, value_col, sites_index, column_var):
+    """Correlation matrix of value_col across the levels of column_var. Rows are the shared sites
+    (sites_index, e.g. ["solute", "site"]); columns are the column_var values (e.g. "solvent" or
+    "sap_nmr_method"). Returns a square DataFrame of Pearson correlations between the columns,
+    computed on the sites where both columns have a value."""
+    pivot = query_df.pivot_table(index=sites_index, columns=column_var, values=value_col)
+    return pivot.corr()
+
+
+# SI Figure S15 "Correlations Between Features": correlations among the composite features within
+# one nucleus and solvent.
+FEATURE_CORRELATION_COLUMNS = ("stationary", "pcm", "desmond", "desmond_vib", "qcd")
+
+
+def feature_correlation_matrix(query_df, nucleus, solvent, cols=FEATURE_CORRELATION_COLUMNS, squared=False):
+    """Pearson correlation matrix (R^2 if squared=True) between the feature columns, within one
+    nucleus and solvent, over sites where both features have a value. query_df must be filtered to a
+    single (sap_nmr_method, sap_basis, sap_geometry_type); si_s15_feature_correlations does that."""
+    sub = query_df[(query_df["nucleus"] == nucleus) & (query_df["solvent"] == solvent)][list(cols)]
+    sub = sub.dropna(how="all")
+    if len(sub) < 2:
+        return pd.DataFrame(np.nan, index=cols, columns=cols)
+    corr = sub.corr(method="pearson")
+    return corr ** 2 if squared else corr
+
+
+def average_feature_correlation_matrix(query_df, nucleus, solvents, cols=FEATURE_CORRELATION_COLUMNS,
+                                       squared=False):
+    """Cell-by-cell nanmean of feature_correlation_matrix over the solvents (SI Figure S15's
+    solvent-averaged matrix). nanmean so a solvent with too few sites for one feature pair drops out
+    of that cell alone, not the whole average. Same single-level-of-theory requirement."""
+    mats = [feature_correlation_matrix(query_df, nucleus, s, cols, squared) for s in solvents]
+    stacked = np.stack([m.to_numpy() for m in mats])
+    with np.errstate(invalid="ignore"):
+        avg = np.nanmean(stacked, axis=0)
+    return pd.DataFrame(avg, index=mats[0].index, columns=mats[0].columns)
+
+
+def si_s15_feature_correlations(query_df_dft, nucleus, solvents=None):
+    """Solvent-averaged feature correlations for one nucleus from the unfiltered
+    load_query_df_dft(...) output: filters to the reference level
+    (MAGNET_PCM_OUTPUT_METHODS[nucleus], pcSseg2, aimnet2), then returns {"r": Pearson R matrix,
+    "r2": R^2 matrix}. The single-level filter is required; pooling levels of theory gives wrong
+    numbers."""
+    solvents = list(solvents) if solvents is not None else sorted(query_df_dft["solvent"].unique())
+    method = MAGNET_PCM_OUTPUT_METHODS[nucleus]
+    sub = query_df_dft[(query_df_dft["nucleus"] == nucleus) & (query_df_dft["sap_nmr_method"] == method)
+                       & (query_df_dft["sap_basis"] == "pcSseg2") & (query_df_dft["sap_geometry_type"] == "aimnet2")]
+    return {"r": average_feature_correlation_matrix(sub, nucleus, solvents),
+            "r2": average_feature_correlation_matrix(sub, nucleus, solvents, squared=True)}
+
+
+def pcm_desmond_correlation_by_solvent(query_df_dft, solvents=None):
+    """SI Figure S15's PCM-vs-Desmond table: per nucleus and solvent, the Pearson correlation
+    between the PCM (implicit) and Desmond (explicit) corrections. Filters to the reference level
+    per nucleus itself (same as si_s15_feature_correlations). Returns a DataFrame indexed by
+    nucleus, one column per solvent."""
+    solvents = list(solvents) if solvents is not None else sorted(query_df_dft["solvent"].unique())
+    rows = {}
+    for nucleus in ("H", "C"):
+        method = MAGNET_PCM_OUTPUT_METHODS[nucleus]
+        nuc_df = query_df_dft[(query_df_dft["nucleus"] == nucleus) & (query_df_dft["sap_nmr_method"] == method)
+                              & (query_df_dft["sap_basis"] == "pcSseg2") & (query_df_dft["sap_geometry_type"] == "aimnet2")]
+        row = {}
+        for solvent in solvents:
+            sub = nuc_df[nuc_df["solvent"] == solvent][["pcm", "desmond"]].dropna()
+            row[solvent] = float(sub["pcm"].corr(sub["desmond"])) if len(sub) > 1 else float("nan")
+        rows[nucleus] = row
+    return pd.DataFrame(rows).T
+
+
+# ---------------------------------------------------------------------------
+# PCM benefit (Figure 3A / Figure S5): how much a correction reduces test error
+
+def median_test_rmse(results_df, by):
+    """Median test_RMSE grouped by the given columns (e.g. ["solvent", "formula"] or
+    ["sap_nmr_method", "formula"]). A small convenience used by several figures."""
+    return results_df.groupby(by)["test_RMSE"].median().reset_index()
+
+
+# ---------------------------------------------------------------------------
+# Figure 3 panel data (Explicit Solvation is Critical for Accuracy)
+
+# the three solvent classes Figure 3 / S6 group by
+SOLVENT_GROUPS = {
+    "Polar Aprotic": ["chloroform", "dichloromethane", "tetrahydrofuran",
+                      "acetonitrile", "dimethylsulfoxide", "acetone"],
+    "Polar Protic": ["methanol", "TIP4P", "trifluoroethanol"],
+    "Aromatic": ["benzene", "toluene", "chlorobenzene"],
+}
+
+
+def fig3a_pcm_benefit(query_df_dft, solvents, n_splits, solutes, nucleus="H", n_test=10):
+    """Panel 3A data: for every DFT method/basis/geometry, the seeded test RMSE with and without the
+    implicit-solvent term ("stationary" vs the single-column "stationary_plus_pcm"). The boxplot
+    shows how much the PCM term lowers error across methods. One nucleus at a time (proton and
+    carbon shieldings sit on different scales, so they are never pooled in a fit). Pass a query_df
+    that has been through add_composite_columns. Returns the tidy run_fits results (one row per
+    method/solvent/formula/seed)."""
+    sub = query_df_dft[query_df_dft["nucleus"] == nucleus]
+    return run_fits(sub, solvents, ["stationary", "stationary_plus_pcm"],
+                    n_splits=n_splits, solutes=solutes,
+                    group_cols=["sap_nmr_method", "sap_basis", "sap_geometry_type"], n_test=n_test)
+
+
+def fig3a_pcm_benefit_by_solvent(query_df_dft, solvents, n_splits, solutes, nucleus="H", n_test=10):
+    """Panel 3A data, split by solvent: for every DFT method and solvent, the per-split percent
+    benefit of adding the implicit-solvent (PCM) term, 100 * (stationary - stationary_plus_pcm) /
+    stationary. A positive value means PCM lowered test RMSE for that split; a negative value means
+    PCM made it worse. Unlike fig3a_pcm_benefit (which reports raw RMSE pooled across solvents),
+    this keeps one row per (method, solvent, seed), so a solvent such as benzene can be shown
+    separately from chloroform and the aromatic-solvent-specific PCM penalty stays visible instead
+    of being averaged away."""
+    results = fig3a_pcm_benefit(query_df_dft, solvents, n_splits, solutes,
+                                nucleus=nucleus, n_test=n_test)
+    pivot = results.pivot_table(
+        index=["sap_nmr_method", "sap_basis", "sap_geometry_type", "solvent", "seed"],
+        columns="formula", values="test_RMSE").reset_index()
+    pivot["percent_benefit"] = (100.0 * (pivot["stationary"] - pivot["stationary_plus_pcm"])
+                                / pivot["stationary"])
+    return pivot
+
+
+def fig3b_shift_differences(query_df_dft, method, basis, geometry, nucleus="H",
+                            x_solvent="methanol", y_solvent="benzene", reference="chloroform"):
+    """Panel 3B data: per site, the experimental solvent-induced shift differences relative to a
+    reference solvent (x = shift in x_solvent minus reference, y = shift in y_solvent minus
+    reference) and the implicit-solvent (PCM) prediction of the same differences. One DFT method.
+    The scatter compares whether PCM reproduces the measured solvent-to-solvent differences."""
+    sub = query_df_dft[(query_df_dft["sap_nmr_method"] == method)
+                       & (query_df_dft["sap_basis"] == basis)
+                       & (query_df_dft["sap_geometry_type"] == geometry)
+                       & (query_df_dft["nucleus"] == nucleus)]
+    cols = ["solute", "site", "experimental", "pcm"]
+    ref = sub[sub["solvent"] == reference][cols]
+    xs = sub[sub["solvent"] == x_solvent][cols]
+    ys = sub[sub["solvent"] == y_solvent][cols]
+    merged = (xs.merge(ref, on=["solute", "site"], suffixes=("_x", "_ref"))
+                .merge(ys, on=["solute", "site"]))
+    merged.columns = ["solute", "site", "exp_x", "pcm_x", "exp_ref", "pcm_ref", "exp_y", "pcm_y"]
+    merged["exp_diff_x"] = merged["exp_x"] - merged["exp_ref"]
+    merged["exp_diff_y"] = merged["exp_y"] - merged["exp_ref"]
+    merged["pcm_diff_x"] = merged["pcm_x"] - merged["pcm_ref"]
+    merged["pcm_diff_y"] = merged["pcm_y"] - merged["pcm_ref"]
+    return merged.dropna(subset=["exp_diff_x", "exp_diff_y"]).reset_index(drop=True)
+
+
+def fig3d_formula_regressions(query_df_dft, method, basis, geometry, formulas, solvents,
+                              n_splits, solutes, nucleus="H", n_test=10):
+    """Panel 3D data: for one DFT method, the seeded test RMSE of an implicit-vs-explicit formula
+    ladder per solvent. Pass a query_df that has been through add_composite_columns. Returns the
+    tidy run_fits results (one row per solvent/formula/seed)."""
+    one = query_df_dft[(query_df_dft["sap_nmr_method"] == method)
+                       & (query_df_dft["sap_basis"] == basis)
+                       & (query_df_dft["sap_geometry_type"] == geometry)
+                       & (query_df_dft["nucleus"] == nucleus)]
+    return run_fits(one, solvents, formulas, n_splits=n_splits, solutes=solutes, n_test=n_test)
+
+
+# ---------------------------------------------------------------------------
+# Figure 4 + SI S16/S17/S18 (solvent-induced shifts: experiment vs prediction)
+
+def add_solvent_mean(query_df, group_cols=("solute", "site", "nucleus")):
+    """Append a 'solvent_mean' pseudo-solvent: for each site, the mean over the real solvents of the
+    numeric columns. Used as the solvent-averaged reference for Figure S18. Returns a new DataFrame
+    with the extra rows."""
+    means = query_df.groupby(list(group_cols)).mean(numeric_only=True).reset_index()
+    means["solvent"] = "solvent_mean"
+    return pd.concat([query_df, means], ignore_index=True)
+
+
+def solvent_pair_differences(query_df, check_solvent, reference_solvent, nucleus="H",
+                             explicit="desmond"):
+    """Per (solute, site), the experimental and predicted correction differences between one solvent
+    and a reference solvent (Figure 4): the experimental difference is reference minus check, and
+    each correction difference (implicit PCM, explicit, and explicit+vibrations) is check minus
+    reference. One DFT method's columns must already be selected."""
+    a = query_df[(query_df["nucleus"] == nucleus) & (query_df["solvent"] == check_solvent)]
+    b = query_df[(query_df["nucleus"] == nucleus) & (query_df["solvent"] == reference_solvent)]
+    merged = a.merge(b, on=["solute", "site"], suffixes=(f"_{check_solvent}", f"_{reference_solvent}"))
+    c, r, e = check_solvent, reference_solvent, explicit
+    out = pd.DataFrame({
+        "solute": merged["solute"], "site": merged["site"],
+        "reference": reference_solvent, "solvent": check_solvent,
+        "exp_diff": merged[f"experimental_{r}"] - merged[f"experimental_{c}"],
+        "implicit_diff": merged[f"pcm_{c}"] - merged[f"pcm_{r}"],
+        "explicit_diff": merged[f"{e}_{c}"] - merged[f"{e}_{r}"],
+        "explicit_vib_diff": ((merged[f"{e}_{c}"] + merged[f"{e}_vib_{c}"])
+                              - (merged[f"{e}_{r}"] + merged[f"{e}_vib_{r}"])),
+    })
+    return out.dropna(subset=["exp_diff"]).reset_index(drop=True)
+
+
+def solvent_induced_shifts(query_df, reference, solvents, nucleus="H", explicit="desmond"):
+    """The full set of solvent-vs-reference differences for one reference solvent (the data behind
+    Figure S16 with chloroform, S17 with benzene, S18 with solvent_mean): every solvent except the
+    reference, stacked. query_df should already include the solvent_mean rows if reference is
+    'solvent_mean' (see add_solvent_mean)."""
+    frames = [solvent_pair_differences(query_df, solvent, reference, nucleus, explicit)
+              for solvent in solvents if solvent != reference]
+    return pd.concat(frames, ignore_index=True) if frames else pd.DataFrame()
+
+
+def fit_differences_to_experimental(diff_df, predicted_col):
+    """Fit exp_diff = slope * predicted_col + intercept (the Figure 4C fit of the experimental
+    solvent-induced shift differences to a correction's predicted differences) and return
+    {slope, intercept, rmse, n}. Rows missing either value are dropped."""
+    data = diff_df.dropna(subset=["exp_diff", predicted_col])
+    x = data[predicted_col].to_numpy(dtype=float)
+    y = data["exp_diff"].to_numpy(dtype=float)
+    intercept, slope = linear_fit_1d(x, y)
+    pred = intercept + slope * x
+    return {"intercept": intercept, "slope": slope, "rmse": rmse(pred, y), "n": int(len(x))}

analysis/code/dft8k_residuals.py

@@ -0,0 +1,216 @@
+"""MagNET-Zero vs DFT residuals on the DFT8K benchmark (~7,000 organics, out of sample). No plotting.
+
+Each nucleus is compared against the level MagNET-Zero reproduces: protons vs WP04/pcSseg-2, carbons
+vs wB97X-D/pcSseg-2, gas phase on AIMNet2 geometries; residual = DFT shielding minus MagNET. Read via
+data/dft8k/dft8k_reader.py. The figure notebook and tests import from here.
+"""
+import numpy as np
+
+from stats import summarize_errors
+
+# the DFT level each MagNET-Zero model reproduces, and the element it predicts
+NUCLEI = {
+    "H": {"atomic_number": 1, "dft": "wp04_pcSseg2", "nn": "nn_wp04_pcSseg2"},
+    "C": {"atomic_number": 6, "dft": "wb97xd_pcSseg2", "nn": "nn_wb97xd_pcSseg2"},
+}
+
+
+def load_shieldings(dft8k_path):
+    """The flat per-atom shielding arrays from the DFT8K AIMNet2 group, through the released
+    reader. Returns the reader's all_shieldings dict (atomic_numbers plus the DFT and MagNET
+    shieldings, each ppm, NaN where not computed)."""
+    from dft8k_reader import DFT8k
+    with DFT8k(dft8k_path) as ds:
+        return ds.aimnet2.all_shieldings()
+
+
+def residuals(shieldings, nucleus):
+    """The DFT-minus-MagNET shielding residuals (ppm) for one nucleus, over the atoms of that
+    element where MagNET makes a prediction. shieldings is the load_shieldings dict; nucleus is
+    "H" or "C". Atoms the model does not predict (its prediction is NaN) are dropped."""
+    spec = NUCLEI[nucleus]
+    z = np.asarray(shieldings["atomic_numbers"])
+    dft = np.asarray(shieldings[spec["dft"]], dtype=float)
+    nn = np.asarray(shieldings[spec["nn"]], dtype=float)
+    keep = (z == spec["atomic_number"]) & np.isfinite(dft) & np.isfinite(nn)
+    return dft[keep] - nn[keep]
+
+
+def residual_stats(errors, small_threshold=0.1):
+    """Summary statistics of a residual array: count, RMSE, MAE, the 95th percentile of the
+    absolute error, and the fraction of atoms whose absolute error is below small_threshold ppm
+    (the paper reports >95% of proton sites below 0.1 ppm)."""
+    errors = np.asarray(errors, dtype=float)
+    abs_err = np.abs(errors)
+    if not errors.size:
+        return {"n": 0, "rmse": float("nan"), "mae": float("nan"), "median_ae": float("nan"),
+                "abs_p95": float("nan"), "frac_below": float("nan")}
+    stats = summarize_errors(errors, np.zeros_like(errors))
+    stats["abs_p95"] = float(np.quantile(abs_err, 0.95))
+    stats["frac_below"] = float(np.mean(abs_err < small_threshold))
+    return stats
+
+
+def residual_summary(dft8k_path, small_threshold=0.1):
+    """Load the DFT8K shieldings and return {nucleus: residual_stats} for both nuclei."""
+    shieldings = load_shieldings(dft8k_path)
+    return {nucleus: residual_stats(residuals(shieldings, nucleus), small_threshold)
+            for nucleus in NUCLEI}
+
+
+# ---------------------------------------------------------------------------------------------
+# Figure 5B's molecule-level panels: the two extreme-residual callouts and the six functional-
+# group boxes. Unlike the whole-dataset histogram above, these need each molecule's SMILES (to
+# find/classify structures), so they read the reader's per-molecule API instead of all_shieldings.
+
+def iter_molecule_residuals(dft8k_path, nucleus):
+    """Yield (molecule_index, molecule_id, smiles, residuals) for every DFT8K molecule, where
+    residuals is that molecule's DFT-minus-MagNET residual array (ppm) for its atoms of the given
+    nucleus (empty if the molecule has none of that element, or MagNET made no prediction there).
+    Molecules with no SMILES on file (the literal "none") are still yielded; callers that need a
+    parseable structure should skip those themselves."""
+    from dft8k_reader import DFT8k
+    spec = NUCLEI[nucleus]
+    with DFT8k(dft8k_path) as ds:
+        g = ds.aimnet2
+        # One bulk read of the flat per-atom arrays, then slice per molecule using the group's own
+        # precomputed atom offsets (g._start). Calling g.molecule(i) per molecule instead does
+        # several small h5py reads each for all ~7,000 molecules, ~260x slower for identical data.
+        flat = g.all_shieldings()
+        z_all, dft_all, nn_all = flat["atomic_numbers"], flat[spec["dft"]], flat[spec["nn"]]
+        start = g._start
+        for i in range(g.n_molecules):
+            sl = slice(int(start[i]), int(start[i + 1]))
+            z, dft, nn = z_all[sl], dft_all[sl], nn_all[sl]
+            keep = (z == spec["atomic_number"]) & np.isfinite(dft) & np.isfinite(nn)
+            yield i, int(g.molecule_ids[i]), g.smiles(i), (dft[keep] - nn[keep])
+
+
+def find_extreme_residual(dft8k_path, nucleus, sign="max"):
+    """The single most extreme DFT-minus-MagNET residual anywhere in the DFT8K set for one
+    nucleus: the largest positive residual if sign="max", or the most negative if sign="min". The
+    exact extremum over the full released test set (not a random sample), so it is deterministic.
+    For 1H, sign="max" reproduces the published figure's positive-tail callout exactly (a porphyrin
+    ring, +17.62 ppm). sign="min" gives a different, chemically unremarkable ester outlier at
+    -4.06 ppm rather than the figure's negative-tail callout (a sulfonium zwitterion, -1.04 ppm):
+    that callout is a hand-picked illustrative example, not the strict extremum. See
+    molecule_by_id/find_extreme_zwitterion_residual for how that specific example was located.
+    Returns a dict with the parent molecule's id/smiles and the residual value, or None if no atom
+    of that nucleus has a prediction anywhere in the set."""
+    if sign not in ("max", "min"):
+        raise ValueError(f"sign must be 'max' or 'min', got {sign!r}")
+    best = None
+    for i, mol_id, smiles, res in iter_molecule_residuals(dft8k_path, nucleus):
+        if res.size == 0:
+            continue
+        val = float(res.max()) if sign == "max" else float(res.min())
+        if best is None or (sign == "max" and val > best["residual"]) or (sign == "min" and val < best["residual"]):
+            best = {"molecule_index": i, "molecule_id": int(mol_id), "smiles": smiles, "residual": val}
+    return best
+
+
+def molecule_by_id(dft8k_path, nucleus, molecule_id):
+    """The residual data for one specific DFT8K molecule, by its integer id. Used to look up
+    Figure 5B's sulfonium-zwitterion callout (id 88779, identified by brute-force scanning the
+    most-negative 1H residuals for a match to the published -1.040 ppm value -- see
+    find_extreme_zwitterion_residual for the automated version of that search) once its id is
+    known, without re-scanning the whole dataset. Returns the same dict shape as
+    find_extreme_residual (residual is that molecule's single most negative residual for the given
+    nucleus), or None if the molecule has no atom of that nucleus with a prediction."""
+    for i, mol_id, smiles, res in iter_molecule_residuals(dft8k_path, nucleus):
+        if mol_id == molecule_id:
+            if res.size == 0:
+                return None
+            return {"molecule_index": i, "molecule_id": int(mol_id), "smiles": smiles,
+                    "residual": float(res.min())}
+    return None
+
+
+def find_extreme_zwitterion_residual(dft8k_path, nucleus, sign="min"):
+    """Like find_extreme_residual, but restricted to molecules containing BOTH a formally
+    positively-charged sulfur ([S+]) and a formally negatively-charged atom ([-]) -- i.e.
+    sulfonium zwitterions. Even this narrower filter does not land on the same molecule as the
+    published callout: a different, more negative sulfonium zwitterion exists in the set. Provided
+    for exploration, not as an exact reproduction. The published example itself is available via
+    molecule_by_id(..., molecule_id=88779)."""
+    from rdkit import Chem
+    if sign not in ("max", "min"):
+        raise ValueError(f"sign must be 'max' or 'min', got {sign!r}")
+    pos_s = Chem.MolFromSmarts("[S+]")
+    neg = Chem.MolFromSmarts("[-]")
+    best = None
+    for i, mol_id, smiles, res in iter_molecule_residuals(dft8k_path, nucleus):
+        if res.size == 0 or smiles in _UNPARSEABLE_SMILES:
+            continue
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None or not (mol.HasSubstructMatch(pos_s) and mol.HasSubstructMatch(neg)):
+            continue
+        val = float(res.max()) if sign == "max" else float(res.min())
+        if best is None or (sign == "max" and val > best["residual"]) or (sign == "min" and val < best["residual"]):
+            best = {"molecule_index": i, "molecule_id": int(mol_id), "smiles": smiles, "residual": val}
+    return best
+
+
+# A handful of DFT8K molecules have no usable SMILES: some store the literal "none" (documented in
+# dft8k_reader.py), and a few others store the literal "nan" (an apparent upstream data quirk). Both
+# are skipped before ever reaching RDKit, which would otherwise print a harmless but noisy
+# parse-error message to stderr for each one.
+_UNPARSEABLE_SMILES = {"none", "nan", ""}
+
+
+# SMARTS patterns for the six functional groups the published figure highlights. These identify
+# whether a molecule CONTAINS the group at all (a molecule-level classification), not which atom
+# specifically carries it: the released dft8k.hdf5 stores atoms in whatever order the original DFT
+# input used, not SMILES parse order, so a matched substructure cannot be reliably mapped back to
+# one specific stored atom without extra data. Averaging over every atom of the nucleus in every
+# matching molecule is the robust alternative used here.
+FUNCTIONAL_GROUP_SMARTS = {
+    "Carbonyls": "[CX3]=[OX1]",
+    "Amines": "[NX3;H2,H1,H0;!$(NC=O);!$(N=[!#6]);!a]",
+    "Sulfonyl": "[SX4](=[OX1])(=[OX1])",
+    "Pyridines": "[n;r6]",
+    "Furans": "[o;r5]",
+    "Nitroso": "[NX2]=[OX1]",
+}
+
+
+def functional_group_errors(dft8k_path, nucleus, patterns=None):
+    """Mean absolute DFT-minus-MagNET residual (ppm), for every atom of `nucleus`, in every DFT8K
+    molecule whose SMILES matches each functional-group SMARTS pattern (default
+    FUNCTIONAL_GROUP_SMARTS). A molecule can match more than one group (e.g. a pyridine bearing a
+    carbonyl counts in both), matching the natural reading of the published panel's per-group boxes
+    as independent slices of the same dataset, not a mutually exclusive partition.
+
+    Returns {group_name: {"mean_abs_error": ppm, "n_molecules": matching molecule count,
+    "n_atoms": total atoms of this nucleus pooled across those molecules}}. A group with no matches
+    in the released set gets NaN/0/0.
+    """
+    from rdkit import Chem
+    if patterns is None:
+        patterns = FUNCTIONAL_GROUP_SMARTS
+    compiled = {name: Chem.MolFromSmarts(smarts) for name, smarts in patterns.items()}
+    for name, patt in compiled.items():
+        if patt is None:
+            raise ValueError(f"invalid SMARTS for group {name!r}: {patterns[name]!r}")
+    pooled_residuals = {name: [] for name in patterns}
+    n_molecules = {name: 0 for name in patterns}
+    for _, _, smiles, res in iter_molecule_residuals(dft8k_path, nucleus):
+        if smiles in _UNPARSEABLE_SMILES or res.size == 0:
+            continue
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            continue
+        for name, patt in compiled.items():
+            if mol.HasSubstructMatch(patt):
+                pooled_residuals[name].append(res)
+                n_molecules[name] += 1
+    out = {}
+    for name in patterns:
+        chunks = pooled_residuals[name]
+        if chunks:
+            all_res = np.concatenate(chunks)
+            out[name] = {"mean_abs_error": float(np.mean(np.abs(all_res))),
+                        "n_molecules": n_molecules[name], "n_atoms": int(all_res.size)}
+        else:
+            out[name] = {"mean_abs_error": float("nan"), "n_molecules": 0, "n_atoms": 0}
+    return out

analysis/code/diels_alder.py

@@ -0,0 +1,137 @@
+"""SI Figure S2: literature Diels-Alder barriers for ethylene + butadiene.
+
+Published quantum-chemistry predictions span ~15-45 kcal/mol (experiment ~23.3). The barriers are a
+literature compilation (paper reference 26); the table below is the data (no dataset/reader). This
+file holds the table and its method-family grouping; the bar chart is drawn inline in the notebook.
+"""
+# experimental activation barrier for the ethylene + butadiene Diels-Alder reaction, kcal/mol
+EXPERIMENTAL_BARRIER = 23.3
+
+# (method // geometry, predicted barrier in kcal/mol), compiled from the literature (reference 26).
+# The label before "//" is the energy method; it sets the bar's group and label.
+DATA = [
+    ("HF/3-21G//HF/3-21G", 35.9),
+    ("HF/6-31G*//HF/6-31G*", 45.0),
+    ("HF/6-31G**//HF/6-31G**", 45.9),
+    ("CASSCF(6,6)/3-21G// CASSCF(6,6)/3-21G", 37.3),
+    ("CASSCF(6,6)/6-31G*//B3LYP/6-31G*", 43.8),
+    ("MRMP2-CAS(6,6)/6-31G*//B3LYP/6-31G*", 28.6),
+    ("CASSCF(6,6)/6-31G*// CASSCF(6,6)/6-31G*", 43.8),
+    ("CASSCF(6,6)/6-31G(d,p)//CASSCF(6,6)/6-31(d,p)", 44.5),
+    ("CASSCF(6,6)-MP2/6-31G(d,p)//CASSCF(6,6)/6-31(d,p)", 39.4),
+    ("CASSCF(6,6)/6-311+G(d,p)//CASSCF(6,6)/6-31(d,p)", 37.0),
+    ("CASSCF(6,6)-MP2/6-311+G(d,p)//CASSCF(6,6)/6-31(d,p)", 40.9),
+    ("CASSCF(6,6)/6-31G*//CASSCF(6,6)/6-31G*", 25.0),
+    ("CASSCF(6,6)/6-31G*//CASSCF(6,6)/6-31G*", 43.6),
+    ("MR-AQCC/6-31G*//MR-AQCC/6-31G*", 25.7),
+    ("MR-AQCC/6-31G**//MR-AQCC/6-31G**", 25.3),
+    ("MR-AQCC/6-31G**//MR-AQCC/6-31G**", 24.2),
+    ("MR-AQCC/6-311G(2d,p)//MR-AQCC/6-311G(2d,p)", 23.7),
+    ("MP2/6-31G*//HF/6-31G*", 16.6),
+    ("MP2/6-31G*//MP2/6-31G*", 18.5),
+    ("MP2/6-31G*//B3LYP/6-31G*", 20.4),
+    ("MP2/6-31G**//MP2/6-31G**", 15.9),
+    ("MP2/6-311G(d,p)//B3LYP/6-31G", 17.5),
+    ("MP3/6-31G*//HF/6-31G*", 26.9),
+    ("MP3/6-311G(d,p)//B3LYP/6-31G", 28.2),
+    ("MP4(SDQ)/6-31G*//HF/6-31G*", 29.0),
+    ("MP4(SDTQ)/6-31G*//HF/6-31G*", 21.9),
+    ("MP4(SDTQ)/6-31G*//MP2/6-31G*", 22.4),
+    ("MP4(SDTQ)/6-311G(d,p)//B3LYP/6-31G", 22.8),
+    ("B3LYP/6-31G*//B3LYP/6-31G*", 23.1),
+    ("B3LYP/6-31G*//B3LYP/6-31G*", 24.9),
+    ("B3LYP/6-31G**//B3LYP/6-31G**", 22.4),
+    ("B3LYP/6-311+G(2d,p)//B3LYP/6-311+G(2d,p)", 27.2),
+    ("BPW91/6-31G*//BPW91/6-31G*", 26.2),
+    ("MPW1K/6-31+G**//MPW1K/6-31+G*", 24.4),
+    ("KMLYP/6-31G*//KMLYP/6-31G*", 21.1),
+    ("KMLYP/6-311G//KMLYP/6-31G", 22.4),
+    ("OLYP/6-31G(d)//OLYP/6-31G(d)", 26.7),
+    ("OLYP/6-311+G(2d,p)//OLYP/6-311+G(2d,p)", 30.1),
+    ("OLYP/6-311G(2df,2pd)// OLYP/6-311G(2df,2pd)", 29.2),
+    ("O3LYP/6-31G(d)//OLYP/6-31G(d)", 26.8),
+    ("O3LYP/6-311+G(2d,p)//OLYP/6-311+G(2d,p)", 30.1),
+    ("O3LYP/6-311G(2df,2pd)// OLYP/6-311G(2df,2pd)", 29.3),
+    ("M06-2x/6-31G(d)", 20.4),
+    ("M06-2x/cc-pVTZ", 23.2),
+    ("ωB97X-D/6-31G(d)", 21.7),
+    ("ωB97X-D/cc-pVTZ", 25.1),
+    ("B2PLYP/cc-pVDZ", 24.0),
+    ("QCISD(T)/6-31G*//CASSCF(6,6)/6-31G*", 25.5),
+    ("QCISD(T)/6-31G*//B3LYP/6-31G*", 25.0),
+    ("CCSD/6-311G(d,p)//B3LYP/6-31G", 30.8),
+    ("CCSD(T)/6-311G(d,p)//B3LYP/6-31G", 25.7),
+    ("CCSD(T)/6-31G**//B3LYP/6-31G**", 27.6),
+    ("G2MP2/6-311+G(3df,2p)", 24.6),
+    ("G2MS/6-311+G(2df,2p)", 23.9),
+    ("CBS-QB3", 22.9),
+]
+
+# the method-family groups, in the left-to-right order the figure lays them out
+CATEGORY_ORDER = ["HF", "MP2", "MP3", "MP4", "DFT", "CASSCF",
+                  "Multireference", "Coupled Cluster", "Composite"]
+
+
+def energy_method(method_string):
+    """The energy method (the part before '//') of a 'method//geometry' string. This is what the
+    bar is labeled with and categorized by; a string with no '//' is returned unchanged."""
+    return method_string.split("//")[0].strip() if "//" in method_string else method_string
+
+
+def categorize_method(method):
+    """The method family ('HF', 'MP2', 'DFT', 'CASSCF', ...) of an energy method, by the same rules
+    as the original figure. Used to group the bars."""
+    m = method.upper()
+    if m.startswith("HF"):
+        return "HF"
+    if m.startswith("MP2"):
+        return "MP2"
+    if m.startswith("MP3"):
+        return "MP3"
+    if m.startswith("MP4"):
+        return "MP4"
+    if any(k in m for k in ("B3LYP", "BPW91", "MPW1K", "KMLYP", "OLYP", "O3LYP", "M06",
+                            "ωB97X", "B2PLYP")):
+        return "DFT"
+    if "CASSCF" in m:
+        return "CASSCF"
+    if "CCSD" in m or "QCISD" in m:
+        return "Coupled Cluster"
+    if m.startswith("G2") or "CBS" in m:
+        return "Composite"
+    if "MR-AQCC" in m or "MRMP2" in m:
+        return "Multireference"
+    return "Other"
+
+
+def ordered_layout(data=DATA, group_spacing=2):
+    """Lay the bars out grouped by method family and sorted by barrier within each group, with a gap
+    of group_spacing between groups. Returns a dict with parallel lists methods (the energy-method
+    labels), energies, positions (x of each bar), and group_labels / group_boundaries (the
+    (start, end) bar position of each group), in CATEGORY_ORDER.
+    """
+    grouped = {}
+    for method_string, energy in data:
+        method = energy_method(method_string)
+        grouped.setdefault(categorize_method(method), []).append((method, energy))
+    for entries in grouped.values():
+        entries.sort(key=lambda me: me[1])
+
+    methods, energies, positions = [], [], []
+    group_labels, group_boundaries = [], []
+    pos = 0
+    for i, category in enumerate(CATEGORY_ORDER):
+        if category not in grouped:
+            continue
+        if i > 0 and methods:   # gap before each group after the first present one
+            pos += group_spacing
+        start = pos
+        for method, energy in grouped[category]:
+            methods.append(method)
+            energies.append(energy)
+            positions.append(pos)
+            pos += 1
+        group_labels.append(category)
+        group_boundaries.append((start, pos - 1))
+    return {"methods": methods, "energies": energies, "positions": positions,
+            "group_labels": group_labels, "group_boundaries": group_boundaries}

analysis/code/fig2b_plots.py

@@ -0,0 +1,52 @@
+"""Correlation-matrix plotting engine for Figure 2B.
+
+The notebook computes the pairwise Pearson correlation table across methods (delta-22 1H stationary
+shieldings, pcSseg-2 basis) and calls `plot_correlation_matrix(...)` to draw the lower-triangle
+heatmap.
+"""
+import numpy as np
+import pandas as pd
+import seaborn as sns
+import matplotlib.pyplot as plt
+
+
+def plot_correlation_matrix(corr_matrix, save_png, figsize=(6, 6)):
+    """Draw and save the Figure 2B lower-triangle heatmap of pairwise Pearson correlations."""
+    arr = corr_matrix.to_numpy(dtype=float, copy=True)
+    np.fill_diagonal(arr, np.nan)                    # ignore the self-correlation diagonal
+    transformed = pd.DataFrame(-np.log10(1 - arr),   # spread values near 1 so differences are visible
+                               index=corr_matrix.index, columns=corr_matrix.columns)
+    # Tight lower triangle: the first row and last method have no pairs below the diagonal, so drop the
+    # first row and last column, then hide the strict upper triangle. Every drawn cell is then a real
+    # method pair, with no empty self-correlation squares on the axes.
+    transformed = transformed.iloc[1:, :-1]
+    mask = np.triu(np.ones(transformed.shape, dtype=bool), k=1)
+
+    sns.set_theme(style="white", context="paper")
+    fig, ax = plt.subplots(figsize=figsize)
+
+    r_ticks = [0.99, 0.999, 0.9999]
+    cb_values = [-np.log10(1 - r) for r in r_ticks]
+    hm = sns.heatmap(transformed, mask=mask, cmap="Reds", vmin=1, vmax=4, square=True,
+                     linewidths=1, cbar_kws={"shrink": 0.65, "pad": -0.05, "ticks": cb_values}, ax=ax)
+
+    cbar = hm.collections[0].colorbar
+    cbar.set_ticks(cb_values)
+    cbar.set_ticklabels([f"{r:g}" for r in r_ticks])  # ticks are Pearson r values
+    cax = cbar.ax
+    pos = cax.get_position()
+    cax.set_position([pos.x0, pos.y0 + 0.04, pos.width, pos.height])
+    cax.text(0.5, 1.08, r"Pearson $\boldsymbol{r}$", transform=cax.transAxes, ha="center", va="bottom",
+             fontsize=11, bbox=dict(boxstyle="round,pad=0.25", facecolor="#EBEBEB", edgecolor="none", alpha=0.3))
+
+    for side in ("left", "bottom"):
+        ax.spines[side].set_visible(True)
+        ax.spines[side].set_linewidth(0.6)
+    ax.tick_params(axis="x", length=1.5, width=0.6, labelsize=10)
+    ax.tick_params(axis="y", length=1.5, width=0.6, labelsize=10)
+    plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="anchor", va="top")
+    plt.setp(ax.get_yticklabels(), rotation=0, va="center")
+    sns.despine(ax=ax, top=True, right=True, left=False, bottom=False)
+
+    fig.savefig(save_png, dpi=300, bbox_inches="tight", pad_inches=0.02)
+    plt.show()

analysis/code/fig3_plots.py

@@ -0,0 +1,160 @@
+"""Plotting engines for Figure 3 panels 3A, 3B, and 3D (delta-22 solvent-correction figures).
+
+Each panel's notebook (fig3a_pcm_benefit.ipynb, fig3b_shifts.ipynb, fig3d_solvent_corrections.ipynb)
+supplies the data table and styling dicts and calls the corresponding function here.
+"""
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import matplotlib.ticker as mticker
+import seaborn as sns
+from matplotlib.ticker import MaxNLocator
+
+
+def plot_pcm_benefit_with_arrows(benefit_df, combos, colors, save_path,
+                                  xlimits=(-20, 35), figsize=(6, 6)):
+    """Figure 3A: implicit-solvent (PCM) benefit by method, chloroform vs benzene."""
+    solvents = ["chloroform", "benzene"]
+    labels = [m for (m, _b, _g) in combos]
+    arrow_label = "__arrow__"
+    ordered_labels = [arrow_label] + labels        # blank spacer row at the top holds the arrows
+
+    rows = []
+    for method, basis, geom in combos:
+        for solvent in solvents:
+            sub = benefit_df[(benefit_df["sap_nmr_method"] == method) &
+                             (benefit_df["sap_basis"] == basis) &
+                             (benefit_df["sap_geometry_type"] == geom) &
+                             (benefit_df["solvent"] == solvent)]
+            for v in sub["percent_benefit"]:
+                rows.append({"Label": method, "Solvent": solvent, "PCM_Benefit": float(v)})
+    long_df = pd.DataFrame(rows)
+    long_df = pd.concat([long_df, pd.DataFrame([{"Label": arrow_label, "Solvent": solvents[0],
+                                                 "PCM_Benefit": float(np.mean(xlimits))}])],
+                        ignore_index=True)
+    long_df["Label"] = pd.Categorical(long_df["Label"], categories=ordered_labels, ordered=True)
+
+    sns.set_theme(context="paper", style="white")
+    fig, ax = plt.subplots(figsize=figsize)
+    sns.boxplot(data=long_df, x="PCM_Benefit", y="Label", order=ordered_labels,
+                hue="Solvent", hue_order=solvents, palette=colors, dodge=False, width=0.3,
+                fliersize=0, linewidth=0.5, whiskerprops=dict(alpha=0.5, linewidth=0.8),
+                capprops=dict(alpha=0.5, linewidth=1), medianprops=dict(alpha=0.9, linewidth=0.4),
+                boxprops=dict(alpha=0.9), ax=ax, orient="h")
+
+    # blank the spacer row's tick label
+    ticks = ax.get_yticks()
+    ax.yaxis.set_major_locator(mticker.FixedLocator(ticks))
+    ax.yaxis.set_major_formatter(mticker.FixedFormatter(
+        ["" if t.get_text() == arrow_label else t.get_text() for t in ax.get_yticklabels()]))
+
+    leg = ax.get_legend()
+    if leg is not None:
+        leg.set_title("")
+        for txt in leg.get_texts():
+            txt.set_text(txt.get_text().capitalize())
+        leg.set_frame_on(True); leg.get_frame().set_alpha(0.95)
+        for t in leg.get_texts():
+            t.set_fontsize(13)
+        leg.set_bbox_to_anchor((0.8, 1.02)); leg._loc = 9
+
+    # arrows on the blank spacer row
+    y_spacer = ordered_labels.index(arrow_label)
+    span = xlimits[1] - xlimits[0]
+    offset = 0.02 * span
+
+    def box_arrow(text, x, ha, facecolor):
+        ax.text(x, y_spacer, text, ha=ha, va="center", fontsize=11, fontweight="bold",
+                color="white", clip_on=False, zorder=10,
+                bbox=dict(boxstyle=f"{'larrow' if ha=='right' else 'rarrow'},pad=0.35",
+                          fc=facecolor, ec="none", alpha=0.85))
+    box_arrow("Improves Acc.", +offset, "left", "#708B75")
+    box_arrow("Decreases Acc.", -offset, "right", "#C3A29E")
+
+    ax.axvline(0, color="black", lw=0.2, alpha=0.4, zorder=0)
+    ax.set_xlim(*xlimits)
+    ax.margins(x=0.01)
+    ax.set_xlabel("Benefit of PCM (%)", fontsize=13)
+    ax.set_ylabel("")
+    sns.despine(ax=ax, left=False, bottom=False)
+    for spine in ax.spines.values():
+        spine.set_linewidth(0.6)
+    ax.tick_params(axis="both", which="major", length=1.5, width=0.6, labelsize=10)
+    fig.tight_layout()
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def plot_shift_vs_pcm(shift_df, save_path, xlimits=(-0.2, 0.2), ylimits=(-1.6, 0.3), figsize=(8, 4)):
+    """Figure 3B: measured vs PCM-predicted solvent-induced shift differences."""
+    sns.set_theme(style="white", context="paper")
+    fig, axes = plt.subplots(1, 2, figsize=figsize)
+    # left: experimental solvent-induced shifts (teal) -- large and uncorrelated
+    sns.scatterplot(data=shift_df, x="exp_diff_x", y="exp_diff_y", ax=axes[0],
+                    color="#44AA99", s=15, alpha=1.0, edgecolor=None)
+    # right: PCM solvent corrections (dark red) -- small and linearly correlated. Same axis limits
+    # as the left panel, so the tiny PCM predictions are visible against the real shifts.
+    pcm = shift_df.dropna(subset=["pcm_diff_x", "pcm_diff_y"])
+    sns.scatterplot(data=pcm, x="pcm_diff_x", y="pcm_diff_y", ax=axes[1],
+                    color="#A72608", s=15, alpha=0.5, edgecolor=None)
+    for ax in axes:
+        ax.set_xlim(*xlimits); ax.set_ylim(*ylimits)
+        ax.xaxis.set_major_locator(MaxNLocator(nbins=6))
+        ax.yaxis.set_major_locator(MaxNLocator(nbins=8))
+        ax.tick_params(axis="both", which="major", direction="out", length=1.5, width=0.6, labelsize=10)
+    axes[0].set_xlabel(r"$\delta_{MeOD} - \delta_{CDCl_3}$ ($^{1}\mathrm{H}$ ppm)", fontsize=13)
+    axes[0].set_ylabel(r"$\delta_{benzene} - \delta_{CDCl_3}$ ($^{1}\mathrm{H}$ ppm)", fontsize=13)
+    axes[1].set_xlabel(r"$\Delta\sigma_{MeOH} - \Delta\sigma_{CHCl_3}$ ($^{1}\mathrm{H}$ ppm)", fontsize=13)
+    axes[1].set_ylabel(r"$\Delta\sigma_{benzene} - \Delta\sigma_{CHCl_3}$ ($^{1}\mathrm{H}$ ppm)", fontsize=13)
+    plt.tight_layout(rect=[0, 0, 1, 0.95])
+    sns.despine()
+    for ax in axes:
+        for spine in ax.spines.values():
+            spine.set_linewidth(0.6)
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def plot_solvent_correction_boxplot(results_df, formula_labels, solvent_order, colors,
+                                    solvent_labels, save_path, box_widths=0.70, figsize=(8, 6)):
+    """Figure 3D: implicit vs explicit solvent corrections, one box per solvent."""
+    labels = list(formula_labels.values())          # implicit, implicit+vib, explicit, explicit+vib
+    palette = dict(zip(labels, colors))
+    df = results_df[results_df["formula"].isin(formula_labels)].copy()
+    df["Model"] = df["formula"].map(formula_labels)
+    df = df[df["solvent"].isin(solvent_order)].copy()
+    df["solvent"] = pd.Categorical(df["solvent"], categories=solvent_order, ordered=True)
+    xpos = {s: float(i) for i, s in enumerate(solvent_order)}
+    df["xpos"] = df["solvent"].map(xpos).astype(float)
+
+    sns.set_theme(context="paper", style="white")
+    fig, ax = plt.subplots(figsize=figsize)
+    sns.boxplot(data=df, x="xpos", y="test_RMSE", hue="Model", hue_order=labels, palette=palette,
+                width=box_widths, showfliers=False, linewidth=0.5,
+                whiskerprops=dict(alpha=0.5, linewidth=0.8), capprops=dict(alpha=0.5, linewidth=1),
+                medianprops=dict(alpha=0.9, linewidth=0.4), boxprops=dict(alpha=0.9), ax=ax)
+
+    # center each solvent's tick under its "Explicit Solvent" box
+    per_box = box_widths / len(labels)
+    offset = (labels.index("Explicit Solvent") - (len(labels) - 1) / 2) * per_box
+    ax.set_xticks([xpos[s] + offset for s in solvent_order])
+    ax.set_xticklabels([solvent_labels.get(s, s) for s in solvent_order], rotation=45, fontsize=8, ha="right")
+
+    ax.set_ylabel(r"Accuracy vs. Benchmark (RMSE, $^{1}\mathrm{H}$ ppm)", fontsize=13)
+    ax.set_xlabel("")
+    ax.set_ylim(bottom=0)
+    ax.set_title("Accuracy Improves Broadly with Explicit Solvation", fontsize=14, fontweight="bold")
+    leg = ax.get_legend()
+    if leg is not None:
+        leg.set_title("")
+        for t in leg.get_texts():
+            t.set_fontsize(13)
+    ax.yaxis.set_major_locator(MaxNLocator(nbins=6))
+    ax.tick_params(axis="both", which="major", direction="out", length=1.5, width=0.6, labelsize=10)
+    ax.grid(False)
+    sns.despine(ax=ax)
+    for spine in ax.spines.values():
+        spine.set_linewidth(0.6)
+    fig.tight_layout()
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()

analysis/code/fig4_plots.py

@@ -0,0 +1,71 @@
+"""Figure 4 panel-drawing engines.
+
+4A/4B share one scatter engine (predicted vs measured solvent-induced shift, implicit vs explicit
+prediction column); 4C is a dumbbell plot of per-solvent fit RMSE vs experimental shift range. The
+notebook supplies the data table and the styling dicts (highlight map, order, palette, markers /
+colors) as arguments and calls `save_panel(...)` to draw and write the PNG.
+"""
+import matplotlib.pyplot as plt
+import seaborn as sns
+from adjustText import adjust_text
+
+
+def draw_shift_scatter(ax, shift_df, pred_col, highlight_map, order, palette, markers, limits=(-0.25, 2.0)):
+    """Figure 4A/4B: scatter predicted vs experimental solvent-induced shift, colored by solvent."""
+    df = shift_df[shift_df["solvent"].isin(highlight_map)].copy()
+    df["Solvent"] = df["solvent"].map(highlight_map)
+    sns.scatterplot(data=df, x="exp_diff", y=pred_col, hue="Solvent", style="Solvent",
+                    hue_order=order, style_order=order, palette=palette,
+                    markers=markers, s=60, alpha=0.6, edgecolor=None, ax=ax)
+    lo, hi = limits
+    ax.plot([lo, hi], [lo, hi], linestyle="--", alpha=0.35, linewidth=0.6, color="k")  # ideal y = x
+    ax.set_xlim(lo, hi); ax.set_ylim(lo, hi); ax.set_box_aspect(1)
+    ax.set_xlabel(r"Experimental Solvent-Induced Shift (Δδ vs. CDCl3, $^{1}\mathrm{H}$ ppm)",
+                  fontsize=13, labelpad=10)
+    ax.set_ylabel(r"Predicted Solvent-Induced Shift (Δδ vs. CDCl3, $^{1}\mathrm{H}$ ppm)",
+                  fontsize=13, labelpad=10)
+    leg = ax.get_legend()
+    if leg is not None:
+        leg.set_title(None)
+    for side in ax.spines:
+        ax.spines[side].set_linewidth(0.6)
+    ax.tick_params(labelsize=10, length=1.5, width=0.6)
+
+
+def save_panel(draw_fn, save_path, figsize=(6, 6)):
+    """Create a figure/axes, call `draw_fn(ax)`, then save and show the panel as a PNG."""
+    sns.set_theme(style="ticks", context="paper")
+    fig, ax = plt.subplots(figsize=figsize)
+    draw_fn(ax)
+    fig.tight_layout()
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def draw_rmse_dumbbell(ax, rows, labels, colors):
+    """Figure 4C: dumbbell plot of implicit vs explicit solvent-correction RMSE per solvent."""
+    rows_sorted = sorted(rows, key=lambda r: r["range"])
+    xs = [r["range"] for r in rows_sorted]
+    for r in rows_sorted:                                   # gray dumbbell connectors
+        ax.plot([r["range"], r["range"]], [r["implicit"], r["explicit"]],
+                color="#cccccc", lw=1.5, zorder=1)
+    ax.scatter(xs, [r["implicit"] for r in rows_sorted], s=60, marker="D",
+               color=colors["Implicit (PCM)"], alpha=0.85, edgecolors="none",
+               label="Implicit (PCM)", zorder=2)
+    ax.scatter(xs, [r["explicit"] for r in rows_sorted], s=60, marker="o",
+               color=colors["Explicit (Desmond)"], alpha=0.85, edgecolors="none",
+               label="Explicit (Desmond)", zorder=2)
+    # frame the data, then label each dumbbell at its midpoint (adjustText avoids overlaps)
+    xr = max(xs) - min(xs); yv = [r[k] for r in rows_sorted for k in ("implicit", "explicit")]
+    ax.set_xlim(min(xs) - 0.06 * xr, max(xs) + 0.10 * xr)
+    ax.set_ylim(min(yv) - 0.06 * (max(yv) - min(yv)), max(yv) + 0.10 * (max(yv) - min(yv)))
+    texts = [ax.text(r["range"], 0.5 * (r["implicit"] + r["explicit"]), labels[r["solvent"]],
+                     fontsize=10, color="black", va="center", ha="left") for r in rows_sorted]
+    adjust_text(texts, ax=ax, arrowprops=dict(arrowstyle="-", color="gray", lw=0.5, alpha=0.8))
+    ax.legend(loc="upper left", framealpha=0.9, fontsize=11)
+    ax.set_box_aspect(1)
+    ax.set_xlabel(r"Range of Solvent-Induced Shifts ($^{1}\mathrm{H}$ ppm)", fontsize=13, labelpad=10)
+    ax.set_ylabel(r"Predicted Solvent-Induced Shift RMSE ($^{1}\mathrm{H}$ ppm)", fontsize=13, labelpad=10)
+    for side in ax.spines:
+        ax.spines[side].set_linewidth(0.6)
+    ax.tick_params(labelsize=10, length=1.5, width=0.6)

analysis/code/fig5b_plots.py

@@ -0,0 +1,76 @@
+"""Plotting engines for Figure 5B (MagNET-Zero vs DFT on the DFT8K benchmark).
+
+The notebook computes residuals / functional-group error tables via `dft8k_residuals` and calls the
+drawing functions here; `show_dft8k_molecule` renders the inline-only RDKit structure callouts (not
+saved to disk).
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+
+def plot_dft8k_residual_histogram(errors, save_path, color="#5D737E", bin_width=0.0025, xlim=0.23,
+                                  grey_box=0.1, x_tick=0.05, figsize=(8, 3)):
+    """Draw the DFT8K residual histogram for one nucleus and save it to save_path (or skip saving
+    if save_path is None)."""
+    # The published panel is slate (#5D737E) with a shaded +/-0.1 ppm window; the "MagNET-Zero Closely
+    # Reproduces DFT" title and "X% of Compounds" bracket are added in the figure compositing step.
+    sns.set_theme(context="paper", style="white")
+    fig, ax = plt.subplots(figsize=figsize)
+    edges = np.arange(-xlim, xlim + bin_width, bin_width)
+    sns.histplot(x=errors, bins=edges, stat="count", element="bars", ax=ax,
+                 color=color, edgecolor=None, alpha=0.9, linewidth=0)
+    ax.set_xlim(-xlim, xlim)
+    ax.set_xlabel("Residual vs. DFT8K (ppm)", fontsize=13)
+    ax.set_ylabel("Counts", fontsize=13)
+    if grey_box is not None:                      # shade the +/- grey_box ppm window
+        ax.axvspan(-grey_box, grey_box, alpha=0.05, facecolor="k", edgecolor="none", zorder=0)
+        for xv in (-grey_box, grey_box):
+            ax.axvline(xv, color="k", ls="--", lw=1.0, alpha=0.35, zorder=1)
+    if x_tick is not None:
+        ax.set_xticks(np.arange(-np.floor(xlim / x_tick) * x_tick, xlim + x_tick / 2, x_tick))
+    sns.despine(ax=ax, top=True, right=True)
+    ax.spines["left"].set_linewidth(0.6); ax.spines["bottom"].set_linewidth(0.6)
+    ax.tick_params(axis="both", which="major", direction="out", length=1.5, width=0.6, labelsize=10)
+    ax.minorticks_off()
+    fig.tight_layout()
+    if save_path:                                 # 13C is an inline-only companion, not saved
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def plot_dft8k_functional_group_errors(group_errors, nucleus, save_path, figsize=(7, 3.5)):
+    """Bar chart of mean |residual| by functional group for one nucleus, saved to save_path."""
+    # The panel composites these mean |residual| values as molecule-icon insets; this bar chart is
+    # the plain reproduction of the numbers.
+    names = list(group_errors.keys())
+    values = [group_errors[n]["mean_abs_error"] for n in names]
+    counts = [group_errors[n]["n_molecules"] for n in names]
+    fig, ax = plt.subplots(figsize=figsize)
+    bars = ax.bar(names, values, color="#4C9B8F")
+    for bar, n in zip(bars, counts):
+        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(), f"n={n:,}",
+                ha="center", va="bottom", fontsize=8)
+    lbl = "¹H" if nucleus == "H" else "¹³C"
+    ax.set_ylabel(f"mean |{lbl} residual| (ppm)")
+    ax.set_title(f"{lbl} error by functional group (whole DFT8K set)")
+    fig.autofmt_xdate(rotation=30)
+    fig.tight_layout()
+    if save_path:                                 # only the 1H panel is a saved deliverable
+        fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.show()
+
+
+def show_dft8k_molecule(smiles, size=(350, 300)):
+    """Display an RDKit structure drawing for the given SMILES inline; nothing is saved to disk."""
+    from rdkit import Chem
+    from rdkit.Chem.Draw import rdMolDraw2D
+    from IPython.display import Image, display
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"RDKit could not parse SMILES: {smiles!r}")
+    drawer = rdMolDraw2D.MolDraw2DCairo(*size)
+    drawer.drawOptions().addStereoAnnotation = False
+    rdMolDraw2D.PrepareAndDrawMolecule(drawer, mol)
+    drawer.FinishDrawing()
+    display(Image(drawer.GetDrawingText()))       # inline only; whole-molecule, no atom highlighting

analysis/code/leveling.py

@@ -0,0 +1,322 @@
+"""Core data loading and analysis for the leveling-effect study (NS372 and delta22). No plotting.
+
+Reads the two datasets, decodes the stored shieldings, and runs the per-nucleus correlation / PCA
+analysis. The figure notebook (built by build_nb_leveling.py) and test_leveling.py import from here.
+
+The claim tested, on each dataset separately (never pooled): once a linear scaling is allowed for,
+any two shielding methods are near-perfectly correlated, PC1 carries almost all between-method
+variance, and the methods fall on a 1D curve in (PC1, PC2).
+"""
+import numpy as np
+import pandas as pd
+import h5py
+import openpyxl
+from sklearn.decomposition import PCA
+from sklearn.preprocessing import StandardScaler
+
+from fixed_point import decode_fixed_point as _decode_fixed_point
+from stats import rmse as _rmse
+
+# canonical functional-family ordering used to sort method columns in every figure
+FAMILY_ORDER = ['LDA', 'GGA', 'mGGA', 'GH', 'RSH', 'LH', 'DH', 'WFT', 'ref']
+
+def sort_by_family(M, methods, families):
+    """Stable-sort method columns of M into FAMILY_ORDER (reference family last)."""
+    rank = {f: i for i, f in enumerate(FAMILY_ORDER)}
+    idx = sorted(range(len(methods)), key=lambda i: (rank.get(families[i], 99), i))
+    return M[:, idx], [methods[i] for i in idx], [families[i] for i in idx]
+
+def log_r(R):
+    """-log10(1-|r|): 1 -> r=0.9, 2 -> 0.99, 3 -> 0.999, ..."""
+    return -np.log10(np.clip(1 - np.abs(R), 1e-6, None))
+
+def scaled_rmse(method, ref):
+    """RMSE of residuals after the per-method linear fit method ~ a*ref + b."""
+    a, b = np.polyfit(ref, method, 1)
+    return _rmse(method, a * ref + b)
+
+def _parabola_fit(xn, yn, n_angles=720):
+    # best degree-2 fit  y_rot ~ x_rot  searched over rotation angle theta
+    best = (0.0, -np.inf, None)
+    for th in np.linspace(0, np.pi, n_angles):
+        c, s = np.cos(th), np.sin(th)
+        xr, yr = c * xn + s * yn, -s * xn + c * yn
+        A = np.column_stack([np.ones_like(xr), xr, xr ** 2])
+        coef, *_ = np.linalg.lstsq(A, yr, rcond=None)
+        ss_res = ((yr - A @ coef) ** 2).sum()
+        ss_tot = ((yr - yr.mean()) ** 2).sum()
+        r2 = 1 - ss_res / ss_tot if ss_tot > 0 else 0.0
+        if r2 > best[1]:
+            best = (th, r2, coef)
+    return best
+
+def analyze_nucleus(M, methods, ref=None):
+    """Correlation, PCA, and parabola-fit diagnostics for one nucleus.
+
+    M is an (n_observations, n_methods) shielding matrix; ref (optional) is the reference method's
+    values, used to compute each method's scaled RMSE against it.
+    """
+    R = np.corrcoef(M, rowvar=False)
+    iu = np.triu_indices_from(R, k=1)
+    rv = R[iu]
+    Z = StandardScaler().fit_transform(M)
+    pca = PCA().fit(Z)
+    ev = pca.explained_variance_ratio_
+    pc1, pc2 = pca.components_[0].copy(), pca.components_[1].copy()
+    if pc1.mean() < 0:                       # cosmetic sign convention
+        pc1 = -pc1
+    mx, my = pc1.mean(), pc2.mean()
+    sx, sy = max(pc1.std(), 1e-12), max(pc2.std(), 1e-12)
+    xn, yn = (pc1 - mx) / sx, (pc2 - my) / sy
+    th, pr2, coef = _parabola_fit(xn, yn)
+    c, s = np.cos(th), np.sin(th)
+    par = dict(theta=th, coef=coef, r2=pr2, mx=mx, sx=sx, my=my, sy=sy,
+               xr=c * xn + s * yn)
+    out = dict(R=R, methods=list(methods), pc1=pc1, pc2=pc2, ev=ev, par=par,
+               r_min=float(rv.min()), r_med=float(np.median(rv)),
+               r_max=float(rv.max()))
+    if ref is not None:
+        out['scaled_rmse'] = {m: scaled_rmse(M[:, i], ref)
+                              for i, m in enumerate(methods)}
+    return out
+
+def parabola_curve(par, npts=300):
+    """Sample points along the fitted parabola (from analyze_nucleus's `par`) in PC1/PC2 space."""
+    th, coef = par['theta'], par['coef']
+    c, s = np.cos(th), np.sin(th)
+    xr = par['xr']
+    vertex = -coef[1] / (2 * coef[2]) if coef[2] != 0 else float(xr.mean())
+    lo, hi = min(xr.min(), vertex), max(xr.max(), vertex)
+    pad = 0.15 * (hi - lo + 1e-9)
+    t = np.linspace(lo - pad, hi + pad, npts)
+    yr = coef[0] + coef[1] * t + coef[2] * t ** 2
+    xn, yn = c * t - s * yr, s * t + c * yr
+    return xn * par['sx'] + par['mx'], yn * par['sy'] + par['my']
+
+def summarize(data, res):
+    """Per-nucleus leveling diagnostics (analyze_nucleus results) as one summary DataFrame."""
+    import pandas as pd
+    rows = []
+    for nuc, r in res.items():
+        rows.append(dict(nucleus=nuc,
+                         n_obs=data[nuc]['M'].shape[0],
+                         n_methods=data[nuc]['M'].shape[1],
+                         r_min=r['r_min'], r_median=r['r_med'],
+                         PC1_pct=r['ev'][0] * 100, PC2_pct=r['ev'][1] * 100,
+                         PC3plus_pct=(1 - r['ev'][0] - r['ev'][1]) * 100,
+                         parabola_R2=r['par']['r2']))
+    return pd.DataFrame(rows)
+
+def dataset_logr_max(res):
+    """Largest -log10(1-|r|) over every nucleus of a dataset (lower triangle of each R)."""
+    vm = 0.0
+    for r in res.values():
+        nm = len(r['methods']); m = nm - 1
+        sub = r['R'][1:, :-1]
+        hide = np.triu(np.ones((m, m), dtype=bool), k=1)
+        vm = max(vm, float(np.nanmax(np.where(hide, np.nan, log_r(sub)))))
+    return float(np.ceil(vm * 10) / 10)
+
+
+# ----------------------------------------------------------------------------
+# NS372 (Kaupp) loader
+# ----------------------------------------------------------------------------
+# (functional, tau-treatment, %HF, family, display label)
+# one tau-variant per functional: tauC (Gaussian default) for tau-dependent
+# meta-GGAs / hybrids, tauD for local hybrids. Match keys use real Unicode.
+METHODS = [
+    ('SVWN', None, 0, 'LDA', 'SVWN'),
+    ('BLYP', None, 0, 'GGA', 'BLYP'), ('BP86', None, 0, 'GGA', 'BP86'),
+    ('PBE', None, 0, 'GGA', 'PBE'), ('B97D', None, 0, 'GGA', 'B97D'),
+    ('HCTH', None, 0, 'GGA', 'HCTH'),
+    ('KT1', None, 0, 'GGA', 'KT1'), ('KT2', None, 0, 'GGA', 'KT2'), ('KT3', None, 0, 'GGA', 'KT3'),
+    ('TPSS', 'τC', 0, 'mGGA', 'TPSS/τC'), ('τ-HCTH', 'τC', 0, 'mGGA', 'τ-HCTH/τC'),
+    ('M06-L', 'τC', 0, 'mGGA', 'M06-L/τC'), ('VSXC', 'τC', 0, 'mGGA', 'VSXC/τC'),
+    ('MN15-L', 'τC', 0, 'mGGA', 'MN15-L/τC'), ('B97M-V', 'τC', 0, 'mGGA', 'B97M-V/τC'),
+    ('SCAN', 'τC', 0, 'mGGA', 'SCAN/τC'), ('rSCAN', 'τC', 0, 'mGGA', 'rSCAN/τC'),
+    ('r2SCAN', 'τC', 0, 'mGGA', 'r2SCAN/τC'),
+    ('TPSSh', 'τC', 10, 'GH', 'TPSSh/τC'), ('B3LYP', None, 20, 'GH', 'B3LYP'),
+    ('B97-2', None, 21, 'GH', 'B97-2'), ('PBE0', None, 25, 'GH', 'PBE0'),
+    ('M06', 'τC', 27, 'GH', 'M06/τC'), ('PW6B95', 'τC', 28, 'GH', 'PW6B95/τC'),
+    ('MN15', 'τC', 44, 'GH', 'MN15/τC'), ('BHLYP', None, 50, 'GH', 'BHLYP'),
+    ('M06-2X', 'τC', 54, 'GH', 'M06-2X/τC'),
+    ('CAM-B3LYP', None, 65, 'RSH', 'CAM-B3LYP'),
+    ('ωB97X-D', None, 100, 'RSH', 'ωB97X-D'), ('ωB97X-V', None, 100, 'RSH', 'ωB97X-V'),
+    ('ωB97M-V', 'τC', 100, 'RSH', 'ωB97M-V/τC'),
+    ('LH07s-SVWN', None, 22, 'LH', 'LH07s-SVWN'),
+    ('MPSTS', 'τD', 30, 'LH', 'mPSTS/τD'),
+    ('LHJ14', 'τD', 10, 'LH', 'LHJ14/τD'),
+    ('LH07t-SVWN', 'τD', 48, 'LH', 'LH07t-SVWN/τD'),
+    ('LH12ct-SsirPW92', 'τD', 65, 'LH', 'LH12ct-SsirPW92/τD'),
+    ('LH12ct-SsifPW92', 'τD', 71, 'LH', 'LH12ct-SsifPW92/τD'),
+    ('LH14t-calPBE', 'τD', 50, 'LH', 'LH14t-calPBE/τD'),
+    ('LH20t', 'τD', 71, 'LH', 'LH20t/τD'),
+    ('B2PLYP', None, 53, 'DH', 'B2PLYP'), ('B2GP-PLYP', None, 65, 'DH', 'B2GP-PLYP'),
+    ('DSD-PBEP86', None, 70, 'DH', 'DSD-PBEP86'),
+    ('Hartree-Fock', None, 100, 'WFT', 'HF'), ('MP2', None, 100, 'WFT', 'MP2'),
+]
+# (display nucleus, source sheet, in-sheet nucleus filter, molecules to drop)
+NUCLEI = [
+    ('1H',  'S6 - 1H Shieldings',  None, ()),
+    ('11B', 'S7 - 11B Shieldings', None, ('BH',)),
+    ('13C', 'S8 - 13C Shieldings', None, ()),
+    ('15N', 'S9 - 15N Shieldings', None, ()),
+    ('17O', 'S10 - 17O Shieldings', None, ('O3',)),
+    ('19F', 'S11 - 19F Shieldings', None, ('F3-',)),
+    ('31P', 'S12 - 31P - 33S Shieldings', '31P', ()),
+    ('33S', 'S12 - 31P - 33S Shieldings', '33S', ()),
+]
+
+def read_nucleus(wb, sheet, nuc_filter=None, exclude=()):
+    """Read the CCSD(T) reference plus each method's shielding vector from one Kaupp SI sheet."""
+    ws = wb[sheet]
+    rows = []; last_nuc = None
+    for r in range(4, ws.max_row + 1):
+        v = ws.cell(r, 1).value
+        if v:
+            last_nuc = v
+        if isinstance(ws.cell(r, 4).value, (int, float)):
+            mol = str(ws.cell(r, 2).value or '').strip()
+            if mol in exclude:
+                continue
+            if nuc_filter is None or last_nuc == nuc_filter:
+                rows.append(r)
+    ref = np.array([ws.cell(r, 4).value for r in rows], dtype=float)
+    last_f = None; cols = []
+    for c in range(7, ws.max_column + 1):
+        h = ws.cell(1, c).value
+        if h is not None:
+            last_f = h
+        cols.append((c, last_f, ws.cell(2, c).value))
+    data = {}
+    for c, f, sub in cols:
+        if f is None:
+            continue
+        vals = []; ok = True
+        for r in rows:
+            v = ws.cell(r, c).value
+            if not isinstance(v, (int, float)):
+                ok = False; break
+            vals.append(v)
+        if not ok or (f, sub) in data:
+            continue
+        data[(f, sub)] = np.array(vals, dtype=float)
+    return ref, data
+
+def load_ns372(path):
+    """Load the NS372 (Kaupp) shielding matrix, sorted by family, for each of the 8 nuclei."""
+    wb = openpyxl.load_workbook(path, data_only=True)
+    out = {}
+    for nuc, sheet, filt, excl in NUCLEI:
+        ref, data = read_nucleus(wb, sheet, filt, excl)
+        used = [m for m in METHODS if (m[0], m[1]) in data]
+        cols = [data[(m[0], m[1])] for m in used]
+        labels = [m[4] for m in used] + ['CCSD(T)']
+        fams = [m[3] for m in used] + ['ref']
+        M = np.column_stack(cols + [ref])
+        M, labels, fams = sort_by_family(M, labels, fams)
+        out[nuc] = dict(M=M, methods=labels, ref=ref, families=fams)
+    return out
+
+
+# ----------------------------------------------------------------------------
+# main-text Figure 2D (convergence to the electronic-structure limit)
+# ----------------------------------------------------------------------------
+# The 1H / pcSseg-3 slice of the Kaupp SI, as a raw-named table, plus the two
+# per-method statistics the 3D panel plots (slope and scaled RMSE vs CCSD(T)).
+# These reuse the NS372 reader above; nothing new is read from disk.
+def load_ns372_1h_convergence(path):
+    """1H CCSD(T) reference plus each benchmarked functional's shieldings, as columns.
+
+    Tau policy differs from the SI S3 leveling analysis: Figure 2D uses the tauD kinetic-energy-density
+    treatment for every tau-dependent functional, whereas load_ns372/METHODS uses tauC for meta-GGAs.
+    So this selects tauD when present, otherwise the functional's single column, keyed by raw
+    functional name ('Hartree-Fock' renamed 'HF' for the labels).
+    """
+    wb = openpyxl.load_workbook(path, data_only=True)
+    ref, data = read_nucleus(wb, 'S6 - 1H Shieldings')
+    names = {m[0] for m in METHODS}                      # the 44 benchmarked functionals
+    variants = {}                                        # functional -> {tau: values}
+    for (func, tau), vals in data.items():
+        if func in names:
+            variants.setdefault(func, {})[tau] = vals
+    cols = {'CCSD(T)': ref}
+    for func, by_tau in variants.items():
+        vals = by_tau['τD'] if 'τD' in by_tau else next(iter(by_tau.values()))
+        cols['HF' if func == 'Hartree-Fock' else func] = vals
+    return pd.DataFrame(cols)
+
+def convergence_statistics(cc_df, reference='CCSD(T)'):
+    """Per method: the slope of method ~ a*CCSD(T) + b, and the scaled RMSE that Figure 2D plots.
+
+    The prediction is rescaled back into reference units, y_scaled = (method - b) / a, and compared
+    to CCSD(T); this equals the residual RMSE divided by the slope, so it differs from scaled_rmse,
+    which the S3 bar plot leaves in the method's own units. These are the x/y coordinates of the
+    Figure 2D points.
+    """
+    ref = cc_df[reference].to_numpy(float)
+    rows = []
+    for method in cc_df.columns:
+        if method == reference:
+            continue
+        y = cc_df[method].to_numpy(float)
+        slope, intercept = np.polyfit(ref, y, 1)
+        y_scaled = (y - intercept) / slope
+        rows.append({'Method': method, 'Slope': slope,
+                     'RMSE (scaled)': _rmse(ref, y_scaled)})
+    return pd.DataFrame(rows)
+
+
+# ----------------------------------------------------------------------------
+# delta22 loader (reads the released file; shieldings are decoded on read)
+# ----------------------------------------------------------------------------
+DELTA22_REF = 'DSD-PBEP86'
+
+def load_delta22(path):
+    """Load the delta22 shielding matrix, sorted by family, for 1H and 13C sites."""
+    f = h5py.File(path, 'r')
+    mn = [m.decode() for m in f['conventional_nmr_method_names'][:]]
+    order = ['hf', 'blyp_d3bj', 'bp86_d3bj', 'b97d3_d3bj', 'tpsstpss_d3bj',
+             'b3lyp_d3bj', 'pbe0_d3bj', 'm062x_d3', 'wb97xd', 'wp04', 'wc04',
+             'B2PLYP', 'mPW2PLYP', 'B2GP_PLYP', 'dsd_pbep86', 'revdsd_pbep86',
+             'mp2', 'dlpno_mp2']
+    disp = {'hf': 'HF', 'blyp_d3bj': 'BLYP-D3', 'bp86_d3bj': 'BP86-D3',
+            'b97d3_d3bj': 'B97D3', 'tpsstpss_d3bj': 'TPSS-D3', 'b3lyp_d3bj': 'B3LYP-D3',
+            'pbe0_d3bj': 'PBE0-D3', 'm062x_d3': 'M06-2X-D3', 'wb97xd': 'ωB97X-D',
+            'wp04': 'WP04', 'wc04': 'WC04', 'B2PLYP': 'B2PLYP', 'mPW2PLYP': 'mPW2PLYP',
+            'B2GP_PLYP': 'B2GP-PLYP', 'dsd_pbep86': 'DSD-PBEP86',
+            'revdsd_pbep86': 'revDSD-PBEP86', 'mp2': 'MP2', 'dlpno_mp2': 'DLPNO-MP2'}
+    fam = {'hf': 'WFT', 'blyp_d3bj': 'GGA', 'bp86_d3bj': 'GGA', 'b97d3_d3bj': 'GGA',
+           'tpsstpss_d3bj': 'mGGA', 'b3lyp_d3bj': 'GH', 'pbe0_d3bj': 'GH',
+           'm062x_d3': 'GH', 'wb97xd': 'RSH', 'wp04': 'GH', 'wc04': 'GH',
+           'B2PLYP': 'DH', 'mPW2PLYP': 'DH', 'B2GP_PLYP': 'DH', 'dsd_pbep86': 'DH',
+           'revdsd_pbep86': 'DH', 'mp2': 'WFT', 'dlpno_mp2': 'WFT'}
+    gas = {}
+    for i, m in enumerate(mn):
+        func, basis, smodel, _ = m.split(',')
+        if smodel == 'gas':
+            gas.setdefault(func, {})[basis] = i
+    sel, labels, families = [], [], []
+    for func in order:
+        b = gas[func]
+        big = 'pcSseg3' if 'pcSseg3' in b else ('pcSseg2' if 'pcSseg2' in b else 'pcSseg1')
+        sel.append(b[big]); labels.append(disp[func]); families.append(fam[func])
+    GEOM = 1  # geometries axis: 0 = AIMNet2, 1 = PBE0/cc-pVTZ
+    rows = {1: [], 6: []}
+    for s in f['solutes']:
+        an = f['solutes'][s]['atomic_numbers'][:]
+        cs = _decode_fixed_point(
+            f['solutes'][s]['stationary_and_pcm']['conventional_shieldings'][:, GEOM, :])
+        for ai, z in enumerate(an):
+            if z in (1, 6):
+                rows[int(z)].append(cs[sel, ai])
+    f.close()
+    out = {}
+    for nuc, z in (('1H', 1), ('13C', 6)):
+        M = np.array(rows[z])
+        M, labs, fams = sort_by_family(M, list(labels), list(families))
+        # use DSD-PBEP86 as the reference vector for delta22 so the scaled RMSE
+        # bar plot mirrors NS372's CCSD(T) role
+        ref_idx = labs.index(DELTA22_REF)
+        out[nuc] = dict(M=M, methods=labs, ref=M[:, ref_idx], families=fams)
+    return out

analysis/code/leveling_plots.py

@@ -0,0 +1,195 @@
+"""Correlation-matrix and PCA-loadings plotting engine for the leveling-effect study (SI Figure S3).
+
+Draws the two figure types built by build_nb_leveling.py: per-nucleus inter-method Pearson
+correlation matrices (log-scaled, with a scaled-RMSE side bar and a colour-scale histogram inset)
+and PC1/PC2 method-loading panels (with a fitted parabola and adjustText-repelled labels). The
+notebook supplies the family-colour map, nucleus display labels, and adjustText tuning as arguments.
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+from adjustText import adjust_text
+
+from leveling import log_r, parabola_curve
+
+
+def _adjust_pca_labels(texts, ax, nuc, pca_adjust, pca_adjust_default):
+    params = pca_adjust.get(nuc, pca_adjust_default)
+    adjust_text(texts, ax=ax, max_move=None, min_arrow_len=0,
+                ensure_inside_axes=True,
+                arrowprops=dict(arrowstyle='-', color='0.55', lw=0.35),
+                **params)
+
+
+def _family_legend(ax, families, family_colors, include_ref=True):
+    present = []
+    for f in families:
+        if f == 'ref' and not include_ref:
+            continue
+        if f not in present:
+            present.append(f)
+    handles = [plt.Line2D([], [], marker='*' if f=='ref' else 'o', ls='',
+                          mfc=family_colors.get(f,'#333333'), mec='k',
+                          ms=12 if f=='ref' else 6.5,
+                          label='CCSD(T)' if f=='ref' else f)
+               for f in present]
+    ax.legend(handles=handles, loc='best', fontsize=6.5, ncol=2,
+              frameon=True, framealpha=0.9)
+
+
+def _pca_panel(ax, name, nuc, r, families, family_colors, nuc_display, pca_adjust, pca_adjust_default):
+    pc1, pc2 = r['pc1'], r['pc2']
+    methods = r['methods']
+    cx, cy = parabola_curve(r['par'])
+    # small padding: the dense central cluster fills more of the plot in inches,
+    # giving adjustText more room per label
+    xpad = 0.08*(pc1.max()-pc1.min()+1e-9)
+    ypad = 0.08*(pc2.max()-pc2.min()+1e-9)
+    xlim = (pc1.min()-xpad, pc1.max()+xpad)
+    ylim = (pc2.min()-ypad, pc2.max()+ypad)
+    m = (cx>=xlim[0])&(cx<=xlim[1])&(cy>=ylim[0])&(cy<=ylim[1])
+    ax.plot(cx[m], cy[m], color='0.75', ls='--', lw=1.2, alpha=0.7, zorder=1)
+    texts = []
+    for i, lab in enumerate(methods):
+        fam = families[i]
+        if fam == 'ref':
+            ax.scatter(pc1[i], pc2[i], s=75, marker='*',
+                       c=family_colors['ref'], edgecolors='white',
+                       linewidths=0.8, zorder=5)
+        else:
+            ax.scatter(pc1[i], pc2[i], s=11,
+                       c=family_colors.get(fam, '#333333'),
+                       edgecolors='black', linewidths=0.3, zorder=4)
+        texts.append(ax.text(pc1[i], pc2[i], lab, fontsize=3.8, zorder=6))
+    ax.set_xlim(*xlim); ax.set_ylim(*ylim)
+    _adjust_pca_labels(texts, ax, nuc, pca_adjust, pca_adjust_default)
+    _family_legend(ax, families, family_colors, include_ref=True)
+    ev = r['ev']
+    ax.set_title(f'{name}  {nuc_display.get(nuc, nuc)} - '
+                 f'method loadings in PC1/PC2 space\n'
+                 f'PC1 = {ev[0]*100:.2f}%,   PC2 = {ev[1]*100:.2f}%,   '
+                 f'parabola R² = {r["par"]["r2"]:.3f}', fontsize=9)
+    ax.set_xlabel('PC1 loading', fontsize=8)
+    ax.set_ylabel('PC2 loading', fontsize=8)
+    ax.tick_params(labelsize=7)
+    ax.axhline(0, color='gray', lw=0.4); ax.axvline(0, color='gray', lw=0.4)
+    ax.grid(True, ls=':', alpha=0.3)
+
+
+def plot_pca_pair(name, nucs, res, fam_lookup, family_colors, nuc_display, pca_adjust, pca_adjust_default):
+    """PC1/PC2 method-loading panels, one nucleus per row, stacked in a single figure."""
+    n = len(nucs)
+    fig, axes = plt.subplots(n, 1, figsize=(7.0, 5.2*n), squeeze=False)
+    for k, nuc in enumerate(nucs):
+        _pca_panel(axes[k][0], name, nuc, res[nuc], fam_lookup[nuc],
+                  family_colors, nuc_display, pca_adjust, pca_adjust_default)
+    fig.tight_layout(h_pad=3.6)
+    return fig
+
+
+def plot_corr_matrix(name, nucs, res, fam_lookup, vmax, ref_name, nuc_display):
+    """Per-nucleus inter-method correlation heatmap with a scaled-RMSE side bar.
+
+    Layout: [ RMSE bar (viridis, bars grow leftward) | heatmap (inferno) ]. Methods are reordered so
+    the reference is first and the rest are sorted by ascending scaled RMSE vs the reference. The
+    heatmap is the lower triangle with the diagonal removed; its upper-right corner carries an inset
+    histogram of all pairwise log-correlations with the colour bar below it, showing two tick sets
+    (bottom = -log10(1-|r|) integers, top = the equivalent raw r values). All matrices share `vmax`.
+    """
+    nm = len(res[nucs[0]]['methods'])
+    m  = nm - 1
+    hm_h, hm_w, bar_w = 0.135*nm + 1.9, 6.0, 1.5
+    fig_w = bar_w + hm_w
+    fig_h = hm_h * len(nucs) + 0.4
+    fig, axes = plt.subplots(len(nucs), 2, sharey='row',
+                             gridspec_kw={'width_ratios': [bar_w, hm_w]},
+                             figsize=(fig_w, fig_h),
+                             layout='constrained', squeeze=False)
+    fig.get_layout_engine().set(wspace=0.01)        # tighten gap to y-labels
+    cmap_corr = plt.cm.inferno.copy(); cmap_corr.set_bad('white')
+    cmap_rmse = plt.cm.viridis
+    fs = 4.3 if m > 24 else 6.0
+    ticks_log = [t for t in [1, 2, 3, 4, 5] if t <= vmax + 0.01]
+    for k, nuc in enumerate(nucs):
+        axR = axes[k, 0]; ax = axes[k, 1]
+        r = res[nuc]; R = r['R']; methods = r['methods']
+        rmse_dict = r['scaled_rmse']
+        # ----- reorder: reference first, others by ascending RMSE vs ref -----
+        ref_idx = methods.index(ref_name)
+        others  = [i for i in range(nm) if i != ref_idx]
+        others_sorted = sorted(others, key=lambda i: rmse_dict[methods[i]])
+        new_order = [ref_idx] + others_sorted
+        methods_s = [methods[i] for i in new_order]
+        R_s = R[np.ix_(new_order, new_order)]
+        rmse_s = np.array([rmse_dict[mth] for mth in methods_s])
+        # ----- heatmap (drop row 0 = ref, col -1 = worst-RMSE method) -----
+        sub  = R_s[1:, :-1]
+        hide = np.triu(np.ones((m, m), dtype=bool), k=1)
+        lr   = log_r(sub)
+        im = ax.imshow(np.ma.masked_array(lr, hide), cmap=cmap_corr,
+                       vmin=0, vmax=vmax, aspect='auto')
+        ax.set_xticks(range(m))
+        ax.set_xticklabels(methods_s[:-1], rotation=90, fontsize=fs)
+        ax.tick_params(left=False, labelleft=False)
+        if m <= 20:
+            for i in range(m):
+                for j in range(i+1):
+                    ax.text(j, i, f'{sub[i,j]:.5f}', ha='center', va='center',
+                            fontsize=3.4,
+                            color='white' if lr[i,j] < vmax*0.5 else 'black')
+        ax.set_title(f'{name}  {nuc_display.get(nuc, nuc)} - '
+                     f'inter-method Pearson correlation', fontsize=9)
+        # ----- histogram + colour bar inset, jammed into the upper-right -----
+        # NS372 (kaupp): scale the inset up 20% so it fills more of the corner.
+        # A small visible gap is left between hist bottom and cb top so the
+        # histogram's tick marks are visible; the dual stacked tick labels
+        # under the cb (top line = -log10(1-|r|), bottom line = Pearson r)
+        # serve both axes since hist and cb share x-extent.
+        all_lr = lr[~hide]
+        counts, edges = np.histogram(all_lr, bins=28, range=(0.0, vmax))
+        centers = 0.5*(edges[:-1] + edges[1:])
+        scale  = 1.2 if name == 'NS372' else 1.0
+        hist_w = 0.50 * scale
+        hist_h = 0.18 * scale
+        cb_h   = 0.014                  # thinner colour bar
+        gap    = 0.012                  # visible gap for hist tick marks
+        hist_top = 0.97
+        x0 = 0.98 - hist_w
+        hist_y0 = hist_top - hist_h
+        cb_y0   = hist_y0 - gap - cb_h
+        hax = ax.inset_axes([x0, hist_y0, hist_w, hist_h])
+        hax.set_facecolor('white')
+        hax.bar(centers, counts, width=(edges[1]-edges[0])*0.92,
+                color=cmap_corr(centers / vmax), edgecolor='black', linewidth=0.3)
+        hax.axvline(float(np.median(all_lr)), color='black', lw=1.6, ls='--')
+        hax.set_xlim(0, vmax)
+        hax.set_xticks(ticks_log)
+        hax.tick_params(axis='x', labelbottom=False, bottom=True, length=3)
+        hax.set_ylabel('count', fontsize=6, labelpad=1)
+        hax.tick_params(axis='y', labelsize=5)
+        hax.grid(axis='y', ls=':', alpha=0.4)
+        # colour bar, slightly below the histogram with a small visible gap
+        cbax = ax.inset_axes([x0, cb_y0, hist_w, cb_h])
+        cb = fig.colorbar(im, cax=cbax, orientation='horizontal')
+        cb.set_ticks(ticks_log)
+        cb.ax.set_xticklabels([f'{t}\n0.{"9"*t}' for t in ticks_log])
+        cb.ax.tick_params(axis='x', labelsize=5.5, pad=1,
+                          top=False, labeltop=False)
+        cb.set_label('-log10(1-|r|)  /  Pearson r', fontsize=6, labelpad=2)
+        # ----- RMSE bar on the LEFT (viridis, bars grow leftward) -----
+        row_methods = methods_s[1:]
+        rmse_row = rmse_s[1:]
+        rmse_max = max(rmse_row.max(), 1e-6)
+        bar_colors = cmap_rmse(rmse_row / rmse_max)
+        axR.barh(range(m), rmse_row, color=bar_colors,
+                 edgecolor='black', linewidth=0.3)
+        axR.yaxis.set_ticks_position('right')
+        axR.yaxis.set_label_position('right')
+        axR.set_yticks(range(m))
+        axR.set_yticklabels(row_methods, fontsize=fs)
+        axR.set_xlim(rmse_max * 1.06, 0)                # reversed: 0 on the right
+        axR.set_xlabel(f'scaled RMSE (ppm)\nvs {ref_name}', fontsize=6.5)
+        axR.tick_params(axis='x', labelsize=5.5)
+        axR.tick_params(axis='y', length=0, pad=2)
+        axR.grid(axis='x', ls=':', alpha=0.4)
+        axR.set_ylim(m - 0.5, -0.5)                     # imshow has inverted y
+    return fig

analysis/code/magnet_benchmark.py

@@ -0,0 +1,327 @@
+"""Tables S3 and S4: MagNET shielding accuracy across stationary, vibrated, and solvated geometries.
+
+How well MagNET reproduces its DFT training reference (PBE0/pcSseg-1) for the foundation model and the
+chloroform MagNET-x, on four test sets:
+
+  stationary_internal   held-out GDB resting geometries         (sigma-shake test split)
+  vibrated_internal     MD frames of those molecules            (gdb_qcd)
+  vibrated_external     perturbed outside-benchmark geometries  (dft8k)
+  isolated_chloroform   isolated-solute chloroform MD snapshots (sigma-fresh)
+
+Two reproduction paths; the shipped notebook uses the first:
+1. EXACT, from `data/magnet_test_predictions/`: raw per-atom predictions over the FULL test sets,
+   computed once from the released checkpoints. `exact_stats_table` maps each (nucleus, model, test
+   set) to its group and calls the reader's `stats()`; median/mean/RMSE reproduce exactly, no
+   checkpoints needed.
+2. FROM SCRATCH: `cases_*` assembles (atomic numbers, geometry, DFT shielding) test cases and
+   `magnet_benchmark_run.py` runs a model over them. Needs the checkpoints and only samples each set.
+
+RMSE note (path 2): median and mean AE reproduce from a modest sample; RMSE does not. For the vibrated
+sets it is dominated by ~1-in-10,000 extreme MD frames (the CO-style breakdown of Figure S12), so the
+published 13C RMSE (7.14 ppm) needs the full ~1.9M-frame set; path 1 sidesteps this. The
+isolated-chloroform row runs in full even under path 2; its foundation 1H RMSE (0.30 ppm) vs MAE
+(0.07 ppm) reflects a broad heavy tail of real deshielded geometries (not DFT failures), kept as-is.
+The Figure S10 stationary-only ablation checkpoint is not released, so that figure is not reproduced.
+"""
+import numpy as np
+
+from stats import summarize_errors
+
+# --------------------------------------------------------------------------- published SI values
+# (median absolute error, mean absolute error, RMSE), per (model, test set). From SI Tables S3/S4.
+# Keys: model in {"MagNET", "MagNET-x"}; test set as named above.
+PUBLISHED = {
+    "1H": {
+        ("MagNET", "stationary_internal"): (0.020946503, 0.030467307, 0.05941977),
+        ("MagNET", "vibrated_internal"): (0.027399063, 0.039091997, 0.10485571),
+        ("MagNET", "vibrated_external"): (0.03241539, 0.04796361, 0.11397923),
+        ("MagNET", "isolated_chloroform"): (0.03216362, 0.071752764, 0.29437021),
+        ("MagNET-x", "stationary_internal"): (0.02163124, 0.031185, 0.06033938),
+        ("MagNET-x", "vibrated_internal"): (0.027683258, 0.039436683, 0.10467214),
+        ("MagNET-x", "vibrated_external"): (0.023479462, 0.034395833, 0.06877531),
+        ("MagNET-x", "isolated_chloroform"): (0.022384644, 0.030986413, 0.04612577),
+    },
+    "13C": {
+        ("MagNET", "stationary_internal"): (0.23955917, 0.36742094, 1.98238429),
+        ("MagNET", "vibrated_internal"): (0.31917763, 0.47851510, 7.13765463),
+        ("MagNET", "vibrated_external"): (0.38428497, 0.57037073, 0.98834883),
+        ("MagNET", "isolated_chloroform"): (0.36051940, 0.58715980, 1.22315871),
+        ("MagNET-x", "stationary_internal"): (0.25761414, 0.39035153, 1.98721069),
+        ("MagNET-x", "vibrated_internal"): (0.33805847, 0.50241500, 7.15089956),
+        ("MagNET-x", "vibrated_external"): (0.32779312, 0.47457147, 0.81721424),
+        ("MagNET-x", "isolated_chloroform"): (0.30544280, 0.43592300, 0.72233414),
+    },
+}
+
+NUCLEUS_Z = {"1H": 1, "13C": 6}
+TEST_SETS = ["stationary_internal", "vibrated_internal", "vibrated_external", "isolated_chloroform"]
+MODELS = ["MagNET", "MagNET-x"]
+
+# Table S5: QCD-correction statistics for the foundation MagNET, one row per nucleus. Computed over
+# the qcdtraj2500 test set (2500 molecules) at PBE0/pcSseg-1. Published (median AE, MAE, RMSE):
+PUBLISHED_S5 = {
+    "1H": (0.01526165, 0.023047674, 0.04020983),
+    "13C": (0.17198181, 0.26223338, 0.50198762),
+}
+QCD_GROUP = {"1H": "qcdtraj2500_H_pretrained_H", "13C": "qcdtraj2500_C_pretrained_C"}
+
+
+def qcd_stats_table(predictions_path, decode_module):
+    """Table S5: (median AE, MAE, RMSE, n) of the QCD-correction errors for the foundation MagNET,
+    one row per nucleus, read from data/magnet_test_predictions/ via qcd_correction_stats."""
+    rows = []
+    for nucleus in ("1H", "13C"):
+        med, mae, rmse, n = decode_module.qcd_correction_stats(predictions_path, QCD_GROUP[nucleus])
+        rows.append(dict(model=f"MagNET ({nucleus})", median_ae=med, mae=mae, rmse=rmse, n=n))
+    return rows
+
+# --------------------------------------------------------------------------- exact reproduction
+# from data/magnet_test_predictions/ (see the module docstring, path 1). That dataset names its
+# groups after its own test-set families (gasphaseinternal, gasphasedft8k, solutesmd500isolated),
+# which don't match the four TEST_SETS names above one-to-one, so this is the mapping between them.
+TEST_PREDICTIONS_MODEL_GROUP = {"MagNET": "pretrained", "MagNET-x": "chloroform"}
+TEST_PREDICTIONS_NUCLEUS = {"1H": "H", "13C": "C"}
+
+
+def test_predictions_group(model, test_set, nucleus):
+    """The data/magnet_test_predictions/ group name holding the raw predictions for one Table S3/S4
+    row. `model` is "MagNET" or "MagNET-x", `test_set` is one of TEST_SETS, `nucleus` is "1H" or
+    "13C". See the TEST_PREDICTIONS_* maps below for the underlying group-naming scheme."""
+    n = TEST_PREDICTIONS_NUCLEUS[nucleus]
+    m = TEST_PREDICTIONS_MODEL_GROUP[model]
+    if test_set == "stationary_internal":
+        return f"gasphaseinternal_{n}_{m}_{n}_vib0"
+    if test_set == "vibrated_internal":
+        return f"gasphaseinternal_{n}_{m}_{n}_vib1"
+    if test_set == "vibrated_external":
+        return f"gasphasedft8k_{n}_{m}_{n}_vib1"
+    if test_set == "isolated_chloroform":
+        return f"solutesmd500isolated_chloroform_{n}_{m}_{n}"
+    raise ValueError(f"unknown test set {test_set!r}")
+
+
+def exact_stats_table(predictions_path, decode_module):
+    """One row per (nucleus, model, test set), with (median_ae, mae, rmse, n) read directly from the
+    full released MagNET predictions -- no sampling, no live inference, so RMSE reproduces exactly
+    too (see the module docstring, path 1). `decode_module` is
+    data/magnet_test_predictions/magnet_test_predictions_reader.py; the caller imports and passes it
+    in, so this module carries no path dependency on data/magnet_test_predictions/."""
+    rows = []
+    for nucleus in NUCLEUS_Z:
+        for model in MODELS:
+            for test_set in TEST_SETS:
+                group = test_predictions_group(model, test_set, nucleus)
+                med, mae, rmse, n = decode_module.stats(predictions_path, group)
+                rows.append({"nucleus": nucleus, "model": model, "test_set": test_set,
+                             "median_ae": med, "mae": mae, "rmse": rmse, "n": n})
+    return rows
+
+# The elements MagNET was trained on (matches SUPPORTED_ELEMENTS in the magnet package). Structures
+# with any other element (e.g. the phosphorus in dft8k) are excluded from the test sets: the model
+# refuses them, and they cannot be fairly scored anyway.
+SUPPORTED_ELEMENTS = frozenset({1, 6, 7, 8, 9, 16, 17})
+
+
+def filter_supported(atomic_numbers_list, geometries_list, dft_list):
+    """Drop any structure whose atoms are not all in SUPPORTED_ELEMENTS. Returns the three lists,
+    filtered in lockstep, plus the number dropped."""
+    keep = [i for i, an in enumerate(atomic_numbers_list)
+            if set(np.unique(an).tolist()) <= SUPPORTED_ELEMENTS]
+    dropped = len(atomic_numbers_list) - len(keep)
+    return ([atomic_numbers_list[i] for i in keep],
+            [geometries_list[i] for i in keep],
+            [dft_list[i] for i in keep], dropped)
+
+
+# --------------------------------------------------------------------------- error statistics
+
+def collect_abs_errors(predicted_list, dft_list, atomic_numbers_list, z):
+    """Concatenate the absolute errors |predicted - dft| over all atoms with atomic number `z`,
+    across a list of test cases. Atoms where either value is NaN (not predicted or not computed) are
+    dropped."""
+    out = []
+    for pred, dft, an in zip(predicted_list, dft_list, atomic_numbers_list):
+        pred = np.asarray(pred, dtype=np.float64)
+        dft = np.asarray(dft, dtype=np.float64)
+        an = np.asarray(an)
+        mask = (an == z) & np.isfinite(pred) & np.isfinite(dft)
+        if mask.any():
+            out.append(np.abs(pred[mask] - dft[mask]))
+    return np.concatenate(out) if out else np.array([])
+
+
+def summarize(abs_errors):
+    """Median absolute error, mean absolute error, RMSE, and count for an array of absolute errors."""
+    e = np.asarray(abs_errors, dtype=np.float64)
+    if e.size == 0:
+        return {"median_ae": float("nan"), "mae": float("nan"), "rmse": float("nan"), "n": 0}
+    return summarize_errors(e, np.zeros_like(e))
+
+
+def error_table(predicted_list, dft_list, atomic_numbers_list):
+    """Summaries for both nuclei: {"1H": {...}, "13C": {...}}."""
+    return {
+        nucleus: summarize(collect_abs_errors(predicted_list, dft_list, atomic_numbers_list, z))
+        for nucleus, z in NUCLEUS_Z.items()
+    }
+
+
+# --------------------------------------------------------------------------- test-case assembly
+# Each function returns three parallel lists: atomic_numbers, geometry (n_atoms, 3), and the DFT
+# reference shielding (n_atoms,), one entry per sampled structure. A model is then run over the
+# geometries and compared with the DFT shieldings.
+
+def cases_stationary_internal(sigma_shake, n_sample=2000, seed=0):
+    """Resting geometries of held-out GDB molecules (sigma-shake test split)."""
+    test_ids = sigma_shake.split("test")
+    rng = np.random.default_rng(seed)
+    chosen = rng.choice(test_ids, size=min(n_sample, len(test_ids)), replace=False)
+    ans, geos, dfts = [], [], []
+    for mid in chosen:
+        m = sigma_shake.by_id(int(mid))
+        ans.append(m["atomic_numbers"])
+        geos.append(m["stationary_coords"])
+        dfts.append(m["shielding_stationary"])
+    return ans, geos, dfts
+
+
+def cases_vibrated_internal(gdb_qcd, n_sample=2000, frames_per_molecule=4, seed=0):
+    """Molecular-dynamics frames of held-out GDB molecules (gdb_qcd)."""
+    rng = np.random.default_rng(seed)
+    mol_idx = rng.choice(gdb_qcd.n_molecules, size=min(n_sample, gdb_qcd.n_molecules), replace=False)
+    ans, geos, dfts = [], [], []
+    for i in mol_idx:
+        m = gdb_qcd.molecule(int(i))
+        nt, nf = m["coordinates"].shape[0], m["coordinates"].shape[1]
+        for _ in range(frames_per_molecule):
+            t, f = int(rng.integers(nt)), int(rng.integers(nf))
+            ans.append(m["atomic_numbers"])
+            geos.append(m["coordinates"][t, f])
+            dfts.append(m["shieldings"][t, f])
+    return ans, geos, dfts
+
+
+def cases_vibrated_external(dft8k, n_sample=2000, seed=0):
+    """Perturbed (non-stationary) geometries of the dft8k benchmark; geometry 0 is the reference, 1
+    and 2 are perturbations, so the vibrated set uses geometries 1 and 2."""
+    rng = np.random.default_rng(seed)
+    n = dft8k.b3lyp.n_molecules
+    chosen = rng.choice(n, size=min(n_sample, n), replace=False)
+    ans, geos, dfts = [], [], []
+    for i in chosen:
+        m = dft8k.b3lyp.molecule(int(i))
+        for g in (1, 2):
+            ans.append(m["atomic_numbers"])
+            geos.append(m["coordinates"][g])
+            dfts.append(m["shielding"][g])
+    return ans, geos, dfts
+
+
+# The isolated-chloroform test set is defined entirely from the released sigma-fresh file, so the
+# table is reproducible without any external file. The foundation model's RMSE on this row is large
+# (0.30 ppm 1H) because the error is heavy-tailed, not because of a few bad frames: the optional
+# scrub in isolated_chloroform_testset removes the most strongly-deshielded frames and the RMSE
+# barely moves (0.31 -> 0.30). Those frames are real geometries, not DFT failures, so the default
+# test set keeps them.
+
+CHLOROFORM_OUTLIER_SIGMA = 5   # cutoff used by the MagNET-x training pipeline (train_MagNET_X.py)
+
+
+def _chloroform_test_solutes(sigma_fresh):
+    """Solute indices flagged as the held-out test split in the released sigma-fresh file. This
+    matches the paper's chloroform test split exactly (534 solutes)."""
+    group = sigma_fresh.f["chloroform"]
+    out = []
+    for name in group.keys():
+        kind = group[name].attrs["type"]
+        if isinstance(kind, bytes):
+            kind = kind.decode()
+        if kind == "test":
+            out.append(int(name.split("_")[1]))
+    return sorted(out)
+
+
+def isolated_chloroform_testset(sigma_fresh, max_frame=10, n_sigma=CHLOROFORM_OUTLIER_SIGMA):
+    """Assemble the chloroform isolated-solute test set for Tables S3/S4, from released data only.
+
+    Each case is one isolated-solute snapshot: the solute geometry from a chloroform MD frame and the
+    isolated (no-solvent) DFT shielding computed for it. The set is defined as:
+
+      * test solutes : attrs['type'] == 'test' (the paper's held-out split; evaluating MagNET-x on
+                       its training solutes would be data leakage).
+      * frames       : frame_index 1..max_frame where the isolated shielding was actually computed.
+
+    The published Tables S3/S4 use the full set (ans_all/geos_all/dfts_all). `flagged` marks the most
+    strongly deshielded frames (~16-17 ppm isolated shielding, far-downfield O-H / N-H), found with the
+    same +/- n_sigma cut the MagNET-x training pipeline applies (train_MagNET_X.py), recomputed here
+    from the dataset's own shieldings. Their shieldings vary smoothly along each MD trajectory and sit
+    on a continuous low-shielding tail, so this split is a diagnostic for documentation, not a
+    correction: ans/geos/dfts (with those frames removed) exist only to show how much of the RMSE they
+    carry.
+
+    Returns a dict:
+      ans, geos, dfts             -- test set with the deshielded-tail frames removed (diagnostic only)
+      ans_all, geos_all, dfts_all -- the full published test set (reproduces Tables S3/S4)
+      flagged                     -- sorted [(solute_index, frame_index), ...] in the deshielded tail
+      bands                       -- {'isolated': (lo, hi), 'solvated': (lo, hi)} n_sigma bands (ppm)
+    """
+    frames = []
+    iso_H, solv_H = [], []
+    for si in _chloroform_test_solutes(sigma_fresh):
+        s = sigma_fresh.solute("chloroform", si)
+        ns = s["n_solute_atoms"]
+        an = s["atomic_numbers"][:ns]                  # solute is the leading ns atoms
+        is_h = an == 1
+        iso = s["shielding_isolated"]                  # (n_frames, n_solute) ppm, NaN where not done
+        solv = s["shielding_solvated"]
+        coords = s["coordinates"]
+        for row in range(min(max_frame, iso.shape[0])):    # row r -> frame_index r+1 (1-based)
+            sh = iso[row]
+            if not np.isfinite(sh).any():
+                continue
+            solv_row = solv[row]
+            frames.append({"solute": si, "frame": row + 1, "an": an,
+                           "geo": coords[row][:ns], "iso": sh, "solv": solv_row, "is_h": is_h})
+            h = sh[is_h]
+            iso_H.append(h[np.isfinite(h)])
+            h = solv_row[is_h]
+            solv_H.append(h[np.isfinite(h)])
+    iso_H = np.concatenate(iso_H)
+    solv_H = np.concatenate(solv_H)
+    iso_band = (iso_H.mean() - n_sigma * iso_H.std(), iso_H.mean() + n_sigma * iso_H.std())
+    solv_band = (solv_H.mean() - n_sigma * solv_H.std(), solv_H.mean() + n_sigma * solv_H.std())
+
+    def is_outlier(fr):
+        h = fr["iso"][fr["is_h"]]
+        h = h[np.isfinite(h)]
+        if ((h < iso_band[0]) | (h > iso_band[1])).any():
+            return True
+        h = fr["solv"][fr["is_h"]]
+        h = h[np.isfinite(h)]
+        return bool(h.size and ((h < solv_band[0]) | (h > solv_band[1])).any())
+
+    out = {"ans": [], "geos": [], "dfts": [], "ans_all": [], "geos_all": [], "dfts_all": [],
+           "flagged": [], "bands": {"isolated": iso_band, "solvated": solv_band}}
+    for fr in frames:
+        out["ans_all"].append(fr["an"])
+        out["geos_all"].append(fr["geo"])
+        out["dfts_all"].append(fr["iso"])
+        if is_outlier(fr):
+            out["flagged"].append((fr["solute"], fr["frame"]))
+        else:
+            out["ans"].append(fr["an"])
+            out["geos"].append(fr["geo"])
+            out["dfts"].append(fr["iso"])
+    out["flagged"].sort()
+    return out
+
+
+def cases_isolated_chloroform(sigma_fresh, scrub=False, max_frame=10):
+    """Isolated-solute snapshots from chloroform MD (sigma-fresh), as (ans, geos, dfts).
+
+    Default (scrub=False) is the full published test set, which reproduces Tables S3/S4. scrub=True
+    drops the strongly-deshielded-proton frames; that is a diagnostic, not a correction (those frames
+    are real geometries, not DFT failures). See isolated_chloroform_testset.
+    """
+    t = isolated_chloroform_testset(sigma_fresh, max_frame=max_frame)
+    return (t["ans"], t["geos"], t["dfts"]) if scrub else (t["ans_all"], t["geos_all"], t["dfts_all"])

analysis/code/magnet_benchmark_run.py

@@ -0,0 +1,121 @@
+"""Produce the magnet_benchmark results table by running MagNET and MagNET-x over the four test sets
+and write it to performance_results.csv.
+
+Needs the checkpoints and datasets, so it does not run in CI; they ship in the repo (Git LFS via the
+Hugging Face checkout) and resolve automatically. One row per (model, test set, nucleus) with median
+AE, MAE, RMSE, and atom count. The three vibrated/stationary sets are sampled; isolated-chloroform is
+run in full. The shipped CSV used --sample 500 --seed 0; re-running reproduces it (sampled rows within
+tolerance, isolated rows exactly).
+
+Usage:
+    python magnet_benchmark_run.py --out performance_results.csv --n-passes 1 --sample 500 --seed 0
+"""
+import argparse
+import csv
+import os
+import sys
+import time
+
+from paths import ensure_on_path, repo_root, dataset_file
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+GH = repo_root(__file__)
+if HERE not in sys.path:
+    sys.path.insert(0, HERE)
+ensure_on_path("analysis", "code", "shared", file=__file__)  # stats.py etc. live here
+ensure_on_path("magnet", file=__file__)
+for _d in ("sigma-shake", "gdb_qcd", "dft8k", "sigma-fresh"):
+    ensure_on_path("data", _d, file=__file__)
+
+import magnet_benchmark as M  # noqa: E402
+from run_magnet import compute_MagNET_foundation_shieldings, _predict_with  # noqa: E402
+
+
+def magnet_x_chloroform_shieldings(ans, geos, n_passes, device):
+    """Absolute shieldings from the MagNET-x chloroform model (run as a plain predictor, no solvent),
+    used to score MagNET-x on the same test sets as the foundation model."""
+    return _predict_with("MagNET-x-chloroform_H", "MagNET-x-chloroform_C", ans, geos,
+                         n_passes=n_passes, device=device)
+
+
+MODEL_RUNNERS = {
+    "MagNET": compute_MagNET_foundation_shieldings,
+    "MagNET-x": magnet_x_chloroform_shieldings,
+}
+
+
+def assemble_all(paths, sample, seed):
+    """Assemble the four test sets, returning {test_set: (ans, geos, dfts)}."""
+    from decode_sigma_shake import SigmaShake
+    from decode_gdb_qcd import GDBQCD
+    from dft8k_reader import DFT8k
+    from decode_sigma_fresh import SigmaFresh
+
+    raw = {}
+    with SigmaShake(paths["sigma_shake"]) as ss:
+        raw["stationary_internal"] = M.cases_stationary_internal(ss, n_sample=sample, seed=seed)
+    with GDBQCD(paths["gdb_qcd"]) as gq:
+        raw["vibrated_internal"] = M.cases_vibrated_internal(gq, n_sample=sample // 3, seed=seed)
+    with DFT8k(paths["dft8k"]) as d8:
+        raw["vibrated_external"] = M.cases_vibrated_external(d8, n_sample=sample, seed=seed)
+    with SigmaFresh(paths["sigma_fresh"]) as sf:
+        # the full published test set (held-out solutes, frames 1..10 with an isolated calculation);
+        # not sampled, so this row reproduces Tables S3/S4 exactly
+        raw["isolated_chloroform"] = M.cases_isolated_chloroform(sf, scrub=False)
+    # exclude molecules with elements outside MagNET's vocabulary (e.g. phosphorus in dft8k); the
+    # model would refuse them anyway, and scoring them would not be a fair test
+    cases = {}
+    for ts, (ans, geos, dfts) in raw.items():
+        a, g, d, dropped = M.filter_supported(ans, geos, dfts)
+        if dropped:
+            print(f"  {ts}: dropped {dropped} structure(s) with unsupported elements", flush=True)
+        cases[ts] = (a, g, d)
+    return cases
+
+
+def main():
+    """Assemble the four test sets, run both models over them, and write performance_results.csv."""
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default=os.path.join(HERE, "..", "..", "data", "magnet_benchmark", "performance_results.csv"))
+    ap.add_argument("--n-passes", type=int, default=1)
+    ap.add_argument("--sample", type=int, default=500)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--sigma-shake", default=dataset_file("sigma-shake", root=GH))
+    ap.add_argument("--gdb-qcd", default=dataset_file("gdb_qcd", root=GH))
+    ap.add_argument("--dft8k", default=dataset_file("dft8k", root=GH))
+    ap.add_argument("--sigma-fresh", default=dataset_file("sigma-fresh", root=GH))
+    args = ap.parse_args()
+
+    paths = {"sigma_shake": args.sigma_shake, "gdb_qcd": args.gdb_qcd,
+             "dft8k": args.dft8k, "sigma_fresh": args.sigma_fresh}
+    print(f"assembling test cases (sample={args.sample}, seed={args.seed})...", flush=True)
+    cases = assemble_all(paths, args.sample, args.seed)
+    for ts, (ans, _, _) in cases.items():
+        print(f"  {ts}: {len(ans)} structures", flush=True)
+
+    rows = []
+    for model, runner in MODEL_RUNNERS.items():
+        for ts in M.TEST_SETS:
+            ans, geos, dfts = cases[ts]
+            t = time.time()
+            pred = runner(ans, geos, args.n_passes, args.device)
+            tab = M.error_table(pred, dfts, ans)
+            dt = time.time() - t
+            for nucleus in ("1H", "13C"):
+                s = tab[nucleus]
+                rows.append(dict(model=model, test_set=ts, nucleus=nucleus,
+                                 median_ae=s["median_ae"], mae=s["mae"], rmse=s["rmse"], n=s["n"]))
+            print(f"{model:9s} {ts:20s} ({dt:.0f}s)  "
+                  f"1H mae {tab['1H']['mae']:.4f}  13C mae {tab['13C']['mae']:.4f}", flush=True)
+
+    with open(args.out, "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=["model", "test_set", "nucleus", "median_ae", "mae", "rmse", "n"])
+        w.writeheader()
+        for r in rows:
+            w.writerow(r)
+    print(f"wrote {args.out} ({len(rows)} rows)")
+
+
+if __name__ == "__main__":
+    main()

analysis/code/nb_build.py

@@ -0,0 +1,28 @@
+"""Shared helpers for the ``build_nb_*.py`` notebook generators.
+
+Each ``build_nb_<x>.py`` assembles a notebook from a list of cells built with :func:`md` and
+:func:`code`, then writes it with :func:`save_notebook`. The notebook metadata written here is
+intentionally minimal -- ``jupyter nbconvert --execute`` fills in the full ``language_info``/kernel
+block when the notebook is run.
+"""
+import nbformat
+
+
+def md(source):
+    """Return a markdown cell from a string, with leading/trailing blank lines stripped."""
+    return nbformat.v4.new_markdown_cell(source.strip("\n"))
+
+
+def code(source):
+    """Return a code cell from a string, with leading/trailing blank lines stripped."""
+    return nbformat.v4.new_code_cell(source.strip("\n"))
+
+
+def save_notebook(cells, path):
+    """Write a notebook made of ``cells`` to ``path`` and print a confirmation line."""
+    nb = nbformat.v4.new_notebook()
+    nb.cells = cells
+    nb.metadata = {"language_info": {"name": "python"}}
+    with open(path, "w") as f:
+        nbformat.write(nb, f)
+    print("wrote", path)

analysis/code/pareto_plot.py

@@ -0,0 +1,198 @@
+"""Broken-axis accuracy-vs-cost Pareto panel used by Figure 2A and SI Figure S1.
+
+Both figures draw the same panel (fitting RMSE vs total compute time, colored by method type, shaped by
+geometry source, sized by basis) and differ only in their config dicts and which panels they show, so
+that shared drawing engine lives here instead of being duplicated in each notebook. Call
+`plot_pareto_panel(...)`; the notebook supplies the points table and the styling dicts.
+"""
+import numpy as np
+import seaborn as sns
+import matplotlib.pyplot as plt
+from matplotlib.lines import Line2D
+from matplotlib.gridspec import GridSpec
+
+# methods colored as double hybrids; dlpno_mp2 is classified ab initio instead (checked first)
+DOUBLE_HYBRID_METHODS = ["dsd_pbep86", "B2GP_PLYP", "B2PLYP", "mPW2PLYP", "revdsd_pbep86", "dlpno_mp2"]
+
+
+def _format_seconds(seconds):
+    if not np.isfinite(seconds) or seconds <= 0:
+        return ""
+    nice = lambda v: str(int(round(v)))
+    if seconds < 60:
+        return f"{nice(seconds)} s"
+    if seconds < 3600:
+        return f"{nice(seconds / 60)} min"
+    if seconds < 86400:
+        return f"{nice(seconds / 3600)} h"
+    return f"{nice(seconds / 86400)} day"
+
+
+def _apply_time_xticklabels(ax, xticks):
+    ax.set_xticks(xticks)
+    ax.set_xticklabels([_format_seconds(10 ** float(t)) for t in xticks], fontsize=10)
+
+
+def _classify_method(method):
+    # returns (label_text, category); category drives color, label_text is the annotation
+    if method in ["hf", "mp2", "dlpno_mp2"]:
+        return method, "ab initio"
+    if method in DOUBLE_HYBRID_METHODS:
+        return method.lower(), "double hybrids"
+    if method.startswith("MagNET"):
+        return "MagNET-Zero", "MagNET-Zero"
+    if method.startswith(("wp", "wc")):
+        return method, "NMR-specific"
+    return method.split("_")[0], "DFT"
+
+
+def _classify_and_plot_points(df, x_range1, ax1, ax2, color_map, geometry_marker_map,
+                              marker_alpha, manual_label_positions, labels_per_category=6):
+    size_dict = {"pcSseg1": 20, "pcSseg2": 50, "pcSseg3": 100, "N/A": 110}
+    color_counts = {k: 0 for k in ["ab initio", "double hybrids", "MagNET-Zero", "NMR-specific", "DFT"]}
+    basis_counts = {k: 0 for k in ["pcSseg1", "pcSseg2", "pcSseg3"]}
+    geometry_counts = {}
+
+    x_all = np.log10(df["total_time"].astype(float).values)
+    x_leftmost = np.nanmin(x_all) if len(x_all) else np.nan
+
+    # label selection: the min-RMSE point per category plus a few spread across time
+    tmp = df.copy()
+    tmp["_log10_time"] = np.log10(tmp["total_time"].astype(float))
+    tmp["_rmse"] = tmp["fitting_RMSE"].astype(float)
+    tmp["_category"] = tmp["nmr_method"].astype(str).apply(lambda m: _classify_method(m)[1])
+    label_indices = set()
+    for cat in ["MagNET-Zero", "ab initio", "double hybrids", "NMR-specific", "DFT"]:
+        sub = tmp[tmp["_category"] == cat]
+        if sub.empty:
+            continue
+        picks = [int(sub["_rmse"].idxmin())]
+        n_extra = max(0, labels_per_category - 1)
+        if n_extra > 0 and len(sub) > 1:
+            tvals = sub["_log10_time"].to_numpy()
+            for q in np.linspace(0.05, 0.95, num=n_extra):
+                picks.append(int((sub["_log10_time"] - np.quantile(tvals, q)).abs().idxmin()))
+        seen, unique = set(), []
+        for p in picks:
+            if p not in seen:
+                unique.append(p); seen.add(p)
+        label_indices.update(unique[:labels_per_category])
+
+    for idx, row in df.iterrows():
+        x_val = float(np.log10(float(row.total_time)))
+        y_val = float(row["fitting_RMSE"])
+        geometry_type = str(row["geometry_type"]); method_raw = str(row["nmr_method"]); basis = str(row["basis"])
+        label_text, category = _classify_method(method_raw)
+        color_counts[category] += 1
+        point_color = color_map[category]
+        marker = geometry_marker_map.get(geometry_type, "o")
+        geometry_counts[geometry_type] = geometry_counts.get(geometry_type, 0) + 1
+        if basis in basis_counts:
+            basis_counts[basis] += 1
+        base_size = size_dict.get(basis, size_dict.get("N/A", 30))
+        size = base_size * (2.8 if abs(x_val - x_leftmost) <= 1e-12 else 1.0)
+        ax = ax1 if (x_range1[0] <= x_val <= x_range1[1]) else ax2
+        unfilled = marker in ["x", "+", "1", "2", "3", "4", "|", "_"]
+        kws = dict(x=[x_val], y=[y_val], ax=ax, color=point_color, s=size, marker=marker,
+                   legend=False, alpha=marker_alpha, zorder=3)
+        kws.update(dict(linewidth=1.2) if unfilled else dict(edgecolor="white", linewidth=0.6))
+        sns.scatterplot(**kws)
+
+        override = None
+        if manual_label_positions:
+            for key in [label_text, method_raw, (method_raw, basis),
+                        (method_raw, basis, geometry_type), (label_text, basis, geometry_type)]:
+                if key in manual_label_positions:
+                    override = manual_label_positions[key]; break
+        if (idx in label_indices) or (override is not None):
+            xytext = tuple(override["xytext"]) if isinstance(override, dict) else (4, 3)
+            ax.annotate(label_text, xy=(x_val, y_val), xytext=xytext, textcoords="offset points",
+                        ha="left", va="bottom", fontsize=8, color=point_color, zorder=4,
+                        annotation_clip=False)
+    return color_counts, basis_counts, geometry_counts
+
+
+def _add_legends(ax1, ax2, color_counts, basis_counts, geometry_counts, color_map,
+                 geometry_marker_map, method_legend_label_map, geometry_legend_label_map):
+    method_order = ["MagNET-Zero", "ab initio", "double hybrids", "NMR-specific", "DFT"]
+    method_elems = [Line2D([0], [0], marker="o", color="w", markerfacecolor=color_map[c],
+                           markeredgecolor="white", markeredgewidth=0.6, markersize=8, linestyle="None",
+                           label=method_legend_label_map.get(c, c))
+                    for c in method_order if color_counts.get(c, 0) > 0]
+    basis_elems = [Line2D([0], [0], marker="o", color="none", markerfacecolor="k", markersize=ms, label=b)
+                   for b, ms in [("pcSseg1", 3), ("pcSseg2", 6), ("pcSseg3", 8)] if basis_counts.get(b, 0) > 0]
+    geometry_elems = []
+    for g in ["aimnet2", "pbe0_tz"]:
+        if geometry_counts.get(g, 0) > 0:
+            gm = geometry_marker_map.get(g, "o")
+            geometry_elems.append(Line2D([0], [0], marker=gm, color="k",
+                                         markerfacecolor="none" if gm == "x" else "k",
+                                         markeredgecolor="k", markersize=7, linestyle="None",
+                                         label=geometry_legend_label_map.get(g, g)))
+    ax1.add_artist(ax1.legend(handles=method_elems, loc="upper left", bbox_to_anchor=(0.005, 0.995),
+                              frameon=True, fontsize=9, title="Method Type", title_fontsize=9,
+                              borderaxespad=0.0, handletextpad=0.4, labelspacing=0.3))
+    ax2.add_artist(ax2.legend(handles=basis_elems, loc="upper right", bbox_to_anchor=(0.995, 0.995),
+                              frameon=True, fontsize=9, title="Basis Set", title_fontsize=9,
+                              borderaxespad=0.0, handletextpad=0.4, labelspacing=0.3))
+    if geometry_elems:
+        ax2.add_artist(ax2.legend(handles=geometry_elems, loc="upper right", bbox_to_anchor=(0.80, 0.995),
+                                  frameon=True, fontsize=9, title="Geometries", title_fontsize=9,
+                                  borderaxespad=0.0, handletextpad=0.4, labelspacing=0.3))
+
+
+def plot_pareto_panel(pareto_df, query_str, nucleus, method_color_map, method_legend_label_map,
+                      geometry_legend_label_map, geometry_marker_map, xlim_left, xlim_right, ylim,
+                      figsize, marker_alpha, manual_label_positions, save_png=None,
+                      suptitle="MagNET-Zero Extends a Flat Pareto Frontier",
+                      supxlabel="Total delta-22 calculation time"):
+    """Draw the broken-axis accuracy-vs-cost Pareto panel and optionally save it as a PNG."""
+    sns.set_theme(style="ticks", context="paper", rc={"axes.grid": False})
+    query_df = pareto_df.query(query_str.format(nucleus=nucleus), engine="python").reset_index(drop=True)
+    default_color_map = {"ab initio": "black", "double hybrids": "#ede6ba", "MagNET-Zero": "#61a89a",
+                         "NMR-specific": "#a72608", "DFT": "black"}
+    color_map = {**default_color_map, **method_color_map}
+
+    fig = plt.figure(figsize=figsize)
+    gs = GridSpec(1, 2, width_ratios=[1, 5], wspace=0.05)
+    ax1 = fig.add_subplot(gs[0]); ax2 = fig.add_subplot(gs[1], sharey=ax1)
+
+    cc, bc, gc = _classify_and_plot_points(query_df, xlim_left, ax1, ax2, color_map,
+                                           geometry_marker_map, marker_alpha, manual_label_positions)
+
+    ax1.set_xlim(xlim_left); ax2.set_xlim(xlim_right)
+    ax1.set_ylim(ylim); ax2.set_ylim(ylim)
+    y_ticks = np.linspace(ylim[0], ylim[1], num=6)
+    ax1.set_yticks(y_ticks); ax1.set_yticklabels([f"{t:.2f}" for t in y_ticks], fontsize=10)
+
+    _apply_time_xticklabels(ax1, np.array([xlim_left[0], (xlim_left[0] + xlim_left[1]) / 2, xlim_left[1]]))
+    _apply_time_xticklabels(ax2, np.arange(xlim_right[0] + 0.25, xlim_right[1] + 0.25, 0.25))
+    ax1.tick_params(axis="both", which="both", length=1.5, width=0.6, labelsize=10)
+    ax2.tick_params(axis="x", which="both", length=1.5, width=0.6, labelsize=10)
+    ax2.tick_params(axis="y", which="both", left=False, right=False, labelleft=False)
+    ax1.grid(False); ax2.grid(False)
+
+    for ax in (ax1, ax2):
+        for side in ("left", "right", "top", "bottom"):
+            ax.spines[side].set_visible(True); ax.spines[side].set_linewidth(0.6)
+    ax1.spines["right"].set_visible(False); ax2.spines["left"].set_visible(False)
+
+    # broken-axis break marks
+    width_ratio = ax2.get_position().width / ax1.get_position().width
+    d1 = 0.01; d2 = d1 / width_ratio; slope2 = width_ratio
+    kw = dict(color="k", clip_on=False, lw=0.6, transform=ax1.transAxes)
+    ax1.plot((1 - d1, 1 + d1), (-d1, d1), **kw); ax1.plot((1 - d1, 1 + d1), (1 - d1, 1 + d1), **kw)
+    kw.update(transform=ax2.transAxes)
+    ax2.plot((-d2, d2), (-d2 * slope2, d2 * slope2), **kw); ax2.plot((-d2, d2), (1 - d2 * slope2, 1 + d2 * slope2), **kw)
+
+    fig.subplots_adjust(left=0.15, right=0.95, top=0.92, bottom=0.12)
+    fig.suptitle(suptitle, fontsize=14)
+    fig.supxlabel(supxlabel, fontsize=13)
+    ax1.set_ylabel(r"RMSE ($^{1}$H ppm, CDCl$_3$)" if nucleus == "H" else r"RMSE ($^{13}$C ppm, CDCl$_3$)",
+                   fontsize=13)
+
+    _add_legends(ax1, ax2, cc, bc, gc, color_map, geometry_marker_map,
+                 method_legend_label_map, geometry_legend_label_map)
+    if save_png:
+        fig.savefig(save_png, dpi=300, bbox_inches="tight")
+    plt.show()

analysis/code/paths.py

@@ -0,0 +1,64 @@
+"""Repo-root and sys.path resolution for the analysis code.
+
+Bootstrap-only: imports nothing from analysis/code/ (it is what makes that folder importable), so it
+stays at analysis/code/ root. repo_root(file) assumes the caller sits two dirs below the repo root
+(analysis/code/x.py -> ../..), which holds for every module here.
+"""
+import os
+import sys
+
+def repo_root(file):
+    """The release repo root, given a module's __file__."""
+    return os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(file)), "..", ".."))
+
+
+def dataset_file(name, filename=None, *, file=None, root=None):
+    """Absolute path to a large released dataset file, in the repo at data/<name>/<filename>.
+
+    `filename` defaults to "<name>.hdf5". Give `root` (the repo root) or `file` (a module __file__,
+    from which the repo root is derived) so the file can be located. Only the large, git-ignored
+    data files go through here; small companion files (readers, spreadsheets) always live in the
+    repo and are addressed directly. The Hugging Face checkout carries these files via Git LFS.
+    """
+    if filename is None:
+        filename = "{}.hdf5".format(name)
+    if root is None:
+        if file is None:
+            raise ValueError(
+                "dataset_file needs root=<repo root> or file=__file__ to locate {!r}".format(filename)
+            )
+        root = repo_root(file)
+    return os.path.join(root, "data", name, filename)
+
+
+def checkpoints_root(*, required=False):
+    """Absolute path to the in-repo model_checkpoints/ folder holding the trained weights.
+
+    Returns the path when it exists, else None so callers can skip checkpoint-dependent work, unless
+    required=True, which raises. The Hugging Face checkout carries the weights via Git LFS; they are
+    git-ignored on GitHub, so this is None in CI.
+    """
+    in_repo = os.path.join(repo_root(__file__), "model_checkpoints")
+    if os.path.exists(in_repo):
+        return in_repo
+    if required:
+        raise RuntimeError(
+            "No model_checkpoints/ found. It ships in the repo via Git LFS in the Hugging Face "
+            "checkout (hf download ekwan16/MagNET --include 'model_checkpoints/*')."
+        )
+    return None
+
+
+def ensure_on_path(*relative_parts, file=None, root=None):
+    """Inserts os.path.join(root, *relative_parts) at the front of sys.path, if not already present.
+
+    Pass either `file` (typically __file__, resolved via repo_root) or `root` directly.
+    """
+    if root is None:
+        if file is None:
+            raise ValueError("ensure_on_path needs either file=__file__ or root=<repo root>")
+        root = repo_root(file)
+    target = os.path.join(root, *relative_parts)
+    if target not in sys.path:
+        sys.path.insert(0, target)
+    return target

analysis/code/scaling_factors.py

@@ -0,0 +1,220 @@
+"""Recommended Linear Scaling Parameters for MagNET-Zero / MagNET-PCM (SI Tables S10 and S11).
+
+Turns MagNET-Zero shieldings into chemical shifts via a per-solvent linear model:
+
+    shift = c_intercept + c_stationary * sigma_zero + c_pcm * delta_pcm
+
+where sigma_zero is the MagNET-Zero gas-phase shielding and delta_pcm the MagNET-PCM chloroform
+correction. Proton (S10) uses a THREE-parameter model (all coefficients free) because for benzene and
+toluene the PCM correction points the wrong way and only a free c_pcm can go positive. Carbon (S11)
+uses a TWO-parameter model: the gas and PCM terms share one per-solvent slope (PCM scaled to solvent
+by a conversion factor first), which fit held-out carbons slightly better; the reported `pcm` is that
+slope times the conversion factor.
+
+Everything reproduces from the released delta-22 data to the precision of its integer encoding
+(`test_scaling_factors.py` locks it down). `build_scaling_tables(..., symmetrized=True)` instead
+computes a separate deployment variant via reflection-symmetrized inference
+(`scaling_factors_symmetrized.py`); it does not match these tables and is not what the SI reports.
+"""
+import os
+import sys
+
+import numpy as np
+import pandas as pd
+
+from paths import repo_root, ensure_on_path, dataset_file
+
+# paths.py is the bootstrap at analysis/code/ root; the other shared utilities (stats, etc.) live in
+# analysis/code/shared/, and the delta-22 reader/harness is reused from data/delta22. Put all of
+# these on sys.path before importing them, so scaling_factors can be imported standalone (as its
+# docstring invites), not only under pytest/notebooks that already pre-wire the paths.
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_REPO = repo_root(__file__)
+if _HERE not in sys.path:
+    sys.path.insert(0, _HERE)
+ensure_on_path("analysis", "code", "shared", file=__file__)
+ensure_on_path("data", "delta22", file=__file__)
+
+from stats import linear_fit_1d  # noqa: E402
+import delta22 as D  # noqa: E402
+
+# The 12 solvents, in the reader's native naming (water is stored as "TIP4P").
+SOLVENTS = tuple(D.DESMOND_SOLVENTS)
+
+# Nitromethane is an outlier for both nuclei (its carbon shift is anomalous); the published tables
+# reproduce only with it excluded throughout.
+EXCLUDE_SOLUTES = ("nitromethane",)
+
+# The model form recommended for each nucleus.
+RECOMMENDED_MODEL = {"H": "three_parameter", "C": "two_parameter"}
+
+_COLUMNS = ["intercept", "stationary", "pcm"]
+
+# The published SI Tables S10 (proton) and S11 (carbon), shipped verbatim so you can convert
+# shieldings to shifts WITHOUT the delta-22 data download. These are exactly what
+# build_scaling_tables() reproduces from delta-22; test_scaling_factors.py asserts the two agree, so
+# they cannot silently drift. Prefer published_scaling_tables() for a quick lookup, and
+# build_scaling_tables() when you want to re-derive them from the raw data.
+_PUBLISHED_TABLE_CSV = {
+    "H": """solvent,intercept,stationary,pcm
+tetrahydrofuran,31.321785,-0.980175,-0.786672
+dichloromethane,31.377310,-0.979870,-0.808573
+chloroform,31.294992,-0.975794,-0.852690
+toluene,31.716957,-0.996733,1.944856
+benzene,31.987569,-1.005289,2.236515
+chlorobenzene,31.681438,-0.993846,0.830236
+acetone,31.512157,-0.987261,-1.238414
+dimethylsulfoxide,31.598698,-0.991103,-1.355248
+acetonitrile,31.496110,-0.985716,-0.974465
+trifluoroethanol,30.876004,-0.959997,-0.958852
+methanol,31.256030,-0.976446,-1.250121
+TIP4P,31.359134,-0.978809,-1.478298
+""",
+    "C": """solvent,intercept,stationary,pcm
+tetrahydrofuran,171.054488,-0.919000,-1.069509
+dichloromethane,171.509389,-0.921167,-1.115015
+chloroform,171.728797,-0.924231,-0.936804
+toluene,171.690907,-0.924987,-0.618410
+benzene,171.967177,-0.927033,-0.593716
+chlorobenzene,171.075910,-0.921238,-0.997641
+acetone,171.308023,-0.919610,-1.239503
+dimethylsulfoxide,170.598425,-0.918674,-1.296501
+acetonitrile,171.879839,-0.922629,-1.287664
+trifluoroethanol,174.237671,-0.939012,-1.290073
+methanol,172.426161,-0.927566,-1.288847
+TIP4P,173.696379,-0.937854,-1.342668
+""",
+}
+
+
+def published_scaling_tables():
+    """The published SI scaling tables, {"H": Table S10, "C": Table S11}, as solvent-indexed
+    DataFrames. No data download needed; use these with predict_shift to turn MagNET-Zero/PCM
+    shieldings into chemical shifts."""
+    import io
+    return {nucleus: pd.read_csv(io.StringIO(csv)).set_index("solvent")
+            for nucleus, csv in _PUBLISHED_TABLE_CSV.items()}
+
+
+def proton_scaling_table(query_df_nn, solvents=SOLVENTS, exclude_solutes=EXCLUDE_SOLUTES):
+    """SI Table S10: per-solvent [intercept, stationary, pcm] for the proton three-parameter model.
+
+    Fits experimental ~ stationary + pcm freely for each solvent on the MagNET-Zero shielding and the
+    MagNET-PCM (chloroform) correction. Returns a DataFrame indexed by solvent.
+    """
+    nn = query_df_nn[(query_df_nn["nucleus"] == "H")
+                     & ~query_df_nn["solute"].isin(set(exclude_solutes))]
+    coeffs = D.full_fit_coefficients(nn, list(solvents), "stationary + pcm")
+    table = coeffs.T.rename(columns={"Intercept": "intercept"})[_COLUMNS]
+    table.index.name = "solvent"
+    return table
+
+
+def carbon_conversion_factors(query_df_dft):
+    """The per-solvent PCM conversion factor for carbon (delta22.pcm_conversion_factors): for each
+    solvent, the through-the-origin slope of that solvent's PCM correction (at the level MagNET-Zero
+    reproduces, wB97X-D) against chloroform's B3LYP-D3(BJ) correction (the level MagNET-PCM
+    reproduces), from the delta-22 DFT PCM data. So one factor converts both chloroform to the solvent
+    and B3LYP to wB97X-D. A Series indexed by solvent."""
+    return D.pcm_conversion_factors(query_df_dft, "C")
+
+
+def carbon_scaling_table(query_df_nn, query_df_dft=None, solvents=SOLVENTS,
+                         exclude_solutes=EXCLUDE_SOLUTES, conversion_factors=None):
+    """SI Table S11: per-solvent [intercept, stationary, pcm] for the carbon two-parameter model.
+
+    For each solvent, fits experimental ~ (stationary + factor * pcm) with a single shared slope,
+    where `factor` is the per-solvent PCM conversion factor. The reported `stationary` is that slope
+    and the reported `pcm` is slope * factor, so the prediction equation in the module docstring
+    applies with MagNET-PCM's chloroform correction. The factors come from query_df_dft
+    (carbon_conversion_factors) unless passed in as `conversion_factors` (a Series or dict). A solvent
+    with fewer than two usable sites, or an undefined conversion factor, gets a row of NaN rather than
+    a fake zero fit. Returns a DataFrame indexed by solvent.
+    """
+    if conversion_factors is None:
+        if query_df_dft is None:
+            raise ValueError("provide query_df_dft or conversion_factors")
+        conversion_factors = carbon_conversion_factors(query_df_dft)
+    nn = query_df_nn[(query_df_nn["nucleus"] == "C")
+                     & ~query_df_nn["solute"].isin(set(exclude_solutes))]
+    rows = {}
+    for solvent in solvents:
+        factor = float(conversion_factors[solvent])
+        sub = nn[nn["solvent"] == solvent]
+        y = pd.to_numeric(sub["experimental"], errors="coerce").to_numpy(float)
+        x = sub["stationary"].to_numpy(float) + factor * sub["pcm"].to_numpy(float)
+        keep = np.isfinite(x) & np.isfinite(y)
+        if not np.isfinite(factor) or keep.sum() < 2:
+            rows[solvent] = {"intercept": np.nan, "stationary": np.nan, "pcm": np.nan}
+            continue
+        intercept, slope = linear_fit_1d(x[keep], y[keep])
+        rows[solvent] = {"intercept": intercept, "stationary": slope, "pcm": slope * factor}
+    table = pd.DataFrame.from_dict(rows, orient="index")[_COLUMNS]
+    table.index.name = "solvent"
+    return table
+
+
+def build_scaling_tables(delta22_path, experimental_path, symmetrized=False, n_passes=10):
+    """Reproduce both recommended-scaling tables from the released delta-22 data.
+
+    symmetrized=False (default) reproduces the published SI Tables S10/S11 exactly, using the
+    HDF5's stored MagNET-Zero/PCM shieldings -- this is what test_scaling_factors.py checks against
+    and what published_scaling_tables() ships.
+
+    symmetrized=True instead computes the shieldings via live, symmetrized inference
+    (predict_shieldings(..., n_passes=n_passes, symmetrize=True), see
+    scaling_factors_symmetrized.py), which corrects a reflection-parity bug in the stored HDF5
+    values (they come from a single unsymmetrized forward pass). Use this for a deployment-quality
+    scaling table -- e.g. serving MagNET-Zero/MagNET-PCM on new "everyday" molecules -- rather than
+    for reproducing the SI. Requires the magnet package and released model checkpoints; slower
+    (re-runs inference on all 22 solutes instead of reading pre-baked values).
+
+    Returns {"H": Table S10 DataFrame, "C": Table S11 DataFrame}.
+    """
+    if symmetrized:
+        import scaling_factors_symmetrized as _S
+        override_df = _S.compute_symmetrized_nn_shieldings_df(delta22_path, n_passes=n_passes,
+                                                               verbose=False)
+        nn = D.load_query_df_nn(delta22_path, experimental_path, verbose=False,
+                                nn_shieldings_override_df=override_df)
+    else:
+        nn = D.load_query_df_nn(delta22_path, experimental_path, verbose=False)
+    dft = D.load_query_df_dft(delta22_path, experimental_path, verbose=False)
+    return {"H": proton_scaling_table(nn), "C": carbon_scaling_table(nn, dft)}
+
+
+def predict_shift(table, solvent, magnet_zero_shielding, magnet_pcm_chloroform_correction):
+    """Apply a scaling table to MagNET-Zero / MagNET-PCM outputs to predict a chemical shift.
+
+    table: a proton or carbon scaling table (DataFrame indexed by solvent). Get one from
+        published_scaling_tables()[nucleus] (no data download) or build_scaling_tables(...)[nucleus]
+        (re-derived from delta-22).
+    solvent: the solvent name (water is "TIP4P").
+    magnet_zero_shielding: the MagNET-Zero gas-phase shielding (scalar or array).
+    magnet_pcm_chloroform_correction: MagNET-PCM's chloroform correction (scalar or array).
+
+    The same equation serves both nuclei because the stored `pcm` coefficient already folds in the
+    per-solvent scaling, so you always pass the chloroform correction.
+
+    End to end, from a geometry (the H/C shieldings each model returns are per-atom arrays):
+
+        from magnet.run_magnet import compute_MagNET_Zero_shieldings, compute_MagNET_PCM_corrections
+        zero = compute_MagNET_Zero_shieldings([Z], [xyz])[0]          # MagNET-Zero shieldings
+        pcm  = compute_MagNET_PCM_corrections([Z], [xyz])[0]          # MagNET-PCM chloroform correction
+        tables = published_scaling_tables()
+        carbons = Z == 6
+        shifts_13C = predict_shift(tables["C"], "benzene", zero[carbons], pcm[carbons])
+    """
+    row = table.loc[solvent]
+    return (row["intercept"]
+            + row["stationary"] * np.asarray(magnet_zero_shielding, dtype=float)
+            + row["pcm"] * np.asarray(magnet_pcm_chloroform_correction, dtype=float))
+
+
+if __name__ == "__main__":
+    h5 = dataset_file("delta22", root=_REPO)
+    xlsx = os.path.join(_REPO, "data", "delta22", "delta22_experimental.xlsx")
+    tables = build_scaling_tables(h5, xlsx)
+    for nucleus, label in (("H", "Table S10 (1H)"), ("C", "Table S11 (13C)")):
+        print(f"\n=== {label}: {RECOMMENDED_MODEL[nucleus]} model ===")
+        print(tables[nucleus].round(6).to_string())

analysis/code/scaling_factors_symmetrized.py

@@ -0,0 +1,113 @@
+"""Computes delta-22's MagNET-Zero / MagNET-PCM shieldings via live, symmetrized inference,
+instead of reading the pre-baked (unsymmetrized, single-pass) values stored in delta22.hdf5.
+
+Why: `data/delta22/delta22_reader.py::load_stationary_and_pcm_nn_shieldings` decodes
+`nn_gas_shieldings`/`nn_pcm_corrections` straight out of the HDF5. Those were computed with a
+single, unsymmetrized forward pass (see `magnet/run_magnet.py`'s module docstring: "The released
+reference shieldings are NOT symmetrized"). Isotropic shielding is parity-even, so the SO(3)-only
+backbone's un-symmetrized predictions carry a spurious reflection error that is not always
+negligible, even on small molecules (delta-22's solutes are 6-22 atoms). `predict_shieldings(...,
+symmetrize=True)` fixes this for free by averaging the prediction on the original geometry and its
+mirror image.
+
+This module produces a scaling-factor table meant for deployment (e.g. serving MagNET-Zero/PCM on
+arbitrary "everyday" molecules), not a replacement for the SI's published Tables S10/S11: it is a
+separate, corrected-methodology table for production use, computed the same way
+`scaling_factors.build_scaling_tables(symmetrized=True)` does.
+"""
+import contextlib
+import os
+import time
+
+import h5py
+import numpy as np
+import pandas as pd
+
+from paths import repo_root, ensure_on_path, checkpoints_root
+
+_REPO = repo_root(__file__)
+# _REPO itself must be on sys.path too: run_magnet.py does `from magnet import ...`, which needs
+# the magnet/ package's PARENT directory on sys.path, not just magnet/ itself. Without this, the
+# import only works by accident when something else (e.g. `pip install -e .`, or pytest's rootdir
+# handling) has already put _REPO on sys.path.
+ensure_on_path(root=_REPO)
+ensure_on_path("magnet", file=__file__)
+ensure_on_path("data", "delta22", file=__file__)
+
+import run_magnet  # noqa: E402
+from run_magnet import compute_MagNET_Zero_shieldings, compute_MagNET_PCM_corrections  # noqa: E402
+from delta22_reader import load_solutes, load_stationary_geometries  # noqa: E402
+
+@contextlib.contextmanager
+def _resolved_checkpoints():
+    """Temporarily points run_magnet.MODEL_CHECKPOINTS at absolute paths, restoring the original
+    (relative-path) dict afterward -- run_magnet.MODEL_CHECKPOINTS is a shared module-level dict,
+    and other code importing run_magnet in the same process expects the original relative paths.
+    Its values already start with "model_checkpoints/", so join each against the repo root."""
+    base = os.path.dirname(checkpoints_root(required=True))
+    original = run_magnet.MODEL_CHECKPOINTS
+    run_magnet.MODEL_CHECKPOINTS = {
+        key: os.path.join(base, rel_path) for key, rel_path in original.items()
+    }
+    try:
+        yield
+    finally:
+        run_magnet.MODEL_CHECKPOINTS = original
+
+_NN_COLUMNS = ["solute", "sap_geometry_type", "sap_nmr_method", "sap_basis",
+               "solvent_model", "solvent", "shieldings", "geometry_time", "nmr_time"]
+
+
+def _load_atomic_numbers(delta22_path, solute):
+    with h5py.File(delta22_path, "r") as f:
+        return np.array(f["solutes"][solute]["atomic_numbers"])
+
+
+def _load_aimnet2_geometry_time(delta22_path, solute):
+    with h5py.File(delta22_path, "r") as f:
+        return float(np.array(f["solutes"][solute]["stationary_and_pcm"]["geometry_optimization_timings"])[0])
+
+
+def compute_symmetrized_nn_shieldings_df(delta22_path, n_passes=10, device=None, verbose=True):
+    """Same row shape as load_stationary_and_pcm_nn_shieldings (one "gas" + one "pcm_correction"
+    row per solute), but the shieldings come from live predict_shieldings(symmetrize=True)
+    inference instead of the HDF5's stored, unsymmetrized values.
+
+    n_passes=10 with symmetrize=True means 10 forward passes on the original geometry and 10 on
+    its mirror image, all 20 averaged together.
+
+    Returns a DataFrame with the same 9 columns load_stationary_and_pcm_nn_shieldings produces;
+    pass it as `nn_shieldings_override_df` to delta22_reader.load_delta22_nn_data.
+    """
+    solutes = load_solutes(delta22_path)
+    atomic_numbers_list = [_load_atomic_numbers(delta22_path, s) for s in solutes]
+    geometries_list = [load_stationary_geometries(delta22_path, s)["aimnet2"] for s in solutes]
+
+    if verbose:
+        print(f"computing symmetrized MagNET-Zero/PCM shieldings for {len(solutes)} solutes "
+              f"(n_passes={n_passes}, symmetrize=True)...", flush=True)
+
+    t0 = time.time()
+    with _resolved_checkpoints():
+        zero_shieldings = compute_MagNET_Zero_shieldings(
+            atomic_numbers_list, geometries_list, n_passes=n_passes, symmetrize=True, device=device)
+        t1 = time.time()
+        pcm_corrections = compute_MagNET_PCM_corrections(
+            atomic_numbers_list, geometries_list, n_passes=n_passes, symmetrize=True, device=device)
+    t2 = time.time()
+
+    if verbose:
+        print(f"  MagNET-Zero: {t1-t0:.1f}s total, MagNET-PCM: {t2-t1:.1f}s total", flush=True)
+
+    zero_time_per_solute = (t1 - t0) / len(solutes)
+    pcm_time_per_solute = (t2 - t1) / len(solutes)
+
+    rows = []
+    for solute, gas, pcm in zip(solutes, zero_shieldings, pcm_corrections):
+        geom_time = _load_aimnet2_geometry_time(delta22_path, solute)
+        rows.append([solute, "aimnet2", "MagNET", "N/A", "gas", "none",
+                     np.asarray(gas, dtype=float), geom_time, zero_time_per_solute])
+        rows.append([solute, "aimnet2", "MagNET", "N/A", "pcm_correction", "chloroform",
+                     np.asarray(pcm, dtype=float), geom_time, pcm_time_per_solute])
+
+    return pd.DataFrame(rows, columns=_NN_COLUMNS)

analysis/code/test_applications_analysis.py

@@ -0,0 +1,424 @@
+"""Tests for analysis/code/applications.py (the no-plotting analysis module).
+
+Builds a tiny synthetic applications.hdf5 + spreadsheet so the correction math can be checked
+in CI without the multi-hundred-megabyte release files.
+"""
+import os
+import sys
+
+import numpy as np
+import pandas as pd
+import h5py
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(HERE, "..", "..", "data", "applications"))
+sys.path.insert(0, HERE)
+import paths
+
+from applications_reader import Applications, SHELL_SIZES, SOLVENTS  # noqa: E402
+import applications as A  # noqa: E402
+
+SCALE = 1e4
+MISS = np.int32(-2147483648)
+
+
+def test_distribution_shift_by_solvent_table_averages_per_solvent():
+    """distribution_shift_by_solvent_table averages the three coefficient-choice RMSEs per solvent,
+    and surfaces the water column as the biggest extrapolation gap (the S15 story)."""
+    remapped = A.FORMULA_REMAP["stationary_plus_qcd + openMM"]
+    per_solute = {"H": pd.DataFrame([
+        {"formula": remapped, "solvent": "chloroform", "solute": "a", "rmse": 0.10},
+        {"formula": remapped, "solvent": "chloroform", "solute": "b", "rmse": 0.20},
+        {"formula": remapped, "solvent": "TIP4P", "solute": "a", "rmse": 0.30},
+        {"formula": remapped, "solvent": "TIP4P", "solute": "b", "rmse": 0.50},
+    ])}
+    per_solvent = {"H": pd.DataFrame([
+        {"formula": remapped, "solvent": "chloroform", "rmse": 0.16},
+        {"formula": remapped, "solvent": "TIP4P", "rmse": 0.45},
+    ])}
+    bootstrap = pd.DataFrame([
+        {"nucleus": "H", "formula": remapped, "solvent": "chloroform", "solute": "a", "seed": 0, "Bootstrap_RMSE": 0.14},
+        {"nucleus": "H", "formula": remapped, "solvent": "chloroform", "solute": "b", "seed": 0, "Bootstrap_RMSE": 0.18},
+        {"nucleus": "H", "formula": remapped, "solvent": "TIP4P", "solute": "a", "seed": 0, "Bootstrap_RMSE": 0.60},
+        {"nucleus": "H", "formula": remapped, "solvent": "TIP4P", "solute": "b", "seed": 0, "Bootstrap_RMSE": 0.80},
+    ])
+    table = A.distribution_shift_by_solvent_table(per_solute, per_solvent, bootstrap, "H",
+                                                  "stationary_plus_qcd + openMM")
+    assert list(table.index) == ["TIP4P", "chloroform"]  # sorted; uppercase sorts before lowercase
+    assert list(table.columns) == ["Extrapolated from delta22", "Scaled to Test Set", "Scaled to Solute"]
+    assert table.loc["chloroform", "Scaled to Solute"] == pytest.approx(0.15)
+    assert table.loc["chloroform", "Scaled to Test Set"] == pytest.approx(0.16)
+    assert table.loc["chloroform", "Extrapolated from delta22"] == pytest.approx(0.16)
+    assert table.loc["TIP4P", "Scaled to Solute"] == pytest.approx(0.40)
+    assert table.loc["TIP4P", "Extrapolated from delta22"] == pytest.approx(0.70)
+    chl_gap = table.loc["chloroform", "Extrapolated from delta22"] - table.loc["chloroform", "Scaled to Solute"]
+    water_gap = table.loc["TIP4P", "Extrapolated from delta22"] - table.loc["TIP4P", "Scaled to Solute"]
+    assert water_gap > chl_gap
+
+
+def _enc(a):
+    a = np.asarray(a, dtype=np.float64)
+    nan = np.isnan(a)
+    out = np.round(np.where(nan, 0.0, a) * SCALE).astype(np.int32)
+    out[nan] = MISS
+    return out
+
+
+@pytest.fixture
+def loader(tmp_path):
+    """One solute (vomicine), 3 atoms, one H site over atoms 1+2, one C site over atom 3.
+    magnet_x is set so the solvated-minus-isolated correction is a known constant per shell."""
+    h5 = tmp_path / "applications.hdf5"
+    xlsx = tmp_path / "exp.xlsx"
+    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], np.int32))
+        g = f.create_group("vomicine")
+        g.create_dataset("atomic_numbers", data=np.array([1, 1, 6], np.int32))
+        mx = g.create_group("magnet_x")
+        for sv in SOLVENTS:
+            svg = mx.create_group(sv)
+            for k, shell in enumerate(SHELL_SIZES):
+                arr = np.zeros((4, 3, 2), dtype=np.float64)
+                arr[:, :, 0] = 10.0                 # isolated
+                arr[:, :, 1] = 10.0 + 0.1 * (k + 1)  # solvated -> correction = 0.1*(k+1)
+                svg.create_dataset(f"shell_{shell}", data=_enc(arr))
+            svg.create_dataset("stationary", data=_enc(np.full(3, 10.0)))
+        mz = g.create_group("magnet_zero")
+        mz.create_dataset("stationary", data=_enc(np.array([100.0, 101.0, 50.0])))
+        mz.create_dataset("pcm_correction", data=_enc(np.array([-1.0, -1.0, -2.0])))
+        q = g.create_group("qcd")
+        q.create_dataset("stationary", data=_enc(np.zeros((3, 4))))
+        q.create_dataset("trajectories", data=_enc(np.zeros((2, 3, 3, 4))))  # corr -> 0
+        cm = f.create_group("composite_model")
+        cmc = cm.create_group("pcm_conversion_factors")
+        # the real conversion table uses 'water' (renamed to TIP4P by build_query_df_nn), never both
+        rows = "solvent,pcm_conversion_factor\n" + "".join(
+            f"{s},1.0\n" for s in ["chloroform", "benzene", "methanol", "water"])
+        cmc.create_dataset("H", data=rows)
+        cmc.create_dataset("C", data=rows)
+    with pd.ExcelWriter(xlsx) as w:
+        pd.DataFrame({
+            "site": ["H", "1 C"], "nucleus": ["H", "C"],
+            "atom_numbers": ["1,2", "3"],
+            "chloroform": [1.0, 2.0], "benzene": [1.0, 2.0],
+            "methanol": [1.0, 2.0], "water": [1.0, 2.0],
+        }).to_excel(w, sheet_name="vomicine", index=False)
+    return Applications(h5, xlsx)
+
+
+def test_shell_convergence_values(loader):
+    df = A.shell_convergence_corrections(loader, "vomicine", "H")
+    # one H site (uuid-prefixed) x 4 solvents
+    assert df.index.get_level_values("solvent").nunique() == len(SOLVENTS)
+    assert df.index.get_level_values("site").unique().tolist() == ["01 H"]
+    # correction for shell k (0-based) is exactly 0.1*(k+1) within the int32 quantum
+    row = df.xs("chloroform", level="solvent").iloc[0]
+    for k, shell in enumerate(SHELL_SIZES):
+        assert row[f"shell_{shell}"] == pytest.approx(0.1 * (k + 1), abs=5e-5)
+
+
+def test_only_requested_nucleus(loader):
+    df = A.shell_convergence_corrections(loader, "vomicine", "H")
+    # C site must not appear in an H-nucleus convergence table
+    sites = df.index.get_level_values("site").unique().tolist()
+    assert all("C" not in s for s in sites)
+
+
+def test_site_atom_indices():
+    assert A.site_atom_indices("1,2") == [0, 1]
+    assert A.site_atom_indices("21") == [20]
+
+
+def test_fit_recovers_known_line():
+    # experimental = 2*stationary + 5 exactly -> intercept ~5, slope ~2, RMSE ~0
+    df = pd.DataFrame({"stationary": [1.0, 2.0, 3.0, 4.0], "experimental": [7.0, 9.0, 11.0, 13.0]})
+    rmse, params = A.fit(df, "stationary")
+    assert rmse < 1e-9
+    assert params["Intercept"] == pytest.approx(5.0, abs=1e-6)
+    assert params["stationary"] == pytest.approx(2.0, abs=1e-6)
+
+
+def test_fit_ignores_rows_with_no_experimental_value():
+    # one site has no experimental measurement (common in the real natural-products spreadsheet);
+    # it must not turn the whole RMSE into NaN, only be left out of the score
+    df = pd.DataFrame({
+        "stationary": [1.0, 2.0, 3.0, 4.0],
+        "experimental": [7.0, 9.0, 11.0, np.nan],
+    })
+    rmse, params = A.fit(df, "stationary")
+    assert rmse < 1e-9
+    assert params["stationary"] == pytest.approx(2.0, abs=1e-6)
+
+
+def test_fit_matches_statsmodels():
+    # the numpy fit must be numerically identical to the original statsmodels OLS oracle, across
+    # multi-term formulas and with missing values present (rows dropped identically by both).
+    rng = np.random.default_rng(7)
+    n = 60
+    df = pd.DataFrame({
+        "stationary": rng.normal(100, 20, n), "pcm": rng.normal(0, 2, n),
+        "qcd": rng.normal(0, 1, n), "openMM": rng.normal(0, 2, n),
+    })
+    df["experimental"] = (2.0 * df["stationary"] + 0.5 * df["pcm"] - 0.3 * df["qcd"]
+                          + 1.0 + rng.normal(0, 1, n))
+    df.loc[[3, 17, 42], "experimental"] = np.nan   # missing experimental values
+    df.loc[[8, 25], "pcm"] = np.nan                # missing predictor values
+    for formula in ["stationary", "stationary + pcm", "stationary + pcm + qcd",
+                    "stationary + pcm + qcd + openMM"]:
+        rmse, params = A.fit(df, formula)
+        s_rmse, s_params = A._fit_statsmodels(df, formula)
+        assert rmse == pytest.approx(s_rmse, abs=1e-8), formula
+        for name in s_params.index:
+            assert params[name] == pytest.approx(s_params[name], abs=1e-6), f"{formula}:{name}"
+
+
+def test_build_query_df_nn(loader):
+    q = A.build_query_df_nn(loader)
+    expected = {"solute", "nucleus", "site", "solvent", "experimental", "stationary", "pcm",
+                "openMM", "openMM_vib", "qcd", "stationary_plus_pcm", "stationary_plus_qcd",
+                "stationary_plus_op_vib"}
+    assert expected.issubset(q.columns)
+    # H site covers atoms 1+2: stationary = mean(100, 101) = 100.5; pcm = mean(-1, -1) = -1;
+    # conversion factor 1.0 -> stationary_plus_pcm = 99.5
+    h = q[(q["nucleus"] == "H") & (q["solvent"] == "chloroform")].iloc[0]
+    assert h["stationary"] == pytest.approx(100.5, abs=1e-3)
+    assert h["stationary_plus_pcm"] == pytest.approx(99.5, abs=1e-3)
+    assert h["qcd"] == pytest.approx(0.0, abs=1e-3)
+
+
+def test_apply_bootstrap_nan_handling():
+    # a missing coeff (openMM=NaN) contributes nothing; a present coeff with missing data -> NaN
+    coeffs = pd.DataFrame([{"solvent": "chloroform", "formula": "pcm2", "seed": 0,
+                            "Intercept": 5.0, "stationary": 2.0, "openMM": np.nan}])
+    data = pd.DataFrame([
+        {"solute": "m", "nucleus": "H", "site": "a", "solvent": "chloroform",
+         "experimental": 7.0, "stationary": 1.0, "openMM": 9.0},                 # -> 5 + 2*1 = 7
+        {"solute": "m", "nucleus": "H", "site": "b", "solvent": "chloroform",
+         "experimental": np.nan, "stationary": np.nan, "openMM": 9.0},           # stationary data NaN -> NaN
+    ])
+    out = A.apply_bootstrap_params_to_full_dataset(data, coeffs, nucleus="H").set_index("site")
+    assert out.loc["a", "predicted"] == pytest.approx(7.0)
+    assert np.isnan(out.loc["b", "predicted"])
+
+
+def test_apply_bootstrap_and_solute_rmse():
+    # one formula/seed/solvent: predicted = 5 + 2*stationary; experimental matches -> RMSE 0
+    coeffs = pd.DataFrame({"solvent": ["chloroform"], "formula": ["pcm2"], "seed": [0],
+                           "Intercept": [5.0], "stationary": [2.0]})
+    data = pd.DataFrame({
+        "solute": ["m", "m"], "nucleus": ["H", "H"], "site": ["01 H", "02 H"],
+        "solvent": ["chloroform", "chloroform"], "experimental": [7.0, 9.0], "stationary": [1.0, 2.0],
+    })
+    preds = A.apply_bootstrap_params_to_full_dataset(data, coeffs, nucleus="H")
+    assert np.allclose(preds["predicted"], preds["experimental"])
+    rmses = A.compute_solute_rmses(preds)
+    assert rmses["Bootstrap_RMSE"].iloc[0] == pytest.approx(0.0, abs=1e-9)
+
+
+# --- SI Figure S15's "Fitting RMSE Comparisons" -----------------------------------------------
+
+def _rmse_comparison_fixture():
+    """Two test-set solutes ("m", "n"), one solvent, formula "stationary + openMM". Each solute has
+    only 2 points against 3 free parameters, so a per-solute-only fit is always exact (RMSE 0) --
+    that part needs no special data. The 4 points are deliberately NOT co-planar in (stationary,
+    openMM, experimental) space (verified: an earlier version of this fixture picked co-planar
+    points by accident, which made the pooled "Scaled to Test Set" fit ALSO exact and silently
+    hid the "formula" column bug this fixture now catches -- see the FORMULA_REMAP note below), so
+    the pooled fit has nonzero RMSE and "Scaled to Test Set" != "Scaled to Solute". The bootstrap
+    coefficients are also chosen far from both fits, so all three columns differ."""
+    query_df_nn = pd.DataFrame([
+        {"solute": "m", "nucleus": "H", "site": "01", "solvent": "chloroform", "experimental": 10.0, "stationary": 5.0, "openMM": 1.0},
+        {"solute": "m", "nucleus": "H", "site": "02", "solvent": "chloroform", "experimental": 12.0, "stationary": 6.0, "openMM": 1.0},
+        {"solute": "n", "nucleus": "H", "site": "01", "solvent": "chloroform", "experimental": 20.0, "stationary": 5.0, "openMM": 2.0},
+        {"solute": "n", "nucleus": "H", "site": "02", "solvent": "chloroform", "experimental": 23.0, "stationary": 7.0, "openMM": 2.0},
+    ])
+    formula = "stationary + openMM"
+    per_solute = {"H": A.fit_formulas_per_solvent_and_solute(query_df_nn, [formula], ["chloroform"],
+                                                             ["m", "n"], A.FORMULA_REMAP)}
+    all_solute = {"H": A.fit_formulas_per_solvent(query_df_nn, [formula], ["chloroform"], A.FORMULA_REMAP)}
+    # The real pipeline (applications.build_bootstrap_seed_coeffs) always stores the REMAPPED
+    # formula name in its "formula" column, same as per_solute_fits/per_solvent_fits above --
+    # fitting_rmse_comparison_table looks up bootstrap_rmses by the remapped name, so a coeffs
+    # frame built with the raw formula string would never match and "Extrapolated from delta22"
+    # would silently come out all-NaN instead of raising. Match that convention here.
+    remapped = A.FORMULA_REMAP[formula]
+    coeffs = pd.DataFrame({"solvent": ["chloroform"], "formula": [remapped], "seed": [0],
+                           "Intercept": [0.0], "stationary": [1.0], "openMM": [1.0]})
+    preds = A.apply_bootstrap_params_to_full_dataset(query_df_nn, coeffs, nucleus="H")
+    bootstrap_rmses = A.compute_solute_rmses(preds)
+    return query_df_nn, per_solute, all_solute, bootstrap_rmses, formula
+
+
+def test_scaled_to_test_set_per_solute_rmse_applies_one_pooled_fit_per_solute():
+    query_df_nn, _, all_solute, _, formula = _rmse_comparison_fixture()
+    out = A.scaled_to_test_set_per_solute_rmse(query_df_nn, all_solute, "H", "chloroform", formula)
+    assert set(out["solute"]) == {"m", "n"}
+    # both solutes score under the SAME pooled coefficients, so their RMSEs need not be equal
+    assert out["Bootstrap_RMSE"].notna().all()
+
+
+def test_fitting_rmse_comparison_table_has_all_three_columns_and_solutes():
+    query_df_nn, per_solute, all_solute, bootstrap_rmses, formula = _rmse_comparison_fixture()
+    table = A.fitting_rmse_comparison_table(query_df_nn, per_solute, all_solute, bootstrap_rmses,
+                                            "H", "chloroform", formula)
+    assert set(table.columns) == {"Scaled to Solute", "Scaled to Test Set", "Extrapolated from delta22"}
+    assert set(table.index) == {"m", "n"}
+    assert table.notna().all().all()   # a formula-remapping mismatch would leave a column all-NaN
+    # "Scaled to Solute" fits each solute alone with 2 free params on 2 points -> exact, RMSE 0
+    assert table["Scaled to Solute"].max() < 1e-8
+    # the fixture's points are deliberately non-coplanar (see _rmse_comparison_fixture), so the
+    # pooled fit is NOT exact and genuinely differs from the per-solute fit
+    assert table["Scaled to Test Set"].max() > 1e-3
+    # the bootstrap coefficients are far from either fit, so this column differs from both
+    assert (table["Extrapolated from delta22"] - table["Scaled to Solute"]).abs().min() > 0.1
+
+
+@pytest.mark.skipif(
+    not os.path.exists(paths.dataset_file("applications", file=__file__)),
+    reason="real applications.hdf5 not present")
+def test_fitting_rmse_comparison_table_reproduces_published_chloroform_values():
+    """Every published "Fitting RMSE Comparisons (chloroform, H)" value (SI Figure S15), read
+    directly off the SI, reproduced from the released data. Cross-checks the delta-22-bootstrap
+    ("Extrapolated") column at a slightly looser tolerance since it averages over random seeds."""
+    real_data_dir = os.path.join(HERE, "..", "..", "data", "applications")
+    loader = Applications(paths.dataset_file("applications", file=__file__),
+                          os.path.join(real_data_dir, "applications_experimental.xlsx"))
+    query_df_nn = A.build_query_df_nn(loader)
+    seed = A.build_bootstrap_seed_coeffs(loader)
+    per_solute = A.per_solute_fits(query_df_nn)
+    all_solute = A.per_solvent_fits(query_df_nn)
+    preds_h = A.apply_bootstrap_params_to_full_dataset(query_df_nn, seed["H"], nucleus="H")
+    bootstrap_rmses_h = A.compute_solute_rmses(preds_h)
+    table = A.fitting_rmse_comparison_table(query_df_nn, per_solute, all_solute, bootstrap_rmses_h,
+                                            "H", "chloroform", "stationary_plus_qcd + openMM")
+    # (solute, column) -> published value, read off the SI Figure S15 bar chart
+    published = {
+        ("isomer_1E", "Scaled to Solute"): 0.138, ("isomer_1E", "Scaled to Test Set"): 0.138,
+        ("vomicine", "Scaled to Solute"): 0.173, ("vomicine", "Scaled to Test Set"): 0.183,
+        ("prednisone", "Scaled to Solute"): 0.130, ("prednisone", "Scaled to Test Set"): 0.153,
+        ("flavone", "Scaled to Solute"): 0.021, ("flavone", "Scaled to Test Set"): 0.062,
+    }
+    for (solute, col), expected in published.items():
+        assert table.loc[solute, col] == pytest.approx(expected, abs=3e-3), f"{solute}/{col}"
+    published_extrapolated = {"isomer_1E": 0.146, "vomicine": 0.180, "prednisone": 0.166, "flavone": 0.060}
+    for solute, expected in published_extrapolated.items():
+        assert table.loc[solute, "Extrapolated from delta22"] == pytest.approx(expected, abs=0.02), solute
+
+
+@pytest.mark.skipif(
+    not os.path.exists(paths.dataset_file("applications", file=__file__)),
+    reason="real applications.hdf5 not present")
+def test_distribution_shift_by_solvent_table_flags_water_on_real_data():
+    """The per-solvent distribution-shift check on real data: for protons, TIP4P (water) must show
+    the largest 'Extrapolated from delta22' minus 'Scaled to Solute' gap -- the exact claim the
+    function's docstring and SI Figure S15's narrative make (the peptide dominates the water column).
+    Also spot-checks two published per-solvent RMSE values."""
+    real_data_dir = os.path.join(HERE, "..", "..", "data", "applications")
+    loader = Applications(paths.dataset_file("applications", file=__file__),
+                          os.path.join(real_data_dir, "applications_experimental.xlsx"))
+    query_df_nn = A.build_query_df_nn(loader)
+    seed = A.build_bootstrap_seed_coeffs(loader)
+    per_solute = A.per_solute_fits(query_df_nn)
+    all_solute = A.per_solvent_fits(query_df_nn)
+    bootstrap_h = A.compute_solute_rmses(
+        A.apply_bootstrap_params_to_full_dataset(query_df_nn, seed["H"], nucleus="H"))
+    table = A.distribution_shift_by_solvent_table(per_solute, all_solute, bootstrap_h, "H",
+                                                  "stationary_plus_qcd + openMM")
+    gap = table["Extrapolated from delta22"] - table["Scaled to Solute"]
+    assert gap.idxmax() == "TIP4P", f"expected water to have the largest extrapolation gap, got {gap.idxmax()}"
+    assert table.loc["TIP4P", "Extrapolated from delta22"] == pytest.approx(0.336, abs=0.01)
+    assert table.loc["chloroform", "Scaled to Solute"] == pytest.approx(0.107, abs=0.01)
+
+
+# --- SI Figure S15's "Feature Space Coverage by Solvent" / "Residuals for Delta22 Fitting Coefficients" ---
+
+def _feature_space_fixture():
+    """A test-set solute ("m") and a delta-22 solute, both nucleus H, in chloroform and benzene
+    (TIP4P/methanol deliberately absent from "m" to check the solvent filter drops rows cleanly).
+    stationary_plus_qcd/openMM are chosen so delta-22's chloroform plane is exact-fittable
+    (3 points, 3 free parameters -> zero residual)."""
+    query_df_nn = pd.DataFrame([
+        {"solute": "m", "nucleus": "H", "solvent": "chloroform", "stationary_plus_qcd": 10.0, "openMM": 1.0, "experimental": 12.0},
+        {"solute": "m", "nucleus": "H", "solvent": "benzene", "stationary_plus_qcd": 11.0, "openMM": 2.0, "experimental": 15.0},
+        {"solute": "m", "nucleus": "C", "solvent": "chloroform", "stationary_plus_op_vib": 100.0, "openMM": 3.0, "experimental": 120.0},
+    ])
+    delta22_query_df_nn = pd.DataFrame([
+        {"solute": "delta22", "nucleus": "H", "solvent": "chloroform", "stationary_plus_qcd": 12.0, "openMM": 0.0, "experimental": 12.0},
+        {"solute": "delta22", "nucleus": "H", "solvent": "chloroform", "stationary_plus_qcd": 14.0, "openMM": 1.0, "experimental": 15.0},
+        {"solute": "delta22", "nucleus": "H", "solvent": "chloroform", "stationary_plus_qcd": 16.0, "openMM": 2.0, "experimental": 18.0},
+        {"solute": "delta22", "nucleus": "H", "solvent": "benzene", "stationary_plus_qcd": 9.0, "openMM": 1.5, "experimental": 13.0},
+    ])
+    return query_df_nn, delta22_query_df_nn
+
+
+def test_feature_space_coverage_table_combines_and_centers_both_datasets():
+    query_df_nn, delta22_query_df_nn = _feature_space_fixture()
+    out = A.feature_space_coverage_table(query_df_nn, delta22_query_df_nn, "H", solvents=("chloroform", "benzene"))
+    assert set(out["dataset"]) == {"Test Set", "Delta22"}
+    assert set(out["solvent"]) == {"chloroform", "benzene"}
+    # global centering: the pooled x column (both datasets, both solvents) averages to 0
+    assert out["x"].mean() == pytest.approx(0.0, abs=1e-9)
+    assert out["y"].mean() == pytest.approx(0.0, abs=1e-9)
+
+
+def test_feature_space_coverage_table_drops_solvents_test_set_lacks():
+    query_df_nn, delta22_query_df_nn = _feature_space_fixture()
+    out = A.feature_space_coverage_table(query_df_nn, delta22_query_df_nn, "H",
+                                         solvents=("chloroform", "benzene", "methanol", "TIP4P"))
+    # "m" (the only test-set H solute) has no methanol/TIP4P rows, so those solvents drop out
+    assert set(out["solvent"]) == {"chloroform", "benzene"}
+
+
+def test_delta22_plane_residuals_table_delta22_residuals_are_near_zero():
+    # delta-22's own 3 chloroform points exactly determine a 3-parameter plane, so delta-22's own
+    # residuals against that plane must be ~0 (an OLS fit's residuals average to 0 with an
+    # intercept; with exactly as many points as parameters, every residual is individually ~0)
+    query_df_nn, delta22_query_df_nn = _feature_space_fixture()
+    out = A.delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, "H", solvents=("chloroform",))
+    d22 = out[out["dataset"] == "Delta22"]
+    assert len(d22) == 3
+    assert np.allclose(d22["residual"], 0.0, atol=1e-6)
+    # the test-set row scores against that SAME plane and is not forced to be exact
+    test_row = out[out["dataset"] == "Test Set"]
+    assert len(test_row) == 1
+
+
+def test_delta22_plane_residuals_table_skips_solvents_with_too_few_delta22_points():
+    # benzene has only 1 delta-22 point (< 3 needed to fit a plane) -> dropped entirely, no crash
+    query_df_nn, delta22_query_df_nn = _feature_space_fixture()
+    out = A.delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, "H", solvents=("chloroform", "benzene"))
+    assert set(out["solvent"]) == {"chloroform"}
+
+
+@pytest.mark.skipif(
+    not os.path.exists(paths.dataset_file("applications", file=__file__)),
+    reason="real applications.hdf5 not present")
+def test_feature_space_and_residuals_tables_sane_on_real_data():
+    """Not an exact-number reproduction (these are scatter plots, not summary statistics) --
+    structural sanity checks: delta-22's own residuals against its own fitted plane average to
+    (near) zero every solvent, and the known H/TIP4P data-scarcity caveat from the SI text ("the
+    available experimental data for water in the test set is limited to one solute") shows up as a
+    small test-set count for that one solvent."""
+    real_data_dir = os.path.join(HERE, "..", "..", "data", "applications")
+    delta22_dir = os.path.join(HERE, "..", "..", "data", "delta22")
+    sys.path.insert(0, delta22_dir)
+    import delta22 as D
+    loader = Applications(paths.dataset_file("applications", file=__file__),
+                          os.path.join(real_data_dir, "applications_experimental.xlsx"))
+    query_df_nn = A.build_query_df_nn(loader)
+    delta22_query_df_nn = D.add_composite_columns(D.load_query_df_nn(
+        paths.dataset_file("delta22", file=__file__), os.path.join(delta22_dir, "delta22_experimental.xlsx"),
+        verbose=False))
+    for nucleus in ["H", "C"]:
+        residuals = A.delta22_plane_residuals_table(query_df_nn, delta22_query_df_nn, nucleus)
+        d22_means = residuals[residuals["dataset"] == "Delta22"].groupby("solvent")["residual"].mean()
+        assert (d22_means.abs() < 1e-6).all()
+        test_counts = residuals[residuals["dataset"] == "Test Set"].groupby("solvent").size()
+        assert test_counts["chloroform"] > 50   # plenty of chloroform test-set data
+        # SI-documented data scarcity: "the available experimental data for water in the test set
+        # is limited to one solute" -- the residuals table (which requires an experimental value,
+        # unlike the feature-space-coverage table) sees far fewer TIP4P test-set rows than chloroform
+        assert 0 < test_counts.get("TIP4P", 0) < test_counts["chloroform"] / 5

analysis/code/test_carbon_monoxide.py

@@ -0,0 +1,41 @@
+"""Tests for analysis/code/carbon_monoxide.py. The figure's claim is that MagNET tracks
+DFT near the equilibrium bond length but drifts in the non-equilibrium tails; that is checked against
+the shipped scan data, along with the loaders and the histogram helper."""
+import os
+import sys
+
+import numpy as np
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+
+import carbon_monoxide as co  # noqa: E402
+
+
+def test_comparison_loads_and_is_sorted():
+    df = co.load_comparison()
+    assert len(df) == 81
+    assert list(df.columns) == ["bond_length", "gaussian_c", "magnet_c", "delta_c",
+                                "gaussian_o", "magnet_o", "delta_o"]
+    assert df["bond_length"].is_monotonic_increasing
+    assert df["bond_length"].min() == pytest.approx(0.70, abs=1e-3)
+    assert df["bond_length"].max() == pytest.approx(1.50, abs=1e-3)
+    # delta is exactly MagNET minus DFT
+    assert np.allclose(df["delta_c"], df["magnet_c"] - df["gaussian_c"])
+
+
+def test_histogram_helper():
+    centers, counts, bin_width = co.load_bond_length_histogram()
+    assert len(centers) == len(counts) == 100
+    assert bin_width > 0
+    assert counts.sum() > 0           # there is real bond-length support behind the error panel
+
+
+def test_error_grows_away_from_equilibrium():
+    """The published claim: the carbon error is small near the equilibrium bond length and much
+    larger in the stretched/compressed tails the model never trained on."""
+    near, tails = co.equilibrium_vs_extreme_error()
+    assert near < 30.0                # near-equilibrium error stays modest (about 19 ppm)
+    assert tails > 100.0              # tails blow up (about 185 ppm)
+    assert tails > 5 * near           # the tails are far worse, which is the point of the figure

analysis/code/test_composite_model.py

@@ -0,0 +1,124 @@
+"""Tests for build_composite_model.py: the composite_model/ group in applications.hdf5 reproduces
+from the shipped delta-22 data, so the natural-products coefficients need no private tree or
+environment variables. Skipped when the large hdf5 files are absent."""
+import os
+import random
+import sys
+
+import numpy as np
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+REPO = os.path.abspath(os.path.join(HERE, "..", ".."))
+for _p in ("analysis/code", "analysis/code/shared", "data/delta22", "data/applications"):
+    sys.path.insert(0, os.path.join(REPO, _p))
+
+import build_composite_model as B  # noqa: E402
+import paths  # noqa: E402
+
+DELTA22 = paths.dataset_file("delta22", root=REPO)
+APPLICATIONS = paths.dataset_file("applications", root=REPO)
+_HAVE_DATA = os.path.exists(DELTA22) and os.path.exists(APPLICATIONS)
+
+# The fixed-point integer encoding (int32 x 1e4) sets the reproduction floor; recipe errors are ~1.
+ENCODING_TOL = 5e-3
+
+
+def _reader():
+    from applications_reader import Applications
+    return Applications(APPLICATIONS)
+
+
+@pytest.mark.skipif(not _HAVE_DATA, reason="delta22.hdf5 / applications.hdf5 not present")
+def test_ols_and_pcm_reproduce():
+    """Every OLS formula (all solvents, both nuclei) and both PCM-conversion tables reproduce."""
+    reader = _reader()
+    full_tab, _fit_tab, _apply_tab, factors = B._load_tables()
+    worst = 0.0
+    for formula in reader.composite_formulas("ols_coefficients"):
+        regressors = B._regressor_columns(formula)
+        for reg in regressors:
+            B._ensure_column(full_tab, reg)
+        for nucleus in ("H", "C"):
+            stored = reader.ols_coefficients(formula, nucleus).set_index("parameter")
+            for solvent in stored.columns:
+                sub = full_tab[(full_tab["nucleus"] == nucleus) & (full_tab["solvent"] == solvent)]
+                got = B._decompose(B._fit(sub, regressors), regressors, solvent, nucleus, factors)
+                worst = max(worst, max(abs(float(stored.loc[p, solvent]) - got[p]) for p in stored.index))
+    assert worst < ENCODING_TOL, f"OLS max diff {worst}"
+
+    for nucleus in ("H", "C"):
+        stored = reader.pcm_conversion_factors(nucleus).set_index("solvent")["pcm_conversion_factor"]
+        got = B._pcm_table(factors, nucleus).set_index("solvent")["pcm_conversion_factor"]
+        assert np.abs(got - stored.reindex(got.index)).max() < 1e-9
+
+
+@pytest.mark.skipif(not _HAVE_DATA, reason="delta22.hdf5 / applications.hdf5 not present")
+def test_bootstrap_and_rmse_reproduce():
+    """Bootstrap coefficients and RMSE distributions reproduce on a sample of seeds (the fixed-seed
+    resampler makes every seed deterministic, so a sample is representative). This is the fast CI
+    check; build_composite_model.main() re-fits all 1000 seeds and is the authoritative full pass."""
+    reader = _reader()
+    _full_tab, fit_tab, apply_tab, factors = B._load_tables()
+    seeds = range(0, 1000, 97)
+
+    worst = 0.0
+    for formula in reader.composite_formulas("bootstrap_coefficients"):
+        regressors = B._regressor_columns(formula)
+        for reg in regressors:
+            B._ensure_column(fit_tab, reg)
+        for nucleus in ("H", "C"):
+            try:
+                stored = reader.bootstrap_coefficients(formula, nucleus)
+            except KeyError:
+                continue
+            params = [c for c in stored.columns if c not in ("solvent", "seed")]
+            for solvent in stored["solvent"].unique():
+                for seed in seeds:
+                    row = stored[(stored["solvent"] == solvent) & (stored["seed"] == seed)]
+                    if not len(row):
+                        continue
+                    got = B._decompose(B._fit(B._bootstrap_sample(fit_tab, nucleus, solvent, seed),
+                                              regressors), regressors, solvent, nucleus, factors)
+                    for p in params:
+                        key = "intercept" if p == "Intercept" else p
+                        worst = max(worst, abs(float(row[p].iloc[0]) - got[key]))
+    assert worst < ENCODING_TOL, f"bootstrap max diff {worst}"
+
+    worst = 0.0
+    for nucleus in ("H", "C"):
+        stored = reader.rmse_distribution(nucleus)
+        for formula in stored["formula"].unique():
+            sub = stored[stored["formula"] == formula]
+            for solvent in sub["solvent"].unique():
+                for seed in seeds:
+                    row = sub[(sub["solvent"] == solvent) & (sub["seed"] == seed)]
+                    if not len(row):
+                        continue
+                    got = B._rmse(fit_tab, apply_tab, formula, nucleus, solvent, seed, factors)
+                    worst = max(worst, abs(float(row["Bootstrap_RMSE"].iloc[0]) - got))
+    assert worst < ENCODING_TOL, f"rmse max diff {worst}"
+
+
+@pytest.mark.skipif(not _HAVE_DATA, reason="delta22.hdf5 / applications.hdf5 not present")
+def test_bootstrap_is_deterministic():
+    """The resampler reseeds internally, so a given seed draws the same solutes no matter the
+    external RNG state going in."""
+    fit_tab = B._load_tables()[1]
+    random.seed(123)
+    a = B._bootstrap_sample(fit_tab, "C", "chloroform", 7)
+    random.seed(456)
+    b = B._bootstrap_sample(fit_tab, "C", "chloroform", 7)
+    assert list(a["solute"]) == list(b["solute"])
+
+
+@pytest.mark.skipif(not _HAVE_DATA, reason="delta22.hdf5 / applications.hdf5 not present")
+def test_write_roundtrips(tmp_path):
+    """write() serializes the regenerated tables into an hdf5 file that reads back unchanged through
+    Applications, so --write cannot silently corrupt the shipped file's composite_model group."""
+    from applications_reader import Applications
+    regenerated = B.regenerate(_reader())
+    scratch = str(tmp_path / "cm.hdf5")
+    B.write(scratch, regenerated)
+    worst = B.verify(Applications(scratch), regenerated)
+    assert worst < 1e-6, f"write/read round-trip diff {worst}"

analysis/code/test_composite_models.py

@@ -0,0 +1,178 @@
+"""Tests for analysis/code/composite_models.py.
+
+The fitting side is exercised on a small synthetic delta-22-shaped query table (composite_models
+adds no new fitting logic of its own -- it calls delta22.py's harness directly, which has its own
+tests). What's specific to this module is the pivot/reindex shape and the openpyxl table-writing
+layer, so those get direct, hermetic tests here: the formula lookup helpers, the %benefit formula
+writer, and the full workbook structure (sheet names, header cells, RMSE cell placement) on
+synthetic data. An opt-in test checks the real data reproduces a plausible RMSE range.
+"""
+import os
+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 composite_models as CM  # noqa: E402
+
+
+def _synthetic_query_df_dft(solutes, lots, seed=0):
+    """A query_df_dft-shaped table with all the columns composite_models' formulas reference, at
+    each of the given (method, basis, geometry) triples, for both nuclei."""
+    rng = np.random.default_rng(seed)
+    rows = []
+    for solute in solutes:
+        for solvent in CM.ORDERED_SOLVENTS:
+            for nucleus in ["H", "C"]:
+                for method, basis, geometry in lots:
+                    stationary = rng.normal(30 if nucleus == "H" else 150, 5)
+                    pcm = rng.normal(0, 0.5)
+                    desmond = rng.normal(0, 0.5)
+                    qcd = rng.normal(0, 0.2)
+                    desmond_vib = rng.normal(0, 0.3)
+                    openMM = rng.normal(0, 0.5)
+                    openMM_vib = rng.normal(0, 0.3)
+                    experimental = stationary + 0.5 * pcm + 0.3 * desmond + rng.normal(0, 0.05)
+                    rows.append(dict(
+                        solute=solute, solvent=solvent, nucleus=nucleus,
+                        sap_nmr_method=method, sap_basis=basis, sap_geometry_type=geometry,
+                        stationary=stationary, pcm=pcm, desmond=desmond, qcd=qcd,
+                        desmond_vib=desmond_vib, openMM=openMM, openMM_vib=openMM_vib,
+                        experimental=experimental))
+    return pd.DataFrame(rows)
+
+
+SOLUTES = [f"m{i}" for i in range(16)]
+DSD_LOT = ("dsd_pbep86", "pcSseg3", "pbe0_tz")
+WP04_LOT = ("wp04", "pcSseg2", "aimnet2")
+WB97XD_LOT = ("wb97xd", "pcSseg2", "aimnet2")
+
+
+def test_canonical_formula_name_ignores_case_spaces_underscores():
+    assert CM._canonical_formula_name("stationary + PCM") == CM._canonical_formula_name("stationary+pcm")
+    assert CM._canonical_formula_name("stationary_plus_pcm") != CM._canonical_formula_name("stationary + pcm")
+
+
+def test_get_rmse_looks_up_by_canonical_name():
+    df = pd.DataFrame({"chloroform": [0.123]}, index=["stationary + pcm"])
+    lookup = CM._build_formula_lookup(df)
+    assert CM._get_rmse(df, lookup, "stationary+PCM", "chloroform") == pytest.approx(0.123)
+
+
+def test_get_rmse_missing_formula_raises_keyerror():
+    df = pd.DataFrame({"chloroform": [0.1]}, index=["stationary"])
+    lookup = CM._build_formula_lookup(df)
+    with pytest.raises(KeyError):
+        CM._get_rmse(df, lookup, "not_a_real_formula", "chloroform")
+
+
+def test_ablation_rmse_table_shape_and_exact_recovery():
+    # experimental = stationary + 0.5*pcm + 0.3*desmond exactly (no noise), so a formula containing
+    # both pcm and qcd should fit near-perfectly (tiny test RMSE) while "stationary" alone
+    # (missing both real predictors) should not. "stationary + pcm + qcd" is a real formula in
+    # CM.FORMULAS; there is no formula combining pcm and desmond together, so the ground truth is
+    # built from pcm and qcd instead.
+    rng = np.random.default_rng(1)
+    rows = []
+    for solute in SOLUTES:
+        for solvent in CM.ORDERED_SOLVENTS:
+            stationary = rng.normal(30, 5)
+            pcm = rng.normal(0, 0.5)
+            desmond = rng.normal(0, 0.5)
+            qcd = rng.normal(0, 0.2)
+            desmond_vib = rng.normal(0, 0.3)
+            openMM = rng.normal(0, 0.5)
+            openMM_vib = rng.normal(0, 0.3)
+            experimental = stationary + 0.5 * pcm + 0.3 * qcd
+            rows.append(dict(solute=solute, solvent=solvent, nucleus="H",
+                             sap_nmr_method="dsd_pbep86", sap_basis="pcSseg3",
+                             sap_geometry_type="pbe0_tz", stationary=stationary, pcm=pcm,
+                             desmond=desmond, qcd=qcd, desmond_vib=desmond_vib, openMM=openMM,
+                             openMM_vib=openMM_vib, experimental=experimental))
+    query_df_dft = pd.DataFrame(rows)
+
+    table = CM.ablation_rmse_table(query_df_dft, "H", "dsd_pbep86", "pcSseg3", "pbe0_tz", SOLUTES,
+                                   n_splits=10)
+    assert list(table.index) == CM.FORMULAS
+    assert list(table.columns) == CM.ORDERED_SOLVENTS + ["Mean Test RMSE"]
+    assert table.loc["stationary + pcm + qcd", "Mean Test RMSE"] < 0.05
+    assert table.loc["stationary", "Mean Test RMSE"] > table.loc["stationary + pcm + qcd", "Mean Test RMSE"]
+
+
+def test_write_nucleus_table_cell_placement():
+    query_df_dft = _synthetic_query_df_dft(SOLUTES, [DSD_LOT])
+    table = CM.ablation_rmse_table(query_df_dft, "H", *DSD_LOT, SOLUTES, n_splits=5)
+
+    from openpyxl import Workbook
+    wb = Workbook()
+    ws = wb.active
+    next_row = CM.write_nucleus_table(
+        ws=ws, start_row=1, nucleus_name="Proton", stationary_desc="DSD-PBEP86/pcSseg3",
+        pcm_desc="B3LYP-D3BJ/pcSseg3", test_rmse_df=table, vmin=0.06, vmax=0.28,
+        formula_config=CM.PROTON_FORMULA_CONFIG)
+
+    assert ws["A1"].value == "Proton"
+    assert ws["A6"].value == "formula"
+    assert ws["A9"].value == "stationary"
+    assert ws["B9"].value == pytest.approx(table.loc["stationary", "chloroform"])
+    # the "Mean Test RMSE" column is the last of the 13 data columns (B..N)
+    assert ws["N5"].value == "Mean Test RMSE"
+    assert ws["N9"].value == pytest.approx(table.loc["stationary", "Mean Test RMSE"])
+    # %benefit rows carry a live Excel formula, not a precomputed number
+    assert str(ws["B14"].value).startswith("=IFERROR(")
+    assert next_row == 1 + 45
+
+
+def test_build_ablations_workbook_single_level_two_sheets():
+    query_df_dft = _synthetic_query_df_dft(SOLUTES, [DSD_LOT])
+    import tempfile
+    with tempfile.TemporaryDirectory() as td:
+        out = os.path.join(td, "ablations.xlsx")
+        CM.build_ablations_workbook(query_df_dft, SOLUTES, out, reference_levels=("dsd",),
+                                    n_splits=5, verbose=False)
+        from openpyxl import load_workbook
+        wb = load_workbook(out)
+        assert wb.sheetnames == ["Proton", "Carbon"]
+
+
+def test_build_ablations_workbook_both_levels_four_sheets_matching_shipped_names():
+    query_df_dft = _synthetic_query_df_dft(SOLUTES, [DSD_LOT, WP04_LOT, WB97XD_LOT])
+    import tempfile
+    with tempfile.TemporaryDirectory() as td:
+        out = os.path.join(td, "ablations.xlsx")
+        CM.build_ablations_workbook(query_df_dft, SOLUTES, out, n_splits=5, verbose=False)
+        from openpyxl import load_workbook
+        wb = load_workbook(out)
+        # matches wolford/ablations_updated 2.xlsx's exact sheet names
+        assert wb.sheetnames == ["Proton (DSD)", "Carbon (DSD)", "Proton (WP04)", "Carbon (wB97XD)"]
+
+
+# --- opt-in: reproduce a plausible ablation table from the real delta-22 data ---------------------
+
+REPO = os.path.abspath(os.path.join(HERE, "..", ".."))
+REAL_H5 = paths.dataset_file("delta22", root=REPO)
+REAL_XLSX = os.path.join(REPO, "data", "delta22", "delta22_experimental.xlsx")
+
+
+@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_ablation_rmse_table_reproduces_plausible_range_from_real_data():
+    import delta22 as D
+    query_df_dft = D.add_composite_columns(D.load_query_df_dft(REAL_H5, REAL_XLSX, verbose=False))
+    solutes = D.delta22_solutes(REAL_H5)
+    level = CM.REFERENCE_LEVELS["dsd"]
+    table = CM.ablation_rmse_table(query_df_dft, "H", level["method_h"], level["basis_h"],
+                                   level["geometry_h"], solutes, n_splits=20)
+    vmin, vmax = CM._VMIN_VMAX["H"]
+    # the shipped workbook's fixed color-scale range is itself a real-data-derived sanity bound:
+    # every formula's mean test RMSE should land inside it
+    assert (table["Mean Test RMSE"] >= vmin * 0.5).all()
+    assert (table["Mean Test RMSE"] <= vmax * 1.5).all()
+    # the fullest model should not be worse than the bare stationary reference
+    assert table.loc["stationary + desmond + qcd", "Mean Test RMSE"] < table.loc["stationary", "Mean Test RMSE"]

analysis/code/test_delta22_analysis.py

@@ -0,0 +1,916 @@
+"""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}"
+
+
+# --- SI Figure S15'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_si_s15_feature_correlations_reproduces_published_matrices():
+    """Every entry of both nuclei's published Pearson R and R^2 "Correlations Between Features"
+    matrices (SI Figure S15), 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.si_s15_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}"

analysis/code/test_dft8k_residuals.py

@@ -0,0 +1,137 @@
+"""Tests for analysis/code/dft8k_residuals.py.
+
+The synthetic tests check the residual selection and statistics on a hand-built shieldings dict.
+One opt-in test runs against the real dft8k.hdf5 and asserts the published DFT8K residual numbers
+(MagNET-Zero is a useful surrogate for DFT): 1H RMSE ~0.11 ppm, 13C RMSE ~0.81 ppm, with >95% of
+proton sites below 0.1 ppm error.
+"""
+import os
+import sys
+
+import numpy as np
+import pytest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, os.path.join(HERE, "..", "..", "data", "dft8k"))  # dft8k_reader
+sys.path.insert(0, HERE)
+import paths
+
+import dft8k_residuals as D  # noqa: E402
+
+
+def make_shieldings():
+    """A flat shieldings dict like the reader's all_shieldings: H, C, C, H, N. MagNET predicts H
+    with the WP04 model and C with the wB97X-D model; other atoms are NaN. The 1H residuals are
+    -0.1 and +0.2; the 13C residuals are 0.1 and 0.1."""
+    return {
+        "atomic_numbers": np.array([1, 6, 6, 1, 7]),
+        "wp04_pcSseg2": np.array([30.5, 150.5, 151.5, 31.5, 200.5]),
+        "wb97xd_pcSseg2": np.array([30.7, 150.7, 151.7, 31.7, 200.7]),
+        "nn_wp04_pcSseg2": np.array([30.6, np.nan, np.nan, 31.3, np.nan]),
+        "nn_wb97xd_pcSseg2": np.array([np.nan, 150.6, 151.6, np.nan, np.nan]),
+    }
+
+
+def test_residuals_select_nucleus_and_drop_unpredicted():
+    s = make_shieldings()
+    h = D.residuals(s, "H")
+    # only the two H atoms (indices 0, 3), residual = DFT - MagNET
+    np.testing.assert_allclose(sorted(h), sorted([30.5 - 30.6, 31.5 - 31.3]))
+    c = D.residuals(s, "C")
+    np.testing.assert_allclose(sorted(c), sorted([150.7 - 150.6, 151.7 - 151.6]))
+    # the nitrogen atom is in neither nucleus' residuals
+    assert len(h) == 2 and len(c) == 2
+
+
+def test_residual_stats():
+    errors = np.array([-0.1, 0.2, 0.05, -0.05])
+    st = D.residual_stats(errors, small_threshold=0.1)
+    assert st["n"] == 4
+    assert st["rmse"] == pytest.approx(np.sqrt(np.mean(errors ** 2)))
+    assert st["mae"] == pytest.approx(np.mean(np.abs(errors)))
+    # two of four are strictly below 0.1 ppm absolute (0.05 and 0.05); 0.1 and 0.2 are not
+    assert st["frac_below"] == pytest.approx(0.5)
+
+
+def test_residual_stats_empty():
+    st = D.residual_stats(np.array([]))
+    assert st["n"] == 0 and np.isnan(st["rmse"])
+
+
+# --- opt-in: reproduce the published DFT8K residual numbers from the released data --------------
+
+REAL = paths.dataset_file("dft8k", file=__file__)
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_reproduces_published_dft8k_residuals():
+    summary = D.residual_summary(REAL)
+    # 1H: RMSE ~0.11 ppm, >95% of sites below 0.1 ppm error (paper Figure 5 / DFT8K text)
+    assert summary["H"]["rmse"] == pytest.approx(0.105, abs=0.01)
+    assert summary["H"]["mae"] == pytest.approx(0.035, abs=0.01)
+    assert summary["H"]["frac_below"] >= 0.95
+    # 13C: RMSE ~0.81 ppm, MAE ~0.41 ppm
+    assert summary["C"]["rmse"] == pytest.approx(0.807, abs=0.02)
+    assert summary["C"]["mae"] == pytest.approx(0.409, abs=0.02)
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_extreme_residual_reproduces_the_published_porphyrin_callout():
+    # Figure 5B's positive-tail outlier callout ("Porphyrin Rings", error +17.620 ppm) is the true
+    # global maximum 1H residual, so this reproduces it exactly.
+    extreme = D.find_extreme_residual(REAL, "H", sign="max")
+    assert extreme["residual"] == pytest.approx(17.6197, abs=1e-3)
+    assert extreme["molecule_id"] == 20170716
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_extreme_residual_min_is_not_the_published_zwitterion_callout():
+    # The negative-tail callout ("Sulfonium Zwitterion", -1.040 ppm) is NOT the true global
+    # minimum -- a different, chemically unremarkable molecule is more negative. This test locks
+    # down that documented discrepancy so it is not silently "fixed" by a future data update.
+    extreme = D.find_extreme_residual(REAL, "H", sign="min")
+    assert extreme["residual"] < -1.5
+    assert extreme["molecule_id"] != 88779
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_molecule_by_id_reproduces_the_published_zwitterion_callout():
+    zwitterion = D.molecule_by_id(REAL, "H", molecule_id=88779)
+    assert zwitterion["residual"] == pytest.approx(-1.0403, abs=1e-3)
+    assert "[S+]" in zwitterion["smiles"] and "[O-]" in zwitterion["smiles"]
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_functional_group_errors_reproduces_published_values():
+    # published (Figure 5B, mean |1H residual| per group): Carbonyls 0.034, Amines 0.036,
+    # Sulfonyl 0.039, Pyridines 0.042, Furans 0.034, Nitroso 0.088 ppm.
+    published = {"Carbonyls": 0.034, "Amines": 0.036, "Sulfonyl": 0.039,
+                 "Pyridines": 0.042, "Furans": 0.034, "Nitroso": 0.088}
+    group_errors = D.functional_group_errors(REAL, "H")
+    for name, expected in published.items():
+        assert group_errors[name]["mean_abs_error"] == pytest.approx(expected, abs=2e-3), name
+        assert group_errors[name]["n_molecules"] > 0
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_functional_group_errors_respects_an_explicit_empty_patterns_dict():
+    # patterns=None means "use the default set"; patterns={} is a real, distinct request for "no
+    # groups" and must return {} rather than silently falling back to the default (an `x or
+    # DEFAULT` check would wrongly treat an empty dict as falsy and do exactly that).
+    assert D.functional_group_errors(REAL, "H", patterns={}) == {}
+
+
+def test_find_extreme_zwitterion_residual_requires_a_valid_sign():
+    with pytest.raises(ValueError):
+        D.find_extreme_zwitterion_residual(REAL, "H", sign="sideways")
+
+
+@pytest.mark.skipif(not os.path.exists(REAL), reason="real dft8k.hdf5 not present")
+def test_find_extreme_zwitterion_residual_finds_a_real_sulfonium_zwitterion():
+    # not asserted to match the published callout exactly (see molecule_by_id's docstring: the
+    # published example is a hand-picked one, not the strict extremum among zwitterions) -- this
+    # just checks the search itself works and returns a genuine sulfonium zwitterion.
+    result = D.find_extreme_zwitterion_residual(REAL, "H", sign="min")
+    assert result is not None
+    assert "[S+]" in result["smiles"] and "[O-]" in result["smiles"]
+    assert result["residual"] < 0

analysis/code/test_leveling.py

@@ -0,0 +1,233 @@
+"""
+Tests for the leveling-effect analysis core (leveling.py).
+
+Most tests build tiny synthetic inputs in the exact on-disk formats (so the fixtures
+double as a spec of what the loaders expect) and need no large data. Two opt-in tests
+run against the real delta22.hdf5 and the Kaupp NS372 spreadsheet if they are present.
+
+Run:  pytest test_leveling.py -q
+Requires: pytest, numpy, h5py, pandas, openpyxl, scikit-learn.
+"""
+import os
+import numpy as np
+import h5py
+import openpyxl
+import pytest
+
+from fixed_point import MISSING_MARKER as _MISSING_MARKER
+from leveling import (
+    _decode_fixed_point,
+    analyze_nucleus,
+    read_nucleus,
+    load_delta22,
+    load_ns372,
+    load_ns372_1h_convergence,
+    convergence_statistics,
+    DELTA22_REF,
+)
+
+SCALE = 1e4
+
+
+def _enc(arr):
+    """Encode floats the way the real file does: int32 = round(value * 1e4), NaN -> marker."""
+    a = np.asarray(arr, dtype=np.float64)
+    mask = np.isnan(a)
+    out = np.zeros(a.shape, dtype=np.int32)
+    out[~mask] = np.rint(a[~mask] * SCALE).astype(np.int32)
+    out[mask] = _MISSING_MARKER
+    return out
+
+
+# ---------------------------------------------------------------------------
+# fixed-point decoding
+# ---------------------------------------------------------------------------
+def test_decode_round_trip_and_marker():
+    values = np.array([0.0, 207.07, -3.5, np.nan, 12345.6789], dtype=np.float64)
+    decoded = _decode_fixed_point(_enc(values))
+    assert np.isnan(decoded[3])
+    np.testing.assert_allclose(decoded[~np.isnan(decoded)],
+                               values[~np.isnan(values)], atol=1e-4)
+
+
+def test_decode_passes_through_float():
+    # an older plain-decimal copy must read 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_fixed_point(values), values)
+
+
+def test_decode_handles_int64_and_no_marker():
+    # a re-encode at int64 must decode like int32 (any integer dtype), and an
+    # integer array with no missing marker must round-trip with no spurious blanks
+    vals = np.array([10000, 20000, -35000], dtype=np.int64)   # = 1.0, 2.0, -3.5
+    out = _decode_fixed_point(vals)
+    np.testing.assert_allclose(out, [1.0, 2.0, -3.5])
+    assert not np.isnan(out).any()
+
+
+# ---------------------------------------------------------------------------
+# leveling analysis on a synthetic, perfectly-leveled matrix
+# ---------------------------------------------------------------------------
+def test_analyze_nucleus_detects_leveling():
+    # build methods that are affine transforms of one base signal plus tiny noise:
+    # after scaling they are interchangeable, so PC1 should dominate and every
+    # pairwise correlation should be near 1.
+    rng = np.random.RandomState(0)
+    base = rng.normal(size=60)
+    slopes = np.linspace(0.8, 1.3, 8)
+    offsets = np.linspace(-5, 5, 8)
+    cols = [a * base + b + 0.01 * rng.normal(size=60)
+            for a, b in zip(slopes, offsets)]
+    M = np.column_stack(cols)
+    methods = [f"m{i}" for i in range(8)]
+    res = analyze_nucleus(M, methods, ref=M[:, 0])
+    assert res['ev'][0] > 0.99            # PC1 carries essentially all variance
+    assert res['r_min'] > 0.99            # worst pairwise correlation still near 1
+    assert set(res['scaled_rmse']) == set(methods)
+    assert 'r2' in res['par'] and 'theta' in res['par']
+
+
+# ---------------------------------------------------------------------------
+# delta22 loader against a synthetic int32 HDF5 in the real layout
+# ---------------------------------------------------------------------------
+DELTA22_FUNCS = ['hf', 'blyp_d3bj', 'bp86_d3bj', 'b97d3_d3bj', 'tpsstpss_d3bj',
+                 'b3lyp_d3bj', 'pbe0_d3bj', 'm062x_d3', 'wb97xd', 'wp04', 'wc04',
+                 'B2PLYP', 'mPW2PLYP', 'B2GP_PLYP', 'dsd_pbep86', 'revdsd_pbep86',
+                 'mp2', 'dlpno_mp2']
+
+
+def _make_synthetic_delta22(path):
+    # one gas method label per functional. most sit at pcSseg3, but mp2 is only at
+    # pcSseg2 and hf only at pcSseg1, so the loader's largest-basis fallback ladder
+    # (pcSseg3 then pcSseg2 then pcSseg1) is exercised.
+    basis_for = {f: "pcSseg3" for f in DELTA22_FUNCS}
+    basis_for["mp2"] = "pcSseg2"
+    basis_for["hf"] = "pcSseg1"
+    names = [f"{f},{basis_for[f]},gas,none" for f in DELTA22_FUNCS]
+    n_methods = len(names)
+    with h5py.File(path, "w") as f:
+        f.create_dataset("conventional_nmr_method_names",
+                         data=np.array(names, dtype=h5py.string_dtype()))
+        # one solute, atoms H, C, H, C, N  (N is ignored: not 1H/13C)
+        z = np.array([1, 6, 1, 6, 7], np.uint8)
+        g = f.create_group("solutes/mol1")
+        g.create_dataset("atomic_numbers", data=z)
+        # conventional_shieldings[method, geometry, atom]; geometry 1 = PBE0/cc-pVTZ
+        base = np.array([30.0, 150.0, 31.0, 151.0, 200.0])
+        cs = np.zeros((n_methods, 2, len(z)))
+        for mi in range(n_methods):
+            cs[mi, 0] = base + 0.5 * mi          # aimnet2 geometry (unused)
+            cs[mi, 1] = base + 0.01 * mi         # pbe0_tz geometry (used)
+        g.create_dataset("stationary_and_pcm/conventional_shieldings", data=_enc(cs))
+    return base
+
+
+def test_load_delta22_synthetic(tmp_path):
+    p = tmp_path / "delta22.hdf5"
+    base = _make_synthetic_delta22(str(p))
+    out = load_delta22(str(p))
+    assert set(out) == {"1H", "13C"}
+    # 1 solute has two H sites (atoms 0,2) and two C sites (atoms 1,3)
+    assert out["1H"]["M"].shape == (2, len(DELTA22_FUNCS))
+    assert out["13C"]["M"].shape == (2, len(DELTA22_FUNCS))
+    assert DELTA22_REF in out["1H"]["methods"]
+    # every functional was selected via the basis-fallback ladder, so no method is missing
+    assert not np.isnan(out["1H"]["M"]).any()
+    assert not np.isnan(out["13C"]["M"]).any()
+    # decoded values must match what we wrote on the pbe0_tz geometry; pick the
+    # DSD-PBEP86 method column for the first H site (atom 0)
+    ref_method_index = DELTA22_FUNCS.index("dsd_pbep86")
+    expected = base[0] + 0.01 * ref_method_index
+    col = out["1H"]["methods"].index(DELTA22_REF)
+    np.testing.assert_allclose(out["1H"]["M"][0, col], expected, atol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# NS372 sheet parsing against a synthetic Kaupp-style worksheet
+# ---------------------------------------------------------------------------
+def _make_synthetic_kaupp_sheet(path):
+    wb = openpyxl.Workbook()
+    ws = wb.active
+    ws.title = "S6 - 1H Shieldings"
+    # row 1: functional headers from column 7; row 2: tau sub-labels
+    # two methods: SVWN (no tau) and TPSS (tau = tauC), matching leveling.METHODS keys
+    ws.cell(1, 7, "SVWN")
+    ws.cell(1, 8, "TPSS"); ws.cell(2, 8, "τC")
+    # data rows from row 4: col1 nucleus (sticky), col2 molecule, col4 CCSD(T) ref
+    rows = [("1H", "molA", 30.0, 30.1, 30.2),
+            ("1H", "molB", 28.0, 28.2, 28.1),
+            ("1H", "molC", 25.0, 25.3, 25.1)]
+    for i, (nuc, mol, ref, svwn, tpss) in enumerate(rows):
+        r = 4 + i
+        ws.cell(r, 1, nuc)
+        ws.cell(r, 2, mol)
+        ws.cell(r, 4, ref)
+        ws.cell(r, 7, svwn)
+        ws.cell(r, 8, tpss)
+    wb.save(path)
+
+
+def test_read_nucleus_synthetic(tmp_path):
+    p = tmp_path / "kaupp.xlsx"
+    _make_synthetic_kaupp_sheet(str(p))
+    wb = openpyxl.load_workbook(str(p), data_only=True)
+    ref, data = read_nucleus(wb, "S6 - 1H Shieldings")
+    np.testing.assert_allclose(ref, [30.0, 28.0, 25.0])
+    assert ("SVWN", None) in data
+    assert ("TPSS", "τC") in data
+    np.testing.assert_allclose(data[("SVWN", None)], [30.1, 28.2, 25.3])
+    np.testing.assert_allclose(data[("TPSS", "τC")], [30.2, 28.1, 25.1])
+
+
+# ---------------------------------------------------------------------------
+# opt-in tests against the real files (skipped when absent, e.g. in CI)
+# ---------------------------------------------------------------------------
+def _find(*candidates):
+    for c in candidates:
+        if c and os.path.exists(c):
+            return c
+    return None
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+import paths
+_REAL_DELTA22 = _find(
+    paths.dataset_file("delta22", file=__file__),   # deposit copy, then in-repo data/delta22/
+    os.path.join(_HERE, "delta22.hdf5"),
+    os.path.join(_HERE, "delta22", "delta22.hdf5"),
+)
+_REAL_KAUPP = _find(
+    os.path.join(_HERE, "..", "..", "data", "ns372", "ct1c00919_si_002.xlsx"),
+)
+
+
+@pytest.mark.skipif(_REAL_DELTA22 is None, reason="real delta22.hdf5 not present")
+def test_real_delta22_reproduces_leveling():
+    out = load_delta22(_REAL_DELTA22)
+    assert out["1H"]["M"].shape == (145, 18)
+    assert out["13C"]["M"].shape == (76, 18)
+    res = analyze_nucleus(out["1H"]["M"], out["1H"]["methods"], out["1H"]["ref"])
+    assert res['ev'][0] > 0.99            # PC1 dominance on the real proton data
+
+
+@pytest.mark.skipif(_REAL_KAUPP is None, reason="real Kaupp NS372 spreadsheet not present")
+def test_real_ns372_loads_all_nuclei():
+    out = load_ns372(_REAL_KAUPP)
+    assert set(out) == {"1H", "11B", "13C", "15N", "17O", "19F", "31P", "33S"}
+    res = analyze_nucleus(out["13C"]["M"], out["13C"]["methods"], out["13C"]["ref"])
+    assert res['ev'][0] > 0.99
+
+
+@pytest.mark.skipif(_REAL_KAUPP is None, reason="real Kaupp NS372 spreadsheet not present")
+def test_real_ns372_1h_convergence_stats():
+    # Figure 2D input: the 1H slice reduces to 44 methods vs the CCSD(T) reference,
+    # and each method's slope/scaled RMSE lands in the expected physical range.
+    cc = load_ns372_1h_convergence(_REAL_KAUPP)
+    assert cc.shape == (124, 45)                       # 44 methods + CCSD(T)
+    assert "CCSD(T)" in cc.columns and "HF" in cc.columns   # 'Hartree-Fock' header renamed
+    res = convergence_statistics(cc).set_index("Method")
+    assert len(res) == 44
+    assert (res["Slope"].between(0.9, 1.15)).all()     # near the y=x limit
+    assert (res["RMSE (scaled)"] > 0).all()
+    # the double hybrid sits closest to the electronic-structure limit
+    assert res["RMSE (scaled)"].idxmin() == "DSD-PBEP86"
+    assert res.loc["DSD-PBEP86", "RMSE (scaled)"] < 0.1

analysis/si_figures/si_figure_s03_leveling.ipynb

@@ -0,0 +1,386 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b61775ea",
+   "metadata": {},
+   "source": [
+    "# SI Figure S3: inter-method correlation and PCA of NS372 and delta-22 shieldings\n",
+    "\n",
+    "Inter-method correlation matrices and PC1/PC2 loadings for NS372 (44 functionals x 8 nuclei) and\n",
+    "delta22 (18 functionals, ¹H/¹³C), plus a combined summary table."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "aeb46df4",
+   "metadata": {},
+   "source": [
+    "Two independent NMR shielding datasets, analyzed separately (never pooled):\n",
+    "\n",
+    "| Dataset | Source | Methods | Nuclei | Reference |\n",
+    "|---|---|---|---|---|\n",
+    "| **NS372** | Schattenberg & Kaupp, *JCTC* **17**, 7602 (2021) | 44 DFT/WFT functionals | ¹H ¹¹B ¹³C ¹⁵N ¹⁷O ¹⁹F ³¹P ³³S | CCSD(T)/pcSseg-3 |\n",
+    "| **delta22** | in-house `delta22.hdf5` | 18 gas-phase functionals, largest basis (pcSseg-3; mp2 → pcSseg-2), PBE0/cc-pVTZ geometry | ¹H ¹³C | DSD-PBEP86 (highest-rung in-set method, stands in for CCSD(T)) |"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d0ce838b",
+   "metadata": {},
+   "source": [
+    "## 1. Configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62157a75",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "REPO = os.path.abspath(\"../..\")\n",
+    "for _p in (\"analysis/code\", \"analysis/code/shared\"):\n",
+    "    sys.path.insert(0, os.path.join(REPO, _p))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8ea46824",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "import pandas as pd\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import paths\n",
+    "import leveling\n",
+    "import leveling_plots"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "34796b33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# inputs:\n",
+    "#   - the Kaupp NS372 spreadsheet is small and ships in the repo's data/ns372/ folder\n",
+    "#   - delta22.hdf5 is large: it resolves from the repo's data/delta22/ folder, carried by the\n",
+    "#     Hugging Face checkout via Git LFS (see analysis/code/paths.py)\n",
+    "# the delta22 file is encoded as whole numbers (real value x 10,000); the loader\n",
+    "# in leveling.py decodes it on read (a plain-decimal copy also works, since the\n",
+    "# decode step is a no-op on floating-point data).\n",
+    "KAUPP_XLSX = os.path.join(REPO, \"data\", \"ns372\", \"ct1c00919_si_002.xlsx\")\n",
+    "DELTA22_H5 = paths.dataset_file(\"delta22\", root=REPO)\n",
+    "SAVE_DPI   = 200        # SI-quality raster output\n",
+    "\n",
+    "def figure_path(name):\n",
+    "    os.makedirs(\"figures\", exist_ok=True)\n",
+    "    return os.path.join(\"figures\", name)\n",
+    "\n",
+    "# self-clean: this notebook builds PNG names dynamically (one per nucleus), so drop any\n",
+    "# previously written figures before regenerating (glob on a missing folder returns [])\n",
+    "for _stale in glob.glob(os.path.join(\"figures\", \"*.png\")):\n",
+    "    os.remove(_stale)\n",
+    "\n",
+    "pd.set_option('display.width', 150)\n",
+    "pd.set_option('display.max_columns', 25)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a402f2fc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# functional-family colours used for every figure (family ordering lives in leveling.py)\n",
+    "FAMILY_COLORS = {'LDA':'#777777','GGA':'#1f77b4','mGGA':'#17becf','GH':'#2ca02c',\n",
+    "                 'RSH':'#9467bd','LH':'#8c564b','DH':'#ff7f0e','WFT':'#d62728',\n",
+    "                 'ref':'#FF1493'}\n",
+    "\n",
+    "# proper Unicode-superscript labels for nuclei in figure titles\n",
+    "NUC_DISPLAY = {'1H':'¹H','11B':'¹¹B','13C':'¹³C',\n",
+    "               '15N':'¹⁵N','17O':'¹⁷O','19F':'¹⁹F',\n",
+    "               '31P':'³¹P','33S':'³³S'}\n",
+    "\n",
+    "# adjustText parameters tuned per nucleus -- 1H and 13C have very dense central\n",
+    "# clusters and need stronger expansion; the paramagnetic-shielding nuclei\n",
+    "# (15N/17O/19F) already spread methods along the parabola and need a gentler\n",
+    "# pass to avoid over-flinging labels.\n",
+    "PCA_ADJUST_DEFAULT = dict(\n",
+    "    force_text=(0.35, 0.55), force_explode=(0.25, 0.40),\n",
+    "    force_static=(0.10, 0.15), force_pull=(0.02, 0.02),\n",
+    "    expand=(1.25, 1.35), time_lim=4,\n",
+    ")\n",
+    "PCA_ADJUST = {\n",
+    "    '1H':  {**PCA_ADJUST_DEFAULT, 'force_text':(0.75, 1.05),\n",
+    "            'force_explode':(0.65, 0.90), 'expand':(1.7, 1.9), 'time_lim':7},\n",
+    "    '11B': {**PCA_ADJUST_DEFAULT, 'force_text':(0.50, 0.75),\n",
+    "            'force_explode':(0.40, 0.60), 'expand':(1.4, 1.55), 'time_lim':5},\n",
+    "    '13C': {**PCA_ADJUST_DEFAULT, 'force_text':(0.70, 1.00),\n",
+    "            'force_explode':(0.60, 0.85), 'expand':(1.6, 1.8), 'time_lim':6},\n",
+    "    '15N': {**PCA_ADJUST_DEFAULT, 'force_text':(0.60, 0.85),\n",
+    "            'force_explode':(0.50, 0.70), 'expand':(1.5, 1.65), 'time_lim':6},\n",
+    "    '17O': {**PCA_ADJUST_DEFAULT, 'force_text':(0.60, 0.85),\n",
+    "            'force_explode':(0.50, 0.70), 'expand':(1.5, 1.65), 'time_lim':6},\n",
+    "    '31P': {**PCA_ADJUST_DEFAULT, 'force_text':(0.65, 0.90),\n",
+    "            'force_explode':(0.55, 0.75), 'expand':(1.55, 1.7), 'time_lim':6},\n",
+    "    '33S': {**PCA_ADJUST_DEFAULT, 'force_text':(0.65, 0.90),\n",
+    "            'force_explode':(0.55, 0.75), 'expand':(1.55, 1.7), 'time_lim':6},\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "52b5cdac",
+   "metadata": {},
+   "source": [
+    "## 2. NS372 (Kaupp) - definitions\n",
+    "\n",
+    "Conventional GIAO shieldings for 44 functionals across 8 main-group nuclei, with a\n",
+    "CCSD(T)/pcSseg-3 reference. Kaupp Reduced-Set exclusions (F₃⁻, O₃, BH - multireference\n",
+    "outliers) are applied. Input: `ct1c00919_si_002.xlsx`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f87c7141",
+   "metadata": {},
+   "source": [
+    "## 3. delta22 - definitions\n",
+    "\n",
+    "Gas-phase conventional GIAO shieldings from `delta22.hdf5`: 18 functionals at their\n",
+    "largest available basis (pcSseg-3; plain `mp2` only to pcSseg-2), at the PBE0/cc-pVTZ\n",
+    "geometry. Observations are pooled ¹H / ¹³C atom sites across all 22 solutes. There is no\n",
+    "CCSD(T) reference in the file; **DSD-PBEP86 is used as the reference** for delta22 - it\n",
+    "is the highest-rung double-hybrid available in this method set and stands in for CCSD(T)\n",
+    "in the same role. The stored whole-number shieldings are decoded on read."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "be87913f",
+   "metadata": {},
+   "source": [
+    "## 4. Load datasets + global colour scale\n",
+    "\n",
+    "Run the loaders, analyse every nucleus, and compute the figure-wide\n",
+    "`-log10(1-|r|)` maximum (`GLOBAL_VMAX`) used as the colour scale on every correlation matrix\n",
+    "below, for NS372 and delta22 alike."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0a24460",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ns372   = leveling.load_ns372(KAUPP_XLSX)\n",
+    "delta22 = leveling.load_delta22(DELTA22_H5)\n",
+    "\n",
+    "res_ns372   = {nuc: leveling.analyze_nucleus(d['M'], d['methods'], d['ref'])\n",
+    "               for nuc, d in ns372.items()}\n",
+    "res_delta22 = {nuc: leveling.analyze_nucleus(d['M'], d['methods'], d['ref'])\n",
+    "               for nuc, d in delta22.items()}\n",
+    "\n",
+    "GLOBAL_VMAX = max(leveling.dataset_logr_max(res_ns372), leveling.dataset_logr_max(res_delta22))\n",
+    "print(f'NS372  : {len(ns372)} nuclei')\n",
+    "print(f'delta22: {len(delta22)} nuclei  (reference = {leveling.DELTA22_REF})')\n",
+    "print(f'global colour scale: 0 -> {GLOBAL_VMAX} on the -log10(1-|r|) axis')\n",
+    "for nuc, d in ns372.items():\n",
+    "    print(f'  NS372 {nuc:4s}: {d[\"M\"].shape[0]:4d} mols  x  {d[\"M\"].shape[1]-1} methods + CCSD(T)')\n",
+    "for nuc, d in delta22.items():\n",
+    "    print(f'  delta22 {nuc:4s}: {d[\"M\"].shape[0]:4d} sites x  {d[\"M\"].shape[1]} methods')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d07150e",
+   "metadata": {},
+   "source": [
+    "## 5. NS372 results"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "747989ff",
+   "metadata": {},
+   "source": [
+    "### 5.1  Leveling diagnostics - NS372"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "52a133ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sum_ns372 = leveling.summarize(ns372, res_ns372)\n",
+    "print('NS372 - leveling diagnostics (CCSD(T) included as a method column):')\n",
+    "sum_ns372.round(5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a385e6f5",
+   "metadata": {},
+   "source": [
+    "### 5.2  Per-method scaled RMSE vs CCSD(T)\n",
+    "\n",
+    "`scaled_rmse` = RMSE of residuals after a per-method linear fit\n",
+    "`sigma_method ~ a*sigma_CCSD(T) + b` - the error that survives empirical linear scaling."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "319abc43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rmse_ns372 = pd.DataFrame({nuc: res_ns372[nuc]['scaled_rmse'] for nuc in res_ns372})\n",
+    "rmse_ns372 = rmse_ns372.drop(index='CCSD(T)', errors='ignore')\n",
+    "rmse_ns372 = rmse_ns372.reindex(ns372['1H']['methods'][:-1])    # family order\n",
+    "print('NS372 - per-method scaled RMSE vs CCSD(T) (ppm):')\n",
+    "rmse_ns372.round(3)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "59cb8672",
+   "metadata": {},
+   "source": [
+    "### 5.3  Inter-method correlation matrices - NS372 (one nucleus per file)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "73c90f1e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NS372_NUCS = ['1H', '11B', '13C', '15N', '17O', '19F', '31P', '33S']\n",
+    "for nuc in NS372_NUCS:\n",
+    "    fams = {nuc: ns372[nuc]['families']}\n",
+    "    fig = leveling_plots.plot_corr_matrix('NS372', [nuc], res_ns372, fams,\n",
+    "                                          GLOBAL_VMAX, ref_name='CCSD(T)',\n",
+    "                                          nuc_display=NUC_DISPLAY)\n",
+    "    fig.savefig(figure_path(f'si_figure_s03_corr_{nuc}.png'),\n",
+    "                dpi=SAVE_DPI, bbox_inches='tight')\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "33da4dd0",
+   "metadata": {},
+   "source": [
+    "### 5.4  PC1/PC2 structure - NS372 (one nucleus per file)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "034bf193",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for nuc in NS372_NUCS:\n",
+    "    fams = {nuc: ns372[nuc]['families']}\n",
+    "    fig = leveling_plots.plot_pca_pair('NS372', [nuc], res_ns372, fams,\n",
+    "                                       family_colors=FAMILY_COLORS, nuc_display=NUC_DISPLAY,\n",
+    "                                       pca_adjust=PCA_ADJUST, pca_adjust_default=PCA_ADJUST_DEFAULT)\n",
+    "    fig.savefig(figure_path(f'si_figure_s03_pca_{nuc}.png'),\n",
+    "                dpi=SAVE_DPI, bbox_inches='tight')\n",
+    "    plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ec202cf",
+   "metadata": {},
+   "source": [
+    "### 5.5  delta22 panels (lead the published figure: corr ¹H/¹³C + PCA)\n",
+    "\n",
+    "The same correlation + PCA layout on the delta22 gas-phase set. DSD-PBEP86 is the reference (delta22\n",
+    "has no CCSD(T)); it appears as an ordinary double-hybrid point in the PCA, with no CCSD(T) star."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc1d57ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# delta22 correlation matrices (canonical S3 A = 1H, B = 13C)\n",
+    "for nuc in ['1H', '13C']:\n",
+    "    fams = {nuc: delta22[nuc]['families']}\n",
+    "    fig = leveling_plots.plot_corr_matrix('delta22', [nuc], res_delta22, fams,\n",
+    "                                          GLOBAL_VMAX, ref_name=leveling.DELTA22_REF,\n",
+    "                                          nuc_display=NUC_DISPLAY)\n",
+    "    fig.savefig(figure_path(f'si_figure_s03_delta22_corr_{nuc}.png'),\n",
+    "                dpi=SAVE_DPI, bbox_inches='tight')\n",
+    "    plt.show()\n",
+    "\n",
+    "# delta22 PCA projection: 1H and 13C stacked in one figure (canonical S3 C)\n",
+    "fams = {nuc: delta22[nuc]['families'] for nuc in ['1H', '13C']}\n",
+    "fig = leveling_plots.plot_pca_pair('delta22', ['1H', '13C'], res_delta22, fams,\n",
+    "                                   family_colors=FAMILY_COLORS, nuc_display=NUC_DISPLAY,\n",
+    "                                   pca_adjust=PCA_ADJUST, pca_adjust_default=PCA_ADJUST_DEFAULT)\n",
+    "fig.savefig(figure_path('si_figure_s03_delta22_pca.png'),\n",
+    "            dpi=SAVE_DPI, bbox_inches='tight')\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4c1e4350",
+   "metadata": {},
+   "source": [
+    "## 6. Combined summary (both datasets)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "185ce298",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# delta22's per-nucleus diagnostics (same summarize() used for NS372 above), computed here\n",
+    "# too so this combined table is self-contained -- it reuses res_delta22 from the shared load\n",
+    "# cell above, so this is cheap (no new correlation/PCA computation).\n",
+    "sum_delta22 = leveling.summarize(delta22, res_delta22)\n",
+    "\n",
+    "combined = pd.concat([sum_ns372.assign(dataset='NS372'),\n",
+    "                      sum_delta22.assign(dataset='delta22')], ignore_index=True)\n",
+    "combined = combined[['dataset','nucleus','n_obs','n_methods','r_min','r_median',\n",
+    "                     'PC1_pct','PC2_pct','PC3plus_pct','parabola_R2']]\n",
+    "print('Leveling effect - both datasets (analyzed separately):')\n",
+    "print(f'  PC1 range : {combined.PC1_pct.min():.2f}% - {combined.PC1_pct.max():.2f}%')\n",
+    "print(f'  min pairwise r : {combined.r_min.min():.5f}  (worst case, all nuclei)')\n",
+    "print(f'  parabola R2 range : {combined.parabola_R2.min():.3f} - {combined.parabola_R2.max():.3f}')\n",
+    "print(f'  global colour scale : 0 -> {GLOBAL_VMAX} on -log10(1-|r|)')\n",
+    "combined.round(5)"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

analysis/si_figures/si_figure_s11.ipynb

@@ -0,0 +1,139 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "934387b2",
+   "metadata": {},
+   "source": [
+    "# SI Figure S11: MagNET vs DFT rovibrational (QCD) corrections\n",
+    "\n",
+    "Four panels: **A** MagNET (NN) vs B3LYP/cc-pVDZ rovibrational (QCD) correction per site, both nuclei;\n",
+    "**B** the NN-minus-DFT error distribution; **C/D** every proton/carbon site's DFT and NN QCD\n",
+    "correction, one column per solute. (QCD is gas-phase, so no solvent axis.)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d004ad45",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "# make the in-repo modules importable (not pip-installed)\n",
+    "REPO = os.path.abspath(\"../..\")\n",
+    "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n",
+    "    sys.path.insert(0, os.path.join(REPO, _p))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "366412b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import delta22\n",
+    "import delta22_plots\n",
+    "import paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d241818b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n",
+    "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n",
+    "\n",
+    "def figure_path(name):\n",
+    "    os.makedirs(\"figures\", exist_ok=True)\n",
+    "    return os.path.join(\"figures\", name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "787bb8c4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n",
+    "nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0f824acd",
+   "metadata": {},
+   "source": [
+    "## Panels A and B: MagNET vs DFT QCD correction, both nuclei"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "31fdd1fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qcd = delta22.compare_dft_nn(dft, nn, \"qcd\", keys=(\"solute\", \"site\", \"nucleus\"))\n",
+    "delta22_plots.plot_qcd_scatter(qcd, save_path=figure_path(\"si_figure_s11a_scatter.png\"))\n",
+    "delta22_plots.plot_qcd_error_histogram(qcd, save_path=figure_path(\"si_figure_s11b_hist.png\"))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "080ca90b",
+   "metadata": {},
+   "source": [
+    "## Panel C: every proton site's QCD correction, one column per solute"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26be5d0a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "delta22_plots.plot_qcd_correction_by_site(delta22.qcd_correction_by_site(dft, nn, nucleus=\"H\"), \"H\", \"upper left\",\n",
+    "                                          save_path=figure_path(\"si_figure_s11c_sites_1H.png\"))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42977cbe",
+   "metadata": {},
+   "source": [
+    "## Panel D: every carbon site's QCD correction, one column per solute"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4430159b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "delta22_plots.plot_qcd_correction_by_site(delta22.qcd_correction_by_site(dft, nn, nucleus=\"C\"), \"C\", \"upper right\",\n",
+    "                                          save_path=figure_path(\"si_figure_s11d_sites_13C.png\"))\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

analysis/si_figures/si_figure_s12_carbon_monoxide.ipynb

@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "cecbae9e",
+   "metadata": {},
+   "source": [
+    "# SI Figure S12: MagNET vs DFT ¹³C shielding of carbon monoxide across bond lengths\n",
+    "\n",
+    "Carbon monoxide stretched/compressed over a wide range of bond lengths, ¹³C shielding by DFT\n",
+    "(PBE1PBE/pcSseg-1, gas) vs the MagNET foundation model. The right panel shades the bond-length\n",
+    "distribution seen in real molecules."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "326ed41a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "# make the in-repo modules importable (not pip-installed)\n",
+    "REPO = os.path.abspath(\"../..\")\n",
+    "for _p in (\"analysis/code\", \"analysis/code/shared\"):\n",
+    "    sys.path.insert(0, os.path.join(REPO, _p))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "19aad17e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import carbon_monoxide"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "42d7d8b0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def figure_path(name):\n",
+    "    os.makedirs(\"figures\", exist_ok=True)\n",
+    "    return os.path.join(\"figures\", name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17186f45",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = carbon_monoxide.load_comparison()\n",
+    "near, tails = carbon_monoxide.equilibrium_vs_extreme_error(df)\n",
+    "print(f'{len(df)} bond lengths from {df.bond_length.min():.2f} to {df.bond_length.max():.2f} angstrom')\n",
+    "print(f'mean |carbon error|: {near:.1f} ppm near equilibrium vs {tails:.1f} ppm in the tails')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "48de095f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n",
+    "\n",
+    "ax = axes[0]\n",
+    "ax.plot(df[\"bond_length\"], df[\"gaussian_c\"], marker=\"o\", linestyle=\"-\", color=\"black\", label=\"Gaussian C\")\n",
+    "ax.plot(df[\"bond_length\"], df[\"magnet_c\"], marker=\"o\", linestyle=\"-\", color=\"gray\", label=\"MagNET C\")\n",
+    "ax.set_xlabel(\"Bond length (Å)\", fontweight=\"bold\")\n",
+    "ax.set_ylabel(\"Shielding (ppm)\", fontweight=\"bold\")\n",
+    "ax.set_title(\"13C Shielding vs. Bond Length\", fontweight=\"bold\")\n",
+    "ax.legend(frameon=False)\n",
+    "\n",
+    "ax = axes[1]\n",
+    "centers, counts, bin_width = carbon_monoxide.load_bond_length_histogram()\n",
+    "ax_hist = ax.twinx()\n",
+    "ax_hist.bar(centers, counts, width=bin_width, color=\"0.75\", edgecolor=\"none\", zorder=0)\n",
+    "ax_hist.set_yticks([])\n",
+    "ax_hist.grid(False)\n",
+    "ax_hist.set_zorder(0)\n",
+    "ax.set_zorder(1)\n",
+    "ax.patch.set_alpha(0)   # let the histogram show through the transparent foreground axes\n",
+    "ax.plot(df[\"bond_length\"], df[\"delta_c\"], marker=\"o\", linestyle=\"-\", color=\"black\")\n",
+    "ax.axhline(0, color=\"gray\", linewidth=1, linestyle=\"--\")\n",
+    "ax.set_xlabel(\"Bond length (Å)\", fontweight=\"bold\")\n",
+    "ax.set_ylabel(\"Δ shielding (MagNET − Gaussian, ppm)\", fontweight=\"bold\")\n",
+    "ax.set_title(\"Shielding Differences vs. Bond Lengths\", fontweight=\"bold\")\n",
+    "\n",
+    "fig.tight_layout()\n",
+    "fig.savefig(figure_path(\"si_figure_s12.png\"), dpi=300, bbox_inches=\"tight\")\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

analysis/si_figures/si_figure_s13.ipynb

@@ -0,0 +1,173 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "1127ad90",
+   "metadata": {},
+   "source": [
+    "# SI Figure S13: MagNET-x vs DFT explicit-solvent corrections\n",
+    "\n",
+    "Four panels, both nuclei: **A** MagNET-x (NN) vs DFT explicit-solvent correction per site/solvent,\n",
+    "Desmond and OpenMM overlaid; **B** the NN-minus-DFT error distribution; **C** every site's DFT and NN\n",
+    "correction in chloroform, one column per solute; **D** the semi-parsimonious composite model's fit accuracy vs experiment across the four OpenMM solvents, split by engine and\n",
+    "DFT-vs-NN source."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "649a6ece",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "# make the in-repo modules importable (not pip-installed)\n",
+    "REPO = os.path.abspath(\"../..\")\n",
+    "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n",
+    "    sys.path.insert(0, os.path.join(REPO, _p))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dd166a12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import delta22\n",
+    "import delta22_plots\n",
+    "import paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da96ff0e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n",
+    "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n",
+    "\n",
+    "def figure_path(name):\n",
+    "    os.makedirs(\"figures\", exist_ok=True)\n",
+    "    return os.path.join(\"figures\", name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0cc06b06",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n",
+    "nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af66d177",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# openMM/Desmond engine hues (used by the scatter/histogram/site/fitting-accuracy panels below)\n",
+    "ENGINE_COLORS = {\"openMM\": \"#2E86AB\", \"desmond\": \"#A23B72\"}\n",
+    "ENGINE_LABELS = {\"desmond\": \"Desmond\", \"openMM\": \"OpenMM\"}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8ddc9cdc",
+   "metadata": {},
+   "source": [
+    "## Panels A and B: MagNET-x vs DFT explicit corrections, Desmond and OpenMM overlaid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "622c1d35",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "explicit_by_engine = delta22.compare_dft_nn_by_engine(dft, nn, keys=(\"solute\", \"site\", \"nucleus\", \"solvent\"))\n",
+    "delta22_plots.plot_dft_nn_scatter_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,\n",
+    "                                            save_path=figure_path(\"si_figure_s13a_scatter.png\"))\n",
+    "delta22_plots.plot_dft_nn_error_histogram_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,\n",
+    "                                                    save_path=figure_path(\"si_figure_s13b_hist.png\"))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2c511445",
+   "metadata": {},
+   "source": [
+    "## Panel C: every proton/carbon site's explicit correction in chloroform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ca4e7871",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for nucleus, label in [(\"H\", \"Proton\"), (\"C\", \"Carbon\")]:\n",
+    "    pairs = delta22.explicit_correction_dft_nn_pairs(dft, nn, \"chloroform\", nucleus)\n",
+    "    delta22_plots.plot_explicit_correction_by_site(\n",
+    "        pairs, f\"All {label} Sites: DFT vs NN Explicit Corrections (chloroform)\", ENGINE_COLORS,\n",
+    "        save_path=figure_path(f\"si_figure_s13c_sites_{'1H' if nucleus == 'H' else '13C'}.png\"))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0813f518",
+   "metadata": {},
+   "source": [
+    "## Panel D: fitting accuracy of the semi-parsimonious composite model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0ebd8702",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# the published panels use 250 seeded train/test splits\n",
+    "N_SPLITS = 250"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8ce39a57",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SOLVENT_ORDER = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]\n",
+    "solutes = sorted(dft[\"solute\"].unique())\n",
+    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
+    "    fitting = delta22.si_s13d_fitting_accuracy(dft, nn, SOLVENT_ORDER, N_SPLITS, solutes, nucleus=nucleus)\n",
+    "    delta22_plots.plot_fitting_accuracy_boxplot(\n",
+    "        fitting, SOLVENT_ORDER, label, ENGINE_COLORS,\n",
+    "        save_path=figure_path(f\"si_figure_s13d_fitting_{'1H' if nucleus == 'H' else '13C'}.png\"))\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

analysis/si_figures/si_figure_s14.ipynb

@@ -0,0 +1,157 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "8cd77576",
+   "metadata": {},
+   "source": [
+    "# SI Figure S14: test RMSE with DFT vs MagNET features, by solvent\n",
+    "\n",
+    "For the four explicit-solvent solvents (chloroform, methanol, TIP4P water, benzene), per-solvent test\n",
+    "RMSE of predicting experimental ¹H/¹³C shifts with **DFT** features (solid) vs end-to-end **MagNET/NN**\n",
+    "features (lightened), each under three conditions (implicit SOTA, explicit, explicit + vibrations).\n",
+    "Two pages, one per MD engine (Desmond, OpenMM); nitromethane dropped."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67f9b9de",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "# make the in-repo modules importable (not pip-installed)\n",
+    "REPO = os.path.abspath(\"../..\")\n",
+    "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n",
+    "    sys.path.insert(0, os.path.join(REPO, _p))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eddfed92",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import delta22\n",
+    "import delta22_plots\n",
+    "import paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c05cd017",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n",
+    "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n",
+    "\n",
+    "def figure_path(name):\n",
+    "    os.makedirs(\"figures\", exist_ok=True)\n",
+    "    return os.path.join(\"figures\", name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "db2e619d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# the published panels use 250 seeded train/test splits\n",
+    "N_SPLITS = 250"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b57840e8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n",
+    "nn = delta22.add_composite_columns(delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False))\n",
+    "print(len(dft), \"DFT rows;\", len(nn), \"NN rows\")\n",
+    "\n",
+    "LABELS = [\"Implicit Solvent (SOTA)\", \"Explicit Solvent\", \"Explicit Solvent + Vibrations\"]\n",
+    "SS_LABELS = {\"TIP4P\": \"Water (TIP4P)\"}\n",
+    "S14_SOLVENTS = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]   # the 4 explicit solvents, SI order"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a0645b27",
+   "metadata": {},
+   "source": [
+    "## Desmond page (DFT vs NN)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a7d9851",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# implicit = 2-term (stationary + pcm); explicit = stationary + desmond;\n",
+    "# explicit + vibrations = stationary_plus_qcd + desmond (1H) / stationary_plus_des_vib + desmond (13C)\n",
+    "DESMOND_FORMULAS = {\n",
+    "    \"H\":  [\"stationary + pcm\", \"stationary + desmond\", \"stationary_plus_qcd + desmond\"],\n",
+    "    \"C\":  [\"stationary + pcm\", \"stationary + desmond\", \"stationary_plus_des_vib + desmond\"],\n",
+    "}\n",
+    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
+    "    formulas = DESMOND_FORMULAS[nucleus]\n",
+    "    delta22_plots.plot_ss_boxplot_dft_vs_nn(\n",
+    "        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),\n",
+    "        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),\n",
+    "        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,\n",
+    "        save_path=figure_path(f\"si_figure_s14_desmond_{label}.png\"))\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c8523187",
+   "metadata": {},
+   "source": [
+    "## OpenMM page (DFT vs NN)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fc91d9af",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# implicit = 1-term composite (stationary_plus_pcm); explicit = stationary + openMM;\n",
+    "# explicit + vibrations = stationary_plus_qcd + openMM (1H) / stationary_plus_op_vib + openMM (13C)\n",
+    "OPENMM_FORMULAS = {\n",
+    "    \"H\":  [\"stationary_plus_pcm\", \"stationary + openMM\", \"stationary_plus_qcd + openMM\"],\n",
+    "    \"C\":  [\"stationary_plus_pcm\", \"stationary + openMM\", \"stationary_plus_op_vib + openMM\"],\n",
+    "}\n",
+    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
+    "    formulas = OPENMM_FORMULAS[nucleus]\n",
+    "    delta22_plots.plot_ss_boxplot_dft_vs_nn(\n",
+    "        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),\n",
+    "        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),\n",
+    "        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,\n",
+    "        save_path=figure_path(f\"si_figure_s14_openmm_{label}.png\"))\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}