#To avoid problems in the future, import everything now
import numpy as np
import scipy.optimize
import scipy.linalg
import time


print("hello")

def mcgratten_weighted_residual(coeff):
    """
    A looped function of the McGratten problem.

    Variables:
        coeff   :   ndarray; a vector of coefficients for the parameterized
                    policy function
    Returns:
        ndarray containing the residuals.

    """
    #Unpack the arguments
    alpha, delta, I, L, k, abcissas, M, a, PIE = args
    coeff = coeff.reshape(I, L)

    #Initialize the residuals
    RES = np.zeros((I*L))

    #Construct the vector of residual funcitons
    for i in range(0, I):
        for l in range(0, L - 1):
            #Compute f at all of the quadrature points on the element
            for m in range(0, M):
                x = abcissas[m]
                kt = scalup(x, k[l+1], k[l])
                u = weights[m]
                bs1 = 0.5*(1 - x)
                bs2 = 0.5*(1 + x)

                #Compute capital next period by projection
                kp = coeff[i, l]*bs1 + coeff[i, l+1]*bs2

                #Compute consumption
                y = a[i]*kt**alpha
                c = y + (1 - delta)*kt - kp

                #Find the element for next periods capital stock
                ii = 0
                for h in range(0, L - 1):
                    if kp >= k[h]:
                        ii = h
                k1p = k[ii]
                k2p = k[ii + 1]

                #Calculate next period capital policy
                x = scaldwn(kp, k2p, k1p)
                bs1p = 0.5*(1 - x)
                bs2p = 0.5*(1 + x)

                #Initialize summation Variables
                sum1 = 0.0

                #Compute the summation and derivative simultaneously
                for h in range(0, I):
                    kpp = coeff[h, ii]*bs1p + coeff[h, ii+1]*bs2p
                    yp = a[h]*kp**alpha
                    cp = yp + (1 - delta)*kp - kpp
                    dudcp = cp**(-gamma)
                    dydkp = alpha*a[h]*kp**(alpha - 1)
                    sum1 += PIE[i, h]*dudcp*(dydkp + (1 - delta))

                ind = l + i*L

                #Generate the full residual entry by mulying qua*res*weight
                res = (beta*sum1 - c**(-gamma))*u
                RES[ind] += res*bs1
                RES[ind+1] += res*bs2

    return RES

"""
Explanation: PyShop Session 3

Advanced Topics in Numerical Methods and Array Manipulation.
This session delves more deeply into the NumPy and SciPy libraries.  Numpy itself is mainly the ndarray object, but this has many features and we will cover its indexing and methods.  SciPy builds on NumPy by adding a host of numerical algorithms.  In order to make the session more fun, we will use NumPy and SciPy to solve an applied problem.  We will construct a program for solving the finite element solution of a simple RBC model.  If time permits, we will then try to vectorize this problem to get some speed up. 
A Quick Introduction to Finite Elements
In economics we are often faced with large, non-linear problems which we would like to solve numerically.  We do this in one of two ways: locally or globally.  A local approximation is the simplest method and is taught in most first year graduate macro courses, in the form of linearizations.  This can also be accomplished through perturbations.  However, we sometimes are interested in the solution far from the steady state.  In this case, we look for global solutions.  Most graduate students are familiar with value- or policy- function iteration, which solve for a policy function over the entire state space.  Another global method is called "Finite Elements", or "Projection Methods".
Projection methods get their name from the assumption of a parameterized policy function.  This reduces the dimensionality of the problem from an infinite dimensional one to finite dimensions.  The name Finite Elements comes from the idea that, given this approximation the solution will be inexact.  However, a Galerkin weighted residual approximation will be exact, and thus we can write an integral equation.  We then break the state space into "elements", over which we integrate.  In this way, we break our problem down to a sequence of smaller problems.
We'll work with the simplest example from McGratten, 1993, the first economics paper on Finite Elements.  Consider an economy with a representative agent who must choose a capital stock each period. Productivity follows a finite state markov chain, capital depreciates at a rate $\delta$, and preferences are CRRA with risk aversion parameter $\gamma$. Their problem can be written as
$$
\begin{equation}
\begin{aligned}
& \underset{{ c_t,k_{t+1} }{t=0}^\infty}{\text{max}} 
& & \mathbb{E} {\sum{t=0}^{\infty}} \beta \frac{c_t^{1-\gamma}}{1-\gamma} \
& \text{s.t.}
& & y_t = a_tk_t^\alpha \
& & & k_{t+1} = (1 - \delta)k_t + i_t \
& & & a_t \in {a_1, ... , a_I} \
& & & \mathbb{P}[a_{t+1} = a_j | a_t = a_i] = \Pi_{i,j}
\end{aligned}
\end{equation}
$$
We can solve this problem to arrive at the euler equation and law of motion for capital (the steps are left to the reader):
$$
\begin{align}
c_t^{-\gamma} &= \beta \sum_{i=1}^{I} \Pi_{ij} \left[ c_{t+1}^{-\gamma} \left(1 - \delta + \alpha a_ik_{t+1}^{\alpha - 1} \right) \right]\
k_{t+1} &= a_ik_t^\alpha +(1 - \delta)k_t - c_t
\end{align}
$$
for all $i, j$.
At this point we cannot solve this problem exactly, as it is highly non-linear.  In order to make it more simple, we can assume that the capital decision rule can be approximated by a parameterized function:
$$
\begin{align}
\hat{k}{t+1} &= \sum{l = 1}^{L} N_l(k_t)\kappa_l
\end{align}
$$
where the $\kappa_l$ are constants
We don't need to make much assumptions about the functions $N_l(\cdot)$, other than that they form a complete basis for the state space.  The easiest and most natural method is to use a polynomial basis, eg chebyshev polynomials.  However, these functions are not very clear to the first time reader, so we will use a linear basis function:
\begin{align}
N_l(k) &= \left{
                \begin{array}{lll}
                    \frac{k - k_{l-1}}{k_l - k_{l-1}}   & k_{l-1}\leq k \leq k_l \
                    \frac{k_{l+1} - k}{k_{l-1} - k_l}   & k_{l}\leq k \leq k_{l + 1}  \
                    0   & else
                \end{array}
          \right.             
\end{align}
This implies that capital can be approximated as a linear interpolant between two constants:
\begin{align}
\hat{k}(k) &= \frac{k_{l+1} - k}{k_{l-1} - k_l}\kappa_l^i + \frac{k - k_{l}}{k_{l + 1} - k_{l}}\kappa_{l+1}^i   & k\in[k_l, k_{l+1}]  
\end{align}
Given this, we can calculate $c_t$, $c_{t+1}$, and $\hat{k_{t+1}}$ conditional on a value for $\kappa$.  However, our euler equation is no longer exact.  But, we can write the Galerkin weak form integral equation (where $R$ is the euler residual):
\begin{align*}
\sum_{i = 1}^{I}\int_{k_0}^{k^L}w(k,i)R(k,i;\kappa)dk =0
\end{align
We define the weighting function to be the same linear interpolator:
\begin{align}
\sum_{i = 1}^{I} \sum_{a = 1}^{L} w_a^i \left {\sum_{l=1}^{L - 1} \int_{k_l}^{k_{l+1}}N_a(k)R(k,i;\kappa)dk\right } =0
\end{align}
If this holds for all values of $w$, then everything inside the braces must be zero.  Finally, we arrive at the system of equations we need to solve:
\begin{align}
\sum_{l=1}^{L - 1} \int_{k_l}^{k_{l+1}}N_a(k)R(k,i;\kappa)dk & & \forall i, a
\end{align}
There are a few things to notice:

The integration is done accross elements.
Since the linear interpolator is zero off of the element and its upper neighbor, the terms of the sum are mainly zeros.
We can calculate the integrals individually and then reconstruct the residual functions.

We'll approximate the integral by gaussian quadrature and use a built in solver to find the solution.
Pre-Programming
Before we set out to solve our problem numerically, we are going to write down EXACTLY what we need to solve.  This will allow us to avoid a lot of head scratching along the way.  Then we'll try to speed up our code by using NumPy's vectorized functions.
First of all, notice that for a value of $a$, $N_a(\cdot)$ is zero for almost all of the terms.  Thus, we have the following system of equations
\begin{align}
\sum_{l=a}^{a + 1} \int_{k_l}^{k_{l+1}}N_a(k)R(k,i;\kappa)dk & & \forall i, a
\end{align}
Next, recall that using gaussian quadrature (which we can generate using a NumPy function) for $M$ quadrature points, we have
\begin{align}
\sum_{l=a}^{a + 1} \sum_{m=1}^{M} N_a(k_l^m)R(k_l^m,i;\kappa)\omega^m & & \forall i, a
\end{align}
This is a much simplified expression of what we need to solve.  Our first shot at the solution will follow the following steps.  We'll loop over all $i\in {1,...,I}$ and all $l\in {1, ...,L-1}$ and do the following:

Scale the abcissas to the state space.
Calculate policy based on a coefficient input.
Calculate consumption and output based on that input.
For all possible outcomes of the state next period, calculate the next period's policy based on the same coefficient input.
Calculate next period's consumption and output based on that calculation.
Calculate the conditional expectation using step 5.
Generate the residual over the current element and add it to the main vector.

If that doesn't seem super clear, it will once we get to programming!
The Looped Application
Here is the looped function that I've used.  I copied it from McGratten's original fortran code.
End of explanation
"""


#Unpack the arguments
alpha, delta, I, L, k, abcissas, M, a, PIE = args
coeff = coeff.reshape(I, L)

#Initialize the residuals
RES = np.zeros((I*L))

"""
Explanation: There are several things you should immediately notice about the structure of this function.  
First, it's a function!  I defined the residual function so that I could pass it to a numerical optimizer.  
Second, you'll notice that it is only a function of the coefficients.  This is in order to simplify the passing to the numerical optimizer.  Often times, we will have functions with additional arguments or keyword arguments.  Many SciPy optimizers can handle this, but some cannot.  To avoid problems, I like to define a tuple object containing those values that do not change during optimization, such as preference parameters, discount rates, etc.
Third, this is not very Pythonic.  There is a looped structure that can be very slow.  Our goal later will be to use vectors to speed this up, but for now we'll talk about what the function is doing.
Let's break the function into sections that we can study individually.
In the next few sections the code I present will not be stand alone, but just repition of parts of the above function to discuss.
Parameter Definition
End of explanation
"""


#Construct the vector of residual funcitons
for i in range(0, I):
    for l in range(0, L - 1):

"""
Explanation: The first thing we do is "unpack" the arguments.  We need to define an args tuple in the global namespace that contains all of these parameters.  alpha is the capital share, delta depreciation rate, I number of discrete states in the markov chain, L the number of points in the partition of capital, k the partition, abcissas a set of gaussian abcissas for estimating the integral, a the different values the markov chain can take, and PIE the markov transition matrix.
Next we initialize a vector that will contain our solution.  Here there are $IL$ equations in $IL$ unknowns.
The Loops
End of explanation
"""


#Compute f at all of the quadrature points on the element
for m in range(0, M):

"""
Explanation: Next, we loop over all elements, calculating the residuals.  There are $I(L-1)$ elements in our problem.
An Unexpected Loop
End of explanation
"""


def scalup(points, upper, lower):
    """
    A linear transformation for state variables: [-1,1]-> [k_low, k_high].

    Variables:
        points  :   ndarray; a vector of points on [-1,1]
        upper   :   ndarray; a vector of upper bounds for the element in which
                    each point is to be projected.
        lower   :   ndarray; a vector of lower bounds for the element in which
                    each point is to be projected.

    Returns:
        ndarray cotaining the transformed points
    """
    return (points + 1.)*(upper - lower)/2 + lower


def scaldwn(points, upper, lower):
    """
    A linear transformation for state variables: [k_low, k_high] -> [-1,1].

    Variables:
        points  :   ndarray; a vector of points on [k_low, k_high]
        upper   :   ndarray; a vector of upper bounds for the element in which
                    each point is to be projected.
        lower   :   ndarray; a vector of lower bounds for the element in which
                    each point is to be projected.
    Returns:
        ndarray cotaining the transformed points

    """
    return 2*(points - lower)/(upper - lower) - 1

"""
Explanation: This isn't as unexpected as it seems.  In order to approximate the integral over each element, we use a gaussian quadrature.  This uses a polynomial approximation rule.  When you tell SciPy you would like quadrature points and weights, you specify the number of points and it returns to you a set of points on $[-1, 1]$ and corresponding weights.  It is then your job to scale these to your state space and calculate the function to be integrated at each point.
To this end, we need to define some functions that will map our quadrature points into the corresponding state space and back.  Here are those functions:
End of explanation
"""


x = abcissas[m]
kt = scalup(x, k[l+1], k[l])
u = weights[m]
bs1 = 0.5*(1 - x)
bs2 = 0.5*(1 + x)

"""
Explanation: Now, we can loop through the $M$ quadrature points on each element to calculate the residual function.
Finding the State
End of explanation
"""


#Compute capital next period by projection
kp = coeff[i, l]*bs1 + coeff[i, l+1]*bs2

#Compute consumption
y = a[i]*kt**alpha
c = y + (1 - delta)*kt - kp

"""
Explanation: Given that we are on element $(i, l)$, studying point $m$ in the list of quadrature points, we need to find our value of the state and calculate the linear interpolators.
These lines do just that, taking x from the list of abcissas, scaling it to the state space on the current element, defining the weight u, and calculating the interpolator.
Calculating Inter- and Intra- temporal
End of explanation
"""


#Find the element for next periods capital stock
ii = 0
for h in range(0, L - 1):
    if kp >= k[h]:
        ii = h
k1p = k[ii]
k2p = k[ii + 1]

"""
Explanation: Now, given the state and the current solution for the coefficients, we can calculate the policy function for next period capital.  Given this, we use the law of motion for capital and the production function to calculate consumption.
Next Period's Element
End of explanation
"""


#Calculate next period capital policy
x = scaldwn(kp, k2p, k1p)
bs1p = 0.5*(1 - x)
bs2p = 0.5*(1 + x)

#Initialize summation Variables
sum1 = 0.0

#Compute the summation and derivative simultaneously
for h in range(0, I):
    kpp = coeff[h, ii]*bs1p + coeff[h, ii+1]*bs2p
    yp = a[h]*kp**alpha
    cp = yp + (1 - delta)*kp - kpp
    dudcp = cp**(-gamma)
    dydkp = alpha*a[h]*kp**(alpha - 1)
    sum1 += PIE[i, h]*dudcp*(dydkp + (1 - delta))


"""
Explanation: In order to calculate the conditional expectation, we need to calculate the policy function for capital next period as well.  That is, in order to find $c_{t+1}$, we need to calculate $k_{t+2}$.  Given our assumption of linear interpolators, we then need to find the element on which to calculate next period's capital.
To do this, we loop through the lower bounds of each element and ask the question: "Is $k_{t+1}$ greater than this lower bound?"  If so, we consider that to be the lower bound of the element for next period's policy function.
Once we are done, we then define the bounds for next period's element.
NOTE:  The assumption of linear interpolators is what makes this necessary.  There are other sets of bases, eg Chebyshev polynomials, which are not functions of the bounds of the element.  Because of this you do not need to do this look up operation.  In particular, were you doing this on several dimensions, this could become quite a pain.  However, linear interpolators are the most conceptually approachable, so they are the easiest to start with.
Calculating Conditional Expectations
End of explanation
"""


ind = l + i*L

#Generate the full residual entry by mulying qua*res*weight
res = (beta*sum1 - c**(-gamma))*u
RES[ind] += res*bs1
RES[ind+1] += res*bs2


"""
Explanation: Given the value of $k_{t+1}$ and the bounds of the element, you can calculate the linear interpolator functions bs1p and bs2p.
Now, you need to calculate the capital policy, production, consumption, marginal utility, and marginal productivity of capital for all possible realizations of the state in $t+1$.  To do this, loop over the indices $i = 0, ... , I-1$ and calculate all of these values.  Then multiply by the probability of this outcome and add it to the sum.
After the loop you will have calculated the conditional expectation:
$$
\sum_{i=1}^{I} \Pi_{ij} \left[ c_{t+1}^{-\gamma} \left(1 - \delta + \alpha a_ik_{t+1}^{\alpha - 1} \right) \right]
$$
Calculate the Residual
End of explanation
"""


def steady(alpha, beta, delta):
    """
    A function to compute the steady state of the simple RBC model.

    Variables:
        alpha   :   float; capital share
        beta    :   float; time discount
        delta   :   float; depreciation rate

    """
    k_star = ((1 - (1 - delta)*beta)/(beta * alpha))**(1/(alpha - 1))
    return k_star, k_star**alpha - delta*k_star


def explicit_transition(states):
    """
    A function to generate an explicit transition matrix.

    Variables:
        states  :   int; number of finite states in the chain

    """
    if states == 5:
        PIE = np.array([[0.1, 0.1, 0.6, 0.1, 0.1],
                        [0.1, 0.1, 0.6, 0.1, 0.1],
                        [0.1, 0.1, 0.6, 0.1, 0.1],
                        [0.1, 0.1, 0.6, 0.1, 0.1],
                        [0.1, 0.1, 0.6, 0.1, 0.1]])
    return PIE

# Step 1: Calibrate Parameters
gamma = 2.0
beta = 0.98
alpha = 0.3
delta = 0.9


# Step 2: Compute steady state
k_star, c_star = steady(alpha, beta, delta)

# Step 3: Generate the state space
#Define upper and lower bound of the region of interest around the steady stat
kmin = 0.5*k_star
kmax = 1.5*k_star

#Define the finite states for the markov chain
I = 5
PIE = explicit_transition(I)
a = np.array([0.95, 0.975, 1.0, 1.025, 1.05])

#Generate a state space vector over k
L = 11
k = np.linspace(kmin, kmax, num=L)

# Step 4: Calculate quadrature weights and abcissas
#Specify number of points in the integral on each element [k_j, k_j+1]
M = 10

#Generate absissas and weights
abcissas, weights = np.polynomial.legendre.leggauss(M)

# Step 5: Guess and solve
#Generate intial guess as half of production at all nodes
coeff0 = np.outer(a, k**alpha)*0.1

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

print(res.shape)

def policy(a_p, k_p, coeff, k, a):
    """
    A function to calculate the capital policy at points.

    Variables:
        points  :   ndarray; a vector of points on
                    [k_low, k_high]x{a1...aI}
        coeff   :   ndarray; a vector of coefficients
        k       :   ndarray; a vector containing the
                    partition of the capital state.
        a       :   ndarray; a vector containing the
                    partition of the productivity state.
    Returns:
        ndarray evaluated funciton

    """
    #Unpack the points
    #k_p = points[:, 0]
    #a_p = points[:, 1]

    #Find the elements for the points
    ll = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                   for point in k_p][0])
    ll[ll > 0] -= 1

    ii = np.array([np.where(point == a)[0] for point in a_p][0])
    #ii[ll > 0] -= 1

    #Define element bounds for next period
    k1p = k[ll]
    k2p = k[ll + 1]

    #Calculate the interpolators
    x = scaldwn(k_p, k2p, k1p)
    bs1 = 0.5*(1 - x)
    bs2 = 0.5*(1 + x)

    #Return the polciy function of the given points
    return coeff[ii, ll]*bs1 + coeff[ii, ll+1]*bs2

#Plot the policy function over the state space
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.gca(projection="3d")

a_points, k_points = np.meshgrid(a, k)

kp = policy(np.reshape(a_points, (I*L)), np.reshape(k_points, (I*L)),
            np.array(res), k, a)

kp = np.reshape(kp, (L, I))

ax.plot_surface(a_points, k_points, kp, alpha=0.3)
#cset = ax.contour(k_points, a_points, kp, zdir='z', offset=-2.0)
#cset = ax.contour(C1, C2, utils, zdir='y', offset=10.0)
#cset = ax.contour(C1, C2, utils, zdir='x', offset=1.0)

plt.show()

"""
Explanation: Now that you have calculate the functional value at the quadrature point $m$ on element $(i, l)$, you can multiply by the weight and the interpolater.  Finally, add it to the residual vector.  You'll notice that you are adding it to severl points in the residual vector simultaneously.  This is because of the linear intorpolater, which is non-zero on two elements simultaneously.
Running the Code
Now that we've gone through how this code works, let's see it in action.  First, we're going to define some functions to generate a steady state and a transition matrix (a seriously sloppy function, but be nice), and define some parameters.  Then we'll calculate a partition over the state and some quadrature points.  Finally, we'll make an initial guess to the value of capital at the nodes (this is the same as the coefficient values) as half of output.  Finally, we package all of our arguments into a tuple and run a jacobian free newton kryolov solver (I like this one both for its speed and its verbosity).
End of explanation
"""


#Increase the number of points
L = 50
k = np.linspace(kmin, kmax, num=L)
coeff0 = np.outer(a, k**alpha)*0.1

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

"""
Explanation: We can see that our policy function is essentially linear.  We knew that, by definition, it would be piecewise linear.  We could also expand the state space and increase the number of points to see whether the shape changes:
End of explanation
"""


#Compute the summation and derivative simultaneously
for h in range(0, I):
    kpp = coeff[h, ii]*bs1p + coeff[h, ii+1]*bs2p
    yp = a[h]*kp**alpha
    cp = yp + (1 - delta)*kp - kpp
    dudcp = cp**(-gamma)
    dydkp = alpha*a[h]*kp**(alpha - 1)
    sum1 += PIE[i, h]*dudcp*(dydkp + (1 - delta))

"""
Explanation: We can see that we don't get much change.  Re-run the earlier cell to get back the old paramteters and we'll move on to vectorizing this function.
Vectorizing
Given our function, we now want to try to speed it up, as you can see from the previous runs, it was quite slow.  Vectorizing it will be a good chance for us to see how much NumPy can increase the speed of your code.
The easiest way (for me) to vectorize something is to work around the loops.  Each loop can be thought of as a set of stacked operations, which we can try to do in a single pass, avoiding the loop.  If you work in this way, it is easiest to start from the inner-most loop and work out.  That is, we have loops over i, l, m, and h.  We'll start with the loop over h and work our way back.
Here is the loop we would like to vectorize:
End of explanation
"""


kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
yp = a*kp**alpha
cp = yp + (1 - delta)*kp*np.ones(I) - kpp
dudcp = cp**(-gamma)
dydkp = alpha*a*kp**(alpha-1)
sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

"""
Explanation: First, we should ask the question, "What is this loop doing?"  It is calculating the conditional expectation of returns on capital next period, discounted by marginal utility.  
Given the markov chain randomness in this model, this lends itself very easily to the dot product.  We can use NumPy's automatic array operations (which are natively element-wise) to calculate a vector of next period returns ($1 - \delta + \alpha a_ik_{t+1}^{\alpha - 1}$) and then take the dot product with the corresponding row from $\Pi$.
Here is the code, take a look and try to figure out what it is doing and then we will discuss it:
End of explanation
"""


temp = 1.
temp2 = np.array([1, 2, 3])
print(temp2 + temp)

"""
Explanation: Got it?  Here's how I think of what this is doing.  First, notice the parts of the loop that do not change.  These are bs1p, bs2p, and kp.  You can treat these as scalars.  So, you can look at the policy function calculation, kpp, as taking the linear combination of the vectors of coefficients weighted by bs1p and bs2p.  
Then you calculate a vector of outputs, each entry corresponding to a different realization of $a$.
Next, calculate consumption, but notice that kp does not vary, so multiply by a vector of ones.  Note: NumPy's broadcasting methods could easily handle this, but I find it clearing to include the vector so those reading your code understand.
End of explanation
"""


def mcgratten_weighted_residual_vec_1(coeff):
    """
    A partially vectorized function of the McGratten problem.

    Variables:
        coeff   :   ndarray; a vector of coefficients for the parameterized
                    policy function
    Returns:
        ndarray containing the residuals.

    """
    #Unpack the arguments
    alpha, delta, I, L, k, abcissas, M, a, PIE = args
    coeff = coeff.reshape(I, L)

    #Initialize the residuals
    RES = np.zeros((I*L))

    #Construct the vector of residual funcitons
    for i in range(0, I):
        for l in range(0, L - 1):
            #Compute f at all of the quadrature points on the element
            for m in range(0, M):
                x = abcissas[m]
                kt = scalup(x, k[l+1], k[l])
                u = weights[m]
                bs1 = 0.5*(1 - x)
                bs2 = 0.5*(1 + x)

                #Compute capital next period by projection
                kp = coeff[i, l]*bs1 + coeff[i, l+1]*bs2

                #Compute consumption
                y = a[i]*kt**alpha
                c = y + (1 - delta)*kt - kp

                #Find the element for next periods capital stock
                ii = 0
                for h in range(0, L - 1):
                    if kp >= k[h]:
                        ii = h
                k1p = k[ii]
                k2p = k[ii + 1]

                #Calculate next period capital policy
                x = scaldwn(kp, k2p, k1p)
                bs1p = 0.5*(1 - x)
                bs2p = 0.5*(1 + x)

                #Initialize summation Variables
                sum1 = 0.0

                #Compute the summation and derivative simultaneously
                kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
                yp = a*kp**alpha
                cp = yp + (1 - delta)*kp*np.ones(I) - kpp
                dudcp = cp**(-gamma)
                dydkp = alpha*a*kp**(alpha-1)
                sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

                ind = l + i*L

                #Generate the full residual entry by mulying qua*res*weight
                res = (beta*sum1 - c**(-gamma))*u
                RES[ind] += res*bs1
                RES[ind+1] += res*bs2

    return RES


"""
Explanation: Finally, you simply calculate $\partial_c u(c_{t+1})$ and  $\partial_k y(k_{t+1})$ as vectors and calculate the expectation by a straight forward dot product.
Here's the full code for the modified function, which I'll call mcgratten_weighted_residual_vec_1:
End of explanation
"""


time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res1 = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec_1, coeff0,
                                    method='lgmres', verbose=True,
                                    maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res1))


"""
Explanation: Let's run the two optimizations side by side and see the difference in speed, as well as take a norm of the difference in the answers to test for accuracy:
End of explanation
"""


for m in range(0, M):
    x = abcissas[m]
    kt = scalup(x, k[l+1], k[l])
    u = weights[m]
    bs1 = 0.5*(1 - x)
    bs2 = 0.5*(1 + x)

    #Compute capital next period by projection
    kp = coeff[i, l]*bs1 + coeff[i, l+1]*bs2

    #Compute consumption
    y = a[i]*kt**alpha
    c = y + (1 - delta)*kt - kp

    #Find the element for next periods capital stock
    ii = 0
    for h in range(0, L - 1):
        if kp >= k[h]:
            ii = h
    k1p = k[ii]
    k2p = k[ii + 1]

    #Calculate next period capital policy
    x = scaldwn(kp, k2p, k1p)
    bs1p = 0.5*(1 - x)
    bs2p = 0.5*(1 + x)

    #Initialize summation Variables
    sum1 = 0.0

    #Compute the summation and derivative simultaneously
    kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
    yp = a*kp**alpha
    cp = yp + (1 - delta)*kp*np.ones(I) - kpp
    dudcp = cp**(-gamma)
    dydkp = alpha*a*kp**(alpha-1)
    sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

    ind = l + i*L

    #Generate the full residual entry by mulying qua*res*weight
    res = (beta*sum1 - c**(-gamma))*u
    RES[ind] += res*bs1
    RES[ind+1] += res*bs2

"""
Explanation: It's actually slower by $50%$!  Why would we ever want to do this!?  What a world...?!
Ok, don't panic.  The reason you get a slow down here is because the vector dot product and broadcasting can be slower for small arrays.  Where you'll get a speed up is when the arrays become large.  The only way in this case to see if this vectorization is effective is to increase the number of discrete states, which is heavy.  We'll have an easier time on furth loops and you'll see the benefits then.
Overall, that's the process!  We'll now go through all the other loops, vectorizing the operations.
The Abcissas Loop
What we'll do is take our previously vectorized function and try to vectorize it further.  Here is the loop we want to vectorize:
End of explanation
"""


#Find the element for next periods capital stock
ii = 0
for h in range(0, L - 1):
    if kp >= k[h]:
        ii = h
k1p = k[ii]
k2p = k[ii + 1]

"""
Explanation: You'll notice here a pesky sub-routine that is trying to find indices for the next period's element:
End of explanation
"""


kp = 0.2
print(np.where(kp >= k))
k

"""
Explanation: This is going to be tough to vectorize, but we can do it!  Let's see what it is doing.  It takes some value, kp, and returns the index for the next lowest value in another object, k.  To do this vectorially we'll use the numpy.where function.  Let's look at an example.
Given a value of kp, we would like to return the index of the next lowest value in k.  If kp$=0.2$, that index is $4$:
End of explanation
"""


print(np.where(kp >= k)[0][-1])

"""
Explanation: Notice though that np.where gives a vector of indices to all of the values of k where the condition is not true.  It also is returning a tuple.  Thus, we need to index to the value we want, the interior of the tuple and the last entry in the vector:
End of explanation
"""


kp = 0.0
print(np.where(kp >= k)[0][-1])

"""
Explanation: However, what if our algorithm steps out of the state space?  For example, what if kp is zero:
End of explanation
"""


kp = np.array(0.2)
print(np.where(kp >= np.append([0], k[:-1]))[0][-1])
kp = np.array(0.0)
print(np.where(kp >= np.append([0], k[:-1]))[0][-1])

"""
Explanation: Oh no!  This is happening because there are no values in k for which the condition is not true.  What we can do is add a zero to the array and cut off the upper point (this will avoid a similar problem on the upper bound).  This gives an answer that is correct for zero, but incorrect for something else:
End of explanation
"""


kp = np.array([0, 0.2])
ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1] 
               for point in kp])
ii[ii > 0] -= 1
print(ii)

"""
Explanation: So, we can just subtract one from all values that are greater than zero.  In order to foresee the needing to make this work for a set of points in kp we can also vectorize the entire calculation using a list comprehension:
End of explanation
"""


#For demonstration purposes
l, i = 0, 0
coeff = coeff0
RES = np.zeros((I*L))

#Compute f at all of the quadrature points on the element
x = abcissas
kt = scalup(x, k[l+1], k[l])
u = weights
bs1 = 0.5*(1 - x)
bs2 = 0.5*(1 + x)

#Compute capital next period by projection
kp = coeff[i, l]*bs1 + coeff[i, l + 1]*bs2
#print bs1
#Compute consumption
y = a[i]*kt**alpha
c = y + (1 - delta)*kt - kp

#Find indices for next period's capital stock for each point
ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
               for point in kp])
ii[ii > 0] -= 1

#Find elements for next period's capital stock
#NOTE: Here I'm using an array as an index
k1p = k[ii]
k2p = k[ii + 1]

#Calculate next period's policy
x = scaldwn(kp, k2p, k1p)
bs1p = 0.5*(1 - x)
bs2p = 0.5*(1 + x)

#Initialize the summation vector
sum1 = np.zeros((M))

#Compute all of the sums simultaneously
#NOTE: You can slice AND index by an array at the same time!
#NOTE: This is a good time to talk about broadcasting
#The resulting matrix contains columns corresponding to each
#abcissa, and rows corresponding to possible states in period
#t+1
kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
yp = np.outer(a, kp**alpha)
cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
dudcp = cp**(-gamma)
dydkp = np.outer(a, kp**(alpha - 1))*alpha
sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

#Calculate the residual functions and fill the vector
ind = l + i*L
res = (beta*sum1 - c**(-gamma))*u
RES[ind] += np.dot(res, bs1)
RES[ind + 1] += np.dot(res, bs2)

"""
Explanation: Now that we've gon over how to use np.where to find the element, we can vectorize the loop in pretty much the same way as before.  Here is the vectorized form of the loop:
End of explanation
"""


kpp

"""
Explanation: The first section is fairly straight forward, but the second section deserves some discussion.  Instead of building a single vector, we want to build M vectors simultaneously.  To this end, kpp is a matrix whose columns correspond to the different quadrature points:
End of explanation
"""


sum1


"""
Explanation: Similarly, yp is a matrix of the same size, as is cp.  Additinally, the sum1 variable now is a vector:
End of explanation
"""


#Calculate the residual functions and fill the vector
ind = l + i*L
res = (beta*sum1 - c**(-gamma))*u
RES[ind] += np.dot(res, bs1)
RES[ind + 1] += np.dot(res, bs2)

"""
Explanation: In the end, we simply use this vector to calculate the integral for each element using a dot product and then add it to the residual:
End of explanation
"""


def mcgratten_weighted_residual_vec_2(coeff):
    """
    A partially vectorized function of the McGratten problem.

    Variables:
        coeff   :   ndarray; a vector of coefficients for the parameterized
                    policy function
    Returns:
        ndarray containing the residuals.

    """
    #Unpack the arguments
    alpha, delta, I, L, k, abcissas, M, a, PIE = args
    coeff = coeff.reshape(I, L)

    #Initialize the residuals
    RES = np.zeros((I*L))

    #Construct the vector of residual funcitons
    for i in range(0, I):
        for l in range(0, L - 1):
            #Compute f at all of the quadrature points on the element
            x = abcissas
            kt = scalup(x, k[l+1], k[l])
            u = weights
            bs1 = 0.5*(1 - x)
            bs2 = 0.5*(1 + x)

            #Compute capital next period by projection
            kp = coeff[i, l]*bs1 + coeff[i, l + 1]*bs2
            #print bs1
            #Compute consumption
            y = a[i]*kt**alpha
            c = y + (1 - delta)*kt - kp

            #Find indices for next period's capital stock for each point
            ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                           for point in kp])
            ii[ii > 0] -= 1

            #Find elements for next period's capital stock
            #NOTE: Here I'm using an array as an index
            k1p = k[ii]
            k2p = k[ii + 1]

            #Calculate next period's policy
            x = scaldwn(kp, k2p, k1p)
            bs1p = 0.5*(1 - x)
            bs2p = 0.5*(1 + x)

            #Initialize the summation vector
            sum1 = np.zeros((M))

            #Compute all of the sums simultaneously
            #NOTE: You can slice AND index by an array at the same time!
            #NOTE: This is a good time to talk about broadcasting
            #The resulting matrix contains columns corresponding to each
            #abcissa, and rows corresponding to possible states in period
            #t+1
            kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
            yp = np.outer(a, kp**alpha)
            cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
            dudcp = cp**(-gamma)
            dydkp = np.outer(a, kp**(alpha - 1))*alpha
            sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

            #Calculate the residual functions and fill the vector
            ind = l + i*L
            res = (beta*sum1 - c**(-gamma))*u
            RES[ind] += np.dot(res, bs1)
            RES[ind + 1] += np.dot(res, bs2)
    return RES

time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res2 = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec_2, coeff0,
                                    method='lgmres', verbose=True,
                                    maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res2))


"""
Explanation: Here is the whole function in one cell, as well as running and comparing our results for the looped and vectorized function:
End of explanation
"""


M = 100

#Generate absissas and weights
abcissas, weights = np.polynomial.legendre.leggauss(M)

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res2 = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec_2, coeff0,
                                    method='lgmres', verbose=True,
                                    maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res2))


"""
Explanation: Huzzah!  We got a speed up!  Now, we can look at how this change affects speed as we vary the number of quadrature points, M.  When we do that we also need to regenerate weights and abcissas, as well as redefine the args tuple:
End of explanation
"""


M = 10

#Generate absissas and weights
abcissas, weights = np.polynomial.legendre.leggauss(M)

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

"""
Explanation: The solution is now achieved in half the time using 100 quadrature points.  We should see similar gains as we vectorize the other loops.  Let's reset the arguments so we don't have problems later, then we'll work on the next loop.
End of explanation
"""


for l in range(0, L - 1):
    #Compute f at all of the quadrature points on the element
    x = abcissas
    kt = scalup(x, k[l+1], k[l])
    u = weights
    bs1 = 0.5*(1 - x)
    bs2 = 0.5*(1 + x)

    #Compute capital next period by projection
    kp = coeff[i, l]*bs1 + coeff[i, l + 1]*bs2
    #print bs1
    #Compute consumption
    y = a[i]*kt**alpha
    c = y + (1 - delta)*kt - kp

    #Find indices for next period's capital stock for each point
    ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                   for point in kp])
    ii[ii > 0] -= 1

    #Find elements for next period's capital stock
    #NOTE: Here I'm using an array as an index
    k1p = k[ii]
    k2p = k[ii + 1]

    #Calculate next period's policy
    x = scaldwn(kp, k2p, k1p)
    bs1p = 0.5*(1 - x)
    bs2p = 0.5*(1 + x)

    #Initialize the summation vector
    sum1 = np.zeros((M))

    #Compute all of the sums simultaneously
    #NOTE: You can slice AND index by an array at the same time!
    #NOTE: This is a good time to talk about broadcasting
    #The resulting matrix contains columns corresponding to each
    #abcissa, and rows corresponding to possible states in period
    #t+1
    kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
    yp = np.outer(a, kp**alpha)
    cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
    dudcp = cp**(-gamma)
    dydkp = np.outer(a, kp**(alpha - 1))*alpha
    sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

    #Calculate the residual functions and fill the vector
    ind = l + i*L
    res = (beta*sum1 - c**(-gamma))*u
    RES[ind] += np.dot(res, bs1)
    RES[ind + 1] += np.dot(res, bs2)

"""
Explanation: The Capital Elements Loop
The next loop we would like to vectorize is the one which loops over capital.  Here it is from the last function:
End of explanation
"""


#Vectorize element loop along k axis
#Scale and calculate capital
x = np.kron(np.ones(L - 1), abcissas)
k1 = np.kron(k[:-1], np.ones(M))
k2 = np.kron(k[1:], np.ones(M))
kt = scalup(x, k2, k1)
u = np.kron(np.ones(L - 1), weights)
bs1 = 0.5*(1 - x)
bs2 = 0.5*(1 + x)
Bs1 = 0.5*(1 - abcissas)
Bs2 = 0.5*(1 + abcissas)

#Compute capital next period by projection
kp = np.kron(coeff[i, :-1], np.ones(M))*bs1\
    + np.kron(coeff[i, 1:], np.ones(M))*bs2

#Compute consumption
y = a[i]*kt**alpha
c = y + (1 - delta)*kt - kp

#Find indices for next period's element bounds
ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
               for point in kp])
ii[ii > 0] -= 1

#Define element bounds for next period
k1p = k[ii]
k2p = k[ii + 1]

#Calculate policy next period
x = scaldwn(kp, k2p, k1p)
bs1p = 0.5*(1 - x)
bs2p = 0.5*(1 + x)

#Initialize the summation
sum1 = np.zeros(((L-1)*M))

#Compute all of the sums simultaneously using linear algebra
kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
yp = np.outer(a, kp**alpha)
cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
dudcp = cp**(-gamma)
dydkp = np.outer(a, kp**(alpha - 1))*alpha
sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

#Calculate the residual functions
res = (beta*sum1 - c**(-gamma))*u
res = res.reshape(L-1, M)
bs1 = bs1.reshape(L-1, M)
bs2 = bs2.reshape(L-1, M)
RES[i*(L-1): (i+1)*(L-1)] += np.dot(res, Bs1)
RES[i*(L-1) + 1: (i+1)*(L-1) + 1] += np.dot(res, Bs2)

"""
Explanation: Now we are going to have to build a series of vectors.  This is going to be slightly more complex, as we'll need to use kronecker products to repeat vectors, but once you get the hang of this it is fairly simple.  Here's the vectorized loop:
End of explanation
"""


x = np.kron(np.ones(L - 1), abcissas)
k1 = np.kron(k[:-1], np.ones(M))
k2 = np.kron(k[1:], np.ones(M))
kt = scalup(x, k2, k1)

print(x[:20])
print(k1[:20])
print(k2[:20])
print(kt[:20])

u = np.kron(np.ones(L - 1), weights)
bs1 = 0.5*(1 - x)
bs2 = 0.5*(1 + x)
Bs1 = 0.5*(1 - abcissas)
Bs2 = 0.5*(1 + abcissas)

"""
Explanation: You'll immediately notice the use kroneckers everywhere.  This is simply staking up all of the loops performed before into one big vector of length $M(L-1)$, so that there is one point for every abcissa in every element on $k$.
Let's look at the first few lines to get a sense of this:
End of explanation
"""


kp = np.kron(coeff[i, :-1], np.ones(M))*bs1\
    + np.kron(coeff[i, 1:], np.ones(M))*bs2

print(bs1)
print(np.kron(coeff[i, :-1], np.ones(M)))
print(kp)

"""
Explanation: So x is just a repetition of the abcissas vector over and over.  Then, k1 and k2 define lower and upper bounds, respectively, for the element.  Finally, kt are the abcissa points.  Notice that these do not repeat, as each point corresponds to a different point on different elements.
We can continue constructing our solution in a similar way throughout.
Next, the computation of kp deserves some attention:
End of explanation
"""


#Find indices for next period's element bounds
ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
               for point in kp])
ii[ii > 0] -= 1

print(ii)

"""
Explanation: This operation is just repeating the coefficients $M$ times in order to carry out an element-wise multiplication with the interpolator.  This is the same calculation as before, but carried out $L-1$ times simultaneously.  
Next let's see how our index finder is doing:
End of explanation
"""


#Define element bounds for next period
k1p = k[ii]
k2p = k[ii + 1]

#Calculate policy next period
x = scaldwn(kp, k2p, k1p)
bs1p = 0.5*(1 - x)
bs2p = 0.5*(1 + x)

#Initialize the summation
sum1 = np.zeros(((L-1)*M))

#Compute all of the sums simultaneously using linear algebra
kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
yp = np.outer(a, kp**alpha)
cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
dudcp = cp**(-gamma)
dydkp = np.outer(a, kp**(alpha - 1))*alpha
sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))


print(kpp.shape)
print(yp.shape)
print(cp.shape)
print(sum1.shape)

"""
Explanation: It seems to be working just fine.  You'll notice that all of the values in kp are below the lower bound on the state space, so we're simply forcing it back into the lower element.
Finally, let's take a look at the conditional expectation:
End of explanation
"""


i = 0
test = np.arange(I*L)
print(test)
print(test[i*L:(i+1)*(L-1)])

#Will this work as we loop through i?
print([test[i*L:(i+1)*(L-1)] for i in range(0,I)])

#NO! Why?  The upper bound is advancing at a rate L-1 while the lower bound moves at L
print([test[i*L:(i+1)*L-1] for i in range(0,I)])

#That gives us the result we wanted!  What about plus one?
print([test[i*L+1:(i+1)*L] for i in range(0,I)])


"""
Explanation: You'll notice that we just repeated the same code as before.  Instead of $M$ columns, the kpp matrix contains $M(L-1)$ columns, one for every quadrature point and element.  Similarly, sum1 now contains all of the conditional expectations at all of the quadrature points on all of the elements in the capital direction.
Finally, we want to calculate the residuals.  First of all, notice that we are going to have some trouble entering adding to the vector of residual, since the loop was only over $0...L-2$, while the length of our RES vector is a function of $L$.  Why is this?  On the lowest and highest elements, we only add once to the residual vector, since there is no overlap with another residual.  In order to deal with this we have to do some fancy slicing.
As a test case, let's create a vector of size $IL$ and try to recove only the entries that we would like to fill.  Given a value for i in ${0, ..., I}$, we would like to fill all of the rows corresponding to $l + i*L$, for all $l$ in ${0...L-2}$, and the same plus one.  Thus, we simply slice the test vector accordingly:
End of explanation
"""


res = (beta*sum1 - c**(-gamma))*u
res = res.reshape(L-1, M)
RES[i*L: (i+1)*L-1] += np.dot(res, Bs1)
RES[i*L + 1: (i+1)*L] += np.dot(res, Bs2)

"""
Explanation: So, on each pass, given a value for i we know how to fill RES, but how do we calculate the integrals?  First, notice that bs1 just repeats over and over.  So, if we reshape the res vector into a matrix, we can carry out a dot product with a single interpolator.  Here's my implimentation:
End of explanation
"""


def mcgratten_weighted_residual_vec_3(coeff):
    """
    A partially vectorized function of the McGratten problem.

    Variables:
        coeff   :   ndarray; a vector of coefficients for the parameterized
                    policy function
    Returns:
        ndarray containing the residuals.

    """
    #Unpack the arguments
    alpha, delta, I, L, k, abcissas, M, a, PIE = args
    coeff = coeff.reshape(I, L)

    #Initialize the residuals
    RES = np.zeros((I*L))

    #Construct the vector of residual funcitons
    for i in range(0, I):
        #Vectorize element loop along k axis
        #Scale and calculate capital
        x = np.kron(np.ones(L - 1), abcissas)
        kt = scalup(x, np.kron(k[1:], np.ones(M)), np.kron(k[:-1], np.ones(M)))
        u = np.kron(np.ones(L - 1), weights)
        bs1 = 0.5*(1 - x)
        bs2 = 0.5*(1 + x)
        Bs1 = 0.5*(1 - abcissas)
        Bs2 = 0.5*(1 + abcissas)

        #Compute capital next period by projection
        kp = np.kron(coeff[i, :-1], np.ones(M))*bs1\
            + np.kron(coeff[i, 1:], np.ones(M))*bs2

        #Compute consumption
        y = a[i]*kt**alpha
        c = y + (1 - delta)*kt - kp

        #Find indices for next period's element bounds
        ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                       for point in kp])
        ii[ii > 0] -= 1

        #Define element bounds for next period
        k1p = k[ii]
        k2p = k[ii + 1]

        #Calculate policy next period
        x = scaldwn(kp, k2p, k1p)
        bs1p = 0.5*(1 - x)
        bs2p = 0.5*(1 + x)

        #Initialize the summation
        sum1 = np.zeros(((L-1)*M))

        #Compute all of the sums simultaneously using linear algebra
        kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
        yp = np.outer(a, kp**alpha)
        cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
        dudcp = cp**(-gamma)
        dydkp = np.outer(a, kp**(alpha - 1))*alpha
        sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

        #Calculate the residual functions
        res = (beta*sum1 - c**(-gamma))*u
        res = res.reshape(L-1, M)
        RES[i*L: (i+1)*L-1] += np.dot(res, Bs1)
        RES[i*L + 1: (i+1)*L] += np.dot(res, Bs2)
    return RES



time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res3 = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec_3, coeff0,
                                    method='lgmres', verbose=True,
                                    maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res3))

"""
Explanation: Ok, here's the full function and a comparison with the older version:
End of explanation
"""


#Generate a state space vector over k
L = 101
k = np.linspace(kmin, kmax, num=L)

#Specify number of points in the integral on each element [k_j, k_j+1]
M = 10

#Generate absissas and weights
abcissas, weights = np.polynomial.legendre.leggauss(M)

#Generate intial guess as half of production at all nodes
coeff0 = np.outer(a, k**alpha)*0.1

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res3 = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec_3, coeff0,
                                    method='lgmres', verbose=True,
                                    maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res3))

"""
Explanation: This is clearly faster, even for lower values of $M$ and $L$, but let's see how it performs as these increase:
End of explanation
"""


#Generate a state space vector over k
L = 11
k = np.linspace(kmin, kmax, num=L)

#Specify number of points in the integral on each element [k_j, k_j+1]
M = 10

#Generate absissas and weights
abcissas, weights = np.polynomial.legendre.leggauss(M)

#Generate intial guess as half of production at all nodes
coeff0 = np.outer(a, k**alpha)*0.1

#Pack up all of the arguments to pass to the solver
args = (alpha, delta, I, L, k, abcissas, M, a, PIE)

"""
Explanation: On my home computer with $L = 101$ and $M = 100$, I got a four times speed up!  Huzzah!
Let's reset our parameters:
End of explanation
"""


for i in range(0, I):
    #Vectorize element loop along k axis
    #Scale and calculate capital
    x = np.kron(np.ones(L - 1), abcissas)
    kt = scalup(x, np.kron(k[1:], np.ones(M)), np.kron(k[:-1], np.ones(M)))
    u = np.kron(np.ones(L - 1), weights)
    bs1 = 0.5*(1 - x)
    bs2 = 0.5*(1 + x)
    Bs1 = 0.5*(1 - abcissas)
    Bs2 = 0.5*(1 + abcissas)

    #Compute capital next period by projection
    kp = np.kron(coeff[i, :-1], np.ones(M))*bs1\
        + np.kron(coeff[i, 1:], np.ones(M))*bs2

    #Compute consumption
    y = a[i]*kt**alpha
    c = y + (1 - delta)*kt - kp

    #Find indices for next period's element bounds
    ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                   for point in kp])
    ii[ii > 0] -= 1

    #Define element bounds for next period
    k1p = k[ii]
    k2p = k[ii + 1]

    #Calculate policy next period
    x = scaldwn(kp, k2p, k1p)
    bs1p = 0.5*(1 - x)
    bs2p = 0.5*(1 + x)

    #Initialize the summation
    sum1 = np.zeros(((L-1)*M))

    #Compute all of the sums simultaneously using linear algebra
    kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
    yp = np.outer(a, kp**alpha)
    cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp
    dudcp = cp**(-gamma)
    dydkp = np.outer(a, kp**(alpha - 1))*alpha
    sum1 = np.dot(PIE[i, :], dudcp*(dydkp + 1 - delta))

    #Calculate the residual functions
    res = (beta*sum1 - c**(-gamma))*u
    res = res.reshape(L-1, M)
    RES[i*(L-1): (i+1)*(L-1)] += np.dot(res, Bs1)
    RES[i*(L-1) + 1: (i+1)*(L-1) + 1] += np.dot(res, Bs2)

"""
Explanation: Finally, let's turn to the last loop.
The Productivity Loop
As mentioned before, you only get a speed up when you vectorize a loop that makes many passes.  That means that this probably won't change much, but we'll give it a try anyways.
One last time, here's the loop we are trying to vectorize:
End of explanation
"""


H = I*(L - 1)
#Vectorize all loops
#Scale and calculate capital
x = np.kron(np.ones(H), abcissas)
kt = scalup(x, np.kron(k[1:], np.ones(I*M)), np.kron(k[:-1], np.ones(I*M)))
u = np.kron(np.ones(H), weights)
bs1 = 0.5*(1 - x)
bs2 = 0.5*(1 + x)
#NOTE: bs is constant across elements, so simplify calculations later
Bs1 = 0.5*(1 - abcissas)
Bs2 = 0.5*(1 + abcissas)

#Compute capital next period by projection
kp = np.kron(coeff[:, :-1].reshape(H,), np.ones((M)))*bs1\
    + np.kron(coeff[:, 1:].reshape(H,), np.ones((M,)))*bs2

#Compute consumption
y = np.kron(a, np.ones((L - 1)*M))*kt**alpha
c = y + (1 - delta)*kt - kp

#Find indices for next periods element bounds
ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
               for point in kp])
ii[ii > 0] -= 1

#Define element bounds for next period
k1p = k[ii]
k2p = k[ii + 1]

#Calculate policy next period
x = scaldwn(kp, k2p, k1p)
bs1p = 0.5*(1 - x)
bs2p = 0.5*(1 + x)

#Initialize the summation
sum1 = np.zeros((H*M))

#Compute the sums
#Calculate next period's policy, output, and consumption
kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
yp = np.outer(a, kp**alpha)
cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp

#Calculate marginal product of capital and marginal utility
dudcp = cp**(-gamma)
dydkp = np.outer(a, kp**(alpha - 1))*alpha
rhs = dudcp*(dydkp + 1 - delta)

#Generate a block diagonal matrix for the expectation
mats = [rhs[:, i*M*(L - 1): (i+1)*M*(L - 1)] for i in range(0, I)]
temp1 = scipy.linalg.block_diag(*mats)
temp2 = PIE.reshape(I*I)
sum1 = np.dot(temp2, temp1)

#Calculate the residual functions
res = (beta*sum1 - c**(-gamma))*u
res = res.reshape((H, M))
RES[[i for i in range(0, I*L) if i not in range((L - 1),I*L,L)]] += np.dot(res, Bs1)
RES[[i for i in range(0, I*L) if i not in range(0, I*L - 1, L)]] += np.dot(res, Bs2)

"""
Explanation: And just to move things along, here is the vectorized version:
End of explanation
"""


kp = np.kron(coeff[:, :-1].reshape(H,), np.ones((M)))*bs1\
    + np.kron(coeff[:, 1:].reshape(H,), np.ones((M,)))*bs2
print(kp.shape)

"""
Explanation: This code looks almost exactly the same, and this is one of the benefits of vectorization.  An increase in dimensions from one to two can be tricky, but from two to three or three to four is usually trivial.  Let's look at exactly what's going on.
First, we build a vector kp the same way as befor, except now it is longer.  One point is the reshpae operation:
End of explanation
"""


test = np.arange(9)
print(test)
print(test.reshape(3,3))
print(test.reshape(3,3).reshape(9,))

"""
Explanation: In order to stack all of the loops as we did befor, we reshape the coefficient matrix.  It is important to think about how this reshape operation works.  The default method is to allow the last (or rightmost) index to vary the fastest.  Let's see an example:
End of explanation
"""


#Generate a block diagonal matrix for the expectation
mats = [rhs[:, i*M*(L - 1): (i+1)*M*(L - 1)] for i in range(0, I)]
temp1 = scipy.linalg.block_diag(*mats)
temp2 = PIE.reshape(I*I)
sum1 = np.dot(temp2, temp1)

"""
Explanation: So as you can see, the reshape method (or corresponding function) will either take items from the first row of a matrix, then the second, and so on.  Or, conversely, if reshaping a vector to a matrix it will fill the first row first, the second row second, and etc.
Looking back to kp, the reshape is just taking the rows of the coeff matrix and stacking them end-to-end to make a long vector.  This is exactly what we want.
The next important point here is the calculation of the integrals.  You'll notice that we construct the rhs matrix, which is $I \times (L-1)I$.  We need to take the dot product of the first row of $\Pi$ with the first $L-1$ columns of rhs, the second with the second, and so forth. This presents a very odd computational challenge, but here is my solution:
We can transform the rhs matrix up into a block diagonal matrix, where each block corresponds to $L-1$ columns from the original matrix.  This matrix will we $I^2\times (L-1)I$.  We can then reshape the $\Pi$ matrix into a vector.  If we use the reshape method, this will stack the rows end to end.  Finally, taking the dot product of these two should give us the desired result!
That's carried out in these lines, where we first create a list of arrays, then pass them to scipy.linalg.block_diag as a star-args argument (allowing you to pass a dynamic number of matrices), and finally reshaping and taking the dot product:
End of explanation
"""


#Here are all the indices
print([i for i in range(0, I*L)])

#Here are all the indices we dont want
#NOTE: The third argument is the step size
print([L-1, I*L, L])

#Here are all the indices we do want!
print([i for i in range(0, I*L) if i not in range((L - 1),I*L,L)])

#And, on the second line, we want all of these indices
print([i for i in range(0, I*L) if i not in range(0, I*L - 1, L)])

"""
Explanation: The last point to make is that we need to fill only a subset of the rows of RES.  In two lines we would like to fill first the rows that aren't on the upper bound of capital, and on the next those that aren't on the lower bound of capital.  This present a need for some truly fancy slicing.
First, notice that in the first case we want the numbers zero through $L-2$, $L$ through $2L-1$, and so on.  Put a different way, we want all of the indices, less those corresponding to $l = L - 1$.  We can do this with a fancy slice and using range.  Check it out:
End of explanation
"""


res = (beta*sum1 - c**(-gamma))*u
res = res.reshape((H, M))
RES[[i for i in range(0, I*L) if i not in range((L - 1),I*L,L)]] += np.dot(res, Bs1)
RES[[i for i in range(0, I*L) if i not in range(0, I*L - 1, L)]] += np.dot(res, Bs2)

"""
Explanation: Thus, we can reshape the res object and fill the corresponding rows of RES using a dot product:
End of explanation
"""


def mcgratten_weighted_residual_vec(coeff):
    """
    A vectorized function of the McGratten problem.

    Variables:
        coeff   :   ndarray; a vector of coefficients for the parameterized
                    policy function
    Returns:
        ndarray containing the residuals.

    """
    #Unpack the arguments
    alpha, delta, I, L, k, abcissas, M, a, PIE = args

    #Reshape coefficients in case they are vector
    coeff = coeff.reshape(I, L)

    #Initialize the residuals
    RES = np.zeros((I*L))

    #For readability define the number of rows
    H = I*(L - 1)

    #Construct the vector of residual funcitons
    #Vectorize all loops
    #Scale and calculate capital
    x = np.kron(np.ones(H), abcissas)
    kt = scalup(x, np.kron(k[1:], np.ones(I*M)), np.kron(k[:-1], np.ones(I*M)))
    u = np.kron(np.ones(H), weights)
    bs1 = 0.5*(1 - x)
    bs2 = 0.5*(1 + x)
    #NOTE: bs is constant across elements, so simplify calculations later
    Bs1 = 0.5*(1 - abcissas)
    Bs2 = 0.5*(1 + abcissas)

    #Compute capital next period by projection
    kp = np.kron(coeff[:, :-1].reshape(H,), np.ones((M)))*bs1\
        + np.kron(coeff[:, 1:].reshape(H,), np.ones((M,)))*bs2

    #Compute consumption
    y = np.kron(a, np.ones((L - 1)*M))*kt**alpha
    c = y + (1 - delta)*kt - kp

    #Find indices for next periods element bounds
    ii = np.array([np.where(point >= np.append([0], k[:-1]))[0][-1]
                   for point in kp])
    ii[ii > 0] -= 1

    #Define element bounds for next period
    k1p = k[ii]
    k2p = k[ii + 1]

    #Calculate policy next period
    x = scaldwn(kp, k2p, k1p)
    bs1p = 0.5*(1 - x)
    bs2p = 0.5*(1 + x)

    #Initialize the summation
    sum1 = np.zeros((H*M))

    #Compute the sums
    #Calculate next period's policy, output, and consumption
    kpp = coeff[:, ii]*bs1p + coeff[:, ii+1]*bs2p
    yp = np.outer(a, kp**alpha)
    cp = yp + (1 - delta)*np.outer(np.ones(I), kp) - kpp

    #Calculate marginal product of capital and marginal utility
    dudcp = cp**(-gamma)
    dydkp = np.outer(a, kp**(alpha - 1))*alpha
    rhs = dudcp*(dydkp + 1 - delta)

    #Generate a block diagonal matrix for the expectation
    mats = [rhs[:, i*M*(L - 1): (i+1)*M*(L - 1)] for i in range(0, I)]
    temp1 = scipy.linalg.block_diag(*mats)
    temp2 = PIE.reshape(I*I)
    sum1 = np.dot(temp2, temp1)

    #Calculate the residual functions
    res = (beta*sum1 - c**(-gamma))*u
    res = res.reshape((H, M))
    RES[[i for i in range(0, I*L) if i not in
        range((L - 1), I*L, L)]] += np.dot(res, Bs1)
    RES[[i for i in range(0, I*L) if i not in
        range(0, I*L - 1, L)]] += np.dot(res, Bs2)

    return RES


time0 = time.time()
res = scipy.optimize.newton_krylov(mcgratten_weighted_residual, coeff0,
                                   method='lgmres', verbose=True, maxiter=1000,
                                   line_search='armijo')
print(time.time() - time0)

time0 = time.time()
res_vec = scipy.optimize.newton_krylov(mcgratten_weighted_residual_vec, coeff0,
                                       method='lgmres', verbose=True,
                                       maxiter=1000, line_search='armijo')
print(time.time() - time0)
print(np.linalg.norm(res - res_vec))


"""
Explanation: Here is the full function with a comparison to the original:
End of explanation
"""


