{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "55d51b05",
   "metadata": {},
   "source": "# Figure 2A: accuracy-vs-compute-cost Pareto frontier (delta-22 ¹H, CDCl3)\n\nFitting RMSE vs total compute time (geometry + shielding), colored by method type, shaped by geometry source, sized by basis."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5fb8cc9",
   "metadata": {},
   "outputs": [],
   "source": "import os, sys\n\n# make the in-repo modules importable (not pip-installed)\nREPO = os.path.abspath(\"../..\")\nfor _p in (\"data/delta22\", \"analysis/code\", \"analysis/code/shared\"):\n    sys.path.insert(0, os.path.join(REPO, _p))"
  },
  {
   "cell_type": "code",
   "id": "24c54da5",
   "source": "import delta22\nimport pareto_plot\nimport paths",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "id": "c2b77287",
   "source": "DELTA22_HDF5 = paths.dataset_file(\"delta22\", root=REPO)\nXLSX = os.path.join(REPO, \"data\", \"delta22\", \"delta22_experimental.xlsx\")\n\ndef figure_path(name):\n    os.makedirs(\"figures\", exist_ok=True)\n    return os.path.join(\"figures\", name)",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8daed8f",
   "metadata": {},
   "outputs": [],
   "source": "dft = delta22.load_query_df_dft(DELTA22_HDF5, XLSX, verbose=False)\nnn = delta22.load_query_df_nn(DELTA22_HDF5, XLSX, verbose=False)\ndft_gas_timings, nn_timings = delta22.load_pareto_timings(DELTA22_HDF5)\n\n# one point per method/basis/geometry/nucleus/solvent (plus solvent-averaged rows); MagNET is a single\n# method (aimnet2, basis \"N/A\"). 100 seeded splits (train on the first 10 shuffled solutes, test on the\n# rest) to match the published panel.\npoints = delta22.fig2a_pareto_points(dft, nn, dft_gas_timings, nn_timings, n_splits=100)\nprint(points[\"nmr_method\"].nunique(), \"methods;\",\n      \"MagNET total time\", float(points.query(\"nmr_method=='MagNET'\")[\"total_time\"].iloc[0]), \"s\")"
  },
  {
   "cell_type": "markdown",
   "id": "556ad0b0",
   "metadata": {},
   "source": [
    "## Proton Pareto panel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c9a93fd",
   "metadata": {},
   "outputs": [],
   "source": "method_legend_label_map = {\n    \"MagNET-Zero\": \"MagNET-Zero\", \"ab initio\": r\"$\\it{ab\\ initio}$\", \"double hybrids\": \"double hybrid\",\n    \"NMR-specific\": \"NMR-specific\", \"DFT\": \"DFT\",\n}\ngeometry_legend_label_map = {\"pbe0_tz\": \"DFT (PBE0/cc-pVTZ)\", \"aimnet2\": \"ML (AIMNet2)\"}\nmethod_color_map = {\"ab initio\": \"#53585f\", \"MagNET-Zero\": \"#000000\", \"double hybrids\": \"#b3a560\",\n                    \"NMR-specific\": \"#a72608\", \"DFT\": \"#4a8075\"}\n\n# which points to draw: MagNET-Zero, all PBE0-geometry DFT methods, plus the wp04 NMR-specific point\nquery_str = (\"nucleus=='{nucleus}' and solvent == 'chloroform' and (\"\n             \" (geometry_type=='aimnet2' and nmr_method.str.startswith('MagNET')) or\"\n             \" (geometry_type=='pbe0_tz' and not nmr_method.str.startswith('MagNET')) or\"\n             \" (geometry_type=='aimnet2' and nmr_method=='wp04' and basis=='pcSseg2' and nucleus=='H'))\")\n\n# hand-tuned label offsets (points, matching the published panel)\nmanual_label_positions = {\n    (\"MagNET-Zero\", \"N/A\", \"aimnet2\"): {\"xytext\": (-28, 14)},\n    (\"wp04\", \"pcSseg2\", \"aimnet2\"): {\"xytext\": (6, 4)},\n    (\"pbe0\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (-22, -2)},\n    (\"b2gp_plyp\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (4, 6)},\n    (\"revdsd_pbep86\", \"pcSseg1\", \"pbe0_tz\"): {\"xytext\": (6, -4)},\n    (\"revdsd_pbep86\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, 6)},\n    (\"tpsstpss\", \"pcSseg2\", \"pbe0_tz\"): {\"xytext\": (6, -2)},\n    (\"b97d3\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (-28, -2)},\n    (\"mpw2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 2)},\n    (\"b2plyp\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (8, 0)},\n    (\"tpsstpss\", \"pcSseg3\", \"pbe0_tz\"): {\"xytext\": (6, -12)},\n}"
  },
  {
   "cell_type": "code",
   "id": "8d87338a",
   "source": "pareto_plot.plot_pareto_panel(\n    points, query_str, nucleus=\"H\",\n    method_color_map=method_color_map,\n    method_legend_label_map=method_legend_label_map,\n    geometry_legend_label_map=geometry_legend_label_map,\n    geometry_marker_map={\"aimnet2\": \"x\", \"pbe0_tz\": \"o\"},\n    xlim_left=(1.7, 1.8), xlim_right=(3.8, 5.5), ylim=(0.08, 0.3),\n    figsize=(12, 6), marker_alpha=0.8,\n    manual_label_positions=manual_label_positions,\n    save_png=figure_path(\"fig2a_pareto_1H.png\"),\n)",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  }
 ],
 "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
}