diff --git a/BCG_complete_mkts.ipynb b/BCG_complete_mkts.ipynb
index a9cda73..5e4fcad 100644
--- a/BCG_complete_mkts.ipynb
+++ b/BCG_complete_mkts.ipynb
@@ -1406,7 +1406,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661940.1660957,
+  "date": 1624431168.2801201,
   "filename": "BCG_complete_mkts.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/BCG_incomplete_mkts.ipynb b/BCG_incomplete_mkts.ipynb
index 5a0e013..8ef152b 100644
--- a/BCG_incomplete_mkts.ipynb
+++ b/BCG_incomplete_mkts.ipynb
@@ -2215,7 +2215,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661940.5953467,
+  "date": 1624431168.5461557,
   "filename": "BCG_incomplete_mkts.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/about_lectures.ipynb b/about_lectures.ipynb
index af82419..d713017 100644
--- a/about_lectures.ipynb
+++ b/about_lectures.ipynb
@@ -28,7 +28,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661940.744541,
+  "date": 1624431168.7001786,
   "filename": "about_lectures.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/additive_functionals.ipynb b/additive_functionals.ipynb
index 90140bf..f290dbd 100644
--- a/additive_functionals.ipynb
+++ b/additive_functionals.ipynb
@@ -1498,7 +1498,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661940.8426106,
+  "date": 1624431168.7973125,
   "filename": "additive_functionals.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/amss.ipynb b/amss.ipynb
index 2d50e53..fd91f22 100644
--- a/amss.ipynb
+++ b/amss.ipynb
@@ -48,7 +48,8 @@
    },
    "outputs": [],
    "source": [
-    "!pip install --upgrade quantecon"
+    "!pip install --upgrade quantecon\n",
+    "!pip install interpolation"
    ]
   },
   {
@@ -70,10 +71,13 @@
    "source": [
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
-    "%matplotlib inline\n",
-    "from scipy.optimize import root, fmin_slsqp\n",
-    "from scipy.interpolate import UnivariateSpline\n",
-    "from quantecon import MarkovChain"
+    "from scipy.optimize import root\n",
+    "from interpolation.splines import eval_linear, UCGrid, nodes\n",
+    "from quantecon import optimize, MarkovChain\n",
+    "from numba import njit, prange, float64\n",
+    "from numba.experimental import jitclass\n",
+    "\n",
+    "%matplotlib inline"
    ]
   },
   {
@@ -481,164 +485,204 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from scipy.optimize import root\n",
-    "from quantecon import MarkovChain\n",
-    "\n",
-    "\n",
-    "class SequentialAllocation:\n",
+    "class SequentialLS:\n",
     "\n",
     "    '''\n",
-    "    Class that takes CESutility or BGPutility object as input returns\n",
-    "    planner's allocation as a function of the multiplier on the\n",
-    "    implementability constraint μ.\n",
+    "    Class that takes a preference object, state transition matrix,\n",
+    "    and state contingent government expenditure plan as inputs, and\n",
+    "    solves the sequential allocation problem described above.\n",
+    "    It returns optimal allocations about consumption and labor supply,\n",
+    "    as well as the multiplier on the implementability constraint Φ.\n",
     "    '''\n",
     "\n",
-    "    def __init__(self, model):\n",
+    "    def __init__(self,\n",
+    "                 pref,\n",
+    "                 π=0.5*np.ones((2, 2)),\n",
+    "                 g=np.array([0.1, 0.2])):\n",
     "\n",
-    "        # Initialize from model object attributes\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.mc, self.Θ = MarkovChain(self.π), model.Θ\n",
-    "        self.S = len(model.π)  # Number of states\n",
-    "        self.model = model\n",
+    "        # Initialize from pref object attributes\n",
+    "        self.β, self.π, self.g = pref.β, π, g\n",
+    "        self.mc = MarkovChain(self.π)\n",
+    "        self.S = len(π)  # Number of states\n",
+    "        self.pref = pref\n",
     "\n",
     "        # Find the first best allocation\n",
     "        self.find_first_best()\n",
     "\n",
+    "    def FOC_first_best(self, c, g):\n",
+    "        '''\n",
+    "        First order conditions that characterize\n",
+    "        the first best allocation.\n",
+    "        '''\n",
+    "\n",
+    "        pref = self.pref\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
+    "\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
+    "\n",
+    "        return Uc(c, l) - Ul(c, l)\n",
+    "\n",
     "    def find_first_best(self):\n",
     "        '''\n",
     "        Find the first best allocation\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, G = self.S, self.Θ, self.G\n",
-    "        Uc, Un = model.Uc, model.Un\n",
-    "\n",
-    "        def res(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:]\n",
-    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
+    "        S, g = self.S, self.g\n",
     "\n",
-    "        res = root(res, 0.5 * np.ones(2 * S))\n",
+    "        res = root(self.FOC_first_best, 0.5 * np.ones(S), args=(g,))\n",
     "\n",
-    "        if not res.success:\n",
+    "        if (res.fun > 1e-10).any():\n",
     "            raise Exception('Could not find first best')\n",
     "\n",
-    "        self.cFB = res.x[:S]\n",
-    "        self.nFB = res.x[S:]\n",
+    "        self.cFB = res.x\n",
+    "        self.nFB = self.cFB + g\n",
     "\n",
-    "        # Multiplier on the resource constraint\n",
-    "        self.ΞFB = Uc(self.cFB, self.nFB)\n",
-    "        self.zFB = np.hstack([self.cFB, self.nFB, self.ΞFB])\n",
+    "    def FOC_time1(self, c, Φ, g):\n",
+    "        '''\n",
+    "        First order conditions that characterize\n",
+    "        optimal time 1 allocation problems.\n",
+    "        '''\n",
+    "\n",
+    "        pref = self.pref\n",
+    "        Uc, Ucc, Ul, Ull, Ulc = pref.Uc, pref.Ucc, pref.Ul, pref.Ull, pref.Ulc\n",
+    "\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
     "\n",
-    "    def time1_allocation(self, μ):\n",
+    "        LHS = (1 + Φ) * Uc(c, l) + Φ * (c * Ucc(c, l) - n * Ulc(c, l))\n",
+    "        RHS = (1 + Φ) * Ul(c, l) + Φ * (c * Ulc(c, l) - n * Ull(c, l))\n",
+    "\n",
+    "        diff = LHS - RHS\n",
+    "\n",
+    "        return diff\n",
+    "\n",
+    "    def time1_allocation(self, Φ):\n",
     "        '''\n",
-    "        Computes optimal allocation for time t >= 1 for a given μ\n",
+    "        Computes optimal allocation for time t >= 1 for a given Φ\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, G = self.S, self.Θ, self.G\n",
-    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
-    "\n",
-    "        def FOC(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:2 * S]\n",
-    "            Ξ = z[2 * S:]\n",
-    "            # FOC of c\n",
-    "            return np.hstack([Uc(c, n) - μ * (Ucc(c, n) * c + Uc(c, n)) - Ξ,\n",
-    "                              Un(c, n) - μ * (Unn(c, n) * n + Un(c, n)) \\\n",
-    "                              + Θ * Ξ,  # FOC of n\n",
-    "                              Θ * n - c - G])\n",
-    "\n",
-    "        # Find the root of the first-order condition\n",
-    "        res = root(FOC, self.zFB)\n",
-    "        if not res.success:\n",
+    "        pref = self.pref\n",
+    "        S, g = self.S, self.g\n",
+    "\n",
+    "        # use the first best allocation as intial guess\n",
+    "        res = root(self.FOC_time1, self.cFB, args=(Φ, g))\n",
+    "\n",
+    "        if (res.fun > 1e-10).any():\n",
     "            raise Exception('Could not find LS allocation.')\n",
-    "        z = res.x\n",
-    "        c, n, Ξ = z[:S], z[S:2 * S], z[2 * S:]\n",
+    "\n",
+    "        c = res.x\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
     "\n",
     "        # Compute x\n",
-    "        I = Uc(c, n) * c + Un(c, n) * n\n",
+    "        I = pref.Uc(c, n) * c - pref.Ul(c, l) * n\n",
     "        x = np.linalg.solve(np.eye(S) - self.β * self.π, I)\n",
     "\n",
-    "        return c, n, x, Ξ\n",
+    "        return c, n, x\n",
     "\n",
-    "    def time0_allocation(self, B_, s_0):\n",
+    "    def FOC_time0(self, c0, Φ, g0, b0):\n",
     "        '''\n",
-    "        Finds the optimal allocation given initial government debt B_ and\n",
-    "        state s_0\n",
+    "        First order conditions that characterize\n",
+    "        time 0 allocation problem.\n",
     "        '''\n",
-    "        model, π, Θ, G, β = self.model, self.π, self.Θ, self.G, self.β\n",
-    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
-    "\n",
-    "        # First order conditions of planner's problem\n",
-    "        def FOC(z):\n",
-    "            μ, c, n, Ξ = z\n",
-    "            xprime = self.time1_allocation(μ)[2]\n",
-    "            return np.hstack([Uc(c, n) * (c - B_) + Un(c, n) * n + β * π[s_0]\n",
-    "                                            @ xprime,\n",
-    "                              Uc(c, n) - μ * (Ucc(c, n)\n",
-    "                                            * (c - B_) + Uc(c, n)) - Ξ,\n",
-    "                              Un(c, n) - μ * (Unn(c, n) * n\n",
-    "                                            + Un(c, n)) + Θ[s_0] * Ξ,\n",
-    "                              (Θ * n - c - G)[s_0]])\n",
-    "\n",
-    "        # Find root\n",
-    "        res = root(FOC, np.array(\n",
-    "            [0, self.cFB[s_0], self.nFB[s_0], self.ΞFB[s_0]]))\n",
-    "        if not res.success:\n",
-    "            raise Exception('Could not find time 0 LS allocation.')\n",
     "\n",
-    "        return res.x\n",
+    "        pref = self.pref\n",
+    "        Ucc, Ulc = pref.Ucc, pref.Ulc\n",
+    "\n",
+    "        n0 = c0 + g0\n",
+    "        l0 = 1 - n0\n",
+    "\n",
+    "        diff = self.FOC_time1(c0, Φ, g0)\n",
+    "        diff -= Φ * (Ucc(c0, l0) - Ulc(c0, l0)) * b0\n",
+    "\n",
+    "        return diff\n",
+    "\n",
+    "    def implementability(self, Φ, b0, s0, cn0_arr):\n",
+    "        '''\n",
+    "        Compute the differences between the RHS and LHS\n",
+    "        of the implementability constraint given Φ,\n",
+    "        initial debt, and initial state.\n",
+    "        '''\n",
+    "\n",
+    "        pref, π, g, β = self.pref, self.π, self.g, self.β\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
+    "        g0 = self.g[s0]\n",
+    "\n",
+    "        c, n, x = self.time1_allocation(Φ)\n",
+    "\n",
+    "        res = root(self.FOC_time0, cn0_arr[0], args=(Φ, g0, b0))\n",
+    "        c0 = res.x\n",
+    "        n0 = c0 + g0\n",
+    "        l0 = 1 - n0\n",
+    "\n",
+    "        cn0_arr[:] = c0, n0\n",
+    "\n",
+    "        LHS = Uc(c0, l0) * b0\n",
+    "        RHS = Uc(c0, l0) * c0 - Ul(c0, l0) * n0 + β * π[s0] @ x\n",
+    "\n",
+    "        return RHS - LHS\n",
     "\n",
-    "    def time1_value(self, μ):\n",
+    "    def time0_allocation(self, b0, s0):\n",
     "        '''\n",
-    "        Find the value associated with multiplier μ\n",
+    "        Finds the optimal time 0 allocation given\n",
+    "        initial government debt b0 and state s0\n",
     "        '''\n",
-    "        c, n, x, Ξ = self.time1_allocation(μ)\n",
-    "        U = self.model.U(c, n)\n",
-    "        V = np.linalg.solve(np.eye(self.S) - self.β * self.π, U)\n",
-    "        return c, n, x, V\n",
     "\n",
-    "    def Τ(self, c, n):\n",
+    "        # use the first best allocation as initial guess\n",
+    "        cn0_arr = np.array([self.cFB[s0], self.nFB[s0]])\n",
+    "\n",
+    "        res = root(self.implementability, 0., args=(b0, s0, cn0_arr))\n",
+    "\n",
+    "        if (res.fun > 1e-10).any():\n",
+    "            raise Exception('Could not find time 0 LS allocation.')\n",
+    "\n",
+    "        Φ = res.x[0]\n",
+    "        c0, n0 = cn0_arr\n",
+    "\n",
+    "        return Φ, c0, n0\n",
+    "\n",
+    "    def τ(self, c, n):\n",
     "        '''\n",
-    "        Computes Τ given c, n\n",
+    "        Computes τ given c, n\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        Uc, Un = model.Uc(c, n), model.Un(c,  n)\n",
+    "        pref = self.pref\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
     "\n",
-    "        return 1 + Un / (self.Θ * Uc)\n",
+    "        return 1 - Ul(c, 1-n) / Uc(c, 1-n)\n",
     "\n",
-    "    def simulate(self, B_, s_0, T, sHist=None):\n",
+    "    def simulate(self, b0, s0, T, sHist=None):\n",
     "        '''\n",
     "        Simulates planners policies for T periods\n",
     "        '''\n",
-    "        model, π, β = self.model, self.π, self.β\n",
-    "        Uc = model.Uc\n",
+    "        pref, π, β = self.pref, self.π, self.β\n",
+    "        Uc = pref.Uc\n",
     "\n",
     "        if sHist is None:\n",
-    "            sHist = self.mc.simulate(T, s_0)\n",
+    "            sHist = self.mc.simulate(T, s0)\n",
     "\n",
-    "        cHist, nHist, Bhist, ΤHist, μHist = np.zeros((5, T))\n",
-    "        RHist = np.zeros(T - 1)\n",
+    "        cHist, nHist, Bhist, τHist, ΦHist = np.empty((5, T))\n",
+    "        RHist = np.empty(T-1)\n",
     "\n",
     "        # Time 0\n",
-    "        μ, cHist[0], nHist[0], _ = self.time0_allocation(B_, s_0)\n",
-    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
-    "        Bhist[0] = B_\n",
-    "        μHist[0] = μ\n",
+    "        Φ, cHist[0], nHist[0] = self.time0_allocation(b0, s0)\n",
+    "        τHist[0] = self.τ(cHist[0], nHist[0])\n",
+    "        Bhist[0] = b0\n",
+    "        ΦHist[0] = Φ\n",
     "\n",
     "        # Time 1 onward\n",
     "        for t in range(1, T):\n",
-    "            c, n, x, Ξ = self.time1_allocation(μ)\n",
-    "            Τ = self.Τ(c, n)\n",
-    "            u_c = Uc(c, n)\n",
+    "            c, n, x = self.time1_allocation(Φ)\n",
+    "            τ = self.τ(c, n)\n",
+    "            u_c = Uc(c, 1-n)\n",
     "            s = sHist[t]\n",
-    "            Eu_c = π[sHist[t - 1]] @ u_c\n",
-    "            cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n[s], x[s] / u_c[s], \\\n",
-    "                                                     Τ[s]\n",
-    "            RHist[t - 1] = Uc(cHist[t - 1], nHist[t - 1]) / (β * Eu_c)\n",
-    "            μHist[t] = μ\n",
+    "            Eu_c = π[sHist[t-1]] @ u_c\n",
+    "            cHist[t], nHist[t], Bhist[t], τHist[t] = c[s], n[s], x[s] / u_c[s], τ[s]\n",
+    "            RHist[t-1] = Uc(cHist[t-1], 1-nHist[t-1]) / (β * Eu_c)\n",
+    "            ΦHist[t] = Φ\n",
+    "\n",
+    "        gHist = self.g[sHist]\n",
+    "        yHist = nHist\n",
     "\n",
-    "        return np.array([cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist])"
+    "        return [cHist, nHist, Bhist, τHist, gHist, yHist, sHist, ΦHist, RHist]"
    ]
   },
   {
@@ -963,304 +1007,218 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from scipy.optimize import fmin_slsqp\n",
-    "from scipy.optimize import root\n",
-    "from quantecon import MarkovChain\n",
-    "\n",
-    "\n",
-    "class RecursiveAllocationAMSS:\n",
-    "\n",
-    "    def __init__(self, model, μgrid, tol_diff=1e-7, tol=1e-7):\n",
-    "\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.mc, self.S = MarkovChain(self.π), len(model.π)  # Number of states\n",
-    "        self.Θ, self.model, self.μgrid = model.Θ, model, μgrid\n",
-    "        self.tol_diff, self.tol = tol_diff, tol\n",
-    "\n",
-    "        # Find the first best allocation\n",
-    "        self.solve_time1_bellman()\n",
-    "        self.T.time_0 = True  # Bellman equation now solves time 0 problem\n",
-    "\n",
-    "    def solve_time1_bellman(self):\n",
-    "        '''\n",
-    "        Solve the time  1 Bellman equation for calibration model and\n",
-    "        initial grid μgrid0\n",
-    "        '''\n",
-    "        model, μgrid0 = self.model, self.μgrid\n",
-    "        π = model.π\n",
-    "        S = len(model.π)\n",
-    "\n",
-    "        # First get initial fit from Lucas Stokey solution.\n",
-    "        # Need to change things to be ex ante\n",
-    "        pp = SequentialAllocation(model)\n",
-    "        interp = interpolator_factory(2, None)\n",
-    "\n",
-    "        def incomplete_allocation(μ_, s_):\n",
-    "            c, n, x, V = pp.time1_value(μ_)\n",
-    "            return c, n, π[s_] @ x, π[s_] @ V\n",
-    "        cf, nf, xgrid, Vf, xprimef = [], [], [], [], []\n",
-    "        for s_ in range(S):\n",
-    "            c, n, x, V = zip(*map(lambda μ: incomplete_allocation(μ, s_), μgrid0))\n",
-    "            c, n = np.vstack(c).T, np.vstack(n).T\n",
-    "            x, V = np.hstack(x), np.hstack(V)\n",
-    "            xprimes = np.vstack([x] * S)\n",
-    "            cf.append(interp(x, c))\n",
-    "            nf.append(interp(x, n))\n",
-    "            Vf.append(interp(x, V))\n",
-    "            xgrid.append(x)\n",
-    "            xprimef.append(interp(x, xprimes))\n",
-    "        cf, nf, xprimef = fun_vstack(cf), fun_vstack(nf), fun_vstack(xprimef)\n",
-    "        Vf = fun_hstack(Vf)\n",
-    "        policies = [cf, nf, xprimef]\n",
-    "\n",
-    "        # Create xgrid\n",
-    "        x = np.vstack(xgrid).T\n",
-    "        xbar = [x.min(0).max(), x.max(0).min()]\n",
-    "        xgrid = np.linspace(xbar[0], xbar[1], len(μgrid0))\n",
-    "        self.xgrid = xgrid\n",
-    "\n",
-    "        # Now iterate on Bellman equation\n",
-    "        T = BellmanEquation(model, xgrid, policies, tol=self.tol)\n",
-    "        diff = 1\n",
-    "        while diff > self.tol_diff:\n",
-    "            PF = T(Vf)\n",
-    "\n",
-    "            Vfnew, policies = self.fit_policy_function(PF)\n",
-    "            diff = np.abs((Vf(xgrid) - Vfnew(xgrid)) / Vf(xgrid)).max()\n",
-    "\n",
-    "            print(diff)\n",
-    "            Vf = Vfnew\n",
-    "\n",
-    "        # Store value function policies and Bellman Equations\n",
-    "        self.Vf = Vf\n",
-    "        self.policies = policies\n",
-    "        self.T = T\n",
-    "\n",
-    "    def fit_policy_function(self, PF):\n",
-    "        '''\n",
-    "        Fits the policy functions\n",
-    "        '''\n",
-    "        S, xgrid = len(self.π), self.xgrid\n",
-    "        interp = interpolator_factory(3, 0)\n",
-    "        cf, nf, xprimef, Tf, Vf = [], [], [], [], []\n",
-    "        for s_ in range(S):\n",
-    "            PFvec = np.vstack([PF(x, s_) for x in self.xgrid]).T\n",
-    "            Vf.append(interp(xgrid, PFvec[0, :]))\n",
-    "            cf.append(interp(xgrid, PFvec[1:1 + S]))\n",
-    "            nf.append(interp(xgrid, PFvec[1 + S:1 + 2 * S]))\n",
-    "            xprimef.append(interp(xgrid, PFvec[1 + 2 * S:1 + 3 * S]))\n",
-    "            Tf.append(interp(xgrid, PFvec[1 + 3 * S:]))\n",
-    "        policies = fun_vstack(cf), fun_vstack(\n",
-    "            nf), fun_vstack(xprimef), fun_vstack(Tf)\n",
-    "        Vf = fun_hstack(Vf)\n",
-    "        return Vf, policies\n",
-    "\n",
-    "    def Τ(self, c, n):\n",
-    "        '''\n",
-    "        Computes Τ given c and n\n",
-    "        '''\n",
-    "        model = self.model\n",
-    "        Uc, Un = model.Uc(c, n), model.Un(c, n)\n",
+    "class AMSS:\n",
+    "    # WARNING: THE CODE IS EXTREMELY SENSITIVE TO CHOCIES OF PARAMETERS.\n",
+    "    # DO NOT CHANGE THE PARAMETERS AND EXPECT IT TO WORK\n",
+    "\n",
+    "    def __init__(self, pref, β, Π, g, x_grid, bounds_v):\n",
+    "        self.β, self.Π, self.g = β, Π, g\n",
+    "        self.x_grid = x_grid\n",
+    "        self.n = x_grid[0][2]\n",
+    "        self.S = len(Π)\n",
+    "        self.bounds = bounds_v\n",
+    "        self.pref = pref\n",
+    "\n",
+    "        self.T_v, self.T_w = bellman_operator_factory(Π, β, x_grid, g,\n",
+    "                                                      bounds_v)\n",
+    "\n",
+    "        self.V_solved = False\n",
+    "        self.W_solved = False\n",
+    "\n",
+    "    def compute_V(self, V, σ_v_star, tol_vfi, maxitr, print_itr):\n",
+    "\n",
+    "        T_v = self.T_v\n",
+    "\n",
+    "        self.success = False\n",
+    "\n",
+    "        V_new = np.zeros_like(V)\n",
+    "\n",
+    "        Δ = 1.0\n",
+    "        for itr in range(maxitr):\n",
+    "            T_v(V, V_new, σ_v_star, self.pref)\n",
+    "\n",
+    "            Δ = np.max(np.abs(V_new - V))\n",
+    "\n",
+    "            if Δ < tol_vfi:\n",
+    "                self.V_solved = True\n",
+    "                print('Successfully completed VFI after %i iterations'\n",
+    "                      % (itr+1))\n",
+    "                break\n",
+    "\n",
+    "            if (itr + 1) % print_itr == 0:\n",
+    "                print('Error at iteration %i : ' % (itr + 1), Δ)\n",
+    "\n",
+    "            V[:] = V_new[:]\n",
+    "\n",
+    "        self.V = V\n",
+    "        self.σ_v_star = σ_v_star\n",
+    "\n",
+    "        return V, σ_v_star\n",
+    "\n",
+    "    def compute_W(self, b_0, W, σ_w_star):\n",
+    "        T_w = self.T_w\n",
+    "        V = self.V\n",
+    "\n",
+    "        T_w(W, σ_w_star, V, b_0, self.pref)\n",
+    "\n",
+    "        self.W = W\n",
+    "        self.σ_w_star = σ_w_star\n",
+    "        self.W_solved = True\n",
+    "        print('Succesfully solved the time 0 problem.')\n",
+    "\n",
+    "        return W, σ_w_star\n",
+    "\n",
+    "    def solve(self, V, σ_v_star,  b_0, W, σ_w_star, tol_vfi=1e-7,\n",
+    "              maxitr=1000, print_itr=10):\n",
+    "        print(\"===============\")\n",
+    "        print(\"Solve time 1 problem\")\n",
+    "        print(\"===============\")\n",
+    "        self.compute_V(V, σ_v_star, tol_vfi, maxitr, print_itr)\n",
+    "        print(\"===============\")\n",
+    "        print(\"Solve time 0 problem\")\n",
+    "        print(\"===============\")\n",
+    "        self.compute_W(b_0, W, σ_w_star)\n",
+    "\n",
+    "    def simulate(self, s_hist, b_0):\n",
+    "        if not (self.V_solved and self.W_solved):\n",
+    "            msg = \"V and W need to be successfully computed before simulation.\"\n",
+    "            raise ValueError(msg)\n",
+    "\n",
+    "        pref = self.pref\n",
+    "        x_grid, g, β, S = self.x_grid, self.g, self.β, self.S\n",
+    "        σ_v_star, σ_w_star = self.σ_v_star, self.σ_w_star\n",
+    "\n",
+    "        T = len(s_hist)\n",
+    "        s_0 = s_hist[0]\n",
+    "\n",
+    "        # Pre-allocate\n",
+    "        n_hist = np.zeros(T)\n",
+    "        x_hist = np.zeros(T)\n",
+    "        c_hist = np.zeros(T)\n",
+    "        τ_hist = np.zeros(T)\n",
+    "        b_hist = np.zeros(T)\n",
+    "        g_hist = np.zeros(T)\n",
+    "\n",
+    "        # Compute t = 0\n",
+    "        l_0, T_0 = σ_w_star[s_0]\n",
+    "        c_0 = (1 - l_0) - g[s_0]\n",
+    "        x_0 = (-pref.Uc(c_0, l_0) * (c_0 - T_0 - b_0) +\n",
+    "               pref.Ul(c_0, l_0) * (1 - l_0))\n",
+    "\n",
+    "        n_hist[0] = (1 - l_0)\n",
+    "        x_hist[0] = x_0\n",
+    "        c_hist[0] = c_0\n",
+    "        τ_hist[0] = 1 - pref.Ul(c_0, l_0) / pref.Uc(c_0, l_0)\n",
+    "        b_hist[0] = b_0\n",
+    "        g_hist[0] = g[s_0]\n",
+    "\n",
+    "        # Compute t > 0\n",
+    "        for t in range(T - 1):\n",
+    "            x_ = x_hist[t]\n",
+    "            s_ = s_hist[t]\n",
+    "            l = np.zeros(S)\n",
+    "            T = np.zeros(S)\n",
+    "            for s in range(S):\n",
+    "                x_arr = np.array([x_])\n",
+    "                l[s] = eval_linear(x_grid, σ_v_star[s_, :, s], x_arr)\n",
+    "                T[s] = eval_linear(x_grid, σ_v_star[s_, :, S+s], x_arr)\n",
     "\n",
-    "        return 1 + Un / (self.Θ * Uc)\n",
+    "            c = (1 - l) - g\n",
+    "            u_c = pref.Uc(c, l)\n",
+    "            Eu_c = Π[s_] @ u_c\n",
     "\n",
-    "    def time0_allocation(self, B_, s0):\n",
-    "        '''\n",
-    "        Finds the optimal allocation given initial government debt B_ and\n",
-    "        state s_0\n",
-    "        '''\n",
-    "        PF = self.T(self.Vf)\n",
-    "        z0 = PF(B_, s0)\n",
-    "        c0, n0, xprime0, T0 = z0[1:]\n",
-    "        return c0, n0, xprime0, T0\n",
+    "            x = u_c * x_ / (β * Eu_c) - u_c * (c - T) + pref.Ul(c, l) * (1 - l)\n",
     "\n",
-    "    def simulate(self, B_, s_0, T, sHist=None):\n",
-    "        '''\n",
-    "        Simulates planners policies for T periods\n",
-    "        '''\n",
-    "        model, π = self.model, self.π\n",
-    "        Uc = model.Uc\n",
-    "        cf, nf, xprimef, Tf = self.policies\n",
+    "            c_next = c[s_hist[t+1]]\n",
+    "            l_next = l[s_hist[t+1]]\n",
     "\n",
-    "        if sHist is None:\n",
-    "            sHist = simulate_markov(π, s_0, T)\n",
+    "            x_hist[t+1] = x[s_hist[t+1]]\n",
+    "            n_hist[t+1] = 1 - l_next\n",
+    "            c_hist[t+1] = c_next\n",
+    "            τ_hist[t+1] = 1 - pref.Ul(c_next, l_next) / pref.Uc(c_next, l_next)\n",
+    "            b_hist[t+1] = x_ / (β * Eu_c)\n",
+    "            g_hist[t+1] = g[s_hist[t+1]]\n",
     "\n",
-    "        cHist, nHist, Bhist, xHist, ΤHist, THist, μHist = np.zeros((7, T))\n",
-    "        # Time 0\n",
-    "        cHist[0], nHist[0], xHist[0], THist[0] = self.time0_allocation(B_, s_0)\n",
-    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
-    "        Bhist[0] = B_\n",
-    "        μHist[0] = self.Vf[s_0](xHist[0])\n",
+    "        return c_hist, n_hist, b_hist, τ_hist, g_hist, n_hist\n",
     "\n",
-    "        # Time 1 onward\n",
-    "        for t in range(1, T):\n",
-    "            s_, x, s = sHist[t - 1], xHist[t - 1], sHist[t]\n",
-    "            c, n, xprime, T = cf[s_, :](x), nf[s_, :](\n",
-    "                x), xprimef[s_, :](x), Tf[s_, :](x)\n",
     "\n",
-    "            Τ = self.Τ(c, n)[s]\n",
-    "            u_c = Uc(c, n)\n",
-    "            Eu_c = π[s_, :] @ u_c\n",
+    "def obj_factory(Π, β, x_grid, g):\n",
+    "    S = len(Π)\n",
     "\n",
-    "            μHist[t] = self.Vf[s](xprime[s])\n",
+    "    @njit\n",
+    "    def obj_V(σ, state, V, pref):\n",
+    "        # Unpack state\n",
+    "        s_, x_ = state\n",
     "\n",
-    "            cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n[s], x / Eu_c, Τ\n",
-    "            xHist[t], THist[t] = xprime[s], T[s]\n",
-    "        return np.array([cHist, nHist, Bhist, ΤHist, THist, μHist, sHist, xHist])\n",
+    "        l = σ[:S]\n",
+    "        T = σ[S:]\n",
     "\n",
+    "        c = (1 - l) - g\n",
+    "        u_c = pref.Uc(c, l)\n",
+    "        Eu_c = Π[s_] @ u_c\n",
+    "        x = u_c * x_ / (β * Eu_c) - u_c * (c - T) + pref.Ul(c, l) * (1 - l)\n",
     "\n",
-    "class BellmanEquation:\n",
-    "    '''\n",
-    "    Bellman equation for the continuation of the Lucas-Stokey Problem\n",
-    "    '''\n",
+    "        V_next = np.zeros(S)\n",
     "\n",
-    "    def __init__(self, model, xgrid, policies0, tol, maxiter=1000):\n",
-    "\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.S = len(model.π)  # Number of states\n",
-    "        self.Θ, self.model, self.tol = model.Θ, model, tol\n",
-    "        self.maxiter = maxiter\n",
-    "\n",
-    "        self.xbar = [min(xgrid), max(xgrid)]\n",
-    "        self.time_0 = False\n",
+    "        for s in range(S):\n",
+    "            V_next[s] = eval_linear(x_grid, V[s], np.array([x[s]]))\n",
     "\n",
-    "        self.z0 = {}\n",
-    "        cf, nf, xprimef = policies0\n",
+    "        out = Π[s_] @ (pref.U(c, l) + β * V_next)\n",
     "\n",
-    "        for s_ in range(self.S):\n",
-    "            for x in xgrid:\n",
-    "                self.z0[x, s_] = np.hstack([cf[s_, :](x),\n",
-    "                                            nf[s_, :](x),\n",
-    "                                            xprimef[s_, :](x),\n",
-    "                                            np.zeros(self.S)])\n",
+    "        return out\n",
     "\n",
-    "        self.find_first_best()\n",
+    "    @njit\n",
+    "    def obj_W(σ, state, V, pref):\n",
+    "        # Unpack state\n",
+    "        s_, b_0 = state\n",
+    "        l, T = σ\n",
     "\n",
-    "    def find_first_best(self):\n",
-    "        '''\n",
-    "        Find the first best allocation\n",
-    "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, Uc, Un, G = self.S, self.Θ, model.Uc, model.Un, self.G\n",
+    "        c = (1 - l) - g[s_]\n",
+    "        x = -pref.Uc(c, l) * (c - T - b_0) + pref.Ul(c, l) * (1 - l)\n",
     "\n",
-    "        def res(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:]\n",
-    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
+    "        V_next = eval_linear(x_grid, V[s_], np.array([x]))\n",
     "\n",
-    "        res = root(res, 0.5 * np.ones(2 * S))\n",
-    "        if not res.success:\n",
-    "            raise Exception('Could not find first best')\n",
+    "        out = pref.U(c, l) + β * V_next\n",
     "\n",
-    "        self.cFB = res.x[:S]\n",
-    "        self.nFB = res.x[S:]\n",
-    "        IFB = Uc(self.cFB, self.nFB) * self.cFB + \\\n",
-    "            Un(self.cFB, self.nFB) * self.nFB\n",
+    "        return out\n",
     "\n",
-    "        self.xFB = np.linalg.solve(np.eye(S) - self.β * self.π, IFB)\n",
+    "    return obj_V, obj_W\n",
     "\n",
-    "        self.zFB = {}\n",
-    "        for s in range(S):\n",
-    "            self.zFB[s] = np.hstack(\n",
-    "                [self.cFB[s], self.nFB[s], self.π[s] @ self.xFB, 0.])\n",
     "\n",
-    "    def __call__(self, Vf):\n",
-    "        '''\n",
-    "        Given continuation value function next period return value function this\n",
-    "        period return T(V) and optimal policies\n",
-    "        '''\n",
-    "        if not self.time_0:\n",
-    "            def PF(x, s): return self.get_policies_time1(x, s, Vf)\n",
-    "        else:\n",
-    "            def PF(B_, s0): return self.get_policies_time0(B_, s0, Vf)\n",
-    "        return PF\n",
+    "def bellman_operator_factory(Π, β, x_grid, g, bounds_v):\n",
+    "    obj_V, obj_W = obj_factory(Π, β, x_grid, g)\n",
+    "    n = x_grid[0][2]\n",
+    "    S = len(Π)\n",
+    "    x_nodes = nodes(x_grid)\n",
     "\n",
-    "    def get_policies_time1(self, x, s_, Vf):\n",
-    "        '''\n",
-    "        Finds the optimal policies \n",
-    "        '''\n",
-    "        model, β, Θ, G, S, π = self.model, self.β, self.Θ, self.G, self.S, self.π\n",
-    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
+    "    @njit(parallel=True)\n",
+    "    def T_v(V, V_new, σ_star, pref):\n",
+    "        for s_ in prange(S):\n",
+    "            for x_i in prange(n):\n",
+    "                state = (s_, x_nodes[x_i])\n",
+    "                x0 = σ_star[s_, x_i]\n",
+    "                res = optimize.nelder_mead(obj_V, x0, bounds=bounds_v,\n",
+    "                                           args=(state, V, pref))\n",
     "\n",
-    "        def objf(z):\n",
-    "            c, n, xprime = z[:S], z[S:2 * S], z[2 * S:3 * S]\n",
+    "                if res.success:\n",
+    "                    V_new[s_, x_i] = res.fun\n",
+    "                    σ_star[s_, x_i] = res.x\n",
+    "                else:\n",
+    "                    print(\"Optimization routine failed.\")\n",
     "\n",
-    "            Vprime = np.empty(S)\n",
-    "            for s in range(S):\n",
-    "                Vprime[s] = Vf[s](xprime[s])\n",
-    "\n",
-    "            return -π[s_] @ (U(c, n) + β * Vprime)\n",
-    "\n",
-    "        def objf_prime(x):\n",
-    "\n",
-    "            epsilon = 1e-7\n",
-    "            x0 = np.asfarray(x)\n",
-    "            f0 = np.atleast_1d(objf(x0))\n",
-    "            jac = np.zeros([len(x0), len(f0)])\n",
-    "            dx = np.zeros(len(x0))\n",
-    "            for i in range(len(x0)):\n",
-    "                dx[i] = epsilon\n",
-    "                jac[i] = (objf(x0+dx) - f0)/epsilon\n",
-    "                dx[i] = 0.0\n",
-    "\n",
-    "            return jac.transpose()\n",
-    "\n",
-    "        def cons(z):\n",
-    "            c, n, xprime, T = z[:S], z[S:2 * S], z[2 * S:3 * S], z[3 * S:]\n",
-    "            u_c = Uc(c, n)\n",
-    "            Eu_c = π[s_] @ u_c\n",
-    "            return np.hstack([\n",
-    "                x * u_c / Eu_c - u_c * (c - T) - Un(c, n) * n - β * xprime,\n",
-    "                Θ * n - c - G])\n",
-    "\n",
-    "        if model.transfers:\n",
-    "            bounds = [(0., 100)] * S + [(0., 100)] * S + \\\n",
-    "                [self.xbar] * S + [(0., 100.)] * S\n",
-    "        else:\n",
-    "            bounds = [(0., 100)] * S + [(0., 100)] * S + \\\n",
-    "                [self.xbar] * S + [(0., 0.)] * S\n",
-    "        out, fx, _, imode, smode = fmin_slsqp(objf, self.z0[x, s_],\n",
-    "                                              f_eqcons=cons, bounds=bounds,\n",
-    "                                              fprime=objf_prime, full_output=True,\n",
-    "                                              iprint=0, acc=self.tol, iter=self.maxiter)\n",
+    "    bounds_w = np.array([[-9.0, 1.0], [0., 10.]])\n",
     "\n",
-    "        if imode > 0:\n",
-    "            raise Exception(smode)\n",
+    "    def T_w(W, σ_star, V, b_0, pref):\n",
+    "        for s_ in prange(S):\n",
+    "            state = (s_, b_0)\n",
+    "            x0 = σ_star[s_]\n",
+    "            res = optimize.nelder_mead(obj_W, x0, bounds=bounds_w,\n",
+    "                                       args=(state, V, pref))\n",
     "\n",
-    "        self.z0[x, s_] = out\n",
-    "        return np.hstack([-fx, out])\n",
+    "            W[s_] = res.fun\n",
+    "            σ_star[s_] = res.x\n",
     "\n",
-    "    def get_policies_time0(self, B_, s0, Vf):\n",
-    "        '''\n",
-    "        Finds the optimal policies \n",
-    "        '''\n",
-    "        model, β, Θ, G = self.model, self.β, self.Θ, self.G\n",
-    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
-    "\n",
-    "        def objf(z):\n",
-    "            c, n, xprime = z[:-1]\n",
-    "\n",
-    "            return -(U(c, n) + β * Vf[s0](xprime))\n",
-    "\n",
-    "        def cons(z):\n",
-    "            c, n, xprime, T = z\n",
-    "            return np.hstack([\n",
-    "                -Uc(c, n) * (c - B_ - T) - Un(c, n) * n - β * xprime,\n",
-    "                (Θ * n - c - G)[s0]])\n",
-    "\n",
-    "        if model.transfers:\n",
-    "            bounds = [(0., 100), (0., 100), self.xbar, (0., 100.)]\n",
-    "        else:\n",
-    "            bounds = [(0., 100), (0., 100), self.xbar, (0., 0.)]\n",
-    "        out, fx, _, imode, smode = fmin_slsqp(objf, self.zFB[s0], f_eqcons=cons,\n",
-    "                                              bounds=bounds, full_output=True,\n",
-    "                                              iprint=0)\n",
-    "\n",
-    "        if imode > 0:\n",
-    "            raise Exception(smode)\n",
-    "\n",
-    "        return np.hstack([-fx, out])"
+    "    return T_v, T_w"
    ]
   },
   {
@@ -1269,88 +1227,7 @@
    "source": [
     "## Examples\n",
     "\n",
-    "We now turn to some examples.\n",
-    "\n",
-    "We will first build some useful functions for solving the model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "hide-output": false
-   },
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "from scipy.interpolate import UnivariateSpline\n",
-    "\n",
-    "\n",
-    "class interpolate_wrapper:\n",
-    "\n",
-    "    def __init__(self, F):\n",
-    "        self.F = F\n",
-    "\n",
-    "    def __getitem__(self, index):\n",
-    "        return interpolate_wrapper(np.asarray(self.F[index]))\n",
-    "\n",
-    "    def reshape(self, *args):\n",
-    "        self.F = self.F.reshape(*args)\n",
-    "        return self\n",
-    "\n",
-    "    def transpose(self):\n",
-    "        self.F = self.F.transpose()\n",
-    "\n",
-    "    def __len__(self):\n",
-    "        return len(self.F)\n",
-    "\n",
-    "    def __call__(self, xvec):\n",
-    "        x = np.atleast_1d(xvec)\n",
-    "        shape = self.F.shape\n",
-    "        if len(x) == 1:\n",
-    "            fhat = np.hstack([f(x) for f in self.F.flatten()])\n",
-    "            return fhat.reshape(shape)\n",
-    "        else:\n",
-    "            fhat = np.vstack([f(x) for f in self.F.flatten()])\n",
-    "            return fhat.reshape(np.hstack((shape, len(x))))\n",
-    "\n",
-    "\n",
-    "class interpolator_factory:\n",
-    "\n",
-    "    def __init__(self, k, s):\n",
-    "        self.k, self.s = k, s\n",
-    "\n",
-    "    def __call__(self, xgrid, Fs):\n",
-    "        shape, m = Fs.shape[:-1], Fs.shape[-1]\n",
-    "        Fs = Fs.reshape((-1, m))\n",
-    "        F = []\n",
-    "        xgrid = np.sort(xgrid)  # Sort xgrid\n",
-    "        for Fhat in Fs:\n",
-    "            F.append(UnivariateSpline(xgrid, Fhat, k=self.k, s=self.s))\n",
-    "        return interpolate_wrapper(np.array(F).reshape(shape))\n",
-    "\n",
-    "\n",
-    "def fun_vstack(fun_list):\n",
-    "\n",
-    "    Fs = [IW.F for IW in fun_list]\n",
-    "    return interpolate_wrapper(np.vstack(Fs))\n",
-    "\n",
-    "\n",
-    "def fun_hstack(fun_list):\n",
-    "\n",
-    "    Fs = [IW.F for IW in fun_list]\n",
-    "    return interpolate_wrapper(np.hstack(Fs))\n",
-    "\n",
-    "\n",
-    "def simulate_markov(π, s_0, T):\n",
-    "\n",
-    "    sHist = np.empty(T, dtype=int)\n",
-    "    sHist[0] = s_0\n",
-    "    S = len(π)\n",
-    "    for t in range(1, T):\n",
-    "        sHist[t] = np.random.choice(np.arange(S), p=π[sHist[t - 1]])\n",
-    "\n",
-    "    return sHist"
+    "We now turn to some examples."
    ]
   },
   {
@@ -1425,44 +1302,52 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "\n",
+    "crra_util_data = [\n",
+    "    ('β', float64),\n",
+    "    ('σ', float64),\n",
+    "    ('γ', float64)\n",
+    "]\n",
     "\n",
+    "@jitclass(crra_util_data)\n",
     "class CRRAutility:\n",
     "\n",
     "    def __init__(self,\n",
     "                 β=0.9,\n",
     "                 σ=2,\n",
-    "                 γ=2,\n",
-    "                 π=0.5*np.ones((2, 2)),\n",
-    "                 G=np.array([0.1, 0.2]),\n",
-    "                 Θ=np.ones(2),\n",
-    "                 transfers=False):\n",
+    "                 γ=2):\n",
     "\n",
     "        self.β, self.σ, self.γ = β, σ, γ\n",
-    "        self.π, self.G, self.Θ, self.transfers = π, G, Θ, transfers\n",
     "\n",
     "    # Utility function\n",
-    "    def U(self, c, n):\n",
+    "    def U(self, c, l):\n",
+    "        # Note: `l` should not be interpreted as labor, it is an auxiliary\n",
+    "        # variable used to conveniently match the code and the equations\n",
+    "        # in the lecture\n",
     "        σ = self.σ\n",
     "        if σ == 1.:\n",
     "            U = np.log(c)\n",
     "        else:\n",
     "            U = (c**(1 - σ) - 1) / (1 - σ)\n",
-    "        return U - n**(1 + self.γ) / (1 + self.γ)\n",
+    "        return U - (1-l) ** (1 + self.γ) / (1 + self.γ)\n",
     "\n",
     "    # Derivatives of utility function\n",
-    "    def Uc(self, c, n):\n",
-    "        return c**(-self.σ)\n",
+    "    def Uc(self, c, l):\n",
+    "        return c ** (-self.σ)\n",
+    "\n",
+    "    def Ucc(self, c, l):\n",
+    "        return -self.σ * c ** (-self.σ - 1)\n",
+    "\n",
+    "    def Ul(self, c, l):\n",
+    "        return (1-l) ** self.γ\n",
     "\n",
-    "    def Ucc(self, c, n):\n",
-    "        return -self.σ * c**(-self.σ - 1)\n",
+    "    def Ull(self, c, l):\n",
+    "        return -self.γ * (1-l) ** (self.γ - 1)\n",
     "\n",
-    "    def Un(self, c, n):\n",
-    "        return -n**self.γ\n",
+    "    def Ucl(self, c, l):\n",
+    "        return 0\n",
     "\n",
-    "    def Unn(self, c, n):\n",
-    "        return -self.γ * n**(self.γ - 1)"
+    "    def Ulc(self, c, l):\n",
+    "        return 0"
    ]
   },
   {
@@ -1487,56 +1372,106 @@
    },
    "outputs": [],
    "source": [
-    "# Initialize μgrid for value function iteration\n",
-    "μ_grid = np.linspace(-0.7, 0.01, 300)\n",
-    "\n",
-    "time_example = CRRAutility()\n",
-    "\n",
-    "time_example.π = np.array([[0, 1, 0,   0,   0,  0],\n",
-    "                           [0, 0, 1,   0,   0,  0],\n",
-    "                           [0, 0, 0, 0.5, 0.5,  0],\n",
-    "                           [0, 0, 0,   0,   0,  1],\n",
-    "                           [0, 0, 0,   0,   0,  1],\n",
-    "                           [0, 0, 0,   0,   0,  1]])\n",
-    "\n",
-    "time_example.G = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
-    "time_example.Θ = np.ones(6)  # Θ can in principle be random\n",
-    "\n",
-    "time_example.transfers = True             # Government can use transfers\n",
-    "# Solve sequential problem\n",
-    "time_sequential = SequentialAllocation(time_example)\n",
-    "# Solve recursive problem\n",
-    "time_bellman = RecursiveAllocationAMSS(time_example, μ_grid)\n",
+    "# WARNING: DO NOT EXPECT THE CODE TO WORK IF YOU CHANGE PARAMETERS\n",
+    "σ = 2\n",
+    "γ = 2\n",
+    "β = 0.9\n",
+    "Π = np.array([[0, 1, 0,   0,   0,  0],\n",
+    "              [0, 0, 1,   0,   0,  0],\n",
+    "              [0, 0, 0, 0.5, 0.5,  0],\n",
+    "              [0, 0, 0,   0,   0,  1],\n",
+    "              [0, 0, 0,   0,   0,  1],\n",
+    "              [0, 0, 0,   0,   0,  1]])\n",
+    "g = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
+    "\n",
+    "x_min = -1.5555\n",
+    "x_max = 17.339\n",
+    "x_num = 300\n",
+    "\n",
+    "x_grid = UCGrid((x_min, x_max, x_num))\n",
+    "\n",
+    "crra_pref = CRRAutility(β=β, σ=σ, γ=γ)\n",
+    "\n",
+    "S = len(Π)\n",
+    "bounds_v = np.vstack([np.hstack([np.ones(S) * -10., np.zeros(S)]),\n",
+    "                      np.hstack([np.ones(S) - g, np.ones(S) * 10])]).T\n",
+    "\n",
+    "amss_model = AMSS(crra_pref, β, Π, g, x_grid, bounds_v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "# WARNING: DO NOT EXPECT THE CODE TO WORK IF YOU CHANGE PARAMETERS\n",
+    "V = np.zeros((len(Π), x_num))\n",
+    "V[:] = -nodes(x_grid).T ** 2\n",
     "\n",
-    "sHist_h = np.array([0, 1, 2, 3, 5, 5, 5])\n",
-    "sHist_l = np.array([0, 1, 2, 4, 5, 5, 5])\n",
+    "σ_v_star = np.ones((S, x_num, S * 2))\n",
+    "σ_v_star[:, :, :S] = 0.0\n",
     "\n",
-    "sim_seq_h = time_sequential.simulate(1, 0, 7, sHist_h)\n",
-    "sim_bel_h = time_bellman.simulate(1, 0, 7, sHist_h)\n",
-    "sim_seq_l = time_sequential.simulate(1, 0, 7, sHist_l)\n",
-    "sim_bel_l = time_bellman.simulate(1, 0, 7, sHist_l)\n",
+    "W = np.empty(len(Π))\n",
+    "b_0 = 1.0\n",
+    "σ_w_star = np.ones((S, 2))\n",
+    "σ_w_star[:, 0] = -0.05"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
     "\n",
-    "# Government spending paths\n",
-    "sim_seq_l[4] = time_example.G[sHist_l]\n",
-    "sim_seq_h[4] = time_example.G[sHist_h]\n",
-    "sim_bel_l[4] = time_example.G[sHist_l]\n",
-    "sim_bel_h[4] = time_example.G[sHist_h]\n",
+    "amss_model.solve(V, σ_v_star, b_0, W, σ_w_star)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "# Solve the LS model\n",
+    "ls_model = SequentialLS(crra_pref, g=g, π=Π)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "# WARNING: DO NOT EXPECT THE CODE TO WORK IF YOU CHANGE PARAMETERS\n",
+    "s_hist_h = np.array([0, 1, 2, 3, 5, 5, 5])\n",
+    "s_hist_l = np.array([0, 1, 2, 4, 5, 5, 5])\n",
     "\n",
-    "# Output paths\n",
-    "sim_seq_l[5] = time_example.Θ[sHist_l] * sim_seq_l[1]\n",
-    "sim_seq_h[5] = time_example.Θ[sHist_h] * sim_seq_h[1]\n",
-    "sim_bel_l[5] = time_example.Θ[sHist_l] * sim_bel_l[1]\n",
-    "sim_bel_h[5] = time_example.Θ[sHist_h] * sim_bel_h[1]\n",
+    "sim_h_amss = amss_model.simulate(s_hist_h, b_0)\n",
+    "sim_l_amss = amss_model.simulate(s_hist_l, b_0)\n",
     "\n",
+    "sim_h_ls = ls_model.simulate(b_0, 0, 7, s_hist_h)\n",
+    "sim_l_ls = ls_model.simulate(b_0, 0, 7, s_hist_l)\n",
     "\n",
     "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
     "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
     "          'Tax Rate', 'Government Spending', 'Output']\n",
     "\n",
-    "for ax, title, sim_l, sim_h, bel_l, bel_h in zip(axes.flatten(), titles,\n",
-    "                                                 sim_seq_l, sim_seq_h,\n",
-    "                                                 sim_bel_l, sim_bel_h):\n",
-    "    ax.plot(sim_l, '-ok', sim_h, '-^k', bel_l, '-or', bel_h, '-^r', alpha=0.7)\n",
+    "for ax, title, ls_l, ls_h, amss_l, amss_h in zip(axes.flatten(), titles,\n",
+    "                                                 sim_l_ls, sim_h_ls,\n",
+    "                                                 sim_l_amss, sim_h_amss):\n",
+    "    ax.plot(ls_l, '-ok', ls_h, '-^k', amss_l, '-or', amss_h, '-^r',\n",
+    "            alpha=0.7)\n",
     "    ax.set(title=title)\n",
     "    ax.grid()\n",
     "\n",
@@ -1618,37 +1553,42 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
+    "log_util_data = [\n",
+    "    ('β', float64),\n",
+    "    ('ψ', float64)\n",
+    "]\n",
     "\n",
+    "@jitclass(log_util_data)\n",
     "class LogUtility:\n",
     "\n",
     "    def __init__(self,\n",
     "                 β=0.9,\n",
-    "                 ψ=0.69,\n",
-    "                 π=0.5*np.ones((2, 2)),\n",
-    "                 G=np.array([0.1, 0.2]),\n",
-    "                 Θ=np.ones(2),\n",
-    "                 transfers=False):\n",
+    "                 ψ=0.69):\n",
     "\n",
-    "        self.β, self.ψ, self.π = β, ψ, π\n",
-    "        self.G, self.Θ, self.transfers = G, Θ, transfers\n",
+    "        self.β, self.ψ = β, ψ\n",
     "\n",
     "    # Utility function\n",
-    "    def U(self, c, n):\n",
-    "        return np.log(c) + self.ψ * np.log(1 - n)\n",
+    "    def U(self, c, l):\n",
+    "        return np.log(c) + self.ψ * np.log(l)\n",
     "\n",
     "    # Derivatives of utility function\n",
-    "    def Uc(self, c, n):\n",
+    "    def Uc(self, c, l):\n",
     "        return 1 / c\n",
     "\n",
-    "    def Ucc(self, c, n):\n",
+    "    def Ucc(self, c, l):\n",
     "        return -c**(-2)\n",
     "\n",
-    "    def Un(self, c, n):\n",
-    "        return -self.ψ / (1 - n)\n",
+    "    def Ul(self, c, l):\n",
+    "        return self.ψ / l\n",
+    "\n",
+    "    def Ull(self, c, l):\n",
+    "        return -self.ψ / l**2\n",
     "\n",
-    "    def Unn(self, c, n):\n",
-    "        return -self.ψ / (1 - n)**2"
+    "    def Ucl(self, c, l):\n",
+    "        return 0\n",
+    "\n",
+    "    def Ulc(self, c, l):\n",
+    "        return 0"
    ]
   },
   {
@@ -1671,34 +1611,83 @@
    },
    "outputs": [],
    "source": [
-    "log_example = LogUtility()\n",
-    "log_example.transfers = True                    # Government can use transfers\n",
-    "log_sequential = SequentialAllocation(log_example)  # Solve sequential problem\n",
-    "log_bellman = RecursiveAllocationAMSS(log_example, μ_grid)\n",
+    "# WARNING: DO NOT EXPECT THE CODE TO WORK IF YOU CHANGE PARAMETERS\n",
+    "ψ = 0.69\n",
+    "Π = 0.5 * np.ones((2, 2))\n",
+    "β = 0.9\n",
+    "g = np.array([0.1, 0.2])\n",
     "\n",
-    "T = 20\n",
-    "sHist = np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1,\n",
-    "                  0, 0, 0, 1, 1, 1, 1, 1, 1, 0])\n",
+    "x_min = -3.4107\n",
+    "x_max = 3.709\n",
+    "x_num = 300\n",
     "\n",
-    "# Simulate\n",
-    "sim_seq = log_sequential.simulate(0.5, 0, T, sHist)\n",
-    "sim_bel = log_bellman.simulate(0.5, 0, T, sHist)\n",
+    "x_grid = UCGrid((x_min, x_max, x_num))\n",
+    "log_pref = LogUtility(β=β, ψ=ψ)\n",
     "\n",
-    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
-    "          'Tax Rate', 'Government Spending', 'Output']\n",
+    "S = len(Π)\n",
+    "bounds_v = np.vstack([np.zeros(2 * S), np.hstack([1 - g, np.ones(S)]) ]).T\n",
+    "\n",
+    "V = np.zeros((len(Π), x_num))\n",
+    "V[:] = -(nodes(x_grid).T + x_max) ** 2 / 14\n",
     "\n",
-    "# Government spending paths\n",
-    "sim_seq[4] = log_example.G[sHist]\n",
-    "sim_bel[4] = log_example.G[sHist]\n",
+    "σ_v_star = 1 - np.ones((S, x_num, S * 2)) * 0.55\n",
     "\n",
-    "# Output paths\n",
-    "sim_seq[5] = log_example.Θ[sHist] * sim_seq[1]\n",
-    "sim_bel[5] = log_example.Θ[sHist] * sim_bel[1]\n",
+    "W = np.empty(len(Π))\n",
+    "b_0 = 0.5\n",
+    "σ_w_star = 1 - np.ones((S, 2)) * 0.55\n",
+    "\n",
+    "amss_model = AMSS(log_pref, β, Π, g, x_grid, bounds_v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "\n",
+    "amss_model.solve(V, σ_v_star, b_0, W, σ_w_star, tol_vfi=3e-5, maxitr=3000,\n",
+    "                 print_itr=100)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "ls_model = SequentialLS(log_pref, g=g, π=Π)  # Solve sequential problem"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "# WARNING: DO NOT EXPECT THE CODE TO WORK IF YOU CHANGE PARAMETERS\n",
+    "s_hist = np.array([0, 0, 0, 0, 0, 0, 0, 0, 1, 1,\n",
+    "                   0, 0, 0, 1, 1, 1, 1, 1, 1, 0])\n",
+    "\n",
+    "T = len(s_hist)\n",
+    "\n",
+    "sim_amss = amss_model.simulate(s_hist, b_0)\n",
+    "sim_ls = ls_model.simulate(0.5, 0, T, s_hist)\n",
+    "\n",
+    "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
+    "          'Tax Rate', 'Government Spending', 'Output']\n",
     "\n",
     "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
     "\n",
-    "for ax, title, seq, bel in zip(axes.flatten(), titles, sim_seq, sim_bel):\n",
-    "    ax.plot(seq, '-ok', bel, '-^b')\n",
+    "for ax, title, ls, amss in zip(axes.flatten(), titles, sim_ls, sim_amss):\n",
+    "    ax.plot(ls, '-ok', amss, '-^b')\n",
     "    ax.set(title=title)\n",
     "    ax.grid()\n",
     "\n",
@@ -1735,27 +1724,33 @@
    },
    "outputs": [],
    "source": [
-    "T = 200  # Set T to 200 periods\n",
-    "sim_seq_long = log_sequential.simulate(0.5, 0, T)\n",
-    "sHist_long = sim_seq_long[-3]\n",
-    "sim_bel_long = log_bellman.simulate(0.5, 0, T, sHist_long)\n",
+    "T = 200\n",
+    "s_0 = 0\n",
+    "mc = MarkovChain(Π)\n",
+    "\n",
+    "s_hist_long = mc.simulate(T, init=s_0, random_state=5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "hide-output": false
+   },
+   "outputs": [],
+   "source": [
+    "sim_amss = amss_model.simulate(s_hist_long, b_0)\n",
+    "sim_ls = ls_model.simulate(0.5, 0, T, s_hist_long)\n",
     "\n",
     "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
     "          'Tax Rate', 'Government Spending', 'Output']\n",
     "\n",
-    "# Government spending paths\n",
-    "sim_seq_long[4] = log_example.G[sHist_long]\n",
-    "sim_bel_long[4] = log_example.G[sHist_long]\n",
-    "\n",
-    "# Output paths\n",
-    "sim_seq_long[5] = log_example.Θ[sHist_long] * sim_seq_long[1]\n",
-    "sim_bel_long[5] = log_example.Θ[sHist_long] * sim_bel_long[1]\n",
     "\n",
     "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
     "\n",
-    "for ax, title, seq, bel in zip(axes.flatten(), titles, sim_seq_long, \\\n",
-    "        sim_bel_long):\n",
-    "    ax.plot(seq, '-k', bel, '-.b', alpha=0.5)\n",
+    "for ax, title, ls, amss in zip(axes.flatten(), titles, sim_ls, \\\n",
+    "        sim_amss):\n",
+    "    ax.plot(ls, '-k', amss, '-.b', alpha=0.5)\n",
     "    ax.set(title=title)\n",
     "    ax.grid()\n",
     "\n",
@@ -1787,7 +1782,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661941.0074658,
+  "date": 1624431168.994837,
   "filename": "amss.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/amss2.ipynb b/amss2.ipynb
index 13fbed9..7eb3c84 100644
--- a/amss2.ipynb
+++ b/amss2.ipynb
@@ -1610,7 +1610,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661941.1767492,
+  "date": 1624431169.1728618,
   "filename": "amss2.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/amss3.ipynb b/amss3.ipynb
index 319c632..a961f31 100644
--- a/amss3.ipynb
+++ b/amss3.ipynb
@@ -1651,7 +1651,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661941.3859596,
+  "date": 1624431169.545487,
   "filename": "amss3.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/arellano.ipynb b/arellano.ipynb
index 884a5de..c4354f3 100644
--- a/arellano.ipynb
+++ b/arellano.ipynb
@@ -1037,7 +1037,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661941.6302557,
+  "date": 1624431169.7652407,
   "filename": "arellano.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/arma.ipynb b/arma.ipynb
index 4a8a52c..f02a038 100644
--- a/arma.ipynb
+++ b/arma.ipynb
@@ -1183,7 +1183,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661942.0405178,
+  "date": 1624431169.9898334,
   "filename": "arma.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/black_litterman.ipynb b/black_litterman.ipynb
index d378556..fe05a33 100644
--- a/black_litterman.ipynb
+++ b/black_litterman.ipynb
@@ -1582,7 +1582,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661942.2040126,
+  "date": 1624431170.2142158,
   "filename": "black_litterman.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/calvo.ipynb b/calvo.ipynb
index 412a7c1..958c87f 100644
--- a/calvo.ipynb
+++ b/calvo.ipynb
@@ -1667,7 +1667,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661942.441115,
+  "date": 1624431170.4078646,
   "filename": "calvo.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/cattle_cycles.ipynb b/cattle_cycles.ipynb
index 5f94ee9..b37d713 100644
--- a/cattle_cycles.ipynb
+++ b/cattle_cycles.ipynb
@@ -548,7 +548,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661942.6324031,
+  "date": 1624431170.758118,
   "filename": "cattle_cycles.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/chang_credible.ipynb b/chang_credible.ipynb
index 94025af..b6c2b41 100644
--- a/chang_credible.ipynb
+++ b/chang_credible.ipynb
@@ -1541,7 +1541,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661942.7952566,
+  "date": 1624431170.923503,
   "filename": "chang_credible.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/chang_ramsey.ipynb b/chang_ramsey.ipynb
index d758512..be678d8 100644
--- a/chang_ramsey.ipynb
+++ b/chang_ramsey.ipynb
@@ -1817,7 +1817,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661943.1950128,
+  "date": 1624431171.1510167,
   "filename": "chang_ramsey.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/classical_filtering.ipynb b/classical_filtering.ipynb
index f5d8244..c001070 100644
--- a/classical_filtering.ipynb
+++ b/classical_filtering.ipynb
@@ -1542,7 +1542,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661943.4328423,
+  "date": 1624431171.3856077,
   "filename": "classical_filtering.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/coase.ipynb b/coase.ipynb
index 783518d..a269f5c 100644
--- a/coase.ipynb
+++ b/coase.ipynb
@@ -854,7 +854,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661943.6288862,
+  "date": 1624431171.5769794,
   "filename": "coase.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/cons_news.ipynb b/cons_news.ipynb
index 8fbe148..713e4a1 100644
--- a/cons_news.ipynb
+++ b/cons_news.ipynb
@@ -1072,7 +1072,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661943.8180664,
+  "date": 1624431171.7677624,
   "filename": "cons_news.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/discrete_dp.ipynb b/discrete_dp.ipynb
index 30910c0..a30a5ae 100644
--- a/discrete_dp.ipynb
+++ b/discrete_dp.ipynb
@@ -1629,7 +1629,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661944.1084075,
+  "date": 1624431172.23872,
   "filename": "discrete_dp.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/dyn_stack.ipynb b/dyn_stack.ipynb
index 4cf6d70..7a0af6e 100644
--- a/dyn_stack.ipynb
+++ b/dyn_stack.ipynb
@@ -1804,7 +1804,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661944.6360881,
+  "date": 1624431172.6109161,
   "filename": "dyn_stack.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/estspec.ipynb b/estspec.ipynb
index 608df2a..b2b13fa 100644
--- a/estspec.ipynb
+++ b/estspec.ipynb
@@ -684,7 +684,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661944.8711634,
+  "date": 1624431172.8013706,
   "filename": "estspec.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/growth_in_dles.ipynb b/growth_in_dles.ipynb
index defe038..bf10a55 100644
--- a/growth_in_dles.ipynb
+++ b/growth_in_dles.ipynb
@@ -880,7 +880,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.019177,
+  "date": 1624431172.9408998,
   "filename": "growth_in_dles.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/hs_invertibility_example.ipynb b/hs_invertibility_example.ipynb
index 627e06f..50af370 100644
--- a/hs_invertibility_example.ipynb
+++ b/hs_invertibility_example.ipynb
@@ -489,7 +489,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.1667056,
+  "date": 1624431173.0863845,
   "filename": "hs_invertibility_example.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/hs_recursive_models.ipynb b/hs_recursive_models.ipynb
index fd33b21..b85134d 100644
--- a/hs_recursive_models.ipynb
+++ b/hs_recursive_models.ipynb
@@ -2767,7 +2767,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.4081616,
+  "date": 1624431173.4950278,
   "filename": "hs_recursive_models.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index.ipynb b/index.ipynb
index 570d43b..4892adb 100644
--- a/index.ipynb
+++ b/index.ipynb
@@ -60,7 +60,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.5646918,
+  "date": 1624431173.6483474,
   "filename": "index.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_asset_pricing.ipynb b/index_asset_pricing.ipynb
index 6786815..93f1bed 100644
--- a/index_asset_pricing.ipynb
+++ b/index_asset_pricing.ipynb
@@ -50,7 +50,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.5911539,
+  "date": 1624431173.6726933,
   "filename": "index_asset_pricing.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_classic_linear_models.ipynb b/index_classic_linear_models.ipynb
index 2b2f275..d20e369 100644
--- a/index_classic_linear_models.ipynb
+++ b/index_classic_linear_models.ipynb
@@ -36,7 +36,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.6128333,
+  "date": 1624431173.6945176,
   "filename": "index_classic_linear_models.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_dynamic_programming_squared.ipynb b/index_dynamic_programming_squared.ipynb
index 0f96032..3f2cd18 100644
--- a/index_dynamic_programming_squared.ipynb
+++ b/index_dynamic_programming_squared.ipynb
@@ -95,7 +95,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.6529777,
+  "date": 1624431173.732319,
   "filename": "index_dynamic_programming_squared.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_hs_recursive_models.ipynb b/index_hs_recursive_models.ipynb
index 0b95442..5af07b4 100644
--- a/index_hs_recursive_models.ipynb
+++ b/index_hs_recursive_models.ipynb
@@ -66,7 +66,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.6871903,
+  "date": 1624431173.7673209,
   "filename": "index_hs_recursive_models.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_lq_control.ipynb b/index_lq_control.ipynb
index 91c165a..ecfab95 100644
--- a/index_lq_control.ipynb
+++ b/index_lq_control.ipynb
@@ -88,7 +88,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.9229577,
+  "date": 1624431173.8061287,
   "filename": "index_lq_control.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_multi_agent_models.ipynb b/index_multi_agent_models.ipynb
index cdeff54..68d4719 100644
--- a/index_multi_agent_models.ipynb
+++ b/index_multi_agent_models.ipynb
@@ -60,7 +60,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.9539936,
+  "date": 1624431173.8364623,
   "filename": "index_multi_agent_models.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_time_series_models.ipynb b/index_time_series_models.ipynb
index 1ec45f5..1169771 100644
--- a/index_time_series_models.ipynb
+++ b/index_time_series_models.ipynb
@@ -72,7 +72,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661945.9898818,
+  "date": 1624431173.8716345,
   "filename": "index_time_series_models.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_toc.ipynb b/index_toc.ipynb
index b77ed03..e8153ce 100644
--- a/index_toc.ipynb
+++ b/index_toc.ipynb
@@ -76,7 +76,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.0735986,
+  "date": 1624431173.9522133,
   "filename": "index_toc.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/index_tools_and_techniques.ipynb b/index_tools_and_techniques.ipynb
index 6550275..5eb414d 100644
--- a/index_tools_and_techniques.ipynb
+++ b/index_tools_and_techniques.ipynb
@@ -63,7 +63,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.1050398,
+  "date": 1624431173.986329,
   "filename": "index_tools_and_techniques.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/irfs_in_hall_model.ipynb b/irfs_in_hall_model.ipynb
index b5cfbf6..9b93d3a 100644
--- a/irfs_in_hall_model.ipynb
+++ b/irfs_in_hall_model.ipynb
@@ -445,7 +445,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.1601908,
+  "date": 1624431174.0410764,
   "filename": "irfs_in_hall_model.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/knowing_forecasts_of_others.ipynb b/knowing_forecasts_of_others.ipynb
index aa78085..3ec5ddd 100644
--- a/knowing_forecasts_of_others.ipynb
+++ b/knowing_forecasts_of_others.ipynb
@@ -1927,7 +1927,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.3871124,
+  "date": 1624431174.2590406,
   "filename": "knowing_forecasts_of_others.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/lqramsey.ipynb b/lqramsey.ipynb
index b50d1b1..4016146 100644
--- a/lqramsey.ipynb
+++ b/lqramsey.ipynb
@@ -1241,7 +1241,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.633319,
+  "date": 1624431174.535595,
   "filename": "lqramsey.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/lu_tricks.ipynb b/lu_tricks.ipynb
index a0b555f..f4c4016 100644
--- a/lu_tricks.ipynb
+++ b/lu_tricks.ipynb
@@ -1525,7 +1525,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661946.8332317,
+  "date": 1624431174.9144583,
   "filename": "lu_tricks.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/lucas_asset_pricing_dles.ipynb b/lucas_asset_pricing_dles.ipynb
index 23f0a16..194409a 100644
--- a/lucas_asset_pricing_dles.ipynb
+++ b/lucas_asset_pricing_dles.ipynb
@@ -456,7 +456,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661947.1463742,
+  "date": 1624431175.0485275,
   "filename": "lucas_asset_pricing_dles.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/lucas_model.ipynb b/lucas_model.ipynb
index e36ea65..eb6f86c 100644
--- a/lucas_model.ipynb
+++ b/lucas_model.ipynb
@@ -732,7 +732,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661947.2823946,
+  "date": 1624431175.1856818,
   "filename": "lucas_model.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/markov_jump_lq.ipynb b/markov_jump_lq.ipynb
index e448c06..49b4ac8 100644
--- a/markov_jump_lq.ipynb
+++ b/markov_jump_lq.ipynb
@@ -1304,7 +1304,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661947.4698737,
+  "date": 1624431175.369055,
   "filename": "markov_jump_lq.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/matsuyama.ipynb b/matsuyama.ipynb
index 2b5ea7c..008335b 100644
--- a/matsuyama.ipynb
+++ b/matsuyama.ipynb
@@ -956,7 +956,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661947.6728525,
+  "date": 1624431175.5666082,
   "filename": "matsuyama.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/muth_kalman.ipynb b/muth_kalman.ipynb
index f4b69c3..9254fd1 100644
--- a/muth_kalman.ipynb
+++ b/muth_kalman.ipynb
@@ -508,7 +508,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661947.7935102,
+  "date": 1624431175.682986,
   "filename": "muth_kalman.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/opt_tax_recur.ipynb b/opt_tax_recur.ipynb
index 60ed935..ec45237 100644
--- a/opt_tax_recur.ipynb
+++ b/opt_tax_recur.ipynb
@@ -48,7 +48,8 @@
    },
    "outputs": [],
    "source": [
-    "!pip install --upgrade quantecon"
+    "!pip install --upgrade quantecon\n",
+    "!pip install interpolation"
    ]
   },
   {
@@ -100,7 +101,13 @@
    "source": [
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
-    "%matplotlib inline"
+    "%matplotlib inline\n",
+    "from scipy.optimize import root\n",
+    "from quantecon import MarkovChain\n",
+    "from quantecon.optimize.nelder_mead import nelder_mead\n",
+    "from interpolation import interp\n",
+    "from numba import njit, prange, float64\n",
+    "from numba.experimental import jitclass"
    ]
   },
   {
@@ -802,7 +809,7 @@
    "source": [
     "### Sequence Implementation\n",
     "\n",
-    "The above steps are implemented in a class called SequentialAllocation"
+    "The above steps are implemented in a class called SequentialLS"
    ]
   },
   {
@@ -813,164 +820,204 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from scipy.optimize import root\n",
-    "from quantecon import MarkovChain\n",
-    "\n",
-    "\n",
-    "class SequentialAllocation:\n",
+    "class SequentialLS:\n",
     "\n",
     "    '''\n",
-    "    Class that takes CESutility or BGPutility object as input returns\n",
-    "    planner's allocation as a function of the multiplier on the\n",
-    "    implementability constraint μ.\n",
+    "    Class that takes a preference object, state transition matrix,\n",
+    "    and state contingent government expenditure plan as inputs, and\n",
+    "    solves the sequential allocation problem described above.\n",
+    "    It returns optimal allocations about consumption and labor supply,\n",
+    "    as well as the multiplier on the implementability constraint Φ.\n",
     "    '''\n",
     "\n",
-    "    def __init__(self, model):\n",
+    "    def __init__(self,\n",
+    "                 pref,\n",
+    "                 π=0.5*np.ones((2, 2)),\n",
+    "                 g=np.array([0.1, 0.2])):\n",
     "\n",
-    "        # Initialize from model object attributes\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.mc, self.Θ = MarkovChain(self.π), model.Θ\n",
-    "        self.S = len(model.π)  # Number of states\n",
-    "        self.model = model\n",
+    "        # Initialize from pref object attributes\n",
+    "        self.β, self.π, self.g = pref.β, π, g\n",
+    "        self.mc = MarkovChain(self.π)\n",
+    "        self.S = len(π)  # Number of states\n",
+    "        self.pref = pref\n",
     "\n",
     "        # Find the first best allocation\n",
     "        self.find_first_best()\n",
     "\n",
+    "    def FOC_first_best(self, c, g):\n",
+    "        '''\n",
+    "        First order conditions that characterize\n",
+    "        the first best allocation.\n",
+    "        '''\n",
+    "\n",
+    "        pref = self.pref\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
+    "\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
+    "\n",
+    "        return Uc(c, l) - Ul(c, l)\n",
+    "\n",
     "    def find_first_best(self):\n",
     "        '''\n",
     "        Find the first best allocation\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, G = self.S, self.Θ, self.G\n",
-    "        Uc, Un = model.Uc, model.Un\n",
-    "\n",
-    "        def res(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:]\n",
-    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
+    "        S, g = self.S, self.g\n",
     "\n",
-    "        res = root(res, 0.5 * np.ones(2 * S))\n",
+    "        res = root(self.FOC_first_best, 0.5 * np.ones(S), args=(g,))\n",
     "\n",
-    "        if not res.success:\n",
+    "        if (res.fun > 1e-10).any():\n",
     "            raise Exception('Could not find first best')\n",
     "\n",
-    "        self.cFB = res.x[:S]\n",
-    "        self.nFB = res.x[S:]\n",
+    "        self.cFB = res.x\n",
+    "        self.nFB = self.cFB + g\n",
+    "\n",
+    "    def FOC_time1(self, c, Φ, g):\n",
+    "        '''\n",
+    "        First order conditions that characterize\n",
+    "        optimal time 1 allocation problems.\n",
+    "        '''\n",
+    "\n",
+    "        pref = self.pref\n",
+    "        Uc, Ucc, Ul, Ull, Ulc = pref.Uc, pref.Ucc, pref.Ul, pref.Ull, pref.Ulc\n",
+    "\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
     "\n",
-    "        # Multiplier on the resource constraint\n",
-    "        self.ΞFB = Uc(self.cFB, self.nFB)\n",
-    "        self.zFB = np.hstack([self.cFB, self.nFB, self.ΞFB])\n",
+    "        LHS = (1 + Φ) * Uc(c, l) + Φ * (c * Ucc(c, l) - n * Ulc(c, l))\n",
+    "        RHS = (1 + Φ) * Ul(c, l) + Φ * (c * Ulc(c, l) - n * Ull(c, l))\n",
     "\n",
-    "    def time1_allocation(self, μ):\n",
+    "        diff = LHS - RHS\n",
+    "\n",
+    "        return diff\n",
+    "\n",
+    "    def time1_allocation(self, Φ):\n",
     "        '''\n",
-    "        Computes optimal allocation for time t >= 1 for a given μ\n",
+    "        Computes optimal allocation for time t >= 1 for a given Φ\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, G = self.S, self.Θ, self.G\n",
-    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
-    "\n",
-    "        def FOC(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:2 * S]\n",
-    "            Ξ = z[2 * S:]\n",
-    "            # FOC of c\n",
-    "            return np.hstack([Uc(c, n) - μ * (Ucc(c, n) * c + Uc(c, n)) - Ξ,\n",
-    "                              Un(c, n) - μ * (Unn(c, n) * n + Un(c, n)) \\\n",
-    "                              + Θ * Ξ,  # FOC of n\n",
-    "                              Θ * n - c - G])\n",
-    "\n",
-    "        # Find the root of the first-order condition\n",
-    "        res = root(FOC, self.zFB)\n",
-    "        if not res.success:\n",
+    "        pref = self.pref\n",
+    "        S, g = self.S, self.g\n",
+    "\n",
+    "        # use the first best allocation as intial guess\n",
+    "        res = root(self.FOC_time1, self.cFB, args=(Φ, g))\n",
+    "\n",
+    "        if (res.fun > 1e-10).any():\n",
     "            raise Exception('Could not find LS allocation.')\n",
-    "        z = res.x\n",
-    "        c, n, Ξ = z[:S], z[S:2 * S], z[2 * S:]\n",
+    "\n",
+    "        c = res.x\n",
+    "        n = c + g\n",
+    "        l = 1 - n\n",
     "\n",
     "        # Compute x\n",
-    "        I = Uc(c, n) * c + Un(c, n) * n\n",
+    "        I = pref.Uc(c, n) * c - pref.Ul(c, l) * n\n",
     "        x = np.linalg.solve(np.eye(S) - self.β * self.π, I)\n",
     "\n",
-    "        return c, n, x, Ξ\n",
+    "        return c, n, x\n",
     "\n",
-    "    def time0_allocation(self, B_, s_0):\n",
+    "    def FOC_time0(self, c0, Φ, g0, b0):\n",
     "        '''\n",
-    "        Finds the optimal allocation given initial government debt B_ and\n",
-    "        state s_0\n",
+    "        First order conditions that characterize\n",
+    "        time 0 allocation problem.\n",
     "        '''\n",
-    "        model, π, Θ, G, β = self.model, self.π, self.Θ, self.G, self.β\n",
-    "        Uc, Ucc, Un, Unn = model.Uc, model.Ucc, model.Un, model.Unn\n",
-    "\n",
-    "        # First order conditions of planner's problem\n",
-    "        def FOC(z):\n",
-    "            μ, c, n, Ξ = z\n",
-    "            xprime = self.time1_allocation(μ)[2]\n",
-    "            return np.hstack([Uc(c, n) * (c - B_) + Un(c, n) * n + β * π[s_0]\n",
-    "                                            @ xprime,\n",
-    "                              Uc(c, n) - μ * (Ucc(c, n)\n",
-    "                                            * (c - B_) + Uc(c, n)) - Ξ,\n",
-    "                              Un(c, n) - μ * (Unn(c, n) * n\n",
-    "                                            + Un(c, n)) + Θ[s_0] * Ξ,\n",
-    "                              (Θ * n - c - G)[s_0]])\n",
-    "\n",
-    "        # Find root\n",
-    "        res = root(FOC, np.array(\n",
-    "            [0, self.cFB[s_0], self.nFB[s_0], self.ΞFB[s_0]]))\n",
-    "        if not res.success:\n",
-    "            raise Exception('Could not find time 0 LS allocation.')\n",
     "\n",
-    "        return res.x\n",
+    "        pref = self.pref\n",
+    "        Ucc, Ulc = pref.Ucc, pref.Ulc\n",
+    "\n",
+    "        n0 = c0 + g0\n",
+    "        l0 = 1 - n0\n",
     "\n",
-    "    def time1_value(self, μ):\n",
+    "        diff = self.FOC_time1(c0, Φ, g0)\n",
+    "        diff -= Φ * (Ucc(c0, l0) - Ulc(c0, l0)) * b0\n",
+    "\n",
+    "        return diff\n",
+    "\n",
+    "    def implementability(self, Φ, b0, s0, cn0_arr):\n",
+    "        '''\n",
+    "        Compute the differences between the RHS and LHS\n",
+    "        of the implementability constraint given Φ,\n",
+    "        initial debt, and initial state.\n",
+    "        '''\n",
+    "\n",
+    "        pref, π, g, β = self.pref, self.π, self.g, self.β\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
+    "        g0 = self.g[s0]\n",
+    "\n",
+    "        c, n, x = self.time1_allocation(Φ)\n",
+    "\n",
+    "        res = root(self.FOC_time0, cn0_arr[0], args=(Φ, g0, b0))\n",
+    "        c0 = res.x\n",
+    "        n0 = c0 + g0\n",
+    "        l0 = 1 - n0\n",
+    "\n",
+    "        cn0_arr[:] = c0, n0\n",
+    "\n",
+    "        LHS = Uc(c0, l0) * b0\n",
+    "        RHS = Uc(c0, l0) * c0 - Ul(c0, l0) * n0 + β * π[s0] @ x\n",
+    "\n",
+    "        return RHS - LHS\n",
+    "\n",
+    "    def time0_allocation(self, b0, s0):\n",
     "        '''\n",
-    "        Find the value associated with multiplier μ\n",
+    "        Finds the optimal time 0 allocation given\n",
+    "        initial government debt b0 and state s0\n",
     "        '''\n",
-    "        c, n, x, Ξ = self.time1_allocation(μ)\n",
-    "        U = self.model.U(c, n)\n",
-    "        V = np.linalg.solve(np.eye(self.S) - self.β * self.π, U)\n",
-    "        return c, n, x, V\n",
     "\n",
-    "    def Τ(self, c, n):\n",
+    "        # use the first best allocation as initial guess\n",
+    "        cn0_arr = np.array([self.cFB[s0], self.nFB[s0]])\n",
+    "\n",
+    "        res = root(self.implementability, 0., args=(b0, s0, cn0_arr))\n",
+    "\n",
+    "        if (res.fun > 1e-10).any():\n",
+    "            raise Exception('Could not find time 0 LS allocation.')\n",
+    "\n",
+    "        Φ = res.x[0]\n",
+    "        c0, n0 = cn0_arr\n",
+    "\n",
+    "        return Φ, c0, n0\n",
+    "\n",
+    "    def τ(self, c, n):\n",
     "        '''\n",
-    "        Computes Τ given c, n\n",
+    "        Computes τ given c, n\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        Uc, Un = model.Uc(c, n), model.Un(c,  n)\n",
+    "        pref = self.pref\n",
+    "        Uc, Ul = pref.Uc, pref.Ul\n",
     "\n",
-    "        return 1 + Un / (self.Θ * Uc)\n",
+    "        return 1 - Ul(c, 1-n) / Uc(c, 1-n)\n",
     "\n",
-    "    def simulate(self, B_, s_0, T, sHist=None):\n",
+    "    def simulate(self, b0, s0, T, sHist=None):\n",
     "        '''\n",
     "        Simulates planners policies for T periods\n",
     "        '''\n",
-    "        model, π, β = self.model, self.π, self.β\n",
-    "        Uc = model.Uc\n",
+    "        pref, π, β = self.pref, self.π, self.β\n",
+    "        Uc = pref.Uc\n",
     "\n",
     "        if sHist is None:\n",
-    "            sHist = self.mc.simulate(T, s_0)\n",
+    "            sHist = self.mc.simulate(T, s0)\n",
     "\n",
-    "        cHist, nHist, Bhist, ΤHist, μHist = np.zeros((5, T))\n",
-    "        RHist = np.zeros(T - 1)\n",
+    "        cHist, nHist, Bhist, τHist, ΦHist = np.empty((5, T))\n",
+    "        RHist = np.empty(T-1)\n",
     "\n",
     "        # Time 0\n",
-    "        μ, cHist[0], nHist[0], _ = self.time0_allocation(B_, s_0)\n",
-    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
-    "        Bhist[0] = B_\n",
-    "        μHist[0] = μ\n",
+    "        Φ, cHist[0], nHist[0] = self.time0_allocation(b0, s0)\n",
+    "        τHist[0] = self.τ(cHist[0], nHist[0])\n",
+    "        Bhist[0] = b0\n",
+    "        ΦHist[0] = Φ\n",
     "\n",
     "        # Time 1 onward\n",
     "        for t in range(1, T):\n",
-    "            c, n, x, Ξ = self.time1_allocation(μ)\n",
-    "            Τ = self.Τ(c, n)\n",
-    "            u_c = Uc(c, n)\n",
+    "            c, n, x = self.time1_allocation(Φ)\n",
+    "            τ = self.τ(c, n)\n",
+    "            u_c = Uc(c, 1-n)\n",
     "            s = sHist[t]\n",
-    "            Eu_c = π[sHist[t - 1]] @ u_c\n",
-    "            cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n[s], x[s] / u_c[s], \\\n",
-    "                                                     Τ[s]\n",
-    "            RHist[t - 1] = Uc(cHist[t - 1], nHist[t - 1]) / (β * Eu_c)\n",
-    "            μHist[t] = μ\n",
+    "            Eu_c = π[sHist[t-1]] @ u_c\n",
+    "            cHist[t], nHist[t], Bhist[t], τHist[t] = c[s], n[s], x[s] / u_c[s], τ[s]\n",
+    "            RHist[t-1] = Uc(cHist[t-1], 1-nHist[t-1]) / (β * Eu_c)\n",
+    "            ΦHist[t] = Φ\n",
+    "\n",
+    "        gHist = self.g[sHist]\n",
+    "        yHist = nHist\n",
     "\n",
-    "        return np.array([cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist])"
+    "        return [cHist, nHist, Bhist, τHist, gHist, yHist, sHist, ΦHist, RHist]"
    ]
   },
   {
@@ -986,7 +1033,7 @@
     "\n",
     "But $ x_t(s^t) $ is  a  natural candidate for a state variable in\n",
     "a recursive formulation of the Ramsey problem, one that records history-dependence and so is\n",
-    "``backward-looking’‘."
+    "`backward-looking`."
    ]
   },
   {
@@ -1316,7 +1363,7 @@
    "source": [
     "### Recursive Implementation\n",
     "\n",
-    "The above steps are implemented in a class called `RecursiveAllocation`."
+    "The above steps are implemented in a class called `RecursiveLS`."
    ]
   },
   {
@@ -1327,295 +1374,250 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from scipy.interpolate import UnivariateSpline\n",
-    "from scipy.optimize import fmin_slsqp\n",
-    "from quantecon import MarkovChain\n",
-    "from scipy.optimize import root\n",
-    "\n",
-    "\n",
-    "class RecursiveAllocation:\n",
+    "class RecursiveLS:\n",
     "\n",
     "    '''\n",
     "    Compute the planner's allocation by solving Bellman\n",
     "    equation.\n",
     "    '''\n",
     "\n",
-    "    def __init__(self, model, μgrid):\n",
+    "    def __init__(self,\n",
+    "                 pref,\n",
+    "                 x_grid,\n",
+    "                 π=0.5*np.ones((2, 2)),\n",
+    "                 g=np.array([0.1, 0.2])):\n",
     "\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.mc, self.S = MarkovChain(self.π), len(model.π)  # Number of states\n",
-    "        self.Θ, self.model, self.μgrid = model.Θ, model, μgrid\n",
+    "        self.π, self.g, self.S = π, g, len(π)\n",
+    "        self.pref, self.x_grid = pref, x_grid\n",
     "\n",
-    "        # Find the first best allocation\n",
-    "        self.solve_time1_bellman()\n",
-    "        self.T.time_0 = True  # Bellman equation now solves time 0 problem\n",
+    "        bounds = np.empty((self.S, 2))\n",
     "\n",
+    "        # bound for n\n",
+    "        bounds[0] = 0, 1\n",
     "\n",
-    "    def solve_time1_bellman(self):\n",
-    "        '''\n",
-    "        Solve the time 1 Bellman equation for calibration model and initial\n",
-    "        grid μgrid0\n",
-    "        '''\n",
-    "        model, μgrid0 = self.model, self.μgrid\n",
-    "        S = len(model.π)\n",
-    "\n",
-    "        # First get initial fit\n",
-    "        pp = SequentialAllocation(model)\n",
-    "        c, n, x, V = map(np.vstack, zip(*map(lambda μ: pp.time1_value(μ), μgrid0)))\n",
-    "\n",
-    "        Vf, cf, nf, xprimef = {}, {}, {}, {}\n",
-    "        for s in range(2):\n",
-    "            ind = np.argsort(x[:, s])   # Sort x\n",
-    "            # Sort arrays according to x\n",
-    "            c, n, x, V = c[ind], n[ind], x[ind], V[ind]\n",
-    "            cf[s] = UnivariateSpline(x[:, s], c[:, s])\n",
-    "            nf[s] = UnivariateSpline(x[:, s], n[:, s])\n",
-    "            Vf[s] = UnivariateSpline(x[:, s], V[:, s])\n",
-    "            for sprime in range(S):\n",
-    "                xprimef[s, sprime] = UnivariateSpline(x[:, s], x[:, s])\n",
-    "        policies = [cf, nf, xprimef]\n",
-    "\n",
-    "        # Create xgrid\n",
-    "        xbar = [x.min(0).max(), x.max(0).min()]\n",
-    "        xgrid = np.linspace(xbar[0], xbar[1], len(μgrid0))\n",
-    "        self.xgrid = xgrid\n",
-    "\n",
-    "        # Now iterate on bellman equation\n",
-    "        T = BellmanEquation(model, xgrid, policies)\n",
-    "        diff = 1\n",
-    "        while diff > 1e-7:\n",
-    "            PF = T(Vf)\n",
-    "            Vfnew, policies = self.fit_policy_function(PF)\n",
-    "            diff = 0\n",
-    "            for s in range(S):\n",
-    "                diff = max(diff, np.abs(\n",
-    "                    (Vf[s](xgrid) - Vfnew[s](xgrid)) / Vf[s](xgrid)).max())\n",
-    "            Vf = Vfnew\n",
-    "\n",
-    "        # Store value function policies and Bellman Equations\n",
-    "        self.Vf = Vf\n",
-    "        self.policies = policies\n",
-    "        self.T = T\n",
-    "\n",
-    "\n",
-    "    def fit_policy_function(self, PF):\n",
+    "        # bound for xprime\n",
+    "        for s in range(self.S-1):\n",
+    "            bounds[s+1] = x_grid.min(), x_grid.max()\n",
+    "\n",
+    "        self.bounds = bounds\n",
+    "\n",
+    "        # initialization of time 1 value function\n",
+    "        self.V = None\n",
+    "\n",
+    "    def time1_allocation(self, V=None, tol=1e-7):\n",
     "        '''\n",
-    "        Fits the policy functions PF using the points xgrid using\n",
-    "        UnivariateSpline\n",
+    "        Solve the optimal time 1 allocation problem\n",
+    "        by iterating Bellman value function.\n",
     "        '''\n",
-    "        xgrid, S = self.xgrid, self.S\n",
     "\n",
-    "        Vf, cf, nf, xprimef = {}, {}, {}, {}\n",
+    "        π, g, S = self.π, self.g, self.S\n",
+    "        pref, x_grid, bounds = self.pref, self.x_grid, self.bounds\n",
+    "\n",
+    "        # initial guess of value function\n",
+    "        if V is None:\n",
+    "            V = np.zeros((len(x_grid), S))\n",
+    "\n",
+    "        # initial guess of policy\n",
+    "        z = np.empty((len(x_grid), S, S+2))\n",
+    "\n",
+    "        # guess of n\n",
+    "        z[:, :, 1] = 0.5\n",
+    "\n",
+    "        # guess of xprime\n",
     "        for s in range(S):\n",
-    "            PFvec = np.vstack(tuple(map(lambda x: PF(x, s), xgrid)))\n",
-    "            Vf[s] = UnivariateSpline(xgrid, PFvec[:, 0], s=0)\n",
-    "            cf[s] = UnivariateSpline(xgrid, PFvec[:, 1], s=0, k=1)\n",
-    "            nf[s] = UnivariateSpline(xgrid, PFvec[:, 2], s=0, k=1)\n",
-    "            for sprime in range(S):\n",
-    "                xprimef[s, sprime] = UnivariateSpline(\n",
-    "                    xgrid, PFvec[:, 3 + sprime], s=0, k=1)\n",
+    "            for i in range(S-1):\n",
+    "                z[:, s, i+2] = x_grid\n",
     "\n",
-    "        return Vf, [cf, nf, xprimef]\n",
+    "        while True:\n",
+    "            # value function iteration\n",
+    "            V_new, z_new = T(V, z, pref, π, g, x_grid, bounds)\n",
     "\n",
+    "            if np.max(np.abs(V - V_new)) < tol:\n",
+    "                break\n",
     "\n",
-    "    def Τ(self, c, n):\n",
+    "            V = V_new\n",
+    "            z = z_new\n",
+    "\n",
+    "        self.V = V_new\n",
+    "        self.z1 = z_new\n",
+    "        self.c1 = z_new[:, :, 0]\n",
+    "        self.n1 = z_new[:, :, 1]\n",
+    "        self.xprime1 = z_new[:, :, 2:]\n",
+    "\n",
+    "        return V_new, z_new\n",
+    "\n",
+    "    def time0_allocation(self, b0, s0):\n",
     "        '''\n",
-    "        Computes Τ given c, n\n",
+    "        Find the optimal time 0 allocation by maximization.\n",
     "        '''\n",
-    "        model = self.model\n",
-    "        Uc, Un = model.Uc(c, n), model.Un(c, n)\n",
     "\n",
-    "        return 1 + Un / (self.Θ * Uc)\n",
+    "        if self.V is None:\n",
+    "            self.time1_allocation()\n",
+    "\n",
+    "        π, g, S = self.π, self.g, self.S\n",
+    "        pref, x_grid, bounds = self.pref, self.x_grid, self.bounds\n",
+    "        V, z1 = self.V, self.z1\n",
+    "\n",
+    "        x = 1. # x is arbitrary\n",
+    "        res = nelder_mead(obj_V,\n",
+    "                          z1[0, s0, 1:-1],\n",
+    "                          args=(x, s0, V, pref, π, g, x_grid, b0),\n",
+    "                          bounds=bounds,\n",
+    "                          tol_f=1e-10)\n",
+    "\n",
+    "        n0, xprime0 = IC(res.x, x, s0, b0, pref, π, g)\n",
+    "        c0 = n0 - g[s0]\n",
+    "        z0 = np.array([c0, n0, *xprime0])\n",
+    "\n",
+    "        self.z0 = z0\n",
+    "        self.n0 = n0\n",
+    "        self.c0 = n0 - g[s0]\n",
+    "        self.xprime0 = xprime0\n",
     "\n",
+    "        return z0\n",
     "\n",
-    "    def time0_allocation(self, B_, s0):\n",
+    "    def τ(self, c, n):\n",
     "        '''\n",
-    "        Finds the optimal allocation given initial government debt B_ and\n",
-    "        state s_0\n",
+    "        Computes τ given c, n\n",
     "        '''\n",
-    "        PF = self.T(self.Vf)\n",
-    "        z0 = PF(B_, s0)\n",
-    "        c0, n0, xprime0 = z0[1], z0[2], z0[3:]\n",
-    "        return c0, n0, xprime0\n",
+    "        pref = self.pref\n",
+    "        uc, ul = pref.Uc(c, 1-n), pref.Ul(c, 1-n)\n",
     "\n",
+    "        return 1 - ul / uc\n",
     "\n",
-    "    def simulate(self, B_, s_0, T, sHist=None):\n",
+    "    def simulate(self, b0, s0, T, sHist=None):\n",
     "        '''\n",
     "        Simulates Ramsey plan for T periods\n",
     "        '''\n",
-    "        model, π = self.model, self.π\n",
-    "        Uc = model.Uc\n",
-    "        cf, nf, xprimef = self.policies\n",
+    "        pref, π = self.pref, self.π\n",
+    "        Uc = pref.Uc\n",
     "\n",
     "        if sHist is None:\n",
-    "            sHist = self.mc.simulate(T, s_0)\n",
+    "            sHist = self.mc.simulate(T, s0)\n",
     "\n",
-    "        cHist, nHist, Bhist, ΤHist, μHist = np.zeros((5, T))\n",
-    "        RHist = np.zeros(T - 1)\n",
+    "        cHist, nHist, Bhist, τHist, xHist = np.empty((5, T))\n",
+    "        RHist = np.zeros(T-1)\n",
     "\n",
     "        # Time 0\n",
-    "        cHist[0], nHist[0], xprime = self.time0_allocation(B_, s_0)\n",
-    "        ΤHist[0] = self.Τ(cHist[0], nHist[0])[s_0]\n",
-    "        Bhist[0] = B_\n",
-    "        μHist[0] = 0\n",
+    "        self.time0_allocation(b0, s0)\n",
+    "        cHist[0], nHist[0], xHist[0] = self.c0, self.n0, self.xprime0[s0]\n",
+    "        τHist[0] = self.τ(cHist[0], nHist[0])\n",
+    "        Bhist[0] = b0\n",
     "\n",
     "        # Time 1 onward\n",
     "        for t in range(1, T):\n",
-    "            s, x = sHist[t], xprime[sHist[t]]\n",
-    "            c, n, xprime = np.empty(self.S), nf[s](x), np.empty(self.S)\n",
-    "            for shat in range(self.S):\n",
-    "                c[shat] = cf[shat](x)\n",
-    "            for sprime in range(self.S):\n",
-    "                xprime[sprime] = xprimef[s, sprime](x)\n",
+    "            s, x = sHist[t], xHist[t-1]\n",
+    "            cHist[t] = interp(self.x_grid, self.c1[:, s], x)\n",
+    "            nHist[t] = interp(self.x_grid, self.n1[:, s], x)\n",
     "\n",
-    "            Τ = self.Τ(c, n)[s]\n",
-    "            u_c = Uc(c, n)\n",
-    "            Eu_c = π[sHist[t - 1]] @ u_c\n",
-    "            μHist[t] = self.Vf[s](x, 1)\n",
+    "            τHist[t] = self.τ(cHist[t], nHist[t])\n",
     "\n",
-    "            RHist[t - 1] = Uc(cHist[t - 1], nHist[t - 1]) / (self.β * Eu_c)\n",
+    "            Bhist[t] = x / Uc(cHist[t], 1-nHist[t])\n",
     "\n",
-    "            cHist[t], nHist[t], Bhist[t], ΤHist[t] = c[s], n, x / u_c[s], Τ\n",
+    "            c, n = np.empty((2, self.S))\n",
+    "            for sprime in range(self.S):\n",
+    "                c[sprime] = interp(x_grid, self.c1[:, sprime], x)\n",
+    "                n[sprime] = interp(x_grid, self.n1[:, sprime], x)\n",
+    "            Euc = π[sHist[t-1]] @ Uc(c, 1-n)\n",
+    "            RHist[t-1] = Uc(cHist[t-1], 1-nHist[t-1]) / (self.pref.β * Euc)\n",
     "\n",
-    "        return np.array([cHist, nHist, Bhist, ΤHist, sHist, μHist, RHist])\n",
+    "            gHist = self.g[sHist]\n",
+    "            yHist = nHist\n",
+    "            \n",
+    "            if t < T-1:\n",
+    "                sprime = sHist[t+1]\n",
+    "                xHist[t] = interp(self.x_grid, self.xprime1[:, s, sprime], x)\n",
     "\n",
+    "        return [cHist, nHist, Bhist, τHist, gHist, yHist, xHist, RHist]\n",
     "\n",
-    "class BellmanEquation:\n",
+    "# Helper functions\n",
     "\n",
+    "@njit(parallel=True)\n",
+    "def T(V, z, pref, π, g, x_grid, bounds):\n",
     "    '''\n",
-    "    Bellman equation for the continuation of the Lucas-Stokey Problem\n",
+    "    One step iteration of Bellman value function.\n",
     "    '''\n",
     "\n",
-    "    def __init__(self, model, xgrid, policies0):\n",
+    "    S = len(π)\n",
     "\n",
-    "        self.β, self.π, self.G = model.β, model.π, model.G\n",
-    "        self.S = len(model.π)  # Number of states\n",
-    "        self.Θ, self.model = model.Θ, model\n",
+    "    V_new = np.empty_like(V)\n",
+    "    z_new = np.empty_like(z)\n",
     "\n",
-    "        self.xbar = [min(xgrid), max(xgrid)]\n",
-    "        self.time_0 = False\n",
+    "    for i in prange(len(x_grid)):\n",
+    "        x = x_grid[i]\n",
+    "        for s in prange(S):\n",
+    "            res = nelder_mead(obj_V,\n",
+    "                              z[i, s, 1:-1],\n",
+    "                              args=(x, s, V, pref, π, g, x_grid),\n",
+    "                              bounds=bounds,\n",
+    "                              tol_f=1e-10)\n",
     "\n",
-    "        self.z0 = {}\n",
-    "        cf, nf, xprimef = policies0\n",
-    "        for s in range(self.S):\n",
-    "            for x in xgrid:\n",
-    "                xprime0 = np.empty(self.S)\n",
-    "                for sprime in range(self.S):\n",
-    "                    xprime0[sprime] = xprimef[s, sprime](x)\n",
-    "                self.z0[x, s] = np.hstack([cf[s](x), nf[s](x), xprime0])\n",
+    "            # optimal policy\n",
+    "            n, xprime = IC(res.x, x, s, None, pref, π, g)\n",
+    "            z_new[i, s, 0] = n - g[s]        # c\n",
+    "            z_new[i, s, 1] = n               # n\n",
+    "            z_new[i, s, 2:] = xprime         # xprime\n",
+    "            \n",
+    "            V_new[i, s] = res.fun\n",
     "\n",
-    "        self.find_first_best()\n",
+    "    return V_new, z_new\n",
     "\n",
+    "@njit\n",
+    "def obj_V(z_sub, x, s, V, pref, π, g, x_grid, b0=None):\n",
+    "    '''\n",
+    "    The objective on the right hand side of the Bellman equation.\n",
+    "    z_sub contains guesses of n and xprime[:-1].\n",
+    "    '''\n",
     "\n",
-    "    def find_first_best(self):\n",
-    "        '''\n",
-    "        Find the first best allocation\n",
-    "        '''\n",
-    "        model = self.model\n",
-    "        S, Θ, Uc, Un, G = self.S, self.Θ, model.Uc, model.Un, self.G\n",
+    "    S = len(π)\n",
+    "    β, U = pref.β, pref.U\n",
     "\n",
-    "        def res(z):\n",
-    "            c = z[:S]\n",
-    "            n = z[S:]\n",
-    "            return np.hstack([Θ * Uc(c, n) + Un(c, n), Θ * n - c - G])\n",
+    "    # find (n, xprime) that satisfies implementability constraint\n",
+    "    n, xprime = IC(z_sub, x, s, b0, pref, π, g)\n",
+    "    c, l = n-g[s], 1-n\n",
     "\n",
-    "        res = root(res, 0.5 * np.ones(2 * S))\n",
-    "        if not res.success:\n",
-    "            raise Exception('Could not find first best')\n",
+    "    # if xprime[-1] violates bound, return large penalty\n",
+    "    if (xprime[-1] < x_grid.min()):\n",
+    "        return -1e9 * (1 + np.abs(xprime[-1] - x_grid.min()))\n",
+    "    elif (xprime[-1] > x_grid.max()):\n",
+    "        return -1e9 * (1 + np.abs(xprime[-1] - x_grid.max()))\n",
     "\n",
-    "        self.cFB = res.x[:S]\n",
-    "        self.nFB = res.x[S:]\n",
-    "        IFB = Uc(self.cFB, self.nFB) * self.cFB + Un(self.cFB, self.nFB) * self.nFB\n",
-    "        self.xFB = np.linalg.solve(np.eye(S) - self.β * self.π, IFB)\n",
-    "        self.zFB = {}\n",
+    "    # prepare Vprime vector\n",
+    "    Vprime = np.empty(S)\n",
+    "    for sprime in range(S):\n",
+    "        Vprime[sprime] = interp(x_grid, V[:, sprime], xprime[sprime])\n",
     "\n",
-    "        for s in range(S):\n",
-    "            self.zFB[s] = np.hstack([self.cFB[s], self.nFB[s], self.xFB])\n",
+    "    # compute the objective value\n",
+    "    obj = U(c, l) + β * π[s] @ Vprime\n",
     "\n",
+    "    return obj\n",
     "\n",
-    "    def __call__(self, Vf):\n",
-    "        '''\n",
-    "        Given continuation value function, next period return value function,\n",
-    "        this period return T(V) and optimal policies\n",
-    "        '''\n",
-    "        if not self.time_0:\n",
-    "            def PF(x, s): return self.get_policies_time1(x, s, Vf)\n",
-    "        else:\n",
-    "            def PF(B_, s0): return self.get_policies_time0(B_, s0, Vf)\n",
-    "        return PF\n",
-    "\n",
-    "\n",
-    "    def get_policies_time1(self, x, s, Vf):\n",
-    "        '''\n",
-    "        Finds the optimal policies\n",
-    "        '''\n",
-    "        model, β, Θ, = self.model, self.β, self.Θ,\n",
-    "        G, S, π = self.G, self.S, self.π\n",
-    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
-    "\n",
-    "        def objf(z):\n",
-    "            c, n, xprime = z[0], z[1], z[2:]\n",
-    "            Vprime = np.empty(S)\n",
-    "            for sprime in range(S):\n",
-    "                Vprime[sprime] = Vf[sprime](xprime[sprime])\n",
-    "\n",
-    "            return -(U(c, n) + β * π[s] @ Vprime)\n",
+    "@njit\n",
+    "def IC(z_sub, x, s, b0, pref, π, g):\n",
+    "    '''\n",
+    "    Find xprime[-1] that satisfies the implementability condition\n",
+    "    given the guesses of n and xprime[:-1].\n",
+    "    '''\n",
     "\n",
-    "        def cons(z):\n",
-    "            c, n, xprime = z[0], z[1], z[2:]\n",
-    "            return np.hstack([x - Uc(c, n) * c - Un(c, n) * n - β * π[s]\n",
-    "                              @ xprime,\n",
-    "                              (Θ * n - c - G)[s]])\n",
+    "    β, Uc, Ul = pref.β, pref.Uc, pref.Ul\n",
     "\n",
-    "        out, fx, _, imode, smode = fmin_slsqp(objf,\n",
-    "                                              self.z0[x, s],\n",
-    "                                              f_eqcons=cons,\n",
-    "                                              bounds=[(0, 100), (0, 100)]\n",
-    "                                              + [self.xbar] * S,\n",
-    "                                              full_output=True,\n",
-    "                                              iprint=0,\n",
-    "                                              acc=1e-10)\n",
+    "    n = z_sub[0]\n",
+    "    xprime = np.empty(len(π))\n",
+    "    xprime[:-1] = z_sub[1:]\n",
     "\n",
-    "        if imode > 0:\n",
-    "            raise Exception(smode)\n",
+    "    c, l = n-g[s], 1-n\n",
+    "    uc = Uc(c, l)\n",
+    "    ul = Ul(c, l)\n",
     "\n",
-    "        self.z0[x, s] = out\n",
-    "        return np.hstack([-fx, out])\n",
+    "    if b0 is None:\n",
+    "        diff = x\n",
+    "    else:\n",
+    "        diff = uc * b0\n",
     "\n",
+    "    diff -= uc * (n - g[s]) - ul * n + β * π[s][:-1] @ xprime[:-1]\n",
+    "    xprime[-1] = diff / (β * π[s][-1])\n",
     "\n",
-    "    def get_policies_time0(self, B_, s0, Vf):\n",
-    "        '''\n",
-    "        Finds the optimal policies\n",
-    "        '''\n",
-    "        model, β, Θ, = self.model, self.β, self.Θ,\n",
-    "        G, S, π = self.G, self.S, self.π\n",
-    "        U, Uc, Un = model.U, model.Uc, model.Un\n",
-    "\n",
-    "        def objf(z):\n",
-    "            c, n, xprime = z[0], z[1], z[2:]\n",
-    "            Vprime = np.empty(S)\n",
-    "            for sprime in range(S):\n",
-    "                Vprime[sprime] = Vf[sprime](xprime[sprime])\n",
-    "\n",
-    "            return -(U(c, n) + β * π[s0] @ Vprime)\n",
-    "\n",
-    "        def cons(z):\n",
-    "            c, n, xprime = z[0], z[1], z[2:]\n",
-    "            return np.hstack([-Uc(c, n) * (c - B_) - Un(c, n) * n - β * π[s0]\n",
-    "                              @ xprime,\n",
-    "                              (Θ * n - c - G)[s0]])\n",
-    "\n",
-    "        out, fx, _, imode, smode = fmin_slsqp(objf, self.zFB[s0], f_eqcons=cons,\n",
-    "                                              bounds=[(0, 100), (0, 100)]\n",
-    "                                              + [self.xbar] * S,\n",
-    "                                              full_output=True, iprint=0,\n",
-    "                                              acc=1e-10)\n",
-    "\n",
-    "        if imode > 0:\n",
-    "            raise Exception(smode)\n",
-    "\n",
-    "        return np.hstack([-fx, out])"
+    "    return n, xprime"
    ]
   },
   {
@@ -1687,44 +1689,52 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "\n",
+    "crra_util_data = [\n",
+    "    ('β', float64),\n",
+    "    ('σ', float64),\n",
+    "    ('γ', float64)\n",
+    "]\n",
     "\n",
+    "@jitclass(crra_util_data)\n",
     "class CRRAutility:\n",
     "\n",
     "    def __init__(self,\n",
     "                 β=0.9,\n",
     "                 σ=2,\n",
-    "                 γ=2,\n",
-    "                 π=0.5*np.ones((2, 2)),\n",
-    "                 G=np.array([0.1, 0.2]),\n",
-    "                 Θ=np.ones(2),\n",
-    "                 transfers=False):\n",
+    "                 γ=2):\n",
     "\n",
     "        self.β, self.σ, self.γ = β, σ, γ\n",
-    "        self.π, self.G, self.Θ, self.transfers = π, G, Θ, transfers\n",
     "\n",
     "    # Utility function\n",
-    "    def U(self, c, n):\n",
+    "    def U(self, c, l):\n",
+    "        # Note: `l` should not be interpreted as labor, it is an auxiliary\n",
+    "        # variable used to conveniently match the code and the equations\n",
+    "        # in the lecture\n",
     "        σ = self.σ\n",
     "        if σ == 1.:\n",
     "            U = np.log(c)\n",
     "        else:\n",
     "            U = (c**(1 - σ) - 1) / (1 - σ)\n",
-    "        return U - n**(1 + self.γ) / (1 + self.γ)\n",
+    "        return U - (1-l) ** (1 + self.γ) / (1 + self.γ)\n",
     "\n",
     "    # Derivatives of utility function\n",
-    "    def Uc(self, c, n):\n",
-    "        return c**(-self.σ)\n",
+    "    def Uc(self, c, l):\n",
+    "        return c ** (-self.σ)\n",
+    "\n",
+    "    def Ucc(self, c, l):\n",
+    "        return -self.σ * c ** (-self.σ - 1)\n",
     "\n",
-    "    def Ucc(self, c, n):\n",
-    "        return -self.σ * c**(-self.σ - 1)\n",
+    "    def Ul(self, c, l):\n",
+    "        return (1-l) ** self.γ\n",
     "\n",
-    "    def Un(self, c, n):\n",
-    "        return -n**self.γ\n",
+    "    def Ull(self, c, l):\n",
+    "        return -self.γ * (1-l) ** (self.γ - 1)\n",
     "\n",
-    "    def Unn(self, c, n):\n",
-    "        return -self.γ * n**(self.γ - 1)"
+    "    def Ucl(self, c, l):\n",
+    "        return 0\n",
+    "\n",
+    "    def Ulc(self, c, l):\n",
+    "        return 0"
    ]
   },
   {
@@ -1747,39 +1757,31 @@
    },
    "outputs": [],
    "source": [
-    "time_π = np.array([[0, 1, 0,   0,   0,  0],\n",
-    "                   [0, 0, 1,   0,   0,  0],\n",
-    "                   [0, 0, 0, 0.5, 0.5,  0],\n",
-    "                   [0, 0, 0,   0,   0,  1],\n",
-    "                   [0, 0, 0,   0,   0,  1],\n",
-    "                   [0, 0, 0,   0,   0,  1]])\n",
-    "\n",
-    "time_G = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
-    "# Θ can in principle be random\n",
-    "time_Θ = np.ones(6)\n",
-    "time_example = CRRAutility(π=time_π, G=time_G, Θ=time_Θ)\n",
+    "π = np.array([[0, 1, 0,   0,   0,  0],\n",
+    "              [0, 0, 1,   0,   0,  0],\n",
+    "              [0, 0, 0, 0.5, 0.5,  0],\n",
+    "              [0, 0, 0,   0,   0,  1],\n",
+    "              [0, 0, 0,   0,   0,  1],\n",
+    "              [0, 0, 0,   0,   0,  1]])\n",
+    "\n",
+    "g = np.array([0.1, 0.1, 0.1, 0.2, 0.1, 0.1])\n",
+    "crra_pref = CRRAutility()\n",
     "\n",
     "# Solve sequential problem\n",
-    "time_allocation = SequentialAllocation(time_example)\n",
+    "seq = SequentialLS(crra_pref, π=π, g=g)\n",
     "sHist_h = np.array([0, 1, 2, 3, 5, 5, 5])\n",
     "sHist_l = np.array([0, 1, 2, 4, 5, 5, 5])\n",
-    "sim_seq_h = time_allocation.simulate(1, 0, 7, sHist_h)\n",
-    "sim_seq_l = time_allocation.simulate(1, 0, 7, sHist_l)\n",
-    "\n",
-    "# Government spending paths\n",
-    "sim_seq_l[4] = time_example.G[sHist_l]\n",
-    "sim_seq_h[4] = time_example.G[sHist_h]\n",
-    "\n",
-    "# Output paths\n",
-    "sim_seq_l[5] = time_example.Θ[sHist_l] * sim_seq_l[1]\n",
-    "sim_seq_h[5] = time_example.Θ[sHist_h] * sim_seq_h[1]\n",
+    "sim_seq_h = seq.simulate(1, 0, 7, sHist_h)\n",
+    "sim_seq_l = seq.simulate(1, 0, 7, sHist_l)\n",
     "\n",
     "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
     "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
     "          'Tax Rate', 'Government Spending', 'Output']\n",
     "\n",
     "for ax, title, sim_l, sim_h in zip(axes.flatten(),\n",
-    "        titles, sim_seq_l, sim_seq_h):\n",
+    "                                   titles,\n",
+    "                                   sim_seq_l[:6],\n",
+    "                                   sim_seq_h[:6]):\n",
     "    ax.set(title=title)\n",
     "    ax.plot(sim_l, '-ok', sim_h, '-or', alpha=0.7)\n",
     "    ax.grid()\n",
@@ -1925,9 +1927,7 @@
    },
    "outputs": [],
    "source": [
-    "tax_sequence = SequentialAllocation(CRRAutility(G=0.15,\n",
-    "                                                π=np.ones((1, 1)),\n",
-    "                                                Θ=np.ones(1)))\n",
+    "tax_seq = SequentialLS(CRRAutility(), g=np.array([0.15]), π=np.ones((1, 1)))\n",
     "\n",
     "n = 100\n",
     "tax_policy = np.empty((n, 2))\n",
@@ -1935,8 +1935,8 @@
     "gov_debt = np.linspace(-1.5, 1, n)\n",
     "\n",
     "for i in range(n):\n",
-    "    tax_policy[i] = tax_sequence.simulate(gov_debt[i], 0, 2)[3]\n",
-    "    interest_rate[i] = tax_sequence.simulate(gov_debt[i], 0, 3)[-1]\n",
+    "    tax_policy[i] = tax_seq.simulate(gov_debt[i], 0, 2)[3]\n",
+    "    interest_rate[i] = tax_seq.simulate(gov_debt[i], 0, 3)[-1]\n",
     "\n",
     "fig, axes = plt.subplots(2, 1, figsize=(10,8), sharex=True)\n",
     "titles = ['Tax Rate', 'Gross Interest Rate']\n",
@@ -2009,9 +2009,7 @@
    },
    "outputs": [],
    "source": [
-    "tax_sequence = SequentialAllocation(CRRAutility(G=0.15,\n",
-    "                                                π=np.ones((1, 1)),\n",
-    "                                                Θ=np.ones(1)))\n",
+    "tax_seq = SequentialLS(CRRAutility(), g=np.array([0.15]), π=np.ones((1, 1)))\n",
     "\n",
     "n = 100\n",
     "tax_policy = np.empty((n, 2))\n",
@@ -2019,8 +2017,8 @@
     "gov_debt = np.linspace(-1.5, 1, n)\n",
     "\n",
     "for i in range(n):\n",
-    "    tax_policy[i] = tax_sequence.simulate(gov_debt[i], 0, 2)[3]\n",
-    "    τ_reset[i] = tax_sequence.simulate(gov_debt[i], 0, 1)[3]\n",
+    "    tax_policy[i] = tax_seq.simulate(gov_debt[i], 0, 2)[3]\n",
+    "    τ_reset[i] = tax_seq.simulate(gov_debt[i], 0, 1)[3]\n",
     "\n",
     "fig, ax = plt.subplots(figsize=(10, 6))\n",
     "ax.plot(gov_debt, tax_policy[:, 1], gov_debt, τ_reset, lw=2)\n",
@@ -2083,37 +2081,42 @@
    },
    "outputs": [],
    "source": [
-    "import numpy as np\n",
+    "log_util_data = [\n",
+    "    ('β', float64),\n",
+    "    ('ψ', float64)\n",
+    "]\n",
     "\n",
+    "@jitclass(log_util_data)\n",
     "class LogUtility:\n",
     "\n",
     "    def __init__(self,\n",
     "                 β=0.9,\n",
-    "                 ψ=0.69,\n",
-    "                 π=0.5*np.ones((2, 2)),\n",
-    "                 G=np.array([0.1, 0.2]),\n",
-    "                 Θ=np.ones(2),\n",
-    "                 transfers=False):\n",
+    "                 ψ=0.69):\n",
     "\n",
-    "        self.β, self.ψ, self.π = β, ψ, π\n",
-    "        self.G, self.Θ, self.transfers = G, Θ, transfers\n",
+    "        self.β, self.ψ = β, ψ\n",
     "\n",
     "    # Utility function\n",
-    "    def U(self, c, n):\n",
-    "        return np.log(c) + self.ψ * np.log(1 - n)\n",
+    "    def U(self, c, l):\n",
+    "        return np.log(c) + self.ψ * np.log(l)\n",
     "\n",
     "    # Derivatives of utility function\n",
-    "    def Uc(self, c, n):\n",
+    "    def Uc(self, c, l):\n",
     "        return 1 / c\n",
     "\n",
-    "    def Ucc(self, c, n):\n",
+    "    def Ucc(self, c, l):\n",
     "        return -c**(-2)\n",
     "\n",
-    "    def Un(self, c, n):\n",
-    "        return -self.ψ / (1 - n)\n",
+    "    def Ul(self, c, l):\n",
+    "        return self.ψ / l\n",
+    "\n",
+    "    def Ull(self, c, l):\n",
+    "        return -self.ψ / l**2\n",
     "\n",
-    "    def Unn(self, c, n):\n",
-    "        return -self.ψ / (1 - n)**2"
+    "    def Ucl(self, c, l):\n",
+    "        return 0\n",
+    "\n",
+    "    def Ulc(self, c, l):\n",
+    "        return 0"
    ]
   },
   {
@@ -2138,38 +2141,32 @@
    "source": [
     "log_example = LogUtility()\n",
     "# Solve sequential problem\n",
-    "seq_log = SequentialAllocation(log_example)\n",
+    "seq_log = SequentialLS(log_example)\n",
     "\n",
     "# Initialize grid for value function iteration and solve\n",
-    "μ_grid = np.linspace(-0.6, 0.0, 200)\n",
+    "x_grid = np.linspace(-3., 3., 200)\n",
+    "\n",
     "# Solve recursive problem\n",
-    "bel_log = RecursiveAllocation(log_example, μ_grid)\n",
+    "rec_log = RecursiveLS(log_example, x_grid)\n",
     "\n",
-    "T = 20\n",
-    "sHist = np.array([0, 0, 0, 0, 0, 0, 0,\n",
-    "                  0, 1, 1, 0, 0, 0, 1,\n",
-    "                  1, 1, 1, 1, 1, 0])\n",
+    "T_length = 20\n",
+    "sHist = np.array([0, 0, 0, 0, 0,\n",
+    "                  0, 0, 0, 1, 1,\n",
+    "                  0, 0, 0, 1, 1,\n",
+    "                  1, 1, 1, 1, 0])\n",
     "\n",
     "# Simulate\n",
-    "sim_seq = seq_log.simulate(0.5, 0, T, sHist)\n",
-    "sim_bel = bel_log.simulate(0.5, 0, T, sHist)\n",
-    "\n",
-    "# Government spending paths\n",
-    "sim_seq[4] = log_example.G[sHist]\n",
-    "sim_bel[4] = log_example.G[sHist]\n",
-    "\n",
-    "# Output paths\n",
-    "sim_seq[5] = log_example.Θ[sHist] * sim_seq[1]\n",
-    "sim_bel[5] = log_example.Θ[sHist] * sim_bel[1]\n",
+    "sim_seq = seq_log.simulate(0.5, 0, T_length, sHist)\n",
+    "sim_rec = rec_log.simulate(0.5, 0, T_length, sHist)\n",
     "\n",
     "fig, axes = plt.subplots(3, 2, figsize=(14, 10))\n",
     "titles = ['Consumption', 'Labor Supply', 'Government Debt',\n",
     "          'Tax Rate', 'Government Spending', 'Output']\n",
     "\n",
-    "for ax, title, sim_s, sim_b in zip(axes.flatten(), titles, sim_seq, sim_bel):\n",
-    "    ax.plot(sim_s, '-ob', sim_b, '-xk', alpha=0.7)\n",
-    "    ax.set(title=title)\n",
-    "    ax.grid()\n",
+    "for ax, title, sim_s, sim_b in zip(axes.flatten(), titles, sim_seq[:6], sim_rec[:6]):\n",
+    "                                   ax.plot(sim_s, '-ob', sim_b, '-xk', alpha=0.7)\n",
+    "                                   ax.set(title=title)\n",
+    "                                   ax.grid()\n",
     "\n",
     "axes.flatten()[0].legend(('Sequential', 'Recursive'))\n",
     "fig.tight_layout()\n",
@@ -2210,7 +2207,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661948.0408337,
+  "date": 1624431175.9157023,
   "filename": "opt_tax_recur.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/orth_proj.ipynb b/orth_proj.ipynb
index 347cc9d..e347420 100644
--- a/orth_proj.ipynb
+++ b/orth_proj.ipynb
@@ -1083,7 +1083,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661948.2942147,
+  "date": 1624431176.3593583,
   "filename": "orth_proj.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/permanent_income_dles.ipynb b/permanent_income_dles.ipynb
index 55895c3..4114203 100644
--- a/permanent_income_dles.ipynb
+++ b/permanent_income_dles.ipynb
@@ -369,7 +369,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661948.6251843,
+  "date": 1624431176.498286,
   "filename": "permanent_income_dles.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/rob_markov_perf.ipynb b/rob_markov_perf.ipynb
index 5bc9fd5..1651068 100644
--- a/rob_markov_perf.ipynb
+++ b/rob_markov_perf.ipynb
@@ -1233,7 +1233,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661948.7616801,
+  "date": 1624431176.6336029,
   "filename": "rob_markov_perf.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/robustness.ipynb b/robustness.ipynb
index 9b31bc8..afa2bd1 100644
--- a/robustness.ipynb
+++ b/robustness.ipynb
@@ -1364,7 +1364,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661948.986384,
+  "date": 1624431176.8532524,
   "filename": "robustness.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/rosen_schooling_model.ipynb b/rosen_schooling_model.ipynb
index 16966e6..df3071b 100644
--- a/rosen_schooling_model.ipynb
+++ b/rosen_schooling_model.ipynb
@@ -459,7 +459,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661949.1206965,
+  "date": 1624431176.9847429,
   "filename": "rosen_schooling_model.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/smoothing.ipynb b/smoothing.ipynb
index 11b9482..7e8a837 100644
--- a/smoothing.ipynb
+++ b/smoothing.ipynb
@@ -1108,7 +1108,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661949.2804203,
+  "date": 1624431177.1357064,
   "filename": "smoothing.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/smoothing_tax.ipynb b/smoothing_tax.ipynb
index 8a32375..1daa614 100644
--- a/smoothing_tax.ipynb
+++ b/smoothing_tax.ipynb
@@ -1194,7 +1194,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661949.4542558,
+  "date": 1624431177.305599,
   "filename": "smoothing_tax.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/stationary_densities.ipynb b/stationary_densities.ipynb
index 45346cd..e766a24 100644
--- a/stationary_densities.ipynb
+++ b/stationary_densities.ipynb
@@ -1264,7 +1264,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661949.8925626,
+  "date": 1624431177.751618,
   "filename": "stationary_densities.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/tax_smoothing_1.ipynb b/tax_smoothing_1.ipynb
index 9943566..bc31d40 100644
--- a/tax_smoothing_1.ipynb
+++ b/tax_smoothing_1.ipynb
@@ -666,7 +666,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.083871,
+  "date": 1624431177.9358778,
   "filename": "tax_smoothing_1.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/tax_smoothing_2.ipynb b/tax_smoothing_2.ipynb
index e7216fb..42efff6 100644
--- a/tax_smoothing_2.ipynb
+++ b/tax_smoothing_2.ipynb
@@ -967,7 +967,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.2123818,
+  "date": 1624431178.0726104,
   "filename": "tax_smoothing_2.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/tax_smoothing_3.ipynb b/tax_smoothing_3.ipynb
index c651e6f..633857a 100644
--- a/tax_smoothing_3.ipynb
+++ b/tax_smoothing_3.ipynb
@@ -415,7 +415,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.316128,
+  "date": 1624431178.1728816,
   "filename": "tax_smoothing_3.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/troubleshooting.ipynb b/troubleshooting.ipynb
index 4bdd2f4..15bd80e 100644
--- a/troubleshooting.ipynb
+++ b/troubleshooting.ipynb
@@ -97,7 +97,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.388517,
+  "date": 1624431178.2283533,
   "filename": "troubleshooting.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/von_neumann_model.ipynb b/von_neumann_model.ipynb
index c3c8ad5..895ffcb 100644
--- a/von_neumann_model.ipynb
+++ b/von_neumann_model.ipynb
@@ -1246,7 +1246,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.5373182,
+  "date": 1624431178.3554227,
   "filename": "von_neumann_model.rst",
   "kernelspec": {
    "display_name": "Python",
diff --git a/zreferences.ipynb b/zreferences.ipynb
index 15e20bb..68660d4 100644
--- a/zreferences.ipynb
+++ b/zreferences.ipynb
@@ -260,7 +260,7 @@
   }
  ],
  "metadata": {
-  "date": 1619661950.6524045,
+  "date": 1624431178.538628,
   "filename": "zreferences.rst",
   "kernelspec": {
    "display_name": "Python",