{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "de02451b",
   "metadata": {},
   "source": [
    "# Tables S3 and S4: MagNET performance across geometries\n",
    "\n",
    "How accurately MagNET reproduces its DFT training reference (PBE0/pcSseg-1) as geometries move from\n",
    "stationary to vibrated to solvated, for the foundation model and the chloroform MagNET-x."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6531dbc",
   "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/magnet_test_predictions\", \"analysis/code\", \"analysis/code/shared\"):\n",
    "    sys.path.insert(0, os.path.join(REPO, _p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab667bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import magnet_test_predictions_reader\n",
    "import magnet_benchmark\n",
    "import paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "979ab00c",
   "metadata": {},
   "outputs": [],
   "source": [
    "PREDICTIONS = paths.dataset_file(\"magnet_test_predictions\", root=REPO)\n",
    "\n",
    "def document_path(name):\n",
    "    # table/spreadsheet outputs go under this notebook's documents/ folder, created on first save\n",
    "    os.makedirs(\"documents\", exist_ok=True)\n",
    "    return os.path.join(\"documents\", name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059bba13",
   "metadata": {},
   "outputs": [],
   "source": [
    "rows = magnet_benchmark.exact_stats_table(PREDICTIONS, magnet_test_predictions_reader)\n",
    "res = pd.DataFrame(rows).set_index([\"nucleus\", \"model\", \"test_set\"])\n",
    "\n",
    "table_rows = []\n",
    "for nucleus in (\"1H\", \"13C\"):\n",
    "    for model in magnet_benchmark.MODELS:\n",
    "        for ts in magnet_benchmark.TEST_SETS:\n",
    "            r = res.loc[(nucleus, model, ts)]\n",
    "            pmed, pmae, prmse = magnet_benchmark.PUBLISHED[nucleus][(model, ts)]\n",
    "            table_rows.append(dict(nucleus=nucleus, model=model, test_set=ts, n=int(r.n),\n",
    "                                   median_repro=round(r.median_ae, 6), median_SI=round(pmed, 6),\n",
    "                                   mae_repro=round(r.mae, 6), mae_SI=round(pmae, 6),\n",
    "                                   rmse_repro=round(r.rmse, 5), rmse_SI=round(prmse, 5)))\n",
    "table = pd.DataFrame(table_rows)\n",
    "table_s3 = table[table.nucleus == \"1H\"].drop(columns=\"nucleus\")\n",
    "table_s4 = table[table.nucleus == \"13C\"].drop(columns=\"nucleus\")\n",
    "print(\"Table S3 (1H):\"); display(table_s3)\n",
    "print(\"Table S4 (13C):\"); display(table_s4)\n",
    "\n",
    "# write the reproduced tables to this notebook's documents/ folder, one sheet per SI table\n",
    "out = document_path(\"si_table_s03_s04_performance.xlsx\")\n",
    "with pd.ExcelWriter(out) as writer:\n",
    "    table_s3.to_excel(writer, sheet_name=\"Table S3 (1H)\", index=False)\n",
    "    table_s4.to_excel(writer, sheet_name=\"Table S4 (13C)\", index=False)\n",
    "print(\"wrote\", os.path.relpath(out, REPO))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d795f1",
   "metadata": {},
   "source": [
    "## Exact-reproduction check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdb61e94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# every row's reproduced median/MAE/RMSE should match the published SI value to a few parts per million\n",
    "max_dev = (table[[\"median_repro\", \"median_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n",
    "          table[[\"mae_repro\", \"mae_SI\"]].diff(axis=1).iloc[:, -1].abs().max(),\n",
    "          table[[\"rmse_repro\", \"rmse_SI\"]].diff(axis=1).iloc[:, -1].abs().max())\n",
    "print(\"largest reproduced-vs-published deviation (median, MAE, RMSE):\", max_dev)\n",
    "assert max(max_dev) < 1e-3, \"a row diverged from the SI by more than float rounding\""
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}