Merge pull request #60 from EcohydrologyTeam/57-improve-timestep-perf…

…ormance-further Adding `examples/dev_sandbox/prof.py`
EcohydrologyTeam · Dec 14, 2023 · 8b53234 · 8b53234
2 parents 999e82d + 9b97882
commit 8b53234
Show file tree

Hide file tree

Showing 2 changed files with 271 additions and 0 deletions.
diff --git a/examples/dev_sandbox/improve_performance_v2.ipynb b/examples/dev_sandbox/improve_performance_v2.ipynb
@@ -0,0 +1,233 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "52b6fe1e-6174-4855-8230-e44c637b754e",
+   "metadata": {},
+   "source": [
+    "# Numba Performance exploration notebook\n",
+    "\n",
+    "**Overview:** After realizing that removing numba speed up performance 8 fold, I wanted to make sure this persists for longer runtimes where JIT is designed to be effective."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a01ac108-0980-4d68-89ac-3b49013fb939",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import numba\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "33e20862-ac39-46be-aa1d-0e1ee51a8994",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "@numba.njit\n",
+    "def jitadd(x, y) -> np.ndarray:\n",
+    "    return (x * 0.2) + (y * 0.3)\n",
+    "\n",
+    "def add(x, y) -> np.ndarray:\n",
+    "    return (x * 0.2) + (y * 0.3)\n",
+    "\n",
+    "@numba.njit\n",
+    "def jitsub(x) -> np.ndarray:\n",
+    "    return x - (0.15 * x)\n",
+    "\n",
+    "def sub(x) -> np.ndarray:\n",
+    "    return x - (0.15 * x)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "3bf02ded-b09c-4791-9d3a-3f299c29bad5",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[0.25289281, 0.3012049 , 0.81652952, ..., 0.08600542, 0.59860995,\n",
+       "        0.7160782 ],\n",
+       "       [0.23049757, 0.7084206 , 0.008605  , ..., 0.12755473, 0.76453943,\n",
+       "        0.70370063],\n",
+       "       [0.34352444, 0.21580305, 0.24908259, ..., 0.54471826, 0.42611405,\n",
+       "        0.9940653 ],\n",
+       "       ...,\n",
+       "       [0.73328428, 0.46605497, 0.29684926, ..., 0.38692169, 0.1791475 ,\n",
+       "        0.24509559],\n",
+       "       [0.22376018, 0.713524  , 0.47561131, ..., 0.0983168 , 0.99624008,\n",
+       "        0.96001501],\n",
+       "       [0.17163253, 0.21930795, 0.68329088, ..., 0.34515248, 0.92341599,\n",
+       "        0.65903902]])"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "input = np.ones(shape=(100, 100)) * np.random.random_sample((100, 100))\n",
+    "input"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "46b28de0-95a5-49b9-8a85-3f468b34993f",
+   "metadata": {},
+   "source": [
+    "**No JIT**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "8760934b-879b-480f-b8f1-6acd8438ef4f",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "61.3 ms ± 5.82 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "TIMESTEPS = 1000\n",
+    "out_dict = {}\n",
+    "for i in range(TIMESTEPS):\n",
+    "    val = add(input, -input)\n",
+    "    out_dict[i] = sub(val)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "00898ada-0ca7-4b5a-baf1-ccad07219545",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "**Just the function being JIT**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "cb8e6fd5-bfc7-4c08-bc3f-a1e4ce0eac06",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "45.2 ms ± 1.65 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "TIMESTEPS = 1000\n",
+    "out_dict = {}\n",
+    "for i in range(TIMESTEPS):\n",
+    "    val = jitadd(input, -input)\n",
+    "    out_dict[i] = jitsub(val)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dde04c99-8b06-47fb-ba8f-838f83501755",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "**The functions and loop being JIT**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "id": "d61c90a2-476d-4769-84ee-40f8c8e57482",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "477 ms ± 24.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%timeit\n",
+    "TIMESTEPS = 1000\n",
+    "\n",
+    "@numba.njit\n",
+    "def loop(input_var, time_steps):\n",
+    "    out_dict = {}\n",
+    "    for i in range(time_steps):\n",
+    "        val = jitadd(input_var, -input_var)\n",
+    "        out_dict[i] = jitsub(val)\n",
+    "    return out_dict\n",
+    "\n",
+    "out_dict = loop(input, TIMESTEPS)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a86bda76-3aec-4855-8101-9dca521bb5e3",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa488e31-835f-424c-b050-a2b9e7d68683",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/dev_sandbox/prof.py b/examples/dev_sandbox/prof.py
@@ -0,0 +1,38 @@
+"""A script to allow for debugging of the TSM module."""
+import clearwater_modules
+import time 
+import sys
+
+def main(iters: int):
+    ti = time.time()
+    # define starting state values
+    state_i = {
+        'water_temp_c': 40.0,
+        'surface_area': 1.0,
+        'volume': 1.0,
+    }
+
+    # instantiate the TSM module
+    tsm = clearwater_modules.tsm.EnergyBudget(
+        initial_state_values=state_i,
+        meteo_parameters={'wind_c': 1.0},
+    )
+    print(tsm.static_variable_values)
+    t2 = time.time()
+    for _ in range(iters):
+        tsm.increment_timestep()
+    print(f'Increment timestep speed (average of {iters}): {(time.time() - t2) / 100}')
+    print(f'Run time: {time.time() - ti}')
+
+if __name__ == '__main__':
+    if len(sys.argv) > 1:
+        try:
+            iters = int(sys.argv[1])
+            print(f'Running {iters} iterations.')
+        except ValueError:
+            raise ValueError('Argument must be an integer # of iterations.')
+    else:
+        print('No argument given, defaulting to 100 iteration.')
+        iters = 100
+
+    main(iters=iters)