Skip to content

Commit

Permalink
Fix notebooks
Browse files Browse the repository at this point in the history
  • Loading branch information
@seibert
seibert committed Aug 12, 2024
1 parent ee09861 commit 1441e57
Show file tree
Hide file tree
Showing 4 changed files with 143 additions and 107 deletions.
127 changes: 66 additions & 61 deletions 1 - Numba basics.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,34 @@
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Notebook requirements\n",
"\n",
"This notebook assumes you have a Python environment with:\n",
"\n",
"* Python 3.9 or later\n",
"* Numba 0.59 or later\n",
"* SciPy\n",
"\n",
"The notebook should work on Windows, Mac, and Linux, but relative performance will depend on the specific hardware you are using."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from numba import __version__ as NUMBA_VERSION\n",
"from scipy import __version__ as SCIPY_VERSION\n",
"\n",
"print('Numba:', NUMBA_VERSION)\n",
"print('SciPy:', SCIPY_VERSION)"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -335,7 +363,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"<h3><span style=\"color:blue\"> Task 3a: Reimplement the above `awkward_sine` function with explicit looping using the template below...</span></h3>"
"<h3><span style=\"color:blue\"> Task 3a: Reimplement the above `awkward_sine` function with a for loop using the template below...</span></h3>"
]
},
{
Expand All @@ -348,7 +376,9 @@
"def awkward_sine_loops(a):\n",
" m, n = a.shape\n",
" result = np.empty_like(a)\n",
" # You write this part!\n",
" for i in range(m):\n",
" for j in range(n):\n",
" # You write this part!\n",
" return result\n",
"\n",
"# check result\n",
Expand Down Expand Up @@ -389,11 +419,8 @@
"outputs": [],
"source": [
"@jit\n",
"def awkward_exp(a):\n",
" return np.cosh(a) + np.sinh(a)\n",
"\n",
"# check result\n",
"np.testing.assert_allclose(awkward_exp(X), np.exp(X))"
"def angle_to_complex(a):\n",
" return np.cos(a) + np.sin(a)*1j"
]
},
{
Expand All @@ -409,7 +436,7 @@
"metadata": {},
"outputs": [],
"source": [
"%timeit awkward_exp.py_func(X)"
"%timeit angle_to_complex.py_func(X)"
]
},
{
Expand All @@ -425,7 +452,7 @@
"metadata": {},
"outputs": [],
"source": [
"%timeit awkward_exp(X)"
"%timeit angle_to_complex(X)"
]
},
{
Expand All @@ -442,51 +469,16 @@
"outputs": [],
"source": [
"@jit\n",
"def awkward_exp_loops(a):\n",
"def angle_to_complex_loops(a):\n",
" m, n = a.shape\n",
" result = np.empty_like(a)\n",
" result = np.empty_like(a, dtype=np.complex128)\n",
" for i in range(m):\n",
" for j in range(n):\n",
" result[i, j] = np.cosh(a[i, j]) + np.sinh(a[i, j])\n",
" result[i, j] = np.cos(a[i, j]) + np.sin(a[i, j]) * 1j\n",
" return result\n",
" \n",
"# time it\n",
"%timeit awkward_exp_loops(X)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's switch on `fastmath`, this plumbs in Intel's Short Vector Math library (SVML) which swaps out `libm` functions for vectorized ones."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"@jit(fastmath=True)\n",
"def awkward_exp_loops_fastmath(a):\n",
" m, n = a.shape\n",
" result = np.empty_like(a)\n",
" for i in range(m):\n",
" for j in range(n):\n",
" result[i, j] = np.cosh(a[i, j]) + np.sinh(a[i, j])\n",
" return result\n",
"\n",
" \n",
"%timeit awkward_exp_loops_fastmath(X)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"That made things a bit quicker but NumPy beat Numba all both cases...\n",
"\n",
"#### <span style=\"color:green\"> Answer: Anaconda NumPy ``ufunc`` behaviour is hugely optimised and uses MKL VML internally. </span>\n"
"%timeit angle_to_complex_loops(X)"
]
},
{
Expand Down Expand Up @@ -530,7 +522,7 @@
"metadata": {},
"outputs": [],
"source": [
"@njit\n",
"@jit\n",
"def jit_cos(x):\n",
" return math.cos(x)\n",
"\n",
Expand All @@ -542,7 +534,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"* Expect Numba to outperform BLAS:"
"* Expect Numba to outperform the optimized linear algebra library (BLAS) used by NumPy and SciPy:"
]
},
{
Expand All @@ -551,15 +543,15 @@
"metadata": {},
"outputs": [],
"source": [
"@njit\n",
"@jit\n",
"def dgemv(A, x):\n",
" m, n = A.shape\n",
" acc = np.zeros((m,))\n",
" for i in range(m):\n",
" for j in range(n):\n",
" acc[i] += A[i, j] * x[j]\n",
" \n",
"@njit\n",
"@jit\n",
"def call_dot(A, x):\n",
" return np.dot(A, x) # Numba will defer this to BLAS::XGEMV\n",
"\n",
Expand All @@ -579,7 +571,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"## The importance of measurement:\n",
"## The power of inlining functions for the compiler to use.\n",
"\n",
"Find the root of `f(x) = cos(x) + 1`. Use the Newton-Raphson algorithm. Use SciPy `optimize.newton` as a performance baseline."
]
Expand Down Expand Up @@ -609,7 +601,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"<h3><span style=\"color:blue\"> Task 4: Write your own Newton-Raphson implementation based on based on pseudo-code from <a href=\"https://en.wikipedia.org/wiki/Newton%27s_method#Pseudocode\">Wikipedia</a>.</span></h3>"
"Let's see how Numba can further optimize this. First we will write our own implementation of Newton-Raphson that Numba can compile:"
]
},
{
Expand All @@ -618,10 +610,23 @@
"metadata": {},
"outputs": [],
"source": [
"@njit\n",
"@jit\n",
"def NR_root(f, dfdx, x0, max_it=max_it, eps=tol):\n",
" converged = False\n",
" # You write this part!\n",
" for _ in range(max_it):\n",
" y = f(x0)\n",
" yprime = dfdx(x0)\n",
"\n",
" if abs(yprime) < 1e-9: # Give up if the denominator is too small\n",
" break\n",
"\n",
" x1 = x0 - y / yprime # Do Newton's computation\n",
"\n",
" if abs(x1 - x0) <= eps: # Stop when the result is within the desired tolerance\n",
" converged = True # x1 is a solution within tolerance and maximum number of iterations\n",
"\n",
" x0 = x1 # Update x0 to start the process again\n",
" \n",
" if converged:\n",
" return x1\n",
" else:\n",
Expand Down Expand Up @@ -660,8 +665,8 @@
"metadata": {},
"outputs": [],
"source": [
"f_jit = njit(f)\n",
"dfdx_jit = njit(dfdx)\n",
"f_jit = jit(f) # We can use jit like a regular function, without the @ decorator symbol.\n",
"dfdx_jit = jit(dfdx)\n",
"\n",
"%timeit NR_root.py_func(f_jit, dfdx_jit, x0)"
]
Expand Down Expand Up @@ -716,7 +721,7 @@
" # Numba will spot this and LLVM will inline them into the machine code\n",
" @njit\n",
" def NR_root(x0, max_it=max_it, eps=tol):\n",
" # Copy your NR function body from above here\n",
" # Copy the contents of NR_root above into this space. Don't forget to indent it properly!\n",
" return NR_root\n",
"\n",
"f_NR_root = specialize(f_jit, dfdx_jit)\n",
Expand All @@ -735,7 +740,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
Expand All @@ -749,7 +754,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.7.4"
"version": "3.12.2"
}
},
"nbformat": 4,
Expand Down
Loading

0 comments on commit 1441e57

Please sign in to comment.