{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hvordan optimere kode med Numba og `@jit`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Første bud for all optimering:\n",
    "**Skriv først kode som virker!** Ikke tenk på ytelsen, bare gjør ting så rett frem som mulig og kjør tester med lite belastning. Deretter kan du gjøre det raskere hvis det absolutt er nødvendig (ofte er det ikke det). Når du modifiserer en funksjon for å få den raskere, sjekk resultatene fortløpende mot den trege versjonen som du vet stemmer.\n",
    "\n",
    "Når det kommer til Numba betyr dette at du **ikke** slenger på `@jit` uten at du først er **sikker på** at funksjonen fungerer og at den går for sakte.\n",
    "\n",
    "Før du bruker Numba, les igjennom dette dokumentet samt Numba-delen av det numeriske kompendiet: https://www.uio.no/studier/emner/matnat/astro/AST2000/h18/undervisningsmateriell_h2018/Numerical%20Compendium/numba.pdf\n",
    "\n",
    "Det vil sannsynligvis ta deg mindre tid enn å løse feilene som typisk oppstår ved å ikke gjøre det."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bruk `@jit` på trege funksjoner med arrayer og løkker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numba as nb\n",
    "from numba import jit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Her er en Python-funksjon med mye beregning, som bare bruker enkle objekter som arrayer og skalarer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vi kan måle tiden den bruker på å regne ut de 1000 første tallene i Fibonacci-følgen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci(1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Men hva hvis vi trenger å gjøre dette raskere?\n",
    "\n",
    "`fibonacci` er et godt eksempel på en funksjon som egner seg til å gjøres raskere med Numba. Ved å sette inn `@jit(nopython=True)` over funksjonen, kan vi be Numba gjøre om hele funksjonen til effektiv maskinkode rett før den skal kalles. JIT står for \"Just-In-Time compilation\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_jit(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci_jit(1000) # ~150 times faster than fibonacci"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unngå `@jit` på raske funksjoner"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merk at `@jit` ikke bør brukes på funksjoner som allerede går raskt nok. De kan fort vise seg å bli treigere grunnet tiden det tar å konvertere til maskinkode. Å legge til `@jit` innebærer også en viss risiko (som vi skal se), så ikke gjør det med mindre det er nødvendig."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_cheat(N):\n",
    "    \"\"\"\n",
    "    Returns the Fibonacci sequence with N=5 elements.\n",
    "    \"\"\"\n",
    "    assert N == 5, 'N must be 5'\n",
    "    return np.array([0.0, 1.0, 1.0, 2.0, 3.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r1 -n1 fibonacci_cheat(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_cheat_jit(N):\n",
    "    \"\"\"\n",
    "    Returns the Fibonacci sequence with N=5 elements.\n",
    "    \"\"\"\n",
    "    assert N == 5, 'N must be 5'\n",
    "    return np.array([0.0, 1.0, 1.0, 2.0, 3.0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit -r1 -n1 fibonacci_cheat_jit(5) # ~50 000 times slower than fibonacci_cheat for first call!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `@jit` har mindre effekt med mer kompliserte datastrukturer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`@jit` fungerer best når funksjonen (og alle funksjoner som kalles inni) kun bruker arrayer og skalarer til å håndtere data, siden de har en fast lengde og kan representeres enkelt på maskinnivå. Lister, for eksempel, er mer kompliserte siden de ikke har en forhåndsdefinert lengde.\n",
    "\n",
    "Følgende variant av `fibonacci` lagrer verdiene i en liste istendefor en array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_list(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    Grows a list instead of inserting into an array.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = [0.0, 1.0]\n",
    "    for i in range(2, N):\n",
    "        a.append(a[i-1] + a[i-2])\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci_list(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_list_jit(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    Grows a list instead of inserting into an array.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = [0.0, 1.0]\n",
    "    for i in range(2, N):\n",
    "        a.append(a[i-1] + a[i-2])\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Den `@jit`ede versjonen går litt raskere, men ikke så veldig mye."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci_list_jit(1000) # ~7 times faster than fibonacci_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Med `@jit(nopython=True)` kan du fort få en feilmelding hvis funksjonen din benytter for kompliserte datastrukturer, som for eksempel dictionaries. `nopython=True` betyr at all koden i funksjonen skal oversettes til rask maskinkode, slik at Python selv kan lene seg tilbake og la maskinen overta når funksjonen kalles. Men kompliserte datastrukturer kan som regel ikke oversettes på denne måten. Med `nopython=True` vil derimot Python selv ta hånd om de delene av funksjonen som ikke kan oversettes direkte, men dette vil typisk gå mye saktere.\n",
    "\n",
    "Her er en variant av `fibonacci` lagrer verdiene i en `dict` istedenfor en array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_dict(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    Builds a dictionary instead of inserting into an array.\n",
    "    (This is a very stupid way of using a dictionary..)\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = {0: 0.0, 1: 1.0}\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci_dict(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_dict_jit(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    Builds a dictionary instead of inserting into an array.\n",
    "    (This is a very stupid way of using a dictionary..)\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = {0: 0.0, 1: 1.0}\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Med `nopython=True` nekter Numba å lystre, siden den ikke vet hvordan man oversetter dictionaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    fibonacci_dict_jit(1000) # UnsupportedError: Failed in nopython mode pipeline\n",
    "except nb.UnsupportedError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Vi kan istedenfor prøve å sette `nopython=False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=False)\n",
    "def fibonacci_dict_jit_python(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    Builds a dictionary instead of inserting into an array.\n",
    "    (This is a very stupid way of using a dictionary..)\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = {0: 0.0, 1: 1.0}\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nå vil funksjonen kjøre, men den går ikke raskere enn den gamle versjonen siden `dict`-objektet er en sentral del av utregningene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fibonacci_dict_jit_python(1000) # Slightly slower than fibonacci_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `@jit` på funksjoner med bugs kan ha katastrofale konsekvenser"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Husk at du **alltid** må teste at funksjonen din fungerer som den skal **før** du bruker `@jit` på den. Dette gjelder også hver gang du endrer noe i funksjonen senere, **kommenter ut `@jit` hver gang du modifiserer funksjonen og ikke aktiver det igjen før du vet at ting fungerer!**\n",
    "\n",
    "Det er i prinsippet ikke noe problem å endre ting mens funksjonen er `@jit`et, neste gang den kalles vil den oversettes på nytt og endringene vil tre i kraft. Problemet er at `@jit`et kode kan gi kryptiske feilmeldinger som følge av problemer som ikke egentlig har noe med `@jit` å gjøre i det hele tatt, eller enda verre, la være å gi feilmeldinger selv om ting går galt!\n",
    "\n",
    "Et skrekkeksempel på dette kan du se her."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_buggy(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    WARNING: THIS CODE HAS A MISTAKE!\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N+1): # Whoops, should be range(2, N)\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I `fibonacci_buggy` har den øvre grensen i `for`-løkken ved en feiltagelse blitt satt til `N+1` istedenfor `N`. Dette medfører at koden vil prøve å skrive til `a[N]` i siste iterasjon, som er ugyldig siden siste element i arrayen er `a[N-1]`. Men dette sier heldigvis Python tydelig ifra om."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    print(fibonacci_buggy(5)) # Saved by bounds check\n",
    "except IndexError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Men hvordan går dette i den `@jit`ede versjonen?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_buggy_jit(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    WARNING: THIS CODE HAS A MISTAKE!\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N+1): # Whoops, should be range(2, N)\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    print(fibonacci_buggy_jit(5)) # Uh oh, anything can happen now\n",
    "except IndexError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ingen feilmelding.. Men det betyr at koden fikk lov å skrive utenfor arrayen. Dette virker kanskje uskyldig nok, men tenk på hva det innebærer. Hva som helst av informasjon som programmet benytter til å kjøre kan ha vært lagret på minneadressen som vi akkurat overskrev, det være seg kodeinstrukser eller variabler. Det betyr at det ikke er noen måte å forutse hva programmet vil finne på! Og utfallet vil variere med tiden, månefasen og hva du lytter til på Spotify. Vi har støtt på enhver feilsøkers verste mareritt: **udefinert oppførsel**!\n",
    "\n",
    "Moralen er at `@jit` skal holdes langt unna kode som du ikke først har testet at fungerer. Og hvis du mistenker at noe er galt bør det første du gjør være å fjerne `@jit` før du feilsøker."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Andre fallgruver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Kalle funksjoner som er definert uten `@jit`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_nested(N):\n",
    "    return fibonacci(N) # fibonacci is not @jit'ed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fibonacci_nested(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_nested_jit(N):\n",
    "    return fibonacci(N) # fibonacci is not @jit'ed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(fibonacci_nested_jit(5)) # TypingError: Failed in nopython mode pipeline\n",
    "except nb.TypingError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_nested_jit_working(N):\n",
    "    return fibonacci_jit(N) # fibonacci_jit is @jit'ed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fibonacci_nested_jit_working(5))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definere en variabel flere ganger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_redefine(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    The array is first defined as a list and then converted into an array.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = [0.0]*N\n",
    "    a = np.array(a) # Variable a is redefined\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fibonacci_redefine(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_redefine_jit(N):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    The array is first defined as a list and then converted into an array.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = [0.0]*N\n",
    "    a = np.array(a) # Variable a is redefined\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = a[i-1] + a[i-2]\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(fibonacci_redefine_jit(5)) # TypingError: Failed in nopython mode pipeline\n",
    "except nb.TypingError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sende inn en funksjon som argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fibonacci_custom_operator(N, operation):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    The argument `operation` is a function that combines two numbers into one.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = operation(a[i-1], a[i-2]) # Uses inputted function to combine elements\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(fibonacci_custom_operator(5, lambda a, b: a + b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def fibonacci_custom_operator_jit(N, operation):\n",
    "    \"\"\"\n",
    "    Computes the Fibonacci sequence with N elements.\n",
    "    The argument `operation` is a function that combines two numbers into one.\n",
    "    \"\"\"\n",
    "    assert N >= 2, 'N cannot be smaller than 2'\n",
    "    a = np.zeros(N)\n",
    "    a[1] = 1.0\n",
    "    for i in range(2, N):\n",
    "        a[i] = operation(a[i-1], a[i-2]) # Uses inputted function to combine elements\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(fibonacci_custom_operator_jit(5, lambda a, b: a + b)) # TypingError: Failed in nopython mode pipeline\n",
    "except nb.TypingError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.8 64-bit ('base': conda)",
   "language": "python",
   "name": "python36864bitbaseconda81def863cce947cfb63776f8088d8820"
  },
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}