[opt_tax_recur+amss] migrate updates from lecture-python-advanced (#76)

* [opt_tax_recur] migrate updates from lecture-python-advanced

* update lecture code files

* [migrate] amss,amss2,amss3 updates from lecture-python-advanced

* [amss] add jupyter-book tags for controlling output and collapse blocks

* [amss2] add jupyter-book tags for controlling output and collapse blocks

Author: mmcky

lectures/_static/lecture_specific/amss/recursive_allocation.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+
+class CRRAutility:
+
+    def __init__(self,
+                 β=0.9,
+                 σ=2,
+                 γ=2,
+                 π=0.5*np.ones((2, 2)),
+                 G=np.array([0.1, 0.2]),
+                 Θ=np.ones(2),
+                 transfers=False):
+
+        self.β, self.σ, self.γ = β, σ, γ
+        self.π, self.G, self.Θ, self.transfers = π, G, Θ, transfers
+
+    # Utility function
+    def U(self, c, n):
+        σ = self.σ
+        if σ == 1.:
+            U = np.log(c)
+        else:
+            U = (c**(1 - σ) - 1) / (1 - σ)
+        return U - n**(1 + self.γ) / (1 + self.γ)
+
+    # Derivatives of utility function
+    def Uc(self, c, n):
+        return c**(-self.σ)
+
+    def Ucc(self, c, n):
+        return -self.σ * c**(-self.σ - 1)
+
+    def Un(self, c, n):
+        return -n**self.γ
+
+    def Unn(self, c, n):
+        return -self.γ * n**(self.γ - 1)

lectures/_static/lecture_specific/amss2/crra_utility.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+class LogUtility:
+
+    def __init__(self,
+                 β=0.9,
+                 ψ=0.69,
+                 π=0.5*np.ones((2, 2)),
+                 G=np.array([0.1, 0.2]),
+                 Θ=np.ones(2),
+                 transfers=False):
+
+        self.β, self.ψ, self.π = β, ψ, π
+        self.G, self.Θ, self.transfers = G, Θ, transfers
+
+    # Utility function
+    def U(self, c, n):
+        return np.log(c) + self.ψ * np.log(1 - n)
+
+    # Derivatives of utility function
+    def Uc(self, c, n):
+        return 1 / c
+
+    def Ucc(self, c, n):
+        return -c**(-2)
+
+    def Un(self, c, n):
+        return -self.ψ / (1 - n)
+
+    def Unn(self, c, n):
+        return -self.ψ / (1 - n)**2

lectures/_static/lecture_specific/amss2/log_utility.py

@@ -0,0 +1,298 @@
+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])