import numpy as np

from scipy.optimize import fmin_slsqp

from scipy.optimize import root

from quantecon import MarkovChain

class RecursiveAllocationAMSS:

def __init__(self, model, μgrid, tol_diff=1e-7, tol=1e-7):

self.β, self.π, self.G = model.β, model.π, model.G

self.mc, self.S = MarkovChain(self.π), len(model.π) # Number of states

self.Θ, self.model, self.μgrid = model.Θ, model, μgrid

self.tol_diff, self.tol = tol_diff, tol

# Find the first best allocation

self.solve_time1_bellman()

self.T.time_0 = True # Bellman equation now solves time 0 problem

def solve_time1_bellman(self):

'''

Solve the time 1 Bellman equation for calibration model and

initial grid μgrid0

'''

model, μgrid0 = self.model, self.μgrid

π = model.π

S = len(model.π)

# First get initial fit from Lucas Stokey solution.

# Need to change things to be ex ante

pp = SequentialAllocation(model)

interp = interpolator_factory(2, None)

def incomplete_allocation(μ_, s_):

c, n, x, V = pp.time1_value(μ_)

return c, n, π[s_] @ x, π[s_] @ V

cf, nf, xgrid, Vf, xprimef = [], [], [], [], []

for s_ in range(S):

c, n, x, V = zip(*map(lambda μ: incomplete_allocation(μ, s_), μgrid0))

c, n = np.vstack(c).T, np.vstack(n).T

x, V = np.hstack(x), np.hstack(V)

xprimes = np.vstack([x] * S)

cf.append(interp(x, c))

nf.append(interp(x, n))

Vf.append(interp(x, V))

xgrid.append(x)

xprimef.append(interp(x, xprimes))

cf, nf, xprimef = fun_vstack(cf), fun_vstack(nf), fun_vstack(xprimef)

Vf = fun_hstack(Vf)

policies = [cf, nf, xprimef]

# Create xgrid

x = np.vstack(xgrid).T

xbar = [x.min(0).max(), x.max(0).min()]

xgrid = np.linspace(xbar[0], xbar[1], len(μgrid0))

self.xgrid = xgrid

# Now iterate on Bellman equation

T = BellmanEquation(model, xgrid, policies, tol=self.tol)

diff = 1

while diff > self.tol_diff:

PF = T(Vf)

Vfnew, policies = self.fit_policy_function(PF)

diff = np.abs((Vf(xgrid) - Vfnew(xgrid)) / Vf(xgrid)).max()

print(diff)

Vf = Vfnew

# Store value function policies and Bellman Equations

self.Vf = Vf

self.policies = policies

self.T = T

def fit_policy_function(self, PF):

'''

Fits the policy functions

'''

S, xgrid = len(self.π), self.xgrid

interp = interpolator_factory(3, 0)

cf, nf, xprimef, Tf, Vf = [], [], [], [], []

for s_ in range(S):

PFvec = np.vstack([PF(x, s_) for x in self.xgrid]).T

Vf.append(interp(xgrid, PFvec[0, :]))

cf.append(interp(xgrid, PFvec[1:1 + S]))

nf.append(interp(xgrid, PFvec[1 + S:1 + 2 * S]))

xprimef.append(interp(xgrid, PFvec[1 + 2 * S:1 + 3 * S]))

Tf.append(interp(xgrid, PFvec[1 + 3 * S:]))

policies = fun_vstack(cf), fun_vstack(

nf), fun_vstack(xprimef), fun_vstack(Tf)

Vf = fun_hstack(Vf)

return Vf, policies

def Τ(self, c, n):

'''

Computes Τ given c and n

'''

model = self.model

Uc, Un = model.Uc(c, n), model.Un(c, n)

return 1 + Un / (self.Θ * Uc)

def time0_allocation(self, B_, s0):

'''

Finds the optimal allocation given initial government debt B_ and

state s_0

'''

PF = self.T(self.Vf)

z0 = PF(B_, s0)

c0, n0, xprime0, T0 = z0[1:]

return c0, n0, xprime0, T0

def simulate(self, B_, s_0, T, sHist=None):

'''

Simulates planners policies for T periods

'''

model, π = self.model, self.π

Uc = model.Uc

cf, nf, xprimef, Tf = self.policies

if sHist is None:

sHist = simulate_markov(π, s_0, T)

cHist, nHist, Bhist, xHist, ΤHist, THist, μHist = np.zeros((7, T))

# Time 0

cHist[0], nHist[0], xHist[0], THist[0] = self.time0_allocation(B_, s_0)

ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]

Bhist[0] = B_

μHist[0] = self.Vf[s_0](xHist[0])

# Time 1 onward

for t in range(1, T):

s_, x, s = sHist[t - 1], xHist[t - 1], sHist[t]

c, n, xprime, T = cf[s_, :](x), nf[s_, :](

x), xprimef[s_, :](x), Tf[s_, :](x)

Τ = self.Τ(c, n)[s]

u_c = Uc(c, n)

Eu_c = π[s_, :] @ u_c

μHist[t] = self.Vf[s](xprime[s])

cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n[s], x / Eu_c, Τ

xHist[t], THist[t] = xprime[s], T[s]

return np.array([cHist, nHist, Bhist, ΤHist, THist, μHist, sHist, xHist])

class BellmanEquation:

'''

Bellman equation for the continuation of the Lucas-Stokey Problem

'''

def __init__(self, model, xgrid, policies0, tol, maxiter=1000):

self.β, self.π, self.G = model.β, model.π, model.G

self.S = len(model.π) # Number of states

self.Θ, self.model, self.tol = model.Θ, model, tol

self.maxiter = maxiter

self.xbar = [min(xgrid), max(xgrid)]

self.time_0 = False

self.z0 = {}

cf, nf, xprimef = policies0

for s_ in range(self.S):

for x in xgrid:

self.z0[x, s_] = np.hstack([cf[s_, :](x),

nf[s_, :](x),

xprimef[s_, :](x),

np.zeros(self.S)])

self.find_first_best()

def find_first_best(self):

'''

Find the first best allocation

'''

model = self.model

S, Θ, Uc, Un, G = self.S, self.Θ, model.Uc, model.Un, self.G

def res(z):

c = z[:S]

n = z[S:]

return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])

res = root(res, 0.5 * np.ones(2 * S))

if not res.success:

raise Exception('Could not find first best')

self.cFB = res.x[:S]

self.nFB = res.x[S:]

IFB = Uc(self.cFB, self.nFB) * self.cFB + \

Un(self.cFB, self.nFB) * self.nFB

self.xFB = np.linalg.solve(np.eye(S) - self.β * self.π, IFB)

self.zFB = {}

for s in range(S):

self.zFB[s] = np.hstack(

[self.cFB[s], self.nFB[s], self.π[s] @ self.xFB, 0.])

def __call__(self, Vf):

'''

Given continuation value function next period return value function this

period return T(V) and optimal policies

'''

if not self.time_0:

def PF(x, s): return self.get_policies_time1(x, s, Vf)

else:

def PF(B_, s0): return self.get_policies_time0(B_, s0, Vf)

return PF

def get_policies_time1(self, x, s_, Vf):

'''

Finds the optimal policies

'''

model, β, Θ, G, S, π = self.model, self.β, self.Θ, self.G, self.S, self.π

U, Uc, Un = model.U, model.Uc, model.Un

def objf(z):

c, n, xprime = z[:S], z[S:2 * S], z[2 * S:3 * S]

Vprime = np.empty(S)

for s in range(S):

Vprime[s] = Vf[s](xprime[s])

return -π[s_] @ (U(c, n) + β * Vprime)

def objf_prime(x):

epsilon = 1e-7

x0 = np.asfarray(x)

f0 = np.atleast_1d(objf(x0))

jac = np.zeros([len(x0), len(f0)])

dx = np.zeros(len(x0))

for i in range(len(x0)):

dx[i] = epsilon

jac[i] = (objf(x0+dx) - f0)/epsilon

dx[i] = 0.0

return jac.transpose()

def cons(z):

c, n, xprime, T = z[:S], z[S:2 * S], z[2 * S:3 * S], z[3 * S:]

u_c = Uc(c, n)

Eu_c = π[s_] @ u_c

return np.hstack([

x * u_c / Eu_c - u_c * (c - T) - Un(c, n) * n - β * xprime,

Θ * n - c - G])

if model.transfers:

bounds = [(0., 100)] * S + [(0., 100)] * S + \

[self.xbar] * S + [(0., 100.)] * S

else:

bounds = [(0., 100)] * S + [(0., 100)] * S + \

[self.xbar] * S + [(0., 0.)] * S

out, fx, _, imode, smode = fmin_slsqp(objf, self.z0[x, s_],

f_eqcons=cons, bounds=bounds,

fprime=objf_prime, full_output=True,

iprint=0, acc=self.tol, iter=self.maxiter)

if imode > 0:

raise Exception(smode)

self.z0[x, s_] = out

return np.hstack([-fx, out])

def get_policies_time0(self, B_, s0, Vf):

'''

Finds the optimal policies

'''

model, β, Θ, G = self.model, self.β, self.Θ, self.G

U, Uc, Un = model.U, model.Uc, model.Un

def objf(z):

c, n, xprime = z[:-1]

return -(U(c, n) + β * Vf[s0](xprime))

def cons(z):

c, n, xprime, T = z

return np.hstack([

-Uc(c, n) * (c - B_ - T) - Un(c, n) * n - β * xprime,

(Θ * n - c - G)[s0]])

if model.transfers:

bounds = [(0., 100), (0., 100), self.xbar, (0., 100.)]

else:

bounds = [(0., 100), (0., 100), self.xbar, (0., 0.)]

out, fx, _, imode, smode = fmin_slsqp(objf, self.zFB[s0], f_eqcons=cons,

bounds=bounds, full_output=True,

iprint=0)

if imode > 0:

raise Exception(smode)

return np.hstack([-fx, out])