#!/usr/bin/env python
# coding: utf-8

# # Quadrature
# 
# This notebook illustrates the quadrature routines available in 
# quantecon. These routines are Python implementations of MATLAB 
# routines originally written by Mario Miranda and Paul Fackler as part 
# of their influential compecon toolkit ([http://www4.ncsu.edu/~pfackler/compecon/toolbox.html](http://www4.ncsu.edu/~pfackler/compecon/toolbox.html)). We are indebted 
# to Mario and Paul for their pioneering work on numerical dynamic 
# programming and their support for the development of Python 
# implementations. For further information on the compecon toolkit see 
# Miranda, Mario J, and Paul L Fackler. Applied Computational Economics and Finance, MIT Press, 2002. 
# 
# The Python versions of the routines are written by Chase Coleman and Spencer Lyon. 
# 
# The examples contained in this document were derived from the examples named `demqua##.m` that are provided with the CompEcon toolbox. Many of them come from the 2005 version of the toolbox, others come from the 2014 version. The year is indiciated next to each reference.

# In[13]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal
from quantecon.quad import *

get_ipython().run_line_magic('matplotlib', 'inline')

np.random.seed(42)  # For reproducability


# ## Plot Equi-Distributed Sequences in 2-D
# 
# ### Based on `demqua01.m` (2005)

# In[14]:


def plotequi(ax, kind, n, a, b, **kwargs):
    """
    This function is to simplify the plotting process.  It takes
    the parameters to qnwequi and plots the output on the axis ax.
    """
    kind_names = {"N":"Neiderreiter", "W":"Weyl", "H":"Haber", "R":"Random"}
    pts, wts = qnwequi(n, a, b, kind)
    
    pt_alph = wts/wts.max()
    
    if n > 1000:
        sze = 3
    else:
        sze = 10
    
    ax.set_title("2-D {} Type Sequence with n={}".format(kind_names[kind], n))
    ax.set_xlabel(r"$x_1$")
    ax.set_ylabel(r"$x_2$")
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1])
    
    ax.scatter(pts[:, 0], pts[:, 1], s=sze, **kwargs)

    return None

# Create a figure and subplots
fig, axess = plt.subplots(2, 2, figsize=(14, 10))
axess = axess.flatten()

# Want to plot these kinds
kinds = ["N", "W", "H", "R"]

n = 4000
a = np.array([0, 0])
b = np.ones(2)

for ind, kind in enumerate(kinds):
  plotequi(axess[ind], kind, n, a, b)

plt.show()


# In[15]:


# Create a figure and subplots
fig, axess = plt.subplots(2, 2, figsize=(14, 10))
axess = axess.flatten()

# Want to plot these kinds
kind = "N"
num_n = [1000, 2000, 4000, 8000]

a = np.array([0, 0])
b = np.ones(2)

for ind, n in enumerate(num_n):
  plotequi(axess[ind], kind, n, a, b)

plt.show()


# ## Montecarlo Integration vs Integration by Quadrature
# 
# ### Based on `demqua02.m` (2014)
# 

# In[16]:


#Set parameters for normal
mu = np.zeros(2)
sigma = np.array([[1., .5], [.5, 1.]])

# Define a function
f = lambda x: x[:, 0]**2 + 2*x[:, 0]*x[:, 1] - 3*x[:, 1]**2

# Setparameters
n = 50000

# Montecarlo Int
mvn = multivariate_normal(cov=sigma)
randsamp = mvn.rvs(n)
mc_int = f(randsamp).sum()/n

# Quadrature Int
n = np.array([3, 3])
pts, wts = qnwnorm(n, mu, sigma)
qnwnorm_int = np.dot(wts.T, f(pts))

# Compute diff
diff_int = mc_int - qnwnorm_int

print("The Montecarlo integration provides the result %.5f" %mc_int)
print("The Quadrature integration provides the result %.5f" %qnwnorm_int)
print("The difference between the two is: %.5f" %diff_int)


# ## Compare Quadrature Methods
# 
# ### Based on `demqua03.m` and `demqua04.m` (2005)
# 
# ## 1d quadrature

# In[17]:


kinds = ["lege", "cheb", "trap", "simp", "N", "W", "H", "R"]

# Define some functions
f1 = lambda x: np.exp(-x)
f2 = lambda x: 1 / (1 + 25 * x**2)
f3 = lambda x: np.abs(x) ** 0.5
f4 = lambda x: np.exp(-x*x / 2)
func_names = ["f1", "f2", "f3", "f4"]

# Integration parameters
n = np.array([3, 5, 11, 21, 31, 51, 101, 401])  # number of nodes
a, b = -1, 1  # endpoints
a4, b4 = -1, 2

# Set up pandas DataFrame to hold results
ind = pd.MultiIndex.from_product([func_names, n])
ind.names=["Function", "Number of Nodes"]
cols = pd.Index(kinds, name="Kind")
res_df = pd.DataFrame(index=ind, columns=cols)

for ind, func in enumerate([f1, f2, f3]):
    func_name = func_names[ind]
    for kind in kinds:
        for num in n:
            res_df.ix[func_name, num][kind] = quadrect(func, num, a, b, kind)

for kind in kinds:
    for num in n:
        res_df.ix["f4", num][kind] = quadrect(f4, num, a4, b4, kind)

res_df


# ## 2d quadrature
# 
# ### Based on `demqua04.m` (2005)

# In[18]:


# Define 2d functions
f1_2 = lambda x: np.exp(x[:, 0] + x[:, 1])
f2_2 = lambda x: np.exp(-x[:, 0] * np.cos(x[:, 1]**2))
func_names_2 = ["f1_2", "f2_2"]

# Set up pandas DataFrame to hold results
a = ([0, 0], [-1, -1])
b = ([1, 2], [1, 1])
ind_2 = pd.MultiIndex.from_product([func_names_2, n**2])
ind_2.names = ["Function", "Number of Nodes"]
res_df_2 = pd.DataFrame(index=ind_2, columns=cols)

for ind, func in enumerate([f1_2, f2_2]):
    func_name = func_names_2[ind]
    for num in n:
        for kind in kinds[:4]:
            res_df_2.ix[func_name, num**2][kind] = quadrect(func, [num, num], a[ind], b[ind], kind);
        for kind in kinds[4:]:
            res_df_2.ix[func_name, num**2][kind] = quadrect(func, num**2, a[ind], b[ind], kind);
        
res_df_2


# ## Compare Chebyshev and Legendre Quadrature Nodes and Weights
# 
# ### Based on `demqua05.m` (2005)

# In[19]:


# Set parameters
n = 15
a = -1
b = 1

pts_cheb, wts_cheb = qnwcheb(n, a, b)
pts_lege, wts_lege = qnwlege(n, a, b)

fig, ax1 = plt.subplots(1, 1, figsize=(10, 8))

ax1.set_title("Quadrature Nodes and Weights")
ax1.set_xlabel("Points")
ax1.set_ylabel("Weights")
ax1.scatter(pts_cheb, wts_cheb, label="Chebyshev", color="k")
ax1.scatter(pts_lege, wts_lege, label="Legendre", color="r")
ax1.legend();


# ## Area under normal pdf using Simpson's rule
# 
# ### Based on `demqua04.m` (2014)
# 
# This example provides a visual for how Simpson's rule calculates the cdf of the standard normal distribution up to the point $z=1$.

# In[20]:


from scipy.stats import norm
# Define parameters
n = 11
a = 0
z = 1

# Compute nodes/weights
x, w = qnwsimp(n,a,z)

# Define f as standard normal pdf
f = norm(0, 1).pdf

prob = 0.5 + w.dot(f(x))

# Plot
b = 4.0
a = -b
n = 500
x = np.linspace(a, b, n)
y = f(x)

fig, ax = plt.subplots(figsize=(8, 5))

ax.plot([a, b], [0.0, 0.0], "k-")
ax.plot([z, z], [0, f(z)], "k-", lw=2)
ax.plot(x, y, lw=2, color="#7F7FFF")
ax.fill_between(x, y, where=x<z, color="#8AC627", alpha=0.2)

# add annotations
ax.annotate(r"Pr$\left(\tilde Z \leq z \right)$", xy=(-0.5, 0.1), 
            xytext=(-2.5, .2), fontsize=16,
            arrowprops=dict(arrowstyle="->"))

ax.set_xticks((z,))
ax.set_yticks(())
ax.set_xticklabels((r'$z$',), fontsize=18)
ax.set_ylim(0, .42)

plt.show()


# ## Willingness to pay, expected utility model
# 
# ### Based on `demqua05.m` (2014)

# In[21]:


n = 100
mu = 0     
var = 0.1
alpha = 2  
ystar = 1
y, w = qnwlogn(n, mu, var)
expectedutility = -w.dot(np.exp(-alpha*y))
certainutility = np.exp(-alpha*ystar)

ystar = -np.log(-expectedutility)/alpha
wtp = w.dot(y)-ystar

print("Expected utility: %.4f" % expectedutility)
print("Certain utility: %.4f" % certainutility)
print("Willingness to pay: %.4f" % wtp)


# ## Area under a curve
# 
# ### Based on `demqua06.m` (2014)
# 
# This example provides a visual for the area that is computed when a function is computed on an interval

# In[22]:


# Define function
f = lambda x: 50 - np.cos(np.pi * x) * (2 * np.pi * x - np.pi + 0.5)**2

xmin, xmax = 0, 1
a, b = 0.25, 0.75
n = 401
x = np.linspace(xmin, xmax, n)
y = f(x)

# plot
fig, ax = plt.subplots(figsize=(8, 5))
ax.plot(x, y, lw=2, color="#7F7FFF")
where_inds = (a <= x) & (x <= b)
ax.fill_between(x, y, 0.0, color="#8AC627", 
                where=where_inds, alpha=0.4)
ax.set_ylim(25, 65)
ax.vlines([a, b], [0, 0], [f(a), f(b)], lw=2, linestyles ="--")

# Annotate the plot
ax.set_xticks((a,b))
ax.set_yticks(())
ax.set_xticklabels((r"$a$", r"$b$"), fontsize=18)

ax.annotate(r"$\int_a^b f(x) dx$", xy=(0.45, 35), fontsize=16)
plt.show()


# ## Illustrating integration using Trapezoidal rule
# 
# ### Based on `demqua07.m` (2014)

# In[23]:


# Define function
c = np.array([2.00, -1.00, 0.50, 0.0])
f = np.poly1d(c)

# Basic Figure Setup
xmin = -1.0
xmax =  1.0
xwid = xmax-xmin
n = 401
x = np.linspace(xmin, xmax, n)
y = f(x)
ymin = min(y)
ymax = max(y)
ywid = ymax - ymin
ymin = ymin - 0.2*ywid
ymax = ymax + 0.1*ywid
fig, axs = plt.subplots(3, 1, figsize=(10, 6))
fig.tight_layout()


def trap_intervals(nint):
    "Split the region defined above into nint intervals"
    nnode = nint + 1
    xnode = np.linspace(xmin, xmax, nnode)
    ynode = f(xnode)
    
    # Calculate bins
    z = np.zeros(n)
    for i in range(1, nnode):
        k = np.where((x >= xnode[i-1]) & (x <= xnode[i]))[0]
        z[k] = ynode[i-1] + ((x[k]-xnode[i-1])*(ynode[i]-ynode[i-1]) 
                             /(xnode[i]-xnode[i-1]))
    
    return z, xnode, ynode


def plot_regions(z, xnode, ynode, ax):
    """
    Take "interval" data z and plot it with the actual function
    on the axes ax.
    """
    nint = xnode.size - 1
    
    # plot
    ax.plot(x, y)
    ax.plot(x, z, "r--", lw=2)
    ax.fill_between(x, z, ymin+0.02, color="#8AC627",
                    alpha=0.4)
    
    # annotate
    # Set ticks
    ax.set_xticks(xnode)
    x_tick_labs = [r"$x_0=a$"]
    x_tick_labs += [r"$x_%i$" % i for i in range(1, nint)]
    x_tick_labs += [r"$x_%i=b$" % nint]
    ax.set_xticklabels(x_tick_labs, fontsize=14)
    ax.xaxis.set_ticks_position('bottom')
    ax.set_yticks(())

    # remove borders
    for d in ["left", "right", "top", "bottom"]:
        ax.spines[d].set_visible(False)
    
    # set plot limits
    ax.set_ylim(ymin, ymax)
    ax.set_xlim(xmin-0.05, xmax+0.05)
    
    # add lines to show bins
    ax.vlines(xnode, ymin, ynode, color="k", linestyles="-", lw=.25)
        
    return

plot_regions(*trap_intervals(2), ax=axs[0])
plot_regions(*trap_intervals(4), ax=axs[1])
plot_regions(*trap_intervals(8), ax=axs[2])


# In[ ]: