{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure @jupyter-widgets/jupyterlab-manager\n",
    "#       and jupyter-matplotlib\n",
    "# are installed and enabled in the extension manager.\n",
    "\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numba\n",
    "\n",
    "# Extra performance libraries\n",
    "import numpy as np\n",
    "\n",
    "# Plotting\n",
    "from ipywidgets import interact"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxiterations = 50\n",
    "\n",
    "# 300 x 400 matrix of complex numbers from [-1.5j, 1.5j] x [-2, 2]\n",
    "c = np.sum(np.broadcast_arrays(*np.ogrid[-1.5j:1.5j:300j, -2:2:400j]), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fractal_numpy(c, maxiterations):\n",
    "    f = np.zeros_like(c, dtype=np.int32)\n",
    "    z = c.copy()\n",
    "\n",
    "    for i in range(1, maxiterations + 1):\n",
    "        z = z**2 + c  # Compute z\n",
    "        diverge = abs(z**2) > 2**2  # Divergence criteria\n",
    "\n",
    "        z[diverge] = 2  # Keep number size small\n",
    "        f[~diverge] = i  # Fill in non-diverged iteration number\n",
    "\n",
    "    return f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@numba.njit((numba.complex128[:, :], numba.int32))\n",
    "def fractal_numba(c, maxiterations):\n",
    "    fractal = np.zeros_like(c, dtype=np.int32)\n",
    "\n",
    "    for yi in range(c.shape[0]):\n",
    "        for xi in range(c.shape[1]):\n",
    "            z = cxy = c[yi, xi]\n",
    "            for i in range(1, maxiterations + 1):\n",
    "                z = z**2 + cxy\n",
    "                if abs(z) > 2:\n",
    "                    break\n",
    "                fractal[yi, xi] = i\n",
    "    return fractal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change the numpy calculation to the numba one. Do you see a difference?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "mesh = ax.imshow(c.real, vmin=0, vmax=1)\n",
    "ax.set_xlabel(\"Re(x)\")\n",
    "ax.set_ylabel(\"Im(y)\")\n",
    "\n",
    "\n",
    "@interact(centerx=(-2.0, 2.0, 0.01), centery=(-2.0, 2.0, 0.1), scale=(-5.0, 2, 0.01))\n",
    "def interactive_fractal(centerx=0.38, centery=-0.6, scale=0.25):\n",
    "    maxiterations = 50\n",
    "    scale = 10**scale\n",
    "\n",
    "    c = np.sum(\n",
    "        np.broadcast_arrays(\n",
    "            *np.ogrid[\n",
    "                (centery - scale) * 1j : (centery + scale) * 1j : 400j,\n",
    "                (centerx - scale) : (centerx + scale) : 400j,\n",
    "            ]\n",
    "        ),\n",
    "        axis=0,\n",
    "    )\n",
    "\n",
    "    f = fractal_numpy(c, maxiterations)\n",
    "    mesh.set_data(f / 50)\n",
    "    mesh.set_extent(\n",
    "        (centerx - scale, centerx + scale, centery - scale, centery + scale)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:performance-minicourse] *",
   "language": "python",
   "name": "conda-env-performance-minicourse-py"
  },
  "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}