Fix AMSS (#315)

* More robust optimization

* Changes to size to eliminate bars

Author: jlperla

lectures/dynamic_programming_squared/amss.md

@@ -1309,12 +1309,18 @@ function get_policies_time0(T::BellmanEquation_Recursive,
     U, Uc, Un = model.U, model.Uc, model.Un
 
     function objf(z, grad)
+        if any(isnan, z)
+            return -Inf
+        end
         c, xprime = z[1], z[2]
         n = (c + G[s0]) / Theta[s0]
         return -(U(c, n) + beta * Vf[s0](xprime))
     end
 
     function cons(z, grad)
+        if any(isnan, z)
+            return -Inf
+        end
         c, xprime, TT = z[1], z[2], z[3]
         n = (c + G[s0]) / Theta[s0]
         return -Uc(c, n) * (c - B_ - TT) - Un(c, n) * n - beta * xprime
@@ -1511,7 +1517,7 @@ sim_seq_h_plot = hcat(sim_seq_h[1:3]..., sim_seq_h[4],
 sim_bel_h_plot = hcat(sim_bel_h[1:3]..., sim_bel_h[5],
                       time_example.G[sHist_h],
                       time_example.Theta[sHist_h] .* sim_bel_h[2])
-p = plot(size = (920, 750), layout = (3, 2),
+p = plot(size = (700, 700), layout = (3, 2),
          xaxis = (0:6), grid = false, titlefont = Plots.font("sans-serif", 10))
 plot!(p, title = titles)
 for i in 1:6
@@ -1626,7 +1632,7 @@ sim_bel_plot = hcat(sim_bel[1:3]...,
                     log_example.Theta[sHist] .* sim_bel[2])
 
 #plot policies
-p = plot(size = (920, 750), layout = grid(3, 2),
+p = plot(size = (700, 700), layout = grid(3, 2),
          xaxis = (0:T), grid = false, titlefont = Plots.font("sans-serif", 10))
 labels = fill(("", ""), 6)
 labels[3] = ("Complete Market", "Incomplete Market")
@@ -1671,7 +1677,7 @@ sim_bel_long_plot = hcat(sim_bel_long[1:3]..., sim_bel_long[5],
                          log_example.G[sHist_long],
                          log_example.Theta[sHist_long] .* sim_bel_long[2])
 
-p = plot(size = (920, 750), layout = (3, 2), xaxis = (0:50:T_long),
+p = plot(size = (700, 700), layout = (3, 2), xaxis = (0:50:T_long),
          grid = false,
          titlefont = Plots.font("sans-serif", 10))
 plot!(p, title = titles)