5. speed.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from chugunov_indicator import skip_chugunov_2009\n",
    "from chugunov_indicator.analytic.predict import _skip_chugunov_2009\n",
    "import chugunov_indicator as chug\n",
    "\n",
    "import keras\n",
    "\n",
    "import pynucastro as pyna\n",
    "import yt\n",
    "\n",
    "fixed_rng = np.random.default_rng(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs = {\"size\": 3*10**5, \"rng\": fixed_rng}\n",
    "\n",
    "train = chug.ScreeningFactorData(**kwargs)\n",
    "validate = chug.ScreeningFactorData(**kwargs)\n",
    "test = chug.ScreeningFactorData(**kwargs)\n",
    "\n",
    "network = chug.ScreeningFactorNetwork(train, validate, test, 0)\n",
    "network.model = keras.saving.load_model(\"chugunov_indicator/keras/model.keras\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = test.T\n",
    "D = test.D\n",
    "y0 = test.y0\n",
    "abar = test.abar\n",
    "z2bar = test.z2bar\n",
    "z1 = test.z1\n",
    "z2 = test.z2\n",
    "\n",
    "inputs = np.vstack((3*np.log10(T) - np.log10(D), abar, np.log10(z2bar), z1, z2)).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "765 ms ± 31.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for i in range(train.size):\n",
    "    F = chug.chugunov_2009(T[i], D[i], abar[i], 1, z2bar[i], z1[i], z2[i], 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.52 s ± 17.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "for i in range(train.size):\n",
    "    if _skip_chugunov_2009(3*np.log10(T[i]) - np.log10(D[i]), abar[i], z2bar[i], z1[i], z2[i]):\n",
    "        F = 1\n",
    "    else:\n",
    "        F = chug.chugunov_2009(T[i], D[i], abar[i], 1, z2bar[i], z1[i], z2[i], 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "yt : [INFO     ] 2024-10-10 01:17:10,147 Parameters: current_time              = 0.800001438968278\n",
      "yt : [INFO     ] 2024-10-10 01:17:10,148 Parameters: domain_dimensions         = [ 640 1280    1]\n",
      "yt : [INFO     ] 2024-10-10 01:17:10,148 Parameters: domain_left_edge          = [0. 0. 0.]\n",
      "yt : [INFO     ] 2024-10-10 01:17:10,149 Parameters: domain_right_edge         = [5.12000000e+09 1.02400000e+10 6.28318531e+00]\n"
     ]
    }
   ],
   "source": [
    "# Double Detonation data: http://groot.astro.sunysb.edu/common/\n",
    "ds = yt.load('./data/subch_plt17526/')\n",
    "ad = ds.all_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "reaclib_library = pyna.ReacLibLibrary()\n",
    "\n",
    "nuclei = [field[2:-1] for field in np.array(ds.field_list)[:,1] if \"X(\" in field]\n",
    "comp = pyna.Composition(nuclei)\n",
    "comp.set_solar_like()\n",
    "\n",
    "mynet = reaclib_library.linking_nuclei(comp.keys())\n",
    "pynet = pyna.PythonNetwork(libraries=[mynet])\n",
    "\n",
    "screen_map = pyna.screening.get_screening_map(\n",
    "    pynet.get_rates(),\n",
    "    symmetric_screening=pynet.symmetric_screening\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "pair = screen_map[5]\n",
    "scn_fac = pyna.make_screen_factors(pair.n1, pair.n2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we consider a case where screening can't be skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13.611971657871864"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plasma = pyna.make_plasma_state(2e7, 1e4, comp.get_molar())\n",
    "pyna.screening.chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.7 µs ± 39.7 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "F = pyna.screening.chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.25 µs ± 29.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "skip = skip_chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we consider a case where it can be skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0012163585951908"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "plasma = pyna.make_plasma_state(1e9, 1e2, comp.get_molar())\n",
    "pyna.screening.chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.67 µs ± 46.4 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "F = pyna.screening.chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.26 µs ± 35.6 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "skip = skip_chugunov_2009(plasma, scn_fac)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that `skip_chugunov_2009` is around 25% faster in all cases.\n",
    "\n",
    "To test how effective this, we can test more double detonation data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "yt : [INFO     ] 2024-10-10 01:18:07,888 Parameters: current_time              = 1.00002871903213\n",
      "yt : [INFO     ] 2024-10-10 01:18:07,889 Parameters: domain_dimensions         = [ 640 1280    1]\n",
      "yt : [INFO     ] 2024-10-10 01:18:07,889 Parameters: domain_left_edge          = [0. 0. 0.]\n",
      "yt : [INFO     ] 2024-10-10 01:18:07,891 Parameters: domain_right_edge         = [5.12000000e+09 1.02400000e+10 6.28318531e+00]\n"
     ]
    }
   ],
   "source": [
    "# Double Detonation data: http://groot.astro.sunysb.edu/common/\n",
    "ds = yt.load('./data/subch_plt21866/')\n",
    "ad = ds.all_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "Xs = (np.array([ad[f\"X({nucleus})\"] for nucleus in nuclei]).T)\n",
    "T, D = ad[\"Temp\"], ad[\"density\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3123643, 3845543, 1062025, ...,  929823,  784546,  171927])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rng = np.random.default_rng(10)\n",
    "\n",
    "total = 10000\n",
    "choices = rng.choice(np.arange(len(ad[\"density\"])), total)\n",
    "choices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "plasma_states = []\n",
    "for i in choices:\n",
    "    comp.set_array(Xs[i])\n",
    "    plasma = pyna.make_plasma_state(T[i], D[i], comp.get_molar())\n",
    "    plasma_states.append(plasma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "18.4 ms ± 712 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for plasma_ in plasma_states:\n",
    "    F = pyna.screening.chugunov_2009(plasma_, scn_fac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.1 ms ± 155 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "\n",
    "for plasma_ in plasma_states:\n",
    "    if skip_chugunov_2009(plasma, scn_fac):\n",
    "        F = 1\n",
    "    else:\n",
    "        F = pyna.screening.chugunov_2009(plasma_, scn_fac)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As such, this can improve screening calculation times by around $25\\%$ when used with `pynucastro`.\n",
    "\n",
    "Note that since these values are the same for given `PlasmaState` and `ScreenFactors` objects, the value $C_{PS} := C_0 + C_1 + C_2$ can be pre-computed when the `PlasmaState` object is initialized, and same with $C_{SF} := C_3$ for the `ScreenFactors` object. Since the same `PlasmaState` is used for several screening calculations, this would speed up the calculation times even further."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}