Refactor optimization test script to streamline steady state calculations and enhance residual function compatibility with NonlinearSolve.jl

Author: Thore Kockerols

test_dynamic_equations_optimization.jl

@@ -30,87 +30,50 @@ end
     β = 0.95
 end
 
-println("Solving model...")
-solve!(RBC)
 
-# Get steady states
-println("Computing steady states...")
-SS_non_stochastic = get_steady_state(RBC, stochastic = false)
-SS_stochastic = get_steady_state(RBC, stochastic = true)
-
-println("Non-stochastic steady state:")
-println(SS_non_stochastic)
-println("\nStochastic steady state:")
-println(SS_stochastic)
-
-# Get auxiliary indices for the steady state input
-aux_idx = RBC.solution.perturbation.auxiliary_indices
-SS_for_func = SS_non_stochastic[aux_idx.dyn_ss_idx]
-
-# Number of variables and shocks
-n_vars = length(RBC.var)
-n_shocks = length(RBC.exo)
-n_dyn_eqs = length(RBC.dyn_equations)
-n_calib_eqs = length(RBC.calibration_equations)
-n_total_eqs = n_dyn_eqs + n_calib_eqs
-
-println("\nModel dimensions:")
-println("  Variables: ", n_vars)
-println("  Shocks: ", n_shocks)
-println("  Dynamic equations: ", n_dyn_eqs)
-println("  Calibration equations: ", n_calib_eqs)
-println("  Total equations: ", n_total_eqs)
-
-# Generate Sobol sequence for shocks
-n_draws = 100
-println("\nGenerating ", n_draws, " Sobol draws for shocks...")
-
-# Generate Sobol sequence in [0,1]
-shock_draws_uniform = QuasiMonteCarlo.sample(n_draws, n_shocks, SobolSample())
+include("./models/Gali_2015_chapter_3_nonlinear.jl")
 
-# Convert to normal distribution using inverse error function
-normal_draws = @. sqrt(2) * erfinv(2 * shock_draws_uniform - 1)
+m = Gali_2015_chapter_3_nonlinear
 
-# Standardize to zero mean and unit variance
-normal_draws .-= mean(normal_draws, dims=2)
-normal_draws ./= Statistics.std(normal_draws, dims=2)
+# Get steady states
+SS_non_stochastic = get_steady_state(m, stochastic = false, derivatives = false) |> collect
+SS_stochastic = get_steady_state(m, stochastic = true, derivatives = false) |> collect
 
-println("Shock draws shape: ", size(normal_draws))
-println("Shock draws mean: ", mean(normal_draws, dims=2))
-println("Shock draws std: ", Statistics.std(normal_draws, dims=2))
 
 # Calibration parameters
-calib_params = zeros(length(RBC.calibration_equations_parameters))
+calib_params = zeros(length(m.calibration_equations_parameters))
 
 # Define residual function that returns per-equation average residuals
 # This is compatible with NonlinearSolve.jl
+# function residual_function(vars_flat, p)
 function residual_function!(residual_avg, vars_flat, p)
-    shock_draws, model, SS_for_func, calib_params = p
+    shock_draws, model, SS_non_stochastic, calib_params = p
 
     n_draws = size(shock_draws, 2)
-    n_vars = length(model.var)
+    # n_vars = length(model.var)
 
     # Unpack variables: [past; present; future]
-    past = vars_flat[1:n_vars]
-    present = vars_flat[n_vars+1:2*n_vars]
-    future = vars_flat[2*n_vars+1:3*n_vars]
+    past = vars_flat#[1:n_vars]
+    present = vars_flat#[n_vars+1:2*n_vars]
+    future = vars_flat#[2*n_vars+1:3*n_vars]
 
     n_eqs = length(model.dyn_equations) + length(model.calibration_equations)
-    residual_temp = zeros(n_eqs)
+    residual_temp = zeros(eltype(vars_flat), n_eqs)
 
     # Initialize average residuals to zero
     fill!(residual_avg, 0.0)
+    # residual_avg = zeros(eltype(vars_flat), n_eqs)
 
     # Sum residuals across all shock draws
     for i in 1:n_draws
         shocks = shock_draws[:, i]
 
         # Evaluate dynamic equations
         get_dynamic_residuals(residual_temp, model.parameter_values, calib_params, 
-                            past, present, future, SS_for_func, shocks, model)
-        
+                            past, present, future, SS_non_stochastic, shocks, model)
+
         # Accumulate residuals
-        residual_avg .+= residual_temp
+        residual_avg .+= abs.(residual_temp)
     end
 
     # Average over draws
@@ -119,75 +82,81 @@ function residual_function!(residual_avg, vars_flat, p)
     return nothing
 end
 
-# Initial guess: use stochastic steady state for all time periods
-println("\nSetting up nonlinear problem...")
-initial_vars = vcat(SS_stochastic, SS_stochastic, SS_stochastic)
+
+# Number of variables and shocks
+n_shocks = length(m.exo)
+
+# Generate Sobol sequence for shocks
+n_draws = 1000
+
+# Generate Sobol sequence in [0,1]
+shock_draws_uniform = QuasiMonteCarlo.sample(n_draws, n_shocks, SobolSample())
+
+# Convert to normal distribution using inverse error function
+normal_draws = @. sqrt(2) * erfinv(2 * shock_draws_uniform - 1)
+
+# Standardize to zero mean and unit variance
+normal_draws .-= mean(normal_draws, dims=2)
+normal_draws ./= Statistics.std(normal_draws, dims=2)
 
 # Package parameters for the residual function
-params = (normal_draws, RBC, SS_for_func, calib_params)
+params = (normal_draws, m, SS_non_stochastic, calib_params)
+
+# Initial guess: use stochastic steady state for all time periods
+initial_vars = copy(SS_stochastic)
 
 # Test initial residuals
-residual_test = zeros(n_total_eqs)
+residual_test = zeros(length(m.dyn_equations) + length(m.calibration_equations))
 residual_function!(residual_test, initial_vars, params)
-println("Initial residual norm: ", norm(residual_test))
+println("Initial residual norm: ", sum(abs2, residual_test))
 println("Initial max|residual|: ", maximum(abs.(residual_test)))
 
 # Set up nonlinear problem
-println("\nSolving nonlinear system...")
-println("This may take a moment...")
-
 prob = NonlinearProblem(residual_function!, initial_vars, params)
 
 # Solve using Newton-Raphson with automatic differentiation
-sol = solve(prob, NewtonRaphson(; autodiff = AutoFiniteDiff()))
+solLM = solve(prob, LevenbergMarquardt(), show_trace = Val(true))
+solTR = solve(prob, TrustRegion(), show_trace = Val(true))
+# sol = solve(prob, show_trace = Val(true))
+
+sum(abs, solTR.u - SS_stochastic)
+sum(abs, solLM.u - SS_stochastic)
+
+sum(abs2, solLM.resid)
+sum(abs2, solTR.resid)
+
+residual_function!(residual_test, solLM.u, params)
+sum(abs2, residual_test)
+
 
 println("\nSolution complete!")
 println("Solution return code: ", sol.retcode)
-println("Final residual norm: ", norm(sol.resid))
+println("Final residual norm: ", sum(abs2, sol.resid))
 println("Final max|residual|: ", maximum(abs.(sol.resid)))
 
 # Extract solved variables
-solved_past = sol.u[1:n_vars]
-solved_present = sol.u[n_vars+1:2*n_vars]
-solved_future = sol.u[2*n_vars+1:3*n_vars]
-
-println("\nSolved variables:")
-println("Past:    ", solved_past)
-println("Present: ", solved_present)
-println("Future:  ", solved_future)
+println("\nSolved variables:", sol.u)
 
 println("\nComparison with steady states:")
-for (i, var) in enumerate(RBC.var)
+for (i, var) in enumerate(m.var)
     println("  ", var, ":")
     println("    Non-stochastic SS: ", SS_non_stochastic[i])
     println("    Stochastic SS:     ", SS_stochastic[i])
-    println("    Solved (past):     ", solved_past[i])
-    println("    Solved (present):  ", solved_present[i])
-    println("    Solved (future):   ", solved_future[i])
+    println("    Solved:     ", sol.u[i])
 end
 
+n_eqs = length(m.dyn_equations) + length(m.calibration_equations)
+
 # Evaluate residuals at the solution for a few shock draws
 println("\nPer-draw residuals at solution for first 5 shock draws:")
-residual = zeros(n_total_eqs)
+residual = zeros(n_eqs)
 for i in 1:min(5, n_draws)
     shocks = normal_draws[:, i]
-    get_dynamic_residuals(residual, RBC.parameter_values, calib_params, 
-                        solved_past, solved_present, solved_future, SS_for_func, shocks, RBC)
+    get_dynamic_residuals(residual, m.parameter_values, calib_params, 
+                        sol.u, sol.u, sol.u, SS_non_stochastic, shocks, m)
     println("Draw ", i, ": max|residual| = ", maximum(abs.(residual)), 
             ", mean|residual| = ", sum(abs.(residual))/length(residual))
 end
 
-println("\nScript completed successfully!")
-
-# Evaluate residuals at the optimal point for a few shock draws
-println("\nResiduals at optimal point for first 5 shock draws:")
-residual = zeros(n_total_eqs)
-for i in 1:min(5, n_draws)
-    shocks = shock_draws[:, i]
-    get_dynamic_residuals(residual, RBC.parameter_values, calib_params, 
-                        opt_past, opt_present, opt_future, SS_for_func, shocks, RBC)
-    println("Draw ", i, ": max|residual| = ", maximum(abs.(residual)), 
-            ", mean|residual| = ", sum(abs.(residual))/length(residual))
-end
 
-println("\nScript completed successfully!")
+m.dyn_equations_func