[WIP] Add Pade approximation for higher order solutions

src/MacroModelling.jl

@@ -550,7 +550,7 @@ function normalize_filtering_options(filter::Symbol,
 
     pruning = algorithm ∈ (:pruned_second_order, :pruned_third_order)
 
-    if shock_decomposition && algorithm ∈ (:second_order, :third_order)
+    if shock_decomposition && algorithm ∈ (:second_order, :third_order, :second_order_pade, :third_order_pade)
         @info "Shock decomposition is not available for $(algorithm) solutions, but is available for first order, pruned second order, and pruned third order solutions. Setting `shock_decomposition = false`." maxlog = maxlog
         shock_decomposition = false
     end
@@ -6963,12 +6963,12 @@ function solve!(𝓂::ℳ;
     # end # timeit_debug
 
     if 𝓂.solution.perturbation.second_order_auxiliary_matrices.𝛔 == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && 
-        algorithm ∈ [:second_order, :pruned_second_order]
+        algorithm ∈ [:second_order, :pruned_second_order, :second_order_pade]
         start_time = time()
         if !silent print("Take symbolic derivatives up to second order:\t\t\t\t") end
         write_functions_mapping!(𝓂, 2)
         if !silent println(round(time() - start_time, digits = 3), " seconds") end
-    elseif 𝓂.solution.perturbation.third_order_auxiliary_matrices.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && algorithm ∈ [:third_order, :pruned_third_order]
+    elseif 𝓂.solution.perturbation.third_order_auxiliary_matrices.𝐂₃ == SparseMatrixCSC{Int, Int64}(ℒ.I,0,0) && algorithm ∈ [:third_order, :pruned_third_order, :third_order_pade]
         start_time = time()
         if !silent print("Take symbolic derivatives up to third order:\t\t\t\t") end
         write_functions_mapping!(𝓂, 3)
@@ -6981,7 +6981,9 @@ function solve!(𝓂::ℳ;
             ((:second_order        == algorithm) && ((:second_order        ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved))) ||
             ((:pruned_second_order == algorithm) && ((:pruned_second_order ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved))) ||
             ((:third_order         == algorithm) && ((:third_order         ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved))) ||
-            ((:pruned_third_order  == algorithm) && ((:pruned_third_order  ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved)))
+            ((:pruned_third_order  == algorithm) && ((:pruned_third_order  ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved))) ||
+            ((:second_order_pade   == algorithm) && ((:second_order_pade   ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved))) ||
+            ((:third_order_pade    == algorithm) && ((:third_order_pade    ∈ 𝓂.solution.outdated_algorithms) || (obc && obc_not_solved)))
 
             # @timeit_debug timer "Solve for NSSS (if necessary)" begin
 
@@ -7183,6 +7185,129 @@ function solve!(𝓂::ℳ;
 
             𝓂.solution.outdated_algorithms = setdiff(𝓂.solution.outdated_algorithms,[:pruned_third_order])
         end
+
+        # Second order Padé approximation (unpruned)
+        # Check if solution matrices are empty (not computed yet) or algorithm is outdated
+        pade_2_not_solved = isempty(𝓂.solution.perturbation.second_order_pade.stochastic_steady_state) || 
+                           size(𝓂.solution.perturbation.second_order_pade.𝐒₁, 1) == 0
+        pade_2_obc_not_solved = obc && size(𝓂.solution.perturbation.second_order_pade.𝐒₁̂, 1) == 0
+
+        if  ((:second_order_pade  == algorithm) && ((:second_order_pade   ∈ 𝓂.solution.outdated_algorithms) || pade_2_not_solved || pade_2_obc_not_solved)) ||
+            ((:third_order_pade  == algorithm) && ((:third_order_pade   ∈ 𝓂.solution.outdated_algorithms) || pade_2_not_solved || pade_2_obc_not_solved))
+
+            stochastic_steady_state, converged, SS_and_pars, solution_error, ∇₁, ∇₂, 𝐒₁, 𝐒₂ = calculate_second_order_stochastic_steady_state(𝓂.parameter_values, 𝓂, opts = opts)
+
+            if !converged  @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end
+
+            # Compute [1,1] Padé coefficient matrix using proper Padé theory
+            # 
+            # For a scalar function f(z) = a₁z + a₂z² + O(z³), the [1,1] Padé is:
+            #   P[1,1](z) = a₁z / (1 - (a₂/a₁)z)
+            # 
+            # For DSGE perturbation solutions, we have vector-valued output and:
+            #   y = S₁*x + (1/2)*S₂*(x⊗x) + O(|x|³)
+            # 
+            # The Padé approximant uses the form:
+            #   y_i = (S₁_i * x) / (1 - d_i * (x⊗x))
+            # 
+            # To match the Taylor expansion to second order, expanding 1/(1-d*z) ≈ 1 + d*z:
+            #   y_i ≈ S₁_i*x * (1 + d_i*(x⊗x)) = S₁_i*x + S₁_i*x*d_i*(x⊗x)
+            #
+            # We want this to match: y_i = S₁_i*x + (1/2)*S₂_i*(x⊗x)
+            # 
+            # For proper matching, we compute d_i element-wise such that the correction
+            # term captures the quadratic contribution relative to the linear response.
+            # The coefficient d_i scales S₂ by the magnitude of the linear response.
+
+            nVars = size(𝐒₁, 1)
+            nAug = size(𝐒₁, 2)
+            nAug² = size(𝐒₂, 2)
+
+            # Compute Padé denominator coefficient matrix with proper normalization
+            # Use the sum of squared linear coefficients as a robust scale factor
+            𝐃 = zeros(eltype(𝐒₂), nVars, nAug²)
+            for i in 1:nVars
+                row_norm_sq = sum(abs2, 𝐒₁[i, :])
+                if row_norm_sq > 1e-14
+                    # Scale factor ensures bounded behavior
+                    𝐃[i, :] = 𝐒₂[i, :] / (2 * row_norm_sq)
+                end
+            end
+
+            # Store OBC matrix and its Padé coefficients if needed
+            if obc
+                𝐒₁̂ = [Ŝ₁[:,1:𝓂.timings.nPast_not_future_and_mixed] zeros(𝓂.timings.nVars) Ŝ₁[:,𝓂.timings.nPast_not_future_and_mixed+1:end]]
+                𝐃̂ = zeros(eltype(𝐒₂), nVars, nAug²)
+                for i in 1:nVars
+                    row_norm_sq = sum(abs2, 𝐒₁̂[i, :])
+                    if row_norm_sq > 1e-14
+                        𝐃̂[i, :] = 𝐒₂[i, :] / (2 * row_norm_sq)
+                    end
+                end
+            else
+                𝐒₁̂ = zeros(0, 0)
+                𝐃̂ = zeros(0, 0)
+            end
+
+            # Store the Padé solution with precomputed coefficient matrices
+            𝓂.solution.perturbation.second_order_pade = second_order_pade_solution(stochastic_steady_state, 𝐒₁, 𝐒₂, 𝐃, 𝐒₁̂, 𝐃̂)
+
+            𝓂.solution.outdated_algorithms = setdiff(𝓂.solution.outdated_algorithms,[:second_order_pade])
+        end
+
+        # Third order Padé approximation (unpruned)
+        pade_3_not_solved = isempty(𝓂.solution.perturbation.third_order_pade.stochastic_steady_state) ||
+                           size(𝓂.solution.perturbation.third_order_pade.𝐒₁, 1) == 0
+        pade_3_obc_not_solved = obc && size(𝓂.solution.perturbation.third_order_pade.𝐒₁̂, 1) == 0
+
+        if  ((:third_order_pade  == algorithm) && ((:third_order_pade   ∈ 𝓂.solution.outdated_algorithms) || pade_3_not_solved || pade_3_obc_not_solved))
+            stochastic_steady_state, converged, SS_and_pars, solution_error, ∇₁, ∇₂, ∇₃, 𝐒₁, 𝐒₂, 𝐒₃ = calculate_third_order_stochastic_steady_state(𝓂.parameter_values, 𝓂, opts = opts)
+
+            if !converged  @warn "Solution does not have a stochastic steady state. Try reducing shock sizes by multiplying them with a number < 1." end
+
+            # Compute [2,1] matrix Padé coefficients for third order
+            # Taylor series: yᵢ = S₁ᵢ*x + S₂ᵢ*(x⊗x)/2 + S₃ᵢ*(x⊗x⊗x)/6
+            # Matrix Padé [2,1]: yᵢ = (S₁ᵢ*x + S₂ᵢ*(x⊗x)/2) / (1 - 𝐃₂ᵢ*(x⊗x) - 𝐃₃ᵢ*(x⊗x⊗x))
+
+            nVars = size(𝐒₁, 1)
+            nAug = size(𝐒₁, 2)
+            nAug² = size(𝐒₂, 2)
+            nAug³ = size(𝐒₃, 2)
+
+            # Compute row-normalized Padé denominator coefficient matrices
+            𝐃₂ = zeros(eltype(𝐒₂), nVars, nAug²)
+            𝐃₃ = zeros(eltype(𝐒₃), nVars, nAug³)
+            for i in 1:nVars
+                row_norm_sq = sum(abs2, 𝐒₁[i, :]) + sum(abs2, 𝐒₂[i, :]) / 4
+                if row_norm_sq > 1e-14
+                    𝐃₂[i, :] = 𝐒₂[i, :] / (2 * row_norm_sq)
+                    𝐃₃[i, :] = 𝐒₃[i, :] / (6 * row_norm_sq)
+                end
+            end
+
+            # Store OBC matrices and Padé coefficients if needed
+            if obc
+                𝐒₁̂ = [Ŝ₁[:,1:𝓂.timings.nPast_not_future_and_mixed] zeros(𝓂.timings.nVars) Ŝ₁[:,𝓂.timings.nPast_not_future_and_mixed+1:end]]
+                𝐃₂̂ = zeros(eltype(𝐒₂), nVars, nAug²)
+                𝐃₃̂ = zeros(eltype(𝐒₃), nVars, nAug³)
+                for i in 1:nVars
+                    row_norm_sq = sum(abs2, 𝐒₁̂[i, :]) + sum(abs2, 𝐒₂[i, :]) / 4
+                    if row_norm_sq > 1e-14
+                        𝐃₂̂[i, :] = 𝐒₂[i, :] / (2 * row_norm_sq)
+                        𝐃₃̂[i, :] = 𝐒₃[i, :] / (6 * row_norm_sq)
+                    end
+                end
+            else
+                𝐒₁̂ = zeros(0, 0)
+                𝐃₂̂ = zeros(0, 0)
+                𝐃₃̂ = zeros(0, 0)
+            end
+
+            # Store the Padé solution with precomputed matrices
+            𝓂.solution.perturbation.third_order_pade = third_order_pade_solution(stochastic_steady_state, 𝐒₁, 𝐒₂, 𝐒₃, 𝐃₂, 𝐃₃, 𝐒₁̂, 𝐃₂̂, 𝐃₃̂)
+
+            𝓂.solution.outdated_algorithms = setdiff(𝓂.solution.outdated_algorithms,[:third_order_pade])
+        end
     end
 
     return nothing
@@ -9692,6 +9817,75 @@ function parse_algorithm_to_state_update(algorithm::Symbol, 𝓂::ℳ, occasiona
         elseif :pruned_third_order == algorithm
             state_update = 𝓂.solution.perturbation.pruned_third_order.state_update_obc
             pruning = true
+        elseif :second_order_pade == algorithm
+            # Matrix Padé [1,1] approximation for OBC
+            pade_sol = 𝓂.solution.perturbation.second_order_pade
+            𝐒₁̂ = pade_sol.𝐒₁̂
+            𝐒₂ = pade_sol.𝐒₂
+            past_not_future_and_mixed_idx = 𝓂.timings.past_not_future_and_mixed_idx
+            state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S}
+                aug_state = [state[past_not_future_and_mixed_idx]; 1; shock]
+                linear_term = 𝐒₁̂ * aug_state
+                kron_aug = ℒ.kron(aug_state, aug_state)
+                quadratic_term = 𝐒₂ * kron_aug / 2
+
+                # Compute Padé [1,1] approximation with stability bounds
+                result = similar(linear_term)
+                for i in eachindex(linear_term)
+                    L = linear_term[i]
+                    Q = quadratic_term[i]
+
+                    if abs(L) < 1e-12
+                        result[i] = L + Q
+                    else
+                        ratio = Q / L
+                        if abs(ratio) < 0.9
+                            result[i] = L / (1 - ratio)
+                        else
+                            result[i] = L + Q
+                        end
+                    end
+                end
+                return result
+            end
+            pruning = false
+        elseif :third_order_pade == algorithm
+            # Matrix Padé [2,1] approximation for OBC
+            pade_sol = 𝓂.solution.perturbation.third_order_pade
+            𝐒₁̂ = pade_sol.𝐒₁̂
+            𝐒₂ = pade_sol.𝐒₂
+            𝐒₃ = pade_sol.𝐒₃
+            past_not_future_and_mixed_idx = 𝓂.timings.past_not_future_and_mixed_idx
+            state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S}
+                aug_state = [state[past_not_future_and_mixed_idx]; 1; shock]
+                linear_term = 𝐒₁̂ * aug_state
+                kron_aug = ℒ.kron(aug_state, aug_state)
+                kron_aug3 = ℒ.kron(kron_aug, aug_state)
+                quadratic_term = 𝐒₂ * kron_aug / 2
+                cubic_term = 𝐒₃ * kron_aug3 / 6
+
+                # Compute Padé [2,1] approximation with stability bounds
+                result = similar(linear_term)
+                numerator = linear_term + quadratic_term
+
+                for i in eachindex(linear_term)
+                    N = numerator[i]
+                    C = cubic_term[i]
+
+                    if abs(N) < 1e-12
+                        result[i] = N + C
+                    else
+                        ratio = C / N
+                        if abs(ratio) < 0.9
+                            result[i] = N / (1 - ratio)
+                        else
+                            result[i] = N + C
+                        end
+                    end
+                end
+                return result
+            end
+            pruning = false
         else
             # @assert false "Provided algorithm not valid. Valid algorithm: $all_available_algorithms"
             state_update = (x,y)->nothing
@@ -9713,6 +9907,92 @@ function parse_algorithm_to_state_update(algorithm::Symbol, 𝓂::ℳ, occasiona
         elseif :pruned_third_order == algorithm
             state_update = 𝓂.solution.perturbation.pruned_third_order.state_update
             pruning = true
+        elseif :second_order_pade == algorithm
+            # Matrix Padé [1,1] approximation using precomputed coefficient matrices
+            # Padé form: y = S₁*x / (1 - d*(x⊗x)) where d captures the quadratic term ratio
+            # This extends better than Taylor when |x| is moderate
+            pade_sol = 𝓂.solution.perturbation.second_order_pade
+            𝐒₁ = pade_sol.𝐒₁
+            𝐒₂ = pade_sol.𝐒₂
+            past_not_future_and_mixed_idx = 𝓂.timings.past_not_future_and_mixed_idx
+            state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S}
+                aug_state = [state[past_not_future_and_mixed_idx]; 1; shock]
+                linear_term = 𝐒₁ * aug_state
+                kron_aug = ℒ.kron(aug_state, aug_state)
+                quadratic_term = 𝐒₂ * kron_aug / 2
+
+                # Compute Padé [1,1] approximation for each element
+                # Standard form: y_i = L_i / (1 - Q_i/L_i) when |Q_i/L_i| < 1
+                # This is equivalent to: y_i = L_i² / (L_i - Q_i)
+                # For stability, we bound the correction and fall back to Taylor when needed
+                result = similar(linear_term)
+                for i in eachindex(linear_term)
+                    L = linear_term[i]
+                    Q = quadratic_term[i]
+
+                    if abs(L) < 1e-12
+                        # Linear term is tiny, use Taylor directly
+                        result[i] = L + Q
+                    else
+                        ratio = Q / L
+                        if abs(ratio) < 0.9  # Safe convergence region
+                            # Padé: y = L / (1 - Q/L) = L / (1 - ratio)
+                            result[i] = L / (1 - ratio)
+                        else
+                            # Outside safe region, use Taylor
+                            result[i] = L + Q
+                        end
+                    end
+                end
+                return result
+            end
+            pruning = false
+        elseif :third_order_pade == algorithm
+            # Matrix Padé [2,1] approximation using precomputed coefficient matrices
+            # Taylor: y = L + Q/2 + C/6 (linear + quadratic + cubic)
+            # Padé [2,1]: y = (L + Q/2) / (1 - correction) matching cubic term
+            pade_sol = 𝓂.solution.perturbation.third_order_pade
+            𝐒₁ = pade_sol.𝐒₁
+            𝐒₂ = pade_sol.𝐒₂
+            𝐒₃ = pade_sol.𝐒₃
+            past_not_future_and_mixed_idx = 𝓂.timings.past_not_future_and_mixed_idx
+            state_update = function(state::Vector{T}, shock::Vector{S}) where {T,S}
+                aug_state = [state[past_not_future_and_mixed_idx]; 1; shock]
+                linear_term = 𝐒₁ * aug_state
+                kron_aug = ℒ.kron(aug_state, aug_state)
+                kron_aug3 = ℒ.kron(kron_aug, aug_state)
+                quadratic_term = 𝐒₂ * kron_aug / 2
+                cubic_term = 𝐒₃ * kron_aug3 / 6
+
+                # Compute Padé [2,1] approximation for each element
+                # Numerator: L + Q (linear + quadratic)
+                # We want the expansion to match L + Q + C (Taylor up to cubic)
+                # Padé form: (L + Q) / (1 - d) ≈ (L + Q)(1 + d) = L + Q + (L+Q)*d
+                # Match: (L+Q)*d = C => d = C / (L + Q)
+                result = similar(linear_term)
+                numerator = linear_term + quadratic_term
+
+                for i in eachindex(linear_term)
+                    N = numerator[i]
+                    C = cubic_term[i]
+
+                    if abs(N) < 1e-12
+                        # Numerator is tiny, use Taylor directly
+                        result[i] = N + C
+                    else
+                        ratio = C / N
+                        if abs(ratio) < 0.9  # Safe convergence region
+                            # Padé [2,1]: y = N / (1 - C/N) = N / (1 - ratio)
+                            result[i] = N / (1 - ratio)
+                        else
+                            # Outside safe region, use Taylor
+                            result[i] = N + C
+                        end
+                    end
+                end
+                return result
+            end
+            pruning = false
         else
             # @assert false "Provided algorithm not valid. Valid algorithm: $all_available_algorithms"
             state_update = (x,y)->nothing
@@ -10270,6 +10550,24 @@ function get_relevant_steady_state_and_state_update(::Val{:first_order},
     return TT, SS_and_pars, 𝐒₁, [state], solved
 end
 
+# Padé versions of get_relevant_steady_state_and_state_update
+# These use the same stochastic steady state calculations as their Taylor counterparts
+function get_relevant_steady_state_and_state_update(::Val{:second_order_pade}, 
+                                                    parameter_values::Vector{S}, 
+                                                    𝓂::ℳ; 
+                                                    opts::CalculationOptions = merge_calculation_options()) where S <: Real
+    # Use the same calculation as second_order
+    return get_relevant_steady_state_and_state_update(Val(:second_order), parameter_values, 𝓂, opts = opts)
+end
+
+function get_relevant_steady_state_and_state_update(::Val{:third_order_pade}, 
+                                                    parameter_values::Vector{S}, 
+                                                    𝓂::ℳ; 
+                                                    opts::CalculationOptions = merge_calculation_options())::Tuple{timings, Vector{S}, Union{Matrix{S},Vector{AbstractMatrix{S}}}, Vector{S}, Bool} where S <: Real
+    # Use the same calculation as third_order
+    return get_relevant_steady_state_and_state_update(Val(:third_order), parameter_values, 𝓂, opts = opts)
+end
+
 end # dispatch_doctor
 
 # @setup_workload begin