| """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 |
|
|
| |
| 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", |
| } |
| |
| 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", |
| ]} |
| |
| 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", |
| } |
| |
| PEPTIDE_HIGHLIGHT_H = {"peptide": "#B0FF40"} |
| PEPTIDE_HIGHLIGHT_C = {"isomer_4N": "#B0FF40", "peptide": "#B0FF40"} |
|
|
| |
| _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) |
| corr = mx[:, idx, 1] - mx[:, idx, 0] |
| corr = np.mean(corr, axis=1) |
| corr = corr[~np.isnan(corr)] |
| 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")] |
|
|
| |
| |
| |
| _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] |
|
|
| |
| 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_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() |
|
|
| |
| 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) |
| 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() |
|
|
| |
| comb = pd.concat([experimental_stack(loader), mz, ex, qcd], axis=1).reset_index() |
| comb = comb.set_index(["solute", "nucleus", "site", "solvent"]) |
| |
| |
| 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() |
|
|
|
|
| |
| |
| |
| |
|
|
| 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: |
| 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]] |
| |
| |
| |
| 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() |
| |
| |
| |
| 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() |
| |
| 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")) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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"] |
| |
| |
| |
| |
| 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]) |
| |
| |
| |
| |
| 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") |
|
|
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| 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) |
|
|
