{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0954a86e",
   "metadata": {},
   "source": [
    "# Figure 3D: implicit vs explicit solvent corrections, one box per solvent\n",
    "\n",
    "All 12 solvents individually (polar aprotic -> polar protic -> aromatic), four schemes each: implicit (PCM), implicit + vibrations, explicit, explicit + vibrations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5172978e",
   "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": "c7e70d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "import delta22\n",
    "import paths\n",
    "import fig3_plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55397f4b",
   "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": "06416ecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the published panels use 250 seeded train/test splits\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e892c346",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n",
    "solutes = sorted(query[\"solute\"].unique())\n",
    "print(len(query), \"rows;\", len(solutes), \"solutes;\", query[\"sap_nmr_method\"].nunique(), \"methods\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32748b6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "FIG3D_LABELS = {\n",
    "    \"stationary + pcm\": \"Implicit Solvent (PCM)\",\n",
    "    \"stationary_plus_qcd + pcm\": \"Implicit Solvent (PCM) + Vibrations\",\n",
    "    \"stationary + desmond\": \"Explicit Solvent\",\n",
    "    \"stationary_plus_qcd + desmond\": \"Explicit Solvent + Vibrations\",\n",
    "}\n",
    "fig3d_results = delta22.fig3d_formula_regressions(\n",
    "    query, \"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\", list(FIG3D_LABELS),\n",
    "    delta22.DESMOND_SOLVENTS, n_splits=N_SPLITS, solutes=solutes)\n",
    "\n",
    "# report implicit vs explicit median test RMSE per solvent (not pooled by class)\n",
    "for solvent in delta22.DESMOND_SOLVENTS:\n",
    "    sub = fig3d_results[fig3d_results[\"solvent\"] == solvent]\n",
    "    mi = sub[sub[\"formula\"] == \"stationary + pcm\"][\"test_RMSE\"].median()\n",
    "    me = sub[sub[\"formula\"] == \"stationary + desmond\"][\"test_RMSE\"].median()\n",
    "    print(f\"{solvent:20s} implicit={mi:.4f}  explicit={me:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e826d46",
   "metadata": {},
   "outputs": [],
   "source": [
    "FIG3D_COLORS = [\"#A72608\", \"#F4BAAD\", \"#5D737E\", \"#D9FFF5\"]   # dark red, pink, gray, mint (as published)\n",
    "FIG3D_SOLVENT_LABELS = {\n",
    "    \"chloroform\": r\"CDCl$_3$\", \"dichloromethane\": \"DCM\", \"tetrahydrofuran\": \"THF\",\n",
    "    \"acetonitrile\": \"MeCN\", \"dimethylsulfoxide\": \"DMSO\", \"methanol\": \"MeOD\",\n",
    "    \"trifluoroethanol\": \"TFE\", \"chlorobenzene\": \"PhCl\",\n",
    "}\n",
    "# solvent order: polar aprotic -> polar protic -> aromatic (group ordering from delta22.SOLVENT_GROUPS)\n",
    "FIG3D_SOLVENT_ORDER = [s for group in delta22.SOLVENT_GROUPS.values() for s in group]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c66631",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3_plots.plot_solvent_correction_boxplot(fig3d_results, FIG3D_LABELS, FIG3D_SOLVENT_ORDER,\n",
    "                                           FIG3D_COLORS, FIG3D_SOLVENT_LABELS,\n",
    "                                           save_path=figure_path(\"fig3d_explicit_solvation_1H.png\"))"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}