setup:


```python
import numpy as np
import matplotlib.pyplot as plt
from numba import njit, prange
import matplotlib.pyplot as plt
import math

plt.style.use("seaborn-talk")
```


```python
def plot_performance(time_dict):
    mean = [t.average for t in time_dict.values()]
    std = [t.stdev for t in time_dict.values()]
    x = range(len(time_dict))
    plt.errorbar(x, mean, yerr=std, lw=3, fmt="o") 
    plt.xticks( np.arange(len(time_dict)), time_dict.keys())
    plt.ylabel("average execution time (s)")
    plt.grid()
```

# Optimization

## Memoization

Remember the fibonnaci problem we had in the debugging lecture?


```python
GOLDEN = (1 + 5 ** 0.5) / 2

def fibonacci(k):

    if k == 0:
        return 0
    if k == 1:
        return 1

    return fibonacci(k - 2) + fibonacci(k - 1)

def compute_golden_ratio(accuracy_level=10):
    return fibonacci(accuracy_level) / fibonacci(accuracy_level - 1)


def plot_golden_ratio_approx(max_k=20):

    ratios = []

    for ii in range(2, max_k):
        ratio = compute_golden_ratio(ii)
        ratios.append(ratio - GOLDEN)

    plt.axhline(0, alpha=0.5, lw=3, color="red")
    plt.scatter(range(2, max_k), ratios)
    plt.ylabel("difference from Golden Ratio")
```


```python

```

We got it to work, but there is still one problem: it's super slow


```python
t = {}
for ii in range(30):
    t[ii] = %timeit -o -n1 -r1 fibonacci(ii)
plot_performance(t)
```

*uh oh!*  We seem to have an algorithm that is $O(k^N)$!

Can we fix it?


```python

```


```python

```

## Speed up a real-world problem!

Here your task is to speed up an algorithm for finding the solution to the *Heat Equation* in 2D using a finite-difference method, given an initial temperature distribution.  

the heat equation is defined as:

$$ \frac{du}{dt}  = \alpha \nabla^2 u$$

Where $u$ is the temperature. This can be approximated simply by iterating in time and approximating the spatial gradient using neighboring array elements. For time-step $k$ of width $\Delta t$ and spatial width $\Delta x$:

\begin{equation}
\frac{u^{k+1}_{ij}-u^k_{ij}}{\Delta t}  = \frac{\alpha}{(\Delta x)^2} \left( u^k_{i,j-1} + u^k_{i-1,j} - 4 u^k_{i,j} + u^k_{i+1,j} + u^k_{i,j+1}\right)
\end{equation}

Below we give a naïeve way to solve this, using for-loops (which are not ideal in python). See if you can speed this up by either:

1. re-writing the code to use numpy to get rid of the spatial loops 
2. using cython or numba to compile the function (may need to experiment also with some of their compile options)

You should also try to see what the memory usage is! (hint: use the memory_profiler module). Is there a memory leak?

#### the setup:
set up the initial condtions (defining the temperature at the boundary, as well as some hot-spots that are initially at a particular temperature)


```python
N=100; M=100  # define the spatial grid
grid = np.zeros(shape=(N,M))
grid[10:40,:] = 100 # a hot-spot on the border
grid[1:-1,1:-1] = 50 # some initial temperature in the middle
grid[60:80, 40:60] = 150 # a hot-spot initially heated at the start, that will cool down

plt.pcolormesh(grid)
plt.colorbar()
plt.title("initial conditions")
```


```python
def solve_heat_equation_loops(init_cond, iterations=100, delta_x=1.0, alpha=1.0):

    delta_t = delta_x**2/(4*alpha)
    prev = np.copy(init_cond)
    cur = np.copy(init_cond)
    N,M = init_cond.shape

    for k in range(iterations):
        for i in range(1,N-1):
            for j in range(1,M-1):
                cur[i,j] = prev[i,j] + alpha*delta_t/delta_x**2 * (
                    prev[i,j-1] + prev[i-1,j] - 4*prev[i,j] + prev[i,j+1] + prev[i+1,j]
                )
        
        prev,cur = cur,prev  #swap pointers

    return prev
```

We'll also define a convenience function to test the results (you can use this same plotter with your own solver)


```python
def plot_heat_equation(solver, iters=(2,10,100,1000)):
    
    fig, axes = plt.subplots(1,len(iters), figsize=(15,3))
    fig.suptitle(solver.__name__)

    for ii, iterations in enumerate(iters):
        result = solver(init_cond=grid, iterations=iterations)
        axes[ii].pcolormesh(result, vmin=0, vmax=100)
        axes[ii].set_title("{} iterations".format(iterations))


```


```python
plot_heat_equation(solver=solve_heat_equation_loops)
```

Note that our code is quite slow... 

### Your turn!

***Write an improved verson***

* how much faster is your version on average?
* how much memory does it use on average? Is it more than the loop version?
* which line is the slowest line? 

(hint: if done right, you should get a factor of about 100 speed increase)


```python
def my_heat_equation_solver(init_cond, iterations=100, delta_x=1.0, alpha=1.0):
    ## your code here
    return init_cond # replace with real return value
```


```python
#plot_heat_equation(solver=my_heat_equation_solver)
```

### SOLUTION

there are many ways to achieve this...


```python
ITERATIONS=50

results = {}
r = %timeit -o solve_heat_equation_loops(grid, iterations=50)
results['loop'] = r
```

#### using Numba:


```python
from numba import jit
solve_heat_equation_numba = njit(solve_heat_equation_loops)
```


```python
r = %timeit -o solve_heat_equation_numba(grid, iterations=50)
results['numba'] = r
```


```python
# (Note: open some views so we can see this always)
plt.figure(figsize=(5,5))
plot_performance(results)
plt.semilogy()
```

#### using numpy

using range-slicing `array[start:end,start:end]` we can get rid of the inner for-loops and turn that part into vector operations


```python
def solve_heat_equation_numpy(init_cond,  iterations=100,  delta_x=1.0, alpha=1.0):

    delta_t = delta_x**2/(4*alpha)
    prev = np.copy(init_cond)
    cur = np.copy(init_cond)

    # define some slices to make it easier to type 
    # just avoids too many things like prev[1:-1,1:-1])
    Z = slice(1,-1) # zero
    P = slice(2,None) # plus 1
    M = slice(0,-2) # minus 1
    
    for k in range(iterations):
        cur[Z,Z] = (
            prev[Z,Z] + alpha*delta_t/delta_x**2 * (
                prev[Z,M] + prev[M,Z] - 4.0*prev[Z,Z] + prev[Z,P] + prev[P,Z]
            )
        )
        prev,cur = cur,prev # swap the pointers

    return prev # since we swapped, prev is the most recent
```


```python
plot_heat_equation(solver=solve_heat_equation_numpy)
```


```python
r = %timeit -o solve_heat_equation_numpy(grid, iterations=ITERATIONS)
results['numpy'] = r
```

#### With numpy and numba


```python
solve_heat_equation_numpy_numba = njit(solve_heat_equation_numpy)

# "prime" it (compile)
plot_heat_equation(solver=solve_heat_equation_numpy_numba)
```


```python
r = %timeit -o solve_heat_equation_numpy_numba(grid, iterations=ITERATIONS)
results['numpy\nnumba'] = r
```

#### using cython:

Cython is a special python-like language that is translated into C-code (or C++ if you request it), and then compiled with your C compiler with a python binding produced automatically.

It has to be explicity compiled (unlike Numba which is "just in time" (JIT) compiled)


```python
%load_ext cython
```


```cython
%%cython  

cimport numpy as cnp
import numpy as np

def solve_heat_equation_cython(init_cond, int iterations=100, double delta_x=1.0, double alpha=1.0):

    cdef int i,j,k, N, M
    cdef float delta_t
    cdef cnp.ndarray[double, mode="c", ndim=2] prev, cur  # this seems to give the biggest improvement

    delta_t = delta_x**2/(4*alpha)
    prev = np.copy(init_cond)
    cur = np.copy(init_cond)
    N,M = init_cond.shape

    for k in range(iterations):
        for i in range(1,N-1):
            for j in range(1,M-1):
                cur[i,j] = prev[i,j] + alpha*delta_t/delta_x**2 * (
                    prev[i,j-1] + prev[i-1,j] - 4*prev[i,j] + prev[i,j+1] + prev[i+1,j]
                )
        
        prev,cur = cur,prev

    return prev
                                                                


```

Try running the previous cell with `%%cython -a`  to get an annotated version to see what it did!


```python
r = %timeit -o solve_heat_equation_cython(grid, iterations=ITERATIONS)
results['cython'] = r
```

#### results


```python
plot_performance(results)
plt.semilogy()
```

### We will come back to this in the next lecture

#### with some parallelization


```python
solve_heat_equation_numpy_numba_parallel = njit(solve_heat_equation_numpy, parallel=True)
```


```python
plot_heat_equation(solver=solve_heat_equation_numpy_numba_parallel)
```


```python
r = %timeit -o solve_heat_equation_numpy_numba_parallel(grid, iterations=50)
results['numpy\nnumba\nparallel'] = r
```

Nice!

*But did it do anything?*


```python
solve_heat_equation_numpy_numba_parallel.parallel_diagnostics()
```

#### with explicit parallelization



```python
@njit(parallel=True)
def solve_heat_equation_numpy_numba_parallel_explicit(init_cond,  iterations=100,  delta_x=1.0, alpha=1.0):

    delta_t = delta_x**2/(4*alpha)
    prev = np.copy(init_cond)
    cur = np.copy(init_cond)

    # define some slices to make it easier to type 
    # just avoids too many things like prev[1:-1,1:-1])
    Z = slice(1,-1) # zero
    P = slice(2,None) # plus 1
    M = slice(0,-2) # minus 1
    
    for k in prange(iterations):
        cur[Z,Z] = (
            prev[Z,Z] + alpha*delta_t/delta_x**2 * (
                prev[Z,M] + prev[M,Z] - 4.0*prev[Z,Z] + prev[Z,P] + prev[P,Z]
            )
        )
        prev,cur = cur,prev # swap the pointers

    return prev # since we swapped, prev is the most recent

answer = solve_heat_equation_numpy_numba_parallel_explicit(grid, iterations=50)

```

<div class="alert alert-block alert-warning">
    WARNING WARNING
</div>


```python
plot_heat_equation(solver=solve_heat_equation_numpy_numba_parallel_explicit)
```


```python
r = %timeit -o solve_heat_equation_numpy_numba_parallel_explicit(grid, iterations=50)
results['numpy\nnumba\nparallel\nexplicit'] = r
```


```python
plot_performance(results)
plt.semilogy()
```


```python
plot_heat_equation(solver=solve_heat_equation_numpy_numba_parallel)
plot_heat_equation(solver=solve_heat_equation_numpy_numba_parallel_explicit)
```


```python

```


```python
Let's look at some code that can be parallelized by numnba
```

ac0e616103002749c8282039c7c2c3649a70a472
ipynb
Jupyter Notebook
optimization-parallelization/optimization.ipynb
asteroidshappen/school2021
7533c8524fa53127c4da1ffe0ad83b09eb59107a
MIT
2021-05-18T11:58:17.000Z
2022-03-12T06:48:52.000Z
optimization-parallelization/optimization.ipynb
asteroidshappen/school2021
7533c8524fa53127c4da1ffe0ad83b09eb59107a
MIT
2021-05-21T14:28:34.000Z
2021-07-12T22:36:06.000Z
optimization-parallelization/optimization.ipynb
asteroidshappen/school2021
7533c8524fa53127c4da1ffe0ad83b09eb59107a
MIT
2021-05-24T18:32:54.000Z
2022-03-26T11:24:16.000Z
Qwen/Qwen-72B
1. YES
2. YES