{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "706410fd",
   "metadata": {},
   "source": [
    "# Figure 5B: MagNET-Zero vs DFT shieldings on the DFT8K benchmark\n",
    "\n",
    "Reproduces Figure 5B: MagNET-Zero gas-phase shielding predictions vs DFT on the DFT8K benchmark (~7,000\n",
    "organics, untrained). Protons vs WP04/pcSseg-2, carbons vs wB97X-D/pcSseg-2; residual = DFT minus\n",
    "MagNET."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e0a3b95",
   "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/dft8k\", \"analysis/code\", \"analysis/code/shared\"):\n",
    "    sys.path.insert(0, os.path.join(REPO, _p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9838352",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dft8k_residuals\n",
    "import paths\n",
    "import fig5b_plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2231e1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "DFT8K_HDF5 = paths.dataset_file(\"dft8k\", root=REPO)\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": "b7d02a23",
   "metadata": {},
   "outputs": [],
   "source": [
    "shieldings = dft8k_residuals.load_shieldings(DFT8K_HDF5)\n",
    "print(\"atoms in DFT8K:\", len(shieldings[\"atomic_numbers\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "867ed2b4",
   "metadata": {},
   "source": [
    "## Residual statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ca1b09",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus in [\"H\", \"C\"]:\n",
    "    errors = dft8k_residuals.residuals(shieldings, nucleus)\n",
    "    stats = dft8k_residuals.residual_stats(errors)\n",
    "    print(f\"{nucleus}: n={stats['n']:,}  RMSE={stats['rmse']:.4f} ppm  MAE={stats['mae']:.4f} ppm  \"\n",
    "          f\"95% abs={stats['abs_p95']:.4f} ppm  fraction < 0.1 ppm={stats['frac_below']:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d95406",
   "metadata": {},
   "source": [
    "## Residual histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf9148b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus in [\"H\", \"C\"]:\n",
    "    errors = dft8k_residuals.residuals(shieldings, nucleus)\n",
    "    stats = dft8k_residuals.residual_stats(errors)\n",
    "    print(f\"{nucleus}: RMSE={stats['rmse']:.3f} ppm  frac within 0.1 ppm={stats['frac_below']:.3f}\")\n",
    "    if nucleus == \"H\":\n",
    "        fig5b_plots.plot_dft8k_residual_histogram(errors, figure_path(\"fig5b_residuals_1H.png\"),\n",
    "                                      bin_width=0.0025, xlim=0.23, grey_box=0.1, x_tick=0.05)\n",
    "    else:   # 13C: wider window, no +/-0.1 box since that threshold is 1H-specific\n",
    "        fig5b_plots.plot_dft8k_residual_histogram(errors, None,\n",
    "                                      bin_width=0.05, xlim=5.0, grey_box=None, x_tick=1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eab006a",
   "metadata": {},
   "source": [
    "## Extreme-residual callouts\n",
    "\n",
    "The largest positive and negative ¹H residuals over the full set. The published negative callout (a\n",
    "sulfonium zwitterion, -1.040 ppm) is not the true global minimum -- see the printed note."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f2315b",
   "metadata": {},
   "outputs": [],
   "source": [
    "extreme_max = dft8k_residuals.find_extreme_residual(DFT8K_HDF5, \"H\", sign=\"max\")\n",
    "print(\"largest positive 1H residual:\", extreme_max)\n",
    "fig5b_plots.show_dft8k_molecule(extreme_max[\"smiles\"])\n",
    "\n",
    "zwitterion = dft8k_residuals.molecule_by_id(DFT8K_HDF5, \"H\", molecule_id=88779)\n",
    "print(\"published sulfonium-zwitterion callout (id 88779):\", zwitterion)\n",
    "fig5b_plots.show_dft8k_molecule(zwitterion[\"smiles\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f4276e7",
   "metadata": {},
   "source": "## Functional-group error breakdown\n\nMean absolute residual for six functional groups (Carbonyls, Amines, Sulfonyl, Pyridines, Furans,\nNitroso), via RDKit SMARTS matching over every molecule's SMILES."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f95f9d58",
   "metadata": {},
   "outputs": [],
   "source": [
    "for nucleus in [\"H\", \"C\"]:\n",
    "    group_errors = dft8k_residuals.functional_group_errors(DFT8K_HDF5, nucleus)\n",
    "    print(f\"--- {nucleus} ---\")\n",
    "    for name, v in group_errors.items():\n",
    "        print(f\"{name:12s} mean|error|={v['mean_abs_error']:.4f} ppm  n_molecules={v['n_molecules']:,}\")\n",
    "    save = figure_path(\"fig5b_functional_groups_1H.png\") if nucleus == \"H\" else None\n",
    "    fig5b_plots.plot_dft8k_functional_group_errors(group_errors, nucleus, save)"
   ]
  }
 ],
 "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
}