# Exercise 1
## JIT the pressure poisson equation
The equation we need to unroll is given by 

\begin{equation}
p_{i,j}^{n} = \frac{1}{4}\left(p_{i+1,j}^{n}+p_{i-1,j}^{n}+p_{i,j+1}^{n}+p_{i,j-1}^{n}\right) - b
\end{equation}

and recall that `b` is already computed, so no need to worry about unrolling that.  We've also filled in the boundary conditions, so don't worry about those.  (don't forget to decorate your function!)


```python
import numpy
from numba import jit
```


```python
def pressure_poisson(p, b, l2_target=1e-4):
    I, J = b.shape
        
    iter_diff = l2_target + 1
    n = 0
    while iter_diff > l2_target and n <= 500:
        pn = p.copy()
        
        #Your code here
        
        #boundary conditions
        for i in range(I):
            p[i, 0] = p[i, 1]
            p[i, -1] = 0

        for j in range(J):
            p[0, j] = p[1, j]
            p[-1, j] = p[-2, j]

        if n % 10 == 0:
            iter_diff = numpy.sqrt(numpy.sum((p - pn)**2)/numpy.sum(pn**2))
            
        n += 1
        
    return p
```


```python
import pickle
from snippets.ns_helper import cavity_flow, velocity_term, quiver_plot
```


```python
def run_cavity():
    nx = 41
    with open('../IC.pickle', 'rb') as f:
        u, v, p, b = pickle.load(f)

    dx = 2 / (nx - 1)
    dt = .005
    nt = 1000
    
    u, v, p = cavity_flow(u, v, p, nt, dt, dx,
                         velocity_term,
                         pressure_poisson,
                         rtol=1e-4)

    return u, v, p
```


```python
un, vn, pn = run_cavity()
```


```python
%timeit run_cavity()
```


```python
with open('../numpy_ans.pickle', 'rb') as f:
    u, v, p = pickle.load(f)
```


```python
assert numpy.allclose(u, un)
assert numpy.allclose(v, vn)
assert numpy.allclose(p, pn)
```

# Exercise 2 (optional)

Finish early?  Just want to try more stuff?  

This line is not super efficient: 

```python
iter_diff = numpy.sqrt(numpy.sum((p - pn)**2)/numpy.sum(pn**2))
```

Try rewriting it using a jitted function and see what kind of performance gain you can get.


```python

```

0e5936f7cf5cec621e0e31c2a188ada3e097a3e4
ipynb
Jupyter Notebook
notebooks/exercises/05.Cavity.Flow.Exercises.ipynb
gforsyth/numba_tutorial_scipy2017
01befd25218783f6d3fb803f55dd9e52f6072ff7
CC-BY-4.0
2017-06-23T10:18:26.000Z
2022-03-27T21:16:56.000Z
notebooks/exercises/05.Cavity.Flow.Exercises.ipynb
gforsyth/numba_tutorial_scipy2017
01befd25218783f6d3fb803f55dd9e52f6072ff7
CC-BY-4.0
2017-06-11T21:20:59.000Z
2018-10-18T13:57:30.000Z
notebooks/exercises/05.Cavity.Flow.Exercises.ipynb
gforsyth/numba_tutorial_scipy2017
01befd25218783f6d3fb803f55dd9e52f6072ff7
CC-BY-4.0
2017-06-26T13:04:48.000Z
2022-01-11T20:36:31.000Z
Qwen/Qwen-72B
1. YES
2. YES