<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Vectorization and JIT compilation in Python
Examples for the QuantEcon 2017 PhD workshops
John Stachurski
Step1:  Problem 1
Step2:  Here's a typical time series
Step3:  Here's a function that simulates for n periods, starting from x0, and returns only the final value
Step4:  Let's see how fast this runs
Step5:  Now let's try this in FORTRAN.  
Note --- this step is intended to be a demo and will only execute if
you have the file fastquad.f90 in your pwd
you have a FORTRAN compiler installed and modify the compilation code below appropriately
Step6:  Now let's do the same thing in Python using Numba's JIT compilation
Step7:  After JIT compilation, function execution speed is about the same as FORTRAN.
But remember, JIT compilation for Python is still limited --- see here
If these limitations frustrate you, then try Julia.
Problem 2
Step8:  Vectorized code
Step9:  JITTed code
A jitted version
Step10:  Numba for vectorization with automatic parallelization - even faster
<ASSISTANT_TASK:>
Python Code:
import numpy as np
from numba import vectorize, jit, float64
from quantecon.util import tic, toc
import matplotlib.pyplot as plt
Explanation: Vectorization and JIT compilation in Python
Examples for the QuantEcon 2017 PhD workshops
John Stachurski
End of explanation
α = 4
Explanation: Problem 1: A Time Series Model
Consider the time series model
$$ x_{t+1} = \alpha x_t (1 - x_t) $$
Let's set $\alpha = 4$
End of explanation
n = 200
x =  np.empty(n)
x[0] = 0.2
for t in range(n-1):
    x[t+1] = α * x[t] * (1 - x[t])
    
plt.plot(x)
plt.show()
Explanation: Here's a typical time series:
End of explanation
def quad(x0, n):
    x = x0
    for i in range(1, n):
        x = α * x * (1 - x)
    return x
Explanation: Here's a function that simulates for n periods, starting from x0, and returns only the final value:
End of explanation
n = 10_000_000
tic()
x = quad(0.2, n)
toc()
Explanation: Let's see how fast this runs:
End of explanation
!cat fastquad.f90
!gfortran -O3 fastquad.f90
!./a.out
Explanation: Now let's try this in FORTRAN.  
Note --- this step is intended to be a demo and will only execute if
you have the file fastquad.f90 in your pwd
you have a FORTRAN compiler installed and modify the compilation code below appropriately
End of explanation
quad_jitted = jit(quad)
tic()
x = quad_jitted(0.2, n)
toc()
tic()
x = quad_jitted(0.2, n)
toc()
Explanation: Now let's do the same thing in Python using Numba's JIT compilation:
End of explanation
def f(x, y):
    return np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib import cm
gridsize = 50
gmin, gmax = -3, 3
xgrid = np.linspace(gmin, gmax, gridsize)
ygrid = xgrid
x, y = np.meshgrid(xgrid, ygrid)
# === plot value function === #
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x,
                y,
                f(x, y),
                rstride=2, cstride=2,
                cmap=cm.jet,
                alpha=0.4,
                linewidth=0.05)
ax.scatter(x, y, c='k', s=0.6)
ax.scatter(x, y, f(x, y), c='k', s=0.6)
ax.view_init(25, -57)
ax.set_zlim(-0, 2.0)
ax.set_xlim(gmin, gmax)
ax.set_ylim(gmin, gmax)
plt.show()
Explanation: After JIT compilation, function execution speed is about the same as FORTRAN.
But remember, JIT compilation for Python is still limited --- see here
If these limitations frustrate you, then try Julia.
Problem 2: Brute Force Optimization
The problem is to maximize the function 
$$ f(x, y) = \frac{\cos \left(x^2 + y^2 \right)}{1 + x^2 + y^2} + 1$$
using brute force --- searching over a grid of $(x, y)$ pairs.
End of explanation
grid = np.linspace(-3, 3, 10000)
x, y = np.meshgrid(grid, grid)
tic()
np.max(f(x, y))
toc()
Explanation: Vectorized code
End of explanation
@jit
def compute_max():
    m = -np.inf
    for x in grid:
        for y in grid:
            z = np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1
            if z > m:
                m = z
    return m
compute_max()
tic()
compute_max()
toc()
Explanation: JITTed code
A jitted version
End of explanation
@vectorize('float64(float64, float64)', target='parallel')
def f_par(x, y):
    return np.cos(x**2 + y**2) / (1 + x**2 + y**2) + 1
x, y = np.meshgrid(grid, grid)
np.max(f_par(x, y))
tic()
np.max(f_par(x, y))
toc()
Explanation: Numba for vectorization with automatic parallelization - even faster:
End of explanation
<END_TASK>