nonlinear cond fcst; pruned sol no longer updating

docs/src/unfinished_docs/todo.md

@@ -5,6 +5,7 @@
 - [ ] add technical details about SS solver, obc solver, and other algorithms
 - [ ] fix translate dynare mod file from file written using write to dynare file (see test models)
 - [ ] rm obc vars from get_SS
+- [ ] write tests/docs for nonlinear obc and forecasting
 - [ ] fix SS solver (failed for backus in guide)
 - [ ] functions to reverse state_update (input: previous shock and current state, output previous state), find shocks corresponding to bringing one state to the next
 - [ ] cover nested case: min(50,a+b+max(c,10))

src/MacroModelling.jl

@@ -303,6 +303,50 @@ function obc_objective_optim_fun(X::Vector{S}, grad::Vector{S}) where S
     sum(abs2, X)
 end
 
+
+
+function match_conditions(res::Vector{S}, X::Vector{S}, jac::Matrix{S}, p) where S
+    Conditions, State_update, Shocks, Cond_var_idx, Free_shock_idx, State, Pruning, 𝒷 = p
+
+    if length(jac) > 0
+        jac .= 𝒜.jacobian(𝒷(), xx -> begin
+                                        Shocks[Free_shock_idx] .= xx
+
+                                        new_State = State_update(State, convert(typeof(xx), Shocks))
+
+                                        cond_vars = Pruning ? sum(new_State) : new_State
+
+                                        return abs2.(Conditions[Cond_var_idx] - cond_vars[Cond_var_idx])
+                                    end, X)[1]'
+    end
+
+    Shocks[Free_shock_idx] .= X
+
+    new_State = State_update(State, convert(typeof(X), Shocks))
+
+    cond_vars = Pruning ? sum(new_State) : new_State
+
+    res .= abs2.(Conditions[Cond_var_idx] - cond_vars[Cond_var_idx])
+end
+
+
+
+# function match_conditions(res::Vector{S}, X::Vector{S}, jac::Matrix{S}, p) where S
+#     conditions, state_update, shocks, cond_var_idx, free_shock_idx, state, 𝒷 = p
+
+#     if length(jac) > 0
+#         jac .= 𝒜.jacobian(𝒷(), xx -> begin
+#                                         shocks[free_shock_idx] .= xx
+#                                         return abs2.(conditions[cond_var_idx] - state_update(state, convert(typeof(xx), shocks))[cond_var_idx])
+#                                     end, X)[1]'
+#     end
+
+#     shocks[free_shock_idx] .= X
+
+#     res .= abs2.(conditions[cond_var_idx] - state_update(state, convert(typeof(X), shocks))[cond_var_idx])
+# end
+
+
 function set_up_obc_violation_function!(𝓂)
     present_varss = collect(reduce(union,match_pattern.(get_symbols.(𝓂.dyn_equations),r"₍₀₎$")))
 
@@ -3264,10 +3308,7 @@ function solve!(𝓂::ℳ;
                 aug_state₁ = [pruned_states[1][𝓂.timings.past_not_future_and_mixed_idx]; 1; shock]
                 aug_state₂ = [pruned_states[2][𝓂.timings.past_not_future_and_mixed_idx]; 0; zero(shock)]
 
-                pruned_states[1] .= 𝐒₁ * aug_state₁
-                pruned_states[2] .= 𝐒₁ * aug_state₂ + 𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2
-
-                return pruned_states[1] + pruned_states[2] # strictly following Andreasen et al. (2018)
+                return [𝐒₁ * aug_state₁, 𝐒₁ * aug_state₂ + 𝐒₂ * ℒ.kron(aug_state₁, aug_state₁) / 2] # strictly following Andreasen et al. (2018)
             end
 
             if obc
@@ -3334,11 +3375,7 @@ function solve!(𝓂::ℳ;
 
                 kron_aug_state₁ = ℒ.kron(aug_state₁, aug_state₁)
 
-                pruned_states[1] .= 𝐒₁ * aug_state₁
-                pruned_states[2] .= 𝐒₁ * aug_state₂ + 𝐒₂ * kron_aug_state₁ / 2
-                pruned_states[3] .= 𝐒₁ * aug_state₃ + 𝐒₂ * ℒ.kron(aug_state₁̂, aug_state₂) + 𝐒₃ * ℒ.kron(kron_aug_state₁,aug_state₁) / 6
-
-                return pruned_states[1] + pruned_states[2] + pruned_states[3]
+                return [𝐒₁ * aug_state₁, 𝐒₁ * aug_state₂ + 𝐒₂ * kron_aug_state₁ / 2, 𝐒₁ * aug_state₃ + 𝐒₂ * ℒ.kron(aug_state₁̂, aug_state₂) + 𝐒₃ * ℒ.kron(kron_aug_state₁,aug_state₁) / 6]
             end
 
             if obc
@@ -4939,22 +4976,30 @@ function irf(state_update::Function,
 
         Y = zeros(T.nVars,periods,1)
 
-        Y[:,1,1] = state_update(initial_state,shock_history[:,1])
+        initial_state = state_update(initial_state,shock_history[:,1])
+
+        Y[:,1,1] = pruning ? sum(initial_state) : initial_state
 
         for t in 1:periods-1
-            Y[:,t+1,1] = state_update(pruning ? initial_state : Y[:,t,1],shock_history[:,t+1])
+            initial_state = state_update(initial_state,shock_history[:,t+1])
+
+            Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state
         end
 
         return KeyedArray(Y[var_idx,:,:] .+ level[var_idx];  Variables = axis1, Periods = 1:periods, Shocks = [:simulate])
     elseif shocks == :none
         Y = zeros(T.nVars,periods,1)
 
         shck = T.nExo == 0 ? Vector{Float64}(undef, 0) : zeros(T.nExo)
-
-        Y[:,1,1] = state_update(initial_state,shck)
+
+        initial_state = state_update(initial_state, shck)
+
+        Y[:,1,1] = pruning ? sum(initial_state) : initial_state
 
         for t in 1:periods-1
-            Y[:,t+1,1] = state_update(pruning ? initial_state : Y[:,t,1], shck)
+            initial_state = state_update(initial_state, shck)
+
+            Y[:,t+1,1] = pruning ? sum(initial_state) : initial_state
         end
 
         return KeyedArray(Y[var_idx,:,:] .+ level[var_idx];  Variables = axis1, Periods = 1:periods, Shocks = [:none])
@@ -4969,10 +5014,14 @@ function irf(state_update::Function,
                 shock_history[ii,1] = negative_shock ? -1 : 1
             end
 
-            Y[:,1,i] = state_update(initial_state_copy, shock_history[:,1])
+            initial_state_copy = state_update(initial_state_copy, shock_history[:,1])
+
+            Y[:,1,1] = pruning ? sum(initial_state_copy) : initial_state_copy
 
             for t in 1:periods-1
-                Y[:,t+1,i] = state_update(pruning ? initial_state_copy : Y[:,t,i], shock_history[:,t+1])
+                initial_state_copy = state_update(initial_state_copy, shock_history[:,t+1])
+
+                Y[:,t+1,1] = pruning ? sum(initial_state_copy) : initial_state_copy
             end
         end
 
@@ -5043,11 +5092,7 @@ function girf(state_update::Function,
             initial_state_copy² = deepcopy(initial_state_copy)
 
             for i in 1:warmup_periods
-                if pruning
-                    state_update(initial_state_copy², randn(T.nExo))
-                else
-                    initial_state_copy² = state_update(initial_state_copy², randn(T.nExo))
-                end
+                initial_state_copy² = state_update(initial_state_copy², randn(T.nExo))
             end
 
             Y₁ = zeros(T.nVars, periods + 1)
@@ -5061,11 +5106,13 @@ function girf(state_update::Function,
             end
 
             if pruning
-                Y₁[:,1] = state_update(initial_state_copy², baseline_noise)
-                Y₂[:,1] = state_update(initial_state_copy², baseline_noise)
+                initial_state_copy² = state_update(initial_state_copy², baseline_noise)
 
                 initial_state₁ = deepcopy(initial_state_copy²)
                 initial_state₂ = deepcopy(initial_state_copy²)
+
+                Y₁[:,1] = initial_state_copy² |> sum
+                Y₂[:,1] = initial_state_copy² |> sum
             else
                 Y₁[:,1] = state_update(initial_state_copy², baseline_noise)
                 Y₂[:,1] = state_update(initial_state_copy², baseline_noise)
@@ -5075,8 +5122,11 @@ function girf(state_update::Function,
                 baseline_noise = randn(T.nExo)
 
                 if pruning
-                    Y₁[:,t+1] = state_update(initial_state₁, baseline_noise)
-                    Y₂[:,t+1] = state_update(initial_state₂, baseline_noise + shock_history[:,t])
+                    initial_state₁ = state_update(initial_state₁, baseline_noise)
+                    initial_state₂ = state_update(initial_state₂, baseline_noise + shock_history[:,t])
+
+                    Y₁[:,t+1] = initial_state₁ |> sum
+                    Y₂[:,t+1] = initial_state₂ |> sum
                 else
                     Y₁[:,t+1] = state_update(Y₁[:,t],baseline_noise)
                     Y₂[:,t+1] = state_update(Y₂[:,t],baseline_noise + shock_history[:,t])