From 71e92cd832d9e700c1fcdbe78e31307ecd6118f8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 15:29:07 +0000 Subject: [PATCH 1/2] Initial plan From db72376aac96c7c4b33d7ec6deac4b45b62f5892 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 2 Jan 2026 15:36:19 +0000 Subject: [PATCH 2/2] Implement Traub's algorithm for Newton-style solvers Co-authored-by: thorek1 <13523097+thorek1@users.noreply.github.com> --- src/algorithms/nonlinear_solver.jl | 218 +++++++++++++++++++ src/filter/find_shocks.jl | 325 +++++++++++++++++++++++++++++ 2 files changed, 543 insertions(+) diff --git a/src/algorithms/nonlinear_solver.jl b/src/algorithms/nonlinear_solver.jl index 4e2af7d8d..93a9f8090 100644 --- a/src/algorithms/nonlinear_solver.jl +++ b/src/algorithms/nonlinear_solver.jl @@ -589,6 +589,224 @@ function newton( end +function traub( + fnj::function_and_jacobian, + initial_guess::Array{T,1}, + parameters_and_solved_vars::Array{T,1}, + lower_bounds::Array{T,1}, + upper_bounds::Array{T,1}, + parameters::solver_parameters; + tol::Tolerances = Tolerances() + )::Tuple{Vector{T}, Tuple{Int, Int, T, T}} where {T <: AbstractFloat} + # Traub's method: A third-order iterative method for root-finding + # Uses two function evaluations per iteration for higher convergence order + # Formula: + # y_n = x_n - f(x_n)/f'(x_n) (Newton step) + # x_{n+1} = y_n - f(y_n)/f'(x_n) (Traub step) + + xtol = tol.NSSS_xtol + ftol = tol.NSSS_ftol + rel_xtol = tol.NSSS_rel_xtol + + iterations = 250 + transformation_level = 0 + + @assert size(lower_bounds) == size(upper_bounds) == size(initial_guess) + + current_guess = copy(initial_guess) + intermediate_guess = similar(current_guess) + guess_update = fnj.lu_buffer.b + + fnj.func(fnj.func_buffer, current_guess, parameters_and_solved_vars) + + current_residuals = fnj.func_buffer + intermediate_residuals = similar(current_residuals) + + ∇ = fnj.jac_buffer + + sol_cache = fnj.lu_buffer + + rel_xtol_reached = 1.0 + rel_ftol_reached = 1.0 + current_residuals_norm = 1.0 + guess_update_norm = 1.0 + + iters = [0, 0] + + for iter in 1:iterations + + if ∇ isa SparseMatrixCSC + ∇.nzval .= 0 + else + ∇ .= 0 + end + + # Compute Jacobian at current point + fnj.jac(∇, current_guess, parameters_and_solved_vars) + + # Compute residuals at current point + fnj.func(current_residuals, current_guess, parameters_and_solved_vars) + + finn = has_nonfinite(current_residuals) + + if finn + rel_xtol_reached = 1.0 + rel_ftol_reached = 1.0 + current_residuals_norm = 1.0 + break + end + + if current_residuals_norm < ftol || rel_xtol_reached < rel_xtol || guess_update_norm < xtol + current_guess_norm = ℒ.norm(current_guess) + + old_residuals_norm = current_residuals_norm + + current_residuals_norm = ℒ.norm(current_residuals) + + # First Traub step: compute intermediate point (Newton step) + if ∇ isa SparseMatrixCSC + sol_cache.A = ∇ + sol_cache.b = current_residuals + 𝒮.solve!(sol_cache) + guess_update .= sol_cache.u + intermediate_guess .= current_guess + ℒ.axpy!(-1, guess_update, intermediate_guess) + else + fact∇ = ℒ.lu!(∇, check = false) + try + if !ℒ.issuccess(fact∇) + fact∇ = ℒ.qr(∇, ℒ.ColumnNorm()) + end + copy!(intermediate_guess, current_residuals) + ℒ.ldiv!(fact∇, intermediate_guess) + guess_update_norm = ℒ.norm(intermediate_guess) + intermediate_guess .= current_guess .- intermediate_guess + catch + rel_xtol_reached = typemax(T) + current_residuals_norm = typemax(T) + break + end + end + + minmax!(intermediate_guess, lower_bounds, upper_bounds) + + # Compute residuals at intermediate point + fnj.func(intermediate_residuals, intermediate_guess, parameters_and_solved_vars) + + # Second Traub step: use same Jacobian with intermediate residuals + if ∇ isa SparseMatrixCSC + sol_cache.A = ∇ + sol_cache.b = intermediate_residuals + 𝒮.solve!(sol_cache) + guess_update .= sol_cache.u + else + try + copy!(guess_update, intermediate_residuals) + ℒ.ldiv!(fact∇, guess_update) + catch + rel_xtol_reached = typemax(T) + current_residuals_norm = typemax(T) + break + end + end + + guess_update_norm = ℒ.norm(guess_update) + current_guess .= intermediate_guess .- guess_update + + iters[1] += 2 # Two function evaluations per iteration + iters[2] += 1 # One Jacobian evaluation + + break + end + + current_guess_norm = ℒ.norm(current_guess) + + old_residuals_norm = current_residuals_norm + + current_residuals_norm = ℒ.norm(current_residuals) + + if iter > 5 && ℒ.norm(rel_xtol_reached) > sqrt(rel_xtol) && current_residuals_norm > old_residuals_norm + break + end + + # First Traub step: compute intermediate point (Newton step) + if ∇ isa SparseMatrixCSC + sol_cache.A = ∇ + sol_cache.b = current_residuals + 𝒮.solve!(sol_cache) + guess_update .= sol_cache.u + intermediate_guess .= current_guess + ℒ.axpy!(-1, guess_update, intermediate_guess) + else + fact∇ = ℒ.lu!(∇, check = false) + try + if !ℒ.issuccess(fact∇) + fact∇ = ℒ.qr(∇, ℒ.ColumnNorm()) + end + copy!(intermediate_guess, current_residuals) + ℒ.ldiv!(fact∇, intermediate_guess) + intermediate_guess .= current_guess .- intermediate_guess + catch + rel_xtol_reached = typemax(T) + current_residuals_norm = typemax(T) + break + end + end + + finn = has_nonfinite(intermediate_guess) + + if finn + rel_xtol_reached = 1.0 + rel_ftol_reached = 1.0 + current_residuals_norm = 1.0 + break + end + + minmax!(intermediate_guess, lower_bounds, upper_bounds) + + # Compute residuals at intermediate point + fnj.func(intermediate_residuals, intermediate_guess, parameters_and_solved_vars) + + # Second Traub step: use same Jacobian with intermediate residuals + if ∇ isa SparseMatrixCSC + sol_cache.A = ∇ + sol_cache.b = intermediate_residuals + 𝒮.solve!(sol_cache) + guess_update .= sol_cache.u + else + try + copy!(guess_update, intermediate_residuals) + ℒ.ldiv!(fact∇, guess_update) + catch + rel_xtol_reached = typemax(T) + current_residuals_norm = typemax(T) + break + end + end + + guess_update_norm = ℒ.norm(guess_update) + current_guess .= intermediate_guess .- guess_update + + finn = has_nonfinite(current_guess) + + if finn + rel_xtol_reached = 1.0 + rel_ftol_reached = 1.0 + current_residuals_norm = 1.0 + break + end + + minmax!(current_guess, lower_bounds, upper_bounds) + + rel_xtol_reached = guess_update_norm / max(current_guess_norm, ℒ.norm(current_guess)) + + iters[1] += 2 # Two function evaluations per iteration + iters[2] += 1 # One Jacobian evaluation + end + + return current_guess, (iters[1], iters[2], rel_xtol_reached, current_residuals_norm) +end + function minmax!(x::Vector{Float64},lb::Vector{Float64},ub::Vector{Float64}) @inbounds for i in eachindex(x) diff --git a/src/filter/find_shocks.jl b/src/filter/find_shocks.jl index ae9f9d7b1..d5abd5d9e 100644 --- a/src/filter/find_shocks.jl +++ b/src/filter/find_shocks.jl @@ -361,6 +361,331 @@ end end # dispatch_doctor +@stable default_mode = "disable" begin + +function find_shocks(::Val{:Traub}, + initial_guess::Vector{Float64}, + kron_buffer::Vector{Float64}, + kron_buffer2::AbstractMatrix{Float64}, + J::ℒ.Diagonal{Bool, Vector{Bool}}, + 𝐒ⁱ::AbstractMatrix{Float64}, + 𝐒ⁱ²ᵉ::AbstractMatrix{Float64}, + shock_independent::Vector{Float64}; + max_iter::Int = 1000, + tol::Float64 = 1e-13) # will fail for higher or lower precision + # Traub's method: A third-order iterative method for shock finding + # Uses two function evaluations per iteration for higher convergence order + x = copy(initial_guess) + y = copy(initial_guess) # Intermediate point + + λ = zeros(size(𝐒ⁱ, 1)) + λy = zeros(size(𝐒ⁱ, 1)) + + xλ = [ x + λ ] + yλ = [ y + λy ] + + Δxλ = copy(xλ) + Δyλ = copy(yλ) + + norm1 = ℒ.norm(shock_independent) + + norm2 = 1.0 + + x̂ = copy(shock_independent) + ŷ = copy(shock_independent) + + x̄ = zeros(size(𝐒ⁱ,2)) + ȳ = zeros(size(𝐒ⁱ,2)) + + ∂x = zero(𝐒ⁱ) + ∂y = zero(𝐒ⁱ) + + fxλ = zeros(length(xλ)) + fyλ = zeros(length(yλ)) + + fxλp = zeros(length(xλ), length(xλ)) + + tmp = zeros(size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + + lI = -2 * vec(ℒ.I(size(𝐒ⁱ, 2))) + + @inbounds for i in 1:max_iter + # Compute Jacobian and residuals at current point x + ℒ.kron!(kron_buffer2, J, x) + + ℒ.mul!(∂x, 𝐒ⁱ²ᵉ, kron_buffer2) + ℒ.axpby!(1, 𝐒ⁱ, 2, ∂x) + + ℒ.mul!(x̄, ∂x', λ) + + ℒ.axpy!(-2, x, x̄) + + copyto!(fxλ, 1, x̄, 1, size(𝐒ⁱ,2)) + copyto!(fxλ, size(𝐒ⁱ,2) + 1, x̂, 1, size(shock_independent,1)) + + ℒ.mul!(tmp, 𝐒ⁱ²ᵉ', λ) + ℒ.axpby!(1, lI, 2, tmp) + + fxλp[1:size(𝐒ⁱ, 2), 1:size(𝐒ⁱ, 2)] = tmp + fxλp[1:size(𝐒ⁱ, 2), size(𝐒ⁱ, 2)+1:end] = ∂x' + + ℒ.rmul!(∂x, -1) + fxλp[size(𝐒ⁱ, 2)+1:end, 1:size(𝐒ⁱ, 2)] = ∂x + + try + f̂xλp = ℒ.factorize(fxλp) + ℒ.ldiv!(Δxλ, f̂xλp, fxλ) + catch + return x, false + end + + if !all(isfinite,Δxλ) break end + + # First step: Newton-like update to intermediate point y + ℒ.axpy!(-1, Δxλ, yλ) + + copyto!(y, 1, yλ, 1, size(𝐒ⁱ,2)) + copyto!(λy, 1, yλ, size(𝐒ⁱ,2) + 1, length(λy)) + + # Compute residuals at intermediate point y + ℒ.kron!(kron_buffer, y, y) + + ℒ.mul!(ŷ, 𝐒ⁱ²ᵉ, kron_buffer) + + ℒ.mul!(ŷ, 𝐒ⁱ, y, 1, 1) + + norm2 = ℒ.norm(ŷ) + + ℒ.axpby!(1, shock_independent, -1, ŷ) + + # Compute Jacobian at intermediate point y (for Lagrange multiplier) + ℒ.kron!(kron_buffer2, J, y) + + ℒ.mul!(∂y, 𝐒ⁱ²ᵉ, kron_buffer2) + ℒ.axpby!(1, 𝐒ⁱ, 2, ∂y) + + ℒ.mul!(ȳ, ∂y', λy) + + ℒ.axpy!(-2, y, ȳ) + + copyto!(fyλ, 1, ȳ, 1, size(𝐒ⁱ,2)) + copyto!(fyλ, size(𝐒ⁱ,2) + 1, ŷ, 1, size(shock_independent,1)) + + # Second step: Use original Jacobian with residuals at y + try + ℒ.ldiv!(Δyλ, f̂xλp, fyλ) + catch + return x, false + end + + if !all(isfinite,Δyλ) break end + + # Update to new point + ℒ.axpy!(-1, Δyλ, xλ) + + copyto!(x, 1, xλ, 1, size(𝐒ⁱ,2)) + copyto!(λ, 1, xλ, size(𝐒ⁱ,2) + 1, length(λ)) + + ℒ.kron!(kron_buffer, x, x) + + ℒ.mul!(x̂, 𝐒ⁱ²ᵉ, kron_buffer) + + ℒ.mul!(x̂, 𝐒ⁱ, x, 1, 1) + + norm2 = ℒ.norm(x̂) + + ℒ.axpby!(1, shock_independent, -1, x̂) + + if ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) + break + end + + # Reset intermediate point for next iteration + copyto!(yλ, xλ) + end + + return x, ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) +end + + +function find_shocks(::Val{:Traub}, + initial_guess::Vector{Float64}, + kron_buffer::Vector{Float64}, + kron_buffer²::Vector{Float64}, + kron_buffer2::AbstractMatrix{Float64}, + kron_buffer3::AbstractMatrix{Float64}, + kron_buffer4::AbstractMatrix{Float64}, + J::ℒ.Diagonal{Bool, Vector{Bool}}, + 𝐒ⁱ::AbstractMatrix{Float64}, + 𝐒ⁱ²ᵉ::AbstractMatrix{Float64}, + 𝐒ⁱ³ᵉ::AbstractMatrix{Float64}, + shock_independent::Vector{Float64}; + max_iter::Int = 1000, + tol::Float64 = 1e-13) # will fail for higher or lower precision + # Traub's method for third-order problems + x = copy(initial_guess) + y = copy(initial_guess) + + λ = zeros(size(𝐒ⁱ, 1)) + λy = zeros(size(𝐒ⁱ, 1)) + + xλ = [ x + λ ] + yλ = [ y + λy ] + + Δxλ = copy(xλ) + Δyλ = copy(yλ) + + norm1 = ℒ.norm(shock_independent) + + norm2 = 1.0 + + x̂ = copy(shock_independent) + ŷ = copy(shock_independent) + + x̄ = zeros(size(𝐒ⁱ,2)) + ȳ = zeros(size(𝐒ⁱ,2)) + + ∂x = zero(𝐒ⁱ) + ∂y = zero(𝐒ⁱ) + + ∂x̂ = zero(𝐒ⁱ) + + fxλ = zeros(length(xλ)) + fyλ = zeros(length(yλ)) + + fxλp = zeros(length(xλ), length(xλ)) + + tmp = zeros(size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + + tmp2 = zeros(size(𝐒ⁱ, 1),size(𝐒ⁱ, 2) * size(𝐒ⁱ, 2)) + + II = sparse(ℒ.I(length(x)^2)) + + lI = -2 * vec(ℒ.I(size(𝐒ⁱ, 2))) + + @inbounds for i in 1:max_iter + # Compute Jacobian and residuals at current point x + ℒ.kron!(kron_buffer2, J, x) + ℒ.kron!(kron_buffer3, J, kron_buffer) + + copy!(∂x, 𝐒ⁱ) + ℒ.mul!(∂x, 𝐒ⁱ²ᵉ, kron_buffer2, 2, 1) + + ℒ.mul!(∂x, 𝐒ⁱ³ᵉ, kron_buffer3, 3, 1) + + ℒ.mul!(x̄, ∂x', λ) + + ℒ.axpy!(-2, x, x̄) + + copyto!(fxλ, 1, x̄, 1, size(𝐒ⁱ,2)) + copyto!(fxλ, size(𝐒ⁱ,2) + 1, x̂, 1, size(shock_independent,1)) + + x_kron_II!(kron_buffer4, x) + ℒ.mul!(tmp2, 𝐒ⁱ³ᵉ, kron_buffer4) + ℒ.mul!(tmp, tmp2', λ) + ℒ.mul!(tmp, 𝐒ⁱ²ᵉ', λ, 2, 6) + ℒ.axpy!(1,lI,tmp) + + fxλp[1:size(𝐒ⁱ, 2), 1:size(𝐒ⁱ, 2)] = tmp + + fxλp[1:size(𝐒ⁱ, 2), size(𝐒ⁱ, 2)+1:end] = ∂x' + + ℒ.rmul!(∂x, -1) + fxλp[size(𝐒ⁱ, 2)+1:end, 1:size(𝐒ⁱ, 2)] = ∂x + + try + f̂xλp = ℒ.factorize(fxλp) + ℒ.ldiv!(Δxλ, f̂xλp, fxλ) + catch + return x, false + end + + if !all(isfinite,Δxλ) break end + + # First step: Newton-like update to intermediate point y + ℒ.axpy!(-1, Δxλ, yλ) + + copyto!(y, 1, yλ, 1, size(𝐒ⁱ,2)) + copyto!(λy, 1, yλ, size(𝐒ⁱ,2) + 1, length(λy)) + + # Compute residuals at intermediate point y + ℒ.kron!(kron_buffer, y, y) + + ℒ.kron!(kron_buffer², y, kron_buffer) + + ℒ.mul!(ŷ, 𝐒ⁱ, y) + + ℒ.mul!(ŷ, 𝐒ⁱ²ᵉ, kron_buffer, 1, 1) + + ℒ.mul!(ŷ, 𝐒ⁱ³ᵉ, kron_buffer², 1, 1) + + norm2 = ℒ.norm(ŷ) + + ℒ.axpby!(1, shock_independent, -1, ŷ) + + # Compute Jacobian at intermediate point y + ℒ.kron!(kron_buffer2, J, y) + ℒ.kron!(kron_buffer3, J, kron_buffer) + + copy!(∂y, 𝐒ⁱ) + ℒ.mul!(∂y, 𝐒ⁱ²ᵉ, kron_buffer2, 2, 1) + + ℒ.mul!(∂y, 𝐒ⁱ³ᵉ, kron_buffer3, 3, 1) + + ℒ.mul!(ȳ, ∂y', λy) + + ℒ.axpy!(-2, y, ȳ) + + copyto!(fyλ, 1, ȳ, 1, size(𝐒ⁱ,2)) + copyto!(fyλ, size(𝐒ⁱ,2) + 1, ŷ, 1, size(shock_independent,1)) + + # Second step: Use original Jacobian with residuals at y + try + ℒ.ldiv!(Δyλ, f̂xλp, fyλ) + catch + return x, false + end + + if !all(isfinite,Δyλ) break end + + # Update to new point + ℒ.axpy!(-1, Δyλ, xλ) + + copyto!(x, 1, xλ, 1, size(𝐒ⁱ,2)) + copyto!(λ, 1, xλ, size(𝐒ⁱ,2) + 1, length(λ)) + + ℒ.kron!(kron_buffer, x, x) + + ℒ.kron!(kron_buffer², x, kron_buffer) + + ℒ.mul!(x̂, 𝐒ⁱ, x) + + ℒ.mul!(x̂, 𝐒ⁱ²ᵉ, kron_buffer, 1, 1) + + ℒ.mul!(x̂, 𝐒ⁱ³ᵉ, kron_buffer², 1, 1) + + norm2 = ℒ.norm(x̂) + + ℒ.axpby!(1, shock_independent, -1, x̂) + + if ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) + break + end + + # Reset intermediate point for next iteration + copyto!(yλ, xλ) + end + + return x, ℒ.norm(x̂) / max(norm1,norm2) < tol && ℒ.norm(Δxλ) / ℒ.norm(xλ) < sqrt(tol) +end + +end # dispatch_doctor + + function rrule(::typeof(find_shocks), ::Val{:LagrangeNewton}, initial_guess::Vector{Float64},