{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2026-04-14T14:53:20.072496Z",
     "start_time": "2026-04-14T14:53:19.712799Z"
    }
   },
   "source": [
    "%config InlineBackend.figure_format = 'retina'\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from plot import discover_runs, load_run_data, get_pynuc_eps\n",
    "\n",
    "from scipy.interpolate import interp1d\n",
    "\n",
    "plt.style.use(\"../ManuscriptFigures/utils/pub.mplstyle\")"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T14:53:22.408612Z",
     "start_time": "2026-04-14T14:53:22.393261Z"
    }
   },
   "cell_type": "code",
   "source": "base_dir = \"GF_Validation_Output\"",
   "id": "6b80320acc86fb9e",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T14:53:23.299257Z",
     "start_time": "2026-04-14T14:53:23.282646Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def collect_species(steps, sp):\n",
    "    abundances = []\n",
    "    for step in steps:\n",
    "        abundances.append(step[\"Composition\"][sp])\n",
    "    return np.array(abundances)\n",
    "\n",
    "def collect_time(steps):\n",
    "    times = []\n",
    "    for step in steps:\n",
    "        times.append(step[\"t\"])\n",
    "    return np.array(times)"
   ],
   "id": "8d4660c043840d00",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T16:13:37.752792Z",
     "start_time": "2026-04-14T16:13:37.735551Z"
    }
   },
   "cell_type": "code",
   "source": "runs = discover_runs(base_dir)",
   "id": "ff332b29e27a5d75",
   "outputs": [],
   "execution_count": 38
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T16:14:42.864567Z",
     "start_time": "2026-04-14T16:14:42.851953Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def plot_run_abundance(run_id, species=(\"H-1\", \"He-3\", \"He-4\", \"C-12\", \"N-14\", \"O-16\")):\n",
    "    fig, axs = plt.subplots(2, 1, figsize=(10, 10), sharex=True, gridspec_kw={'hspace': 0})\n",
    "    data = load_run_data(base_dir, run_id)\n",
    "\n",
    "    for spID, sp in enumerate(species):\n",
    "        gf_t, gf_y = collect_time(data[0][\"Steps\"]), collect_species(data[0][\"Steps\"], sp)\n",
    "        py_t, py_y = collect_time(data[1][\"Steps\"]), collect_species(data[1][\"Steps\"], sp)\n",
    "\n",
    "        axs[0].loglog(gf_t, gf_y, color=f\"C{spID}\", label=f\"GF {sp}\")\n",
    "        axs[0].loglog(py_t, py_y, color=f\"C{spID}\", linestyle=\"dashed\")\n",
    "\n",
    "        el = sp.split('-')[0]\n",
    "        A = sp.split('-')[1]\n",
    "        axs[0].text(gf_t[0], gf_y[0]*0.45, rf\"$^{{{A}}}${el}\", color=f\"C{spID}\", fontsize=12)\n",
    "\n",
    "        f_pync = interp1d(py_t, py_y, fill_value=\"extrapolate\", bounds_error=False)\n",
    "        f_gf = interp1d(gf_t, gf_y, fill_value=\"extrapolate\", bounds_error=False)\n",
    "\n",
    "        t_common = np.logspace(np.log10(max(min(gf_t), min(py_t))), np.log10(min(max(gf_t), max(py_t))), 100)\n",
    "        py_interp = f_pync(t_common)\n",
    "        gf_interp = f_gf(t_common)\n",
    "\n",
    "        log_diff = np.abs(np.log10(py_interp) - np.log10(gf_interp))\n",
    "\n",
    "        axs[1].loglog(t_common, log_diff, color=f\"C{spID}\")\n",
    "        # axs[1].set_xlim(t_common[0], 3.1e17)\n",
    "\n",
    "    axs[0].set_ylabel(\"Abundance [mol/g]\")\n",
    "    axs[1].set_xlabel(\"Time [s]\")\n",
    "    axs[1].set_ylabel(r\"$\\left|\\Delta\\log_{10}\\right|$ [dex]\")\n",
    "\n",
    "    return fig, axs\n"
   ],
   "id": "32114ccdfbdc0c30",
   "outputs": [],
   "execution_count": 40
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T16:22:25.941637Z",
     "start_time": "2026-04-14T16:22:23.679453Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for run in runs:\n",
    "    if \"No_QSE\" in run:\n",
    "        fig, axs = plot_run_abundance(run);\n",
    "        axs[0].set_ylim(1e-8)\n",
    "        plt.savefig(f\"{run}_abundance.pdf\", bbox_inches='tight')\n",
    "        plt.close()"
   ],
   "id": "bd04488fead04880",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/v3/xfq17d692dg_gzdwv4yzmrkw0000gp/T/ipykernel_42214/2916075565.py:23: RuntimeWarning: divide by zero encountered in log10\n",
      "  log_diff = np.abs(np.log10(py_interp) - np.log10(gf_interp))\n",
      "/var/folders/v3/xfq17d692dg_gzdwv4yzmrkw0000gp/T/ipykernel_42214/2916075565.py:23: RuntimeWarning: invalid value encountered in log10\n",
      "  log_diff = np.abs(np.log10(py_interp) - np.log10(gf_interp))\n"
     ]
    }
   ],
   "execution_count": 44
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T15:30:39.211161Z",
     "start_time": "2026-04-14T15:30:38.434430Z"
    }
   },
   "cell_type": "code",
   "source": [
    "num_species = []\n",
    "py_time = []\n",
    "gf_time = []\n",
    "for run in runs:\n",
    "    if \"Depth\" in run:\n",
    "        data = load_run_data(base_dir, run)\n",
    "        num_species.append(data[0][\"Metadata\"][\"NumSpecies\"])\n",
    "        gf_time.append(data[0][\"Metadata\"][\"ElapsedTime\"])\n",
    "        py_time.append(data[1][\"Metadata\"][\"RunTime1\"])\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 7))\n",
    "ax.semilogy(num_species, gf_time, 'o', label='GridFire', color='C2')\n",
    "ax.semilogy(num_species, py_time, 'D', label='pynucastro', color='black')\n",
    "\n",
    "ax.legend(fontsize=17, frameon=False)\n",
    "ax.set_xlabel(\"Number of Species\", fontsize=25)\n",
    "ax.set_ylabel(\"Wall Time [s]\", fontsize=25)\n",
    "\n",
    "plt.savefig(\"GF_PyNuc_WallTime.pdf\", bbox_inches='tight')\n",
    "\n"
   ],
   "id": "e64c82995a57c114",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Figure size 1000x700 with 1 Axes>"
      ],
      "image/png": "[encoded data, 61528 characters]"
     },
     "metadata": {
      "image/png": {
       "width": 896,
       "height": 622
      }
     },
     "output_type": "display_data",
     "jetTransient": {
      "display_id": null
     }
    }
   ],
   "execution_count": 26
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2026-04-14T15:30:42.930128Z",
     "start_time": "2026-04-14T15:30:42.888964Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for ns, t in zip(num_species, py_time):\n",
    "    print(f\"pynucastro: {ns} species, {t:.2e} seconds\")"
   ],
   "id": "b0b9ceebfa5538b3",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pynucastro: 357 species, 8.47e-03 seconds\n",
      "pynucastro: 457 species, 1.07e-02 seconds\n",
      "pynucastro: 568 species, 1.17e-02 seconds\n",
      "pynucastro: 61 species, 8.80e-03 seconds\n",
      "pynucastro: 580 species, 1.05e-02 seconds\n",
      "pynucastro: 98 species, 8.55e-03 seconds\n",
      "pynucastro: 136 species, 8.49e-03 seconds\n",
      "pynucastro: 177 species, 7.85e-03 seconds\n",
      "pynucastro: 219 species, 1.21e-02 seconds\n",
      "pynucastro: 261 species, 1.21e-02 seconds\n"
     ]
    }
   ],
   "execution_count": 27
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
