{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "29f74212",
   "metadata": {},
   "source": [
    "# SI Figure S6: per-solvent test RMSE of six solvent-correction models\n",
    "\n",
    "Per nucleus, the per-solvent test RMSE of six solvent-correction models: SotA implicit (single-slope\n",
    "PCM), the paper's implicit fit (stationary + PCM as separate terms), and explicit (Desmond), each\n",
    "with and without a vibrational correction (QCD for ¹H, Desmond vibration for ¹³C). Level of theory:\n",
    "dsd_pbep86/pcSseg3 on pbe0_tz geometries, PCM substituted from b3lyp_d3bj/pcSseg3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6871225d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "\n",
    "# make the in-repo modules importable (not pip-installed)\n",
    "REPO = os.path.abspath(\"../..\")\n",
    "for _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n",
    "    sys.path.insert(0, os.path.join(REPO, _p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bad055e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import delta22\n",
    "import delta22_plots\n",
    "import paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a55272ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\n",
    "XLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n",
    "\n",
    "def figure_path(name):\n",
    "    os.makedirs(\"figures\", exist_ok=True)\n",
    "    return os.path.join(\"figures\", name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "221a152c",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = delta22.add_composite_columns(delta22.load_query_df_dft(\n",
    "    DELTA22_HDF5, XLSX, pcm_reference_method=\"b3lyp_d3bj\", pcm_reference_basis=\"pcSseg3\", verbose=False))\n",
    "solutes = sorted(query[\"solute\"].unique())\n",
    "\n",
    "METHOD, BASIS, GEOM = \"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"\n",
    "\n",
    "# same six categories for both nuclei; the \"+ Vibrations\" term differs (qcd for H, desmond_vib for C)\n",
    "LADDER = {\n",
    "    \"H\": {\n",
    "        \"stationary_plus_pcm\": \"SotA Implicit Solvent\",\n",
    "        \"stationary_plus_pcm_plus_qcd\": \"SotA Implicit Solvent + Vibrations\",\n",
    "        \"stationary + pcm\": \"Implicit Solvent\",\n",
    "        \"stationary_plus_qcd + pcm\": \"Implicit Solvent + Vibrations\",\n",
    "        \"stationary + desmond\": \"Explicit Solvent\",\n",
    "        \"stationary_plus_qcd + desmond\": \"Explicit Solvent + Vibrations\"},\n",
    "    \"C\": {\n",
    "        \"stationary_plus_pcm\": \"SotA Implicit Solvent\",\n",
    "        \"stationary_plus_pcm_plus_des_vib\": \"SotA Implicit Solvent + Vibrations\",\n",
    "        \"stationary + pcm\": \"Implicit Solvent\",\n",
    "        \"stationary_plus_des_vib + pcm\": \"Implicit Solvent + Vibrations\",\n",
    "        \"stationary + desmond\": \"Explicit Solvent\",\n",
    "        \"stationary_plus_des_vib + desmond\": \"Explicit Solvent + Vibrations\"},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1b678b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# one dark/light pair per solvent-treatment family (SotA implicit / implicit / explicit), light = +Vibrations\n",
    "S6_COLORS = [\"#C4B037\", \"#F5EDA0\", \"#A72608\", \"#E4A0A0\", \"#5F7C8A\", \"#B9E7DF\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53817ec5",
   "metadata": {},
   "source": [
    "## Formula ladder per nucleus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cd062c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the published panels use 250 seeded train/test splits\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d57a9788",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus, labels in LADDER.items():\n",
    "    results = delta22.fig3d_formula_regressions(query, METHOD, BASIS, GEOM, list(labels),\n",
    "                                                delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS,\n",
    "                                                solutes=solutes, nucleus=nucleus)\n",
    "    med = results.groupby(\"formula\")[\"test_RMSE\"].median()\n",
    "    print(f\"--- {nucleus} ({METHOD}) median test RMSE ---\")\n",
    "    for formula in labels:\n",
    "        print(f\"  {labels[formula]:38s} {med[formula]:.4f}\")\n",
    "    delta22_plots.plot_formula_ladder_boxplot(\n",
    "        results, labels, delta22.SOLVENT_GROUPS, nucleus=nucleus, colors=S6_COLORS,\n",
    "        title=f\"Explicit Solvation + Vibrational Effects Improve Predictions Across Solvents ({nucleus} Nucleus)\",\n",
    "        save_path=figure_path(f\"si_figure_s06_ladder_{'1H' if nucleus == 'H' else '13C'}.png\"))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}