Update to optimal growth lecture (#145)

* In progress on change to type-safety + named tuples

* Commented out old one

* Fixed exercise solution code

* Fixed typo

Author: jlperla

lectures/dynamic_programming/optgrowth.md

@@ -464,8 +464,8 @@ The next figure illustrates piecewise linear interpolation of an arbitrary funct
 tags: [hide-output]
 ---
 using LinearAlgebra, Statistics
-using Plots, QuantEcon, Interpolations, NLsolve, Optim, Random
-
+using Plots, QuantEcon, Interpolations, NLsolve, Optim, Random, Parameters
+using Optim: maximum, maximizer
 ```
 
 ```{code-cell} julia
@@ -498,21 +498,19 @@ Another advantage of piecewise linear interpolation is that it preserves useful
 Here's a function that implements the Bellman operator using linear interpolation
 
 ```{code-cell} julia
-using Optim
-
-function T(w, grid, β, u, f, shocks; compute_policy = false)
-    w_func = LinearInterpolation(grid, w)
-    # objective for each grid point
-    objectives = (c -> u(c) + β * mean(w_func.(f(y - c) .* shocks)) for y in grid)
-    results = maximize.(objectives, 1e-10, grid) # solver result for each grid point
-
-    Tw = Optim.maximum.(results)
-    if compute_policy
-        σ = Optim.maximizer.(results)
-        return Tw, σ
+function T(w;p, tol = 1e-10)
+    @unpack β, u, f, ξ, y = p # unpack parameters
+    w_func = LinearInterpolation(y, w)
+
+    Tw = similar(w)
+    σ = similar(w)
+    for (i, y_val) in enumerate(y)
+        # solve maximization for each point in y, using y itself as initial condition.
+        results = maximize(c -> u(c;p) + β * mean(w_func.(f(y_val - c;p) .* ξ)), tol, y_val)
+        Tw[i] = maximum(results)
+        σ[i] = maximizer(results)
     end
-
-    return Tw
+    return (;w = Tw, σ) # returns named tuple of results
 end
 ```
 
@@ -561,76 +559,34 @@ $$
 
 Let's code this up now so we can test against it below
 
-```{code-cell} julia
-α = 0.4
-β = 0.96
-μ = 0
-s = 0.1
-
-c1 = log(1 - α * β) / (1 - β)
-c2 = (μ + α * log(α * β)) / (1 - α)
-c3 = 1 / (1 - β)
-c4 = 1 / (1 - α * β)
-
-# Utility
-u(c) = log(c)
-
-∂u∂c(c) = 1 / c
-
-# Deterministic part of production function
-f(k) = k^α
-
-f′(k) = α * k^(α - 1)
-
-# True optimal policy
-c_star(y) = (1 - α * β) * y
-
-# True value function
-v_star(y) = c1 + c2 * (c3 - c4) + c4 * log(y)
-```
+In addition to the model parameters, we need a grid and some shock draws for Monte Carlo integration.
 
 ```{code-cell} julia
----
-tags: [remove-cell]
----
-@testset "Primitives Tests" begin
-    @test [c1, c2, c3, c4] ≈ [-12.112707886215421, -0.6380751509296068, 24.99999999999998,
-                              1.6233766233766234]
-    @test ∂u∂c(c1) ≈ -0.08255792258789846
-    @test v_star(3) ≈ -25.245288867900843
+Random.seed!(42) # for reproducible results
+u(c;p) = log(c) # utility
+f(k;p) = k^p.α # deterministic part of production function
+OptimalGrowthModel = @with_kw (α = 0.4, β = 0.96, μ = 0.0, s = 0.1,
+                  u = u, f = f, # defaults defined above
+                  y = range(1e-5, 4.0, length = 200), # grid on y
+                  ξ = exp.(μ .+ s * randn(250)) # monte carlo shocks
+) # named tuples defaults
+
+# True value and policy function
+function v_star(y;p)
+    @unpack α, μ, β = p
+    c1 = log(1 - α * β) / (1 - β)
+    c2 = (μ + α * log(α * β)) / (1 - α)
+    c3 = 1 / (1 - β)
+    c4 = 1 / (1 - α * β)
+    return c1 + c2 * (c3 - c4) + c4 * log(y)
 end
+c_star(y;p) = (1 - p.α * p.β) * y
 ```
 
 ### A First Test
 
 To test our code, we want to see if we can replicate the analytical solution numerically, using fitted value function iteration.
 
-We need a grid and some shock draws for Monte Carlo integration.
-
-```{code-cell} julia
----
-tags: [hide-output]
----
-using Random
-Random.seed!(42) # For reproducible results.
-
-grid_max = 4         # Largest grid point
-grid_size = 200      # Number of grid points
-shock_size = 250     # Number of shock draws in Monte Carlo integral
-
-grid_y = range(1e-5,  grid_max, length = grid_size)
-shocks = exp.(μ .+ s * randn(shock_size))
-```
-
-```{code-cell} julia
----
-tags: [remove-cell]
----
-@testset "Shock Invariance Tests" begin
-    #test shocks[4] ≈ 0.9704956010607036 && length(shocks) == 250
-end
-```
-
 Now let's do some tests.
 
 As one preliminary test, let's see what happens when we apply our Bellman operator to the exact solution $v^*$.
@@ -640,11 +596,14 @@ In theory, the resulting function should again be $v^*$.
 In practice we expect some small numerical error.
 
 ```{code-cell} julia
-w = T(v_star.(grid_y), grid_y, β, log, k -> k^α, shocks)
+p = OptimalGrowthModel() # use all default parameters from named tuple
+w_star = v_star.(p.y; p)  # evaluate closed form value along grid
+
+w = T(w_star; p).w # evaluate operator, access Tw results
 
 plt = plot(ylim = (-35,-24))
-plot!(plt, grid_y, w, linewidth = 2, alpha = 0.6, label = "T(v_star)")
-plot!(plt, v_star, grid_y, linewidth = 2, alpha=0.6, label = "v_star")
+plot!(plt, p.y, w, linewidth = 2, alpha = 0.6, label = "T(v_star)")
+plot!(plt, p.y, w_star, linewidth = 2, alpha=0.6, label = "v_star")
 plot!(plt, legend = :bottomright)
 ```
 
@@ -656,20 +615,20 @@ from an arbitrary initial condition.
 The initial condition we'll start with is $w(y) = 5 \ln (y)$
 
 ```{code-cell} julia
-w = 5 * log.(grid_y)  # An initial condition -- fairly arbitrary
+w = 5 * log.(p.y)  # An initial condition -- fairly arbitrary
 n = 35
 
-plot(xlim = (extrema(grid_y)), ylim = (-50, 10))
+plot(xlim = (extrema(p.y)), ylim = (-50, 10))
 lb = "initial condition"
-plt = plot(grid_y, w, color = :black, linewidth = 2, alpha = 0.8, label = lb)
+plt = plot(p.y, w, color = :black, linewidth = 2, alpha = 0.8, label = lb)
 for i in 1:n
-    w = T(w, grid_y, β, log, k -> k^α, shocks)
-    plot!(grid_y, w, color = RGBA(i/n, 0, 1 - i/n, 0.8), linewidth = 2, alpha = 0.6,
+    w = T(w; p).w
+    plot!(p.y, w, color = RGBA(i/n, 0, 1 - i/n, 0.8), linewidth = 2, alpha = 0.6,
           label = "")
 end
 
 lb = "true value function"
-plot!(plt, v_star, grid_y, color = :black, linewidth = 2, alpha = 0.8, label = lb)
+plot!(plt, p.y, v_star.(p.y; p), color = :black, linewidth = 2, alpha = 0.8, label = lb)
 plot!(plt, legend = :bottomright)
 ```
 
@@ -678,7 +637,7 @@ plot!(plt, legend = :bottomright)
 tags: [remove-cell]
 ---
 @testset begin
-    #test v_star(grid_y[2]) ≈ -33.370496456772266
+    #test v_star(grid_y[2];OptimalGrowthModel()) ≈ -33.370496456772266
 end
 ```
 
@@ -694,25 +653,32 @@ We are clearly getting closer.
 We can write a function that computes the exact fixed point
 
 ```{code-cell} julia
-function solve_optgrowth(initial_w; tol = 1e-6, max_iter = 500)
-    fixedpoint(w -> T(w, grid_y, β, u, f, shocks), initial_w).zero # gets returned
+function solve_optgrowth(initial_w; p, iterations = 500, m = 3, show_trace = false) 
+    results = fixedpoint(w -> T(w;p).w, initial_w; iterations, m, show_trace) # Anderson iteration
+    v_star = results.zero
+    σ = T(results.zero;p).σ
+    return (;v_star, σ, results)
 end
 ```
 
 We can check our result by plotting it against the true value
 
 ```{code-cell} julia
-initial_w = 5 * log.(grid_y)
-v_star_approx = solve_optgrowth(initial_w)
+initial_w = 5 * log.(p.y)
+sol = solve_optgrowth(initial_w;p)
+v_star_approx = sol.v_star
+println("Converged in $(sol.results.iterations) to an ||residuals||_∞ of $(sol.results.residual_norm)")
 
 plt = plot(ylim = (-35, -24))
-plot!(plt, grid_y, v_star_approx, linewidth = 2, alpha = 0.6,
+plot!(plt, p.y, v_star_approx, linewidth = 2, alpha = 0.6,
       label = "approximate value function")
-plot!(plt, v_star, grid_y, linewidth = 2, alpha = 0.6, label = "true value function")
+plot!(plt, p.y, v_star.(p.y;p), linewidth = 2, alpha = 0.6, label = "true value function")
 plot!(plt, legend = :bottomright)
 ```
 
-The figure shows that we are pretty much on the money.
+The figure shows that we are pretty much on the money
+
+Note that this converges in fewer than the 36 iterations printed above because it is using Anderson iteration - where the $m=0$ parameter is naive fixed point iteration 
 
 ### The Policy Function
 
@@ -726,12 +692,8 @@ The next figure compares the result to the exact solution, which, as mentioned
 above, is $\sigma(y) = (1 - \alpha \beta) y$.
 
 ```{code-cell} julia
-Tw, σ = T(v_star_approx, grid_y, β, log, k -> k^α, shocks;
-                         compute_policy = true)
-cstar = (1 - α * β) * grid_y
-
-plt = plot(grid_y, σ, lw=2, alpha=0.6, label = "approximate policy function")
-plot!(plt, grid_y, cstar, lw = 2, alpha = 0.6, label = "true policy function")
+plt = plot(p.y, T(v_star_approx; p).σ, lw=2, alpha=0.6, label = "approximate policy function")
+plot!(plt, p.y, c_star.(p.y; p), lw = 2, alpha = 0.6, label = "true policy function")
 plot!(plt, legend = :bottomright)
 ```
 
@@ -740,7 +702,7 @@ plot!(plt, legend = :bottomright)
 tags: [remove-cell]
 ---
 @testset begin
-    #test cstar[102] ≈ 1.2505758978894472
+    #test c_star.(p.y; p)[102] ≈ 1.2505758978894472
 end
 ```
 
@@ -776,17 +738,8 @@ Random.seed!(42);
 ---
 tags: [hide-output]
 ---
-s = 0.05
-shocks = exp.(μ .+ s * randn(shock_size))
-```
-
-```{code-cell} julia
----
-tags: [remove-cell]
----
-@testset begin
-    #test shocks[25] ≈ 0.8281356207080574
-end
+p = OptimalGrowthModel(s = 0.05)
+p.ξ
 ```
 
 Otherwise, the parameters and primitives are the same as the log linear model discussed earlier in the lecture.
@@ -802,57 +755,28 @@ Replicate the figure modulo randomness.
 Here's one solution (assuming as usual that you've executed everything above)
 
 ```{code-cell} julia
-function simulate_og(σ, y0 = 0.1, ts_length=100)
+function simulate_og(σ, p, y0, ts_length)
     y = zeros(ts_length)
-    ξ = randn(ts_length-1)
     y[1] = y0
     for t in 1:(ts_length-1)
-        y[t+1] = (y[t] - σ(y[t]))^α * exp(μ + s * ξ[t])
+        y[t+1] = (y[t] - σ(y[t]))^p.α * exp(p.μ + p.s * randn())
     end
     return y
 end
 
+β_vals = [0.9 0.94 0.98]
+ts_length = 100
+y0 = 0.1
 plt = plot()
 
-for β in (0.9, 0.94, 0.98)
-    Tw = similar(grid_y)
-    initial_w = 5 * log.(grid_y)
-
-    v_star_approx = fixedpoint(w -> T(w, grid_y, β, u, f, shocks),
-                               initial_w).zero
-    Tw, σ = T(v_star_approx, grid_y, β, log, k -> k^α, shocks,
-                             compute_policy = true)
-    σ_func = LinearInterpolation(grid_y, σ)
-    y = simulate_og(σ_func)
+for β in β_vals
+    p = OptimalGrowthModel(;β) # change β from default
+    initial_w = 5 * log.(p.y)
+    sol = solve_optgrowth(initial_w;p)
+    σ_func = LinearInterpolation(p.y, sol.σ)
+    y = simulate_og(σ_func, p,y0, ts_length)
 
-    plot!(plt, y, lw = 2, alpha = 0.6, label = label = "beta = $β")
+    plot!(plt, 0:(ts_length-1), y, lw = 2, alpha = 0.6, label = "beta = $β")
 end
-
-plot!(plt, legend = :bottomright)
-```
-
-```{code-cell} julia
----
-tags: [remove-cell]
----
-using Random
-Random.seed!(42)
-β = 0.94
-
-Tw = similar(grid_y)
-initial_w = 5 * log.(grid_y)
-v_star_approx = fixedpoint(w -> T(w, grid_y, β, u, f, shocks),
-                           initial_w).zero
-Tw, σ = T(v_star_approx, grid_y, β, log, k -> k^α, shocks,
-                         compute_policy = true)
-σ_func = LinearInterpolation(grid_y, σ)
-y = simulate_og(σ_func);
-
-@testset begin
-    ##test y[5] ≈ 0.4896390574930345
-    ##test σ[3] ≈ 0.025084788719798193 atol = 1e-6
-    ##test Tw[4] ≈ -22.24681487036426
-    ##test v_star_approx[50] ≈ -17.76952387641302
-end
-```
-
+plt
+```