"""Reproduces the delta-22 composite-formula ablation study (the ablations.xlsx workbook). 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 this ablation analysis 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