{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8cd77576",
   "metadata": {},
   "source": [
    "# SI Figure S14: test RMSE with DFT vs MagNET features, by solvent\n",
    "\n",
    "For the four explicit-solvent solvents (chloroform, methanol, TIP4P water, benzene), per-solvent test\n",
    "RMSE of predicting experimental ¹H/¹³C shifts with **DFT** features (solid) vs end-to-end **MagNET/NN**\n",
    "features (lightened), each under three conditions (implicit SOTA, explicit, explicit + vibrations).\n",
    "Two pages, one per MD engine (Desmond, OpenMM); nitromethane dropped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f9b9de",
   "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": "eddfed92",
   "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": "c05cd017",
   "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": "db2e619d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the published panels use 250 seeded train/test splits\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b57840e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n",
    "nn = delta22.add_composite_columns(delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False))\n",
    "print(len(dft), \"DFT rows;\", len(nn), \"NN rows\")\n",
    "\n",
    "LABELS = [\"Implicit Solvent (SOTA)\", \"Explicit Solvent\", \"Explicit Solvent + Vibrations\"]\n",
    "SS_LABELS = {\"TIP4P\": \"Water (TIP4P)\"}\n",
    "S14_SOLVENTS = [\"chloroform\", \"methanol\", \"TIP4P\", \"benzene\"]   # the 4 explicit solvents, SI order"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0645b27",
   "metadata": {},
   "source": [
    "## Desmond page (DFT vs NN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a7d9851",
   "metadata": {},
   "outputs": [],
   "source": [
    "# implicit = 2-term (stationary + pcm); explicit = stationary + desmond;\n",
    "# explicit + vibrations = stationary_plus_qcd + desmond (1H) / stationary_plus_des_vib + desmond (13C)\n",
    "DESMOND_FORMULAS = {\n",
    "    \"H\":  [\"stationary + pcm\", \"stationary + desmond\", \"stationary_plus_qcd + desmond\"],\n",
    "    \"C\":  [\"stationary + pcm\", \"stationary + desmond\", \"stationary_plus_des_vib + desmond\"],\n",
    "}\n",
    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
    "    formulas = DESMOND_FORMULAS[nucleus]\n",
    "    delta22_plots.plot_ss_boxplot_dft_vs_nn(\n",
    "        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),\n",
    "        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),\n",
    "        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,\n",
    "        save_path=figure_path(f\"si_figure_s14_desmond_{label}.png\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8523187",
   "metadata": {},
   "source": [
    "## OpenMM page (DFT vs NN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc91d9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# implicit = 1-term composite (stationary_plus_pcm); explicit = stationary + openMM;\n",
    "# explicit + vibrations = stationary_plus_qcd + openMM (1H) / stationary_plus_op_vib + openMM (13C)\n",
    "OPENMM_FORMULAS = {\n",
    "    \"H\":  [\"stationary_plus_pcm\", \"stationary + openMM\", \"stationary_plus_qcd + openMM\"],\n",
    "    \"C\":  [\"stationary_plus_pcm\", \"stationary + openMM\", \"stationary_plus_op_vib + openMM\"],\n",
    "}\n",
    "for nucleus, label in [(\"H\", \"1H\"), (\"C\", \"13C\")]:\n",
    "    formulas = OPENMM_FORMULAS[nucleus]\n",
    "    delta22_plots.plot_ss_boxplot_dft_vs_nn(\n",
    "        delta22_plots.ss_fits(dft, nucleus, formulas, S14_SOLVENTS, N_SPLITS, dft=True),\n",
    "        delta22_plots.ss_fits(nn, nucleus, formulas, S14_SOLVENTS, N_SPLITS),\n",
    "        S14_SOLVENTS, formulas, LABELS, label, solvent_labels=SS_LABELS,\n",
    "        save_path=figure_path(f\"si_figure_s14_openmm_{label}.png\"))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}