"""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() # Composite-formula ablation workbook, "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); ablations_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 (the ablation workbook'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 ablations_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): """The ablation workbook'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 ablations_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))}