{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1127ad90",
   "metadata": {},
   "source": [
    "# SI Figure S13: MagNET-x vs DFT explicit-solvent corrections\n",
    "\n",
    "Four panels, both nuclei: **A** MagNET-x (NN) vs DFT explicit-solvent correction per site/solvent,\n",
    "Desmond and OpenMM overlaid; **B** the NN-minus-DFT error distribution; **C** every site's DFT and NN\n",
    "correction in chloroform, one column per solute; **D** the semi-parsimonious composite model's fit accuracy vs experiment across the four OpenMM solvents, split by engine and\n",
    "DFT-vs-NN source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649a6ece",
   "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": "dd166a12",
   "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": "da96ff0e",
   "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": "0cc06b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\n",
    "nn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af66d177",
   "metadata": {},
   "outputs": [],
   "source": [
    "# openMM/Desmond engine hues (used by the scatter/histogram/site/fitting-accuracy panels below)\n",
    "ENGINE_COLORS = {\"openMM\": \"#2E86AB\", \"desmond\": \"#A23B72\"}\n",
    "ENGINE_LABELS = {\"desmond\": \"Desmond\", \"openMM\": \"OpenMM\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddc9cdc",
   "metadata": {},
   "source": [
    "## Panels A and B: MagNET-x vs DFT explicit corrections, Desmond and OpenMM overlaid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "622c1d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "explicit_by_engine = delta22.compare_dft_nn_by_engine(dft, nn, keys=(\"solute\", \"site\", \"nucleus\", \"solvent\"))\n",
    "delta22_plots.plot_dft_nn_scatter_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,\n",
    "                                            save_path=figure_path(\"si_figure_s13a_scatter.png\"))\n",
    "delta22_plots.plot_dft_nn_error_histogram_by_engine(explicit_by_engine, ENGINE_COLORS, ENGINE_LABELS,\n",
    "                                                    save_path=figure_path(\"si_figure_s13b_hist.png\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c511445",
   "metadata": {},
   "source": [
    "## Panel C: every proton/carbon site's explicit correction in chloroform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4e7871",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus, label in [(\"H\", \"Proton\"), (\"C\", \"Carbon\")]:\n",
    "    pairs = delta22.explicit_correction_dft_nn_pairs(dft, nn, \"chloroform\", nucleus)\n",
    "    delta22_plots.plot_explicit_correction_by_site(\n",
    "        pairs, f\"All {label} Sites: DFT vs NN Explicit Corrections (chloroform)\", ENGINE_COLORS,\n",
    "        save_path=figure_path(f\"si_figure_s13c_sites_{'1H' if nucleus == 'H' else '13C'}.png\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0813f518",
   "metadata": {},
   "source": [
    "## Panel D: fitting accuracy of the semi-parsimonious composite model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ebd8702",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the published panels use 250 seeded train/test splits\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ce39a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "SOLVENT_ORDER = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]\n",
    "solutes = sorted(dft[\"solute\"].unique())\n",
    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
    "    fitting = delta22.si_s13d_fitting_accuracy(dft, nn, SOLVENT_ORDER, N_SPLITS, solutes, nucleus=nucleus)\n",
    "    delta22_plots.plot_fitting_accuracy_boxplot(\n",
    "        fitting, SOLVENT_ORDER, label, ENGINE_COLORS,\n",
    "        save_path=figure_path(f\"si_figure_s13d_fitting_{'1H' if nucleus == 'H' else '13C'}.png\"))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}