{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "77f6c813",
   "metadata": {},
   "source": [
    "# Composite-formula ablation workbook\n",
    "\n",
    "Rebuilds the composite-formula ablation grid: ~25 composite-formula variants (stationary geometry\n",
    "plus some combination of implicit PCM, explicit Desmond, and rovibrational QCD corrections) across\n",
    "all 12 delta-22 solvents, at two reference levels: DSD-PBEP86/pcSseg-3 (geometry PBE0/tz) and the\n",
    "MagNET-Zero training reference (WP04/pcSseg-2 for ¹H, ωB97X-D/pcSseg-2 for ¹³C). Also reports the\n",
    "solvent-averaged correlations between the composite-model features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e00f8b39",
   "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": "b3acc33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import delta22\n",
    "import composite_models\n",
    "import composite_plots\n",
    "import paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a909ba74",
   "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)\n",
    "\n",
    "def document_path(name):\n",
    "    os.makedirs(\"documents\", exist_ok=True)\n",
    "    return os.path.join(\"documents\", name)\n",
    "\n",
    "# the shipped ablations.xlsx used 250 seeded train/test splits per (formula, solvent)\n",
    "N_SPLITS = 250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5df608e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_df_dft = delta22.add_composite_columns(delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False))\n",
    "solutes = delta22.delta22_solutes(DELTA22_HDF5)\n",
    "print(len(query_df_dft), \"rows;\", len(solutes), \"solutes\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7bc63a7",
   "metadata": {},
   "source": [
    "## Preview: one formula's mean test RMSE at the DSD-PBEP86 reference level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad79f75e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# preview before the full build (all ~25 formulas x 12 solvents x 250 splits x 2 nuclei x 2\n",
    "# reference levels)\n",
    "level = composite_models.REFERENCE_LEVELS[\"dsd\"]\n",
    "preview = composite_models.ablation_rmse_table(query_df_dft, \"H\", level[\"method_h\"], level[\"basis_h\"],\n",
    "                                 level[\"geometry_h\"], solutes, n_splits=20)\n",
    "preview[[\"chloroform\", \"benzene\", \"Mean Test RMSE\"]].round(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ae7d282",
   "metadata": {},
   "source": [
    "## Build the full ablations workbook (both reference levels, both nuclei)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9acdb9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = document_path(\"ablations.xlsx\")\n",
    "composite_models.build_ablations_workbook(query_df_dft, solutes, output_path, n_splits=N_SPLITS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66102dfb",
   "metadata": {},
   "source": [
    "## Correlations Between Features\n",
    "\n",
    "Solvent-averaged Pearson r correlation matrix between the five composite-model features (stationary\n",
    "shielding, PCM, Desmond, its vibrational analogue, QCD), both nuclei, plus a per-solvent\n",
    "PCM-vs-Desmond table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "208d343b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus, label in [(\"H\", \"Proton\"), (\"C\", \"Carbon\")]:\n",
    "    corr = delta22.ablations_feature_correlations(query_df_dft, nucleus)\n",
    "    nuc_label = \"1H\" if nucleus == \"H\" else \"13C\"\n",
    "    nuc_title = \"$^{1}$H\" if nucleus == \"H\" else \"$^{13}$C\"\n",
    "    composite_plots.plot_feature_correlation_heatmap(\n",
    "        corr[\"r\"], vmin=-1, vmax=1, cmap=\"RdBu\",\n",
    "        title=f\"Solvent-Averaged Pearson $r$ Correlation Matrix for {nuc_title}\",\n",
    "        save_path=figure_path(f\"ablations_feature_corr_r_{nuc_label}.png\"))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3701cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "pcm_desmond_corr = delta22.pcm_desmond_correlation_by_solvent(query_df_dft)\n",
    "print(\"PCM vs. Desmond Pearson R by nucleus and solvent:\")\n",
    "display(pcm_desmond_corr.round(3))\n",
    "composite_plots.plot_pcm_desmond_correlation_table(pcm_desmond_corr,\n",
    "                                   save_path=figure_path(\"ablations_pcm_desmond_corr_table.png\"))\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}