Done update along with formatting

lectures/multi_agent_models/lucas_model.md

@@ -386,69 +386,54 @@ using Test
 ```
 
 ```{code-cell} julia
-using LinearAlgebra, Statistics
-using Distributions, Interpolations, LaTeXStrings, Parameters, Plots, Random
+using LinearAlgebra, Statistics, Random
+using Distributions, Interpolations, LaTeXStrings, Plots, NLsolve
 
 ```
 
 ```{code-cell} julia
-# model
-function LucasTree(;gamma = 2.0,
-                    beta = 0.95,
-                    alpha = 0.9,
-                    sigma = 0.1,
-                    grid_size = 100)
-
+function LucasTree(; gamma = 2.0,
+                   beta = 0.95,
+                   alpha = 0.9,
+                   sigma = 0.1,
+                   grid_size = 100,
+                   num_z = 500)
     phi = LogNormal(0.0, sigma)
-    shocks = rand(phi, 500)
+    z = rand(phi, num_z)
 
     # build a grid with mass around stationary distribution
     ssd = sigma / sqrt(1 - alpha^2)
-    grid_min, grid_max = exp(-4ssd), exp(4ssd)
+    grid_min = exp(-4 * ssd)
+    grid_max = exp(4 * ssd)
     grid = range(grid_min, grid_max, length = grid_size)
 
     # set h(y) = beta * int u'(G(y,z)) G(y,z) phi(dz)
     h = similar(grid)
     for (i, y) in enumerate(grid)
-        h[i] = beta * mean((y^alpha .* shocks).^(1 - gamma))
+        h[i] = beta * mean((y^alpha .* z) .^ (1 - gamma))
     end
 
-    return (;gamma, beta, alpha, sigma, phi, grid, shocks, h)
-end
-
-# approximate Lucas operator, which returns the updated function Tf on the grid
-function lucas_operator(lt, f)
-
-    # unpack input
-    (;grid, alpha, beta, h) = lt
-    z = lt.shocks
-
-    Af = linear_interpolation(grid, f, extrapolation_bc=Line()) 
-
-    Tf = [ h[i] + beta * mean(Af.(grid[i]^alpha .* z)) for i in 1:length(grid) ]
-    return Tf
+    return (; gamma, beta, alpha, sigma, phi, grid, z, h)
 end
 
 # get equilibrium price for Lucas tree
-function solve_lucas_model(lt;
-                           tol = 1e-6,
-                           max_iter = 500)
-
-    (;grid, gamma) = lt
-
-    i = 0
-    f = zero(grid)  # Initial guess of f
-    error = tol + 1
-
-    while (error > tol) && (i < max_iter)
-        f_new = lucas_operator(lt, f)
-        error = maximum(abs, f_new - f)
-        f = f_new
-        i += 1
+function solve_lucas_model(lt; ftol = 1e-8, iterations = 500)
+    (; grid, gamma, alpha, beta, h, z) = lt
+
+    # approximate Lucas operator, which returns the updated function Tf on the grid
+    function T(f)
+        Af = linear_interpolation(grid, f, extrapolation_bc = Line())
+        # Using z for monte-carlo integration
+        Tf = [h[i] + beta * mean(Af.(grid[i]^alpha .* z))
+              for i in 1:length(grid)]
+        return Tf
     end
 
-    # p(y) = f(y) * y ^ gamma
-    price = f .* grid.^gamma
+    sol = fixedpoint(T, zero(grid); ftol, iterations)
+    converged(sol) || error("Failed to converge in $(sol.iterations) iter")
+    f = sol.zero
+
+    price = f .* grid .^ gamma # f(y)*y^gamma
 
     return price
 end
@@ -459,7 +444,7 @@ An example of usage is given in the docstring and repeated here
 ```{code-cell} julia
 Random.seed!(42) # For reproducible results.
 
-tree = LucasTree(gamma = 2.0, beta = 0.95, alpha = 0.90, sigma = 0.1)
+tree = LucasTree(; gamma = 2.0, beta = 0.95, alpha = 0.90, sigma = 0.1)
 price_vals = solve_lucas_model(tree);
 ```
 
@@ -468,9 +453,9 @@ price_vals = solve_lucas_model(tree);
 tags: [remove-cell]
 ---
 @testset begin
-    #test price_vals[57] ≈ 44.5077566916004
-    #test price_vals[78] ≈ 68.42956586308563
-    #test price_vals[13] ≈ 9.880376662058682
+    @test price_vals[57] ≈ 41.35601991726328
+    @test price_vals[78] ≈ 63.474988006734925
+    @test price_vals[13] ≈ 9.24450665849126
 end
 ```
 
@@ -519,8 +504,8 @@ Random.seed!(42);
 
 ```{code-cell} julia
 plot()
-for beta in (.95, 0.98)
-    tree = LucasTree(;beta = beta)
+for beta in (0.95, 0.98)
+    tree = LucasTree(; beta)
     grid = tree.grid
     price_vals = solve_lucas_model(tree)
     plot!(grid, price_vals, lw = 2, label = L"\beta = %$beta")
@@ -536,8 +521,8 @@ tags: [remove-cell]
 @testset begin # For the 0.98, since the other one is overwritten.
     Random.seed!(42)
     price_vals = solve_lucas_model(LucasTree(beta = 0.98))
-    #test price_vals[20] ≈ 35.00073581199659
-    #test price_vals[57] ≈ 124.32987344509688
+    @test price_vals[20] ≈ 32.17514173843709
+    @test price_vals[57] ≈ 113.89100588660085
 end
 ```