{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d13d8aa5",
   "metadata": {},
   "source": [
    "# Figure 3A: implicit-solvent (PCM) benefit by method, chloroform vs benzene\n",
    "\n",
    "Per-split percent benefit of adding a PCM solvent correction, for each method, in chloroform and benzene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80ceb0d1",
   "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": "88b3d113",
   "metadata": {},
   "outputs": [],
   "source": [
    "import delta22\n",
    "import paths\n",
    "import fig3_plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e830f642",
   "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": "df43c551",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the published panels use 250 seeded train/test splits\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a780b7",
   "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": "6c68c523",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3a_by_solvent = delta22.fig3a_pcm_benefit_by_solvent(query, [\"chloroform\", \"benzene\"],\n",
    "                                                        n_splits=N_SPLITS, solutes=solutes)\n",
    "\n",
    "# per-method median percent benefit, then mean +/- std across methods\n",
    "chloro = fig3a_by_solvent[fig3a_by_solvent[\"solvent\"] == \"chloroform\"]\n",
    "benz = fig3a_by_solvent[fig3a_by_solvent[\"solvent\"] == \"benzene\"]\n",
    "chloro_by_method = chloro.groupby(\"sap_nmr_method\")[\"percent_benefit\"].median()\n",
    "benz_by_method = benz.groupby(\"sap_nmr_method\")[\"percent_benefit\"].median()\n",
    "print(f\"PCM benefit in chloroform: {chloro_by_method.mean():+.1f} +/- {chloro_by_method.std():.1f}%\"\n",
    "      f\"  (positive = PCM helps)\")\n",
    "print(f\"PCM benefit in benzene:    {benz_by_method.mean():+.1f} +/- {benz_by_method.std():.1f}%\"\n",
    "      f\"  (negative = PCM hurts)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f282780f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fixed method order, grouped by family (ab initio -> conventional DFT -> NMR-specific -> double\n",
    "# hybrid), each as (method, basis, geometry). This ordering is what gives the panel its family\n",
    "# structure; sorting by benefit value would scramble it.\n",
    "FIG3A_COMBOS = [\n",
    "    (\"hf\", \"pcSseg3\", \"pbe0_tz\"), (\"mp2\", \"pcSseg2\", \"pbe0_tz\"), (\"dlpno_mp2\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"b3lyp_d3bj\", \"pcSseg3\", \"pbe0_tz\"), (\"pbe0_d3bj\", \"pcSseg2\", \"pbe0_tz\"),\n",
    "    (\"m062x_d3\", \"pcSseg3\", \"pbe0_tz\"), (\"b97d3_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"wb97xd\", \"pcSseg3\", \"pbe0_tz\"), (\"tpsstpss_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"bp86_d3bj\", \"pcSseg3\", \"pbe0_tz\"), (\"blyp_d3bj\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"wp04\", \"pcSseg3\", \"pbe0_tz\"), (\"wc04\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"dsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"), (\"B2GP_PLYP\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"B2PLYP\", \"pcSseg3\", \"pbe0_tz\"), (\"mPW2PLYP\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "    (\"revdsd_pbep86\", \"pcSseg3\", \"pbe0_tz\"),\n",
    "]\n",
    "FIG3A_COLORS = {\"chloroform\": \"#44AA99\", \"benzene\": \"#DDCC77\"}   # teal / tan, as published"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fa6355b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig3_plots.plot_pcm_benefit_with_arrows(fig3a_by_solvent, FIG3A_COMBOS, FIG3A_COLORS,\n",
    "                                        save_path=figure_path(\"fig3a_pcm_benefit_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
}