Merge pull request #24 from lassepe/feature/krylov

Use Kyrlov linear solver for the Newton step.

Project.toml

@@ -10,6 +10,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SymbolicTracingUtils = "77ddf47f-b2ab-4ded-95ee-54f4fa148129"
 TrajectoryGamesBase = "ac1ac542-73eb-4349-ae1b-660ab3609574"
@@ -22,6 +23,7 @@ DataStructures = "0.18.20"
 FiniteDiff = "2.26.2"
 ForwardDiff = "0.10.38"
 LinearAlgebra = "1.11.0"
+LinearSolve = "2.38.0"
 SparseArrays = "1.11.0"
 SymbolicTracingUtils = "0.1.1"
 TrajectoryGamesBase = "0.3.10"

src/MixedComplementarityProblems.jl

@@ -5,6 +5,7 @@ using LinearAlgebra: LinearAlgebra, I, norm, eigvals
 using BlockArrays: blocks, blocksizes
 using TrajectoryGamesBase: to_blockvector
 using SymbolicTracingUtils: SymbolicTracingUtils as SymbolicTracingUtils
+using LinearSolve: LinearProblem, init, solve!, KrylovJL_GMRES
 
 include("mcp.jl")
 include("solver.jl")

src/solver.jl

@@ -14,6 +14,23 @@ and for `y` it chooses step size `α_s` such that
 A typical value of τ is 0.995. Once we converge to ||F(z; \epsilon)|| ≤ ϵ,
 we typically decrease ϵ by a factor of 0.1 or 0.2, with smaller values chosen
 when the previous subproblem is solved in fewer iterations.
+
+Positional arguments:
+    - `mcp::PrimalDualMCP`: the mixed complementarity problem to solve.
+    - `θ::AbstractVector{<:Real}`: the parameter vector.
+    - `x₀::AbstractVector{<:Real}`: the initial primal variable.
+    - `y₀::AbstractVector{<:Real}`: the initial dual variable.
+    - `s₀::AbstractVector{<:Real}`: the initial slack variable.
+
+Keyword arguments:
+    - `tol::Real = 1e-4`: the tolerance for the KKT error.
+    - `max_inner_iters::Int = 20`: the maximum number of inner iterations.
+    - `max_outer_iters::Int = 50`: the maximum number of outer iterations.
+    - `tightening_rate::Real = 0.1`: the rate at which to tighten the tolerance.
+    - `loosening_rate::Real = 0.5`: the rate at which to loosen the tolerance.
+    - `min_stepsize::Real = 1e-2`: the minimum step size for the linesearch.
+    - `verbose::Bool = false`: whether to print debug information.
+    - `linear_solve_algorithm::KrylovJL_GMRES()`: the linear solve algorithm to use. Any solver from `LinearSolve.jl` can be used.
 """
 function solve(
     ::InteriorPoint,
@@ -29,6 +46,7 @@ function solve(
     loosening_rate = 0.5,
     min_stepsize = 1e-2,
     verbose = false,
+    linear_solve_algorithm = KrylovJL_GMRES(),
 )
     # Set up common memory.
     ∇F = mcp.∇F_z!.result_buffer
@@ -39,6 +57,8 @@ function solve(
         @view δz[(mcp.unconstrained_dimension + 1):(mcp.unconstrained_dimension + mcp.constrained_dimension)]
     δs = @view δz[(mcp.unconstrained_dimension + mcp.constrained_dimension + 1):end]
 
+    linsolve = init(LinearProblem(∇F, δz), linear_solve_algorithm)
+
     # Main solver loop.
     x = x₀
     y = y₀
@@ -53,10 +73,14 @@ function solve(
 
         while kkt_error > ϵ && inner_iters < max_inner_iters
             # Compute the Newton step.
-            # TODO! Can add some adaptive regularization.
+            # TODO: Can add some adaptive regularization.
+            # TODO: use a linear operator with a lazy gradient computation here.
             mcp.F!(F, x, y, s; θ, ϵ)
             mcp.∇F_z!(∇F, x, y, s; θ, ϵ)
-            δz .= ((∇F + tol * I) \ F) .* -1
+            linsolve.A = ∇F + tol * I
+            linsolve.b = -F
+            solution = solve!(linsolve)
+            δz .= solution.u
 
             # Fraction to the boundary linesearch.
             α_s = fraction_to_the_boundary_linesearch(s, δs; tol = min_stepsize)
@@ -78,7 +102,7 @@ function solve(
         end
 
         ϵ *=
-            (status == :solved) ? 1 - exp(-tightening_rate * inner_iters) :
+            (status === :solved) ? 1 - exp(-tightening_rate * inner_iters) :
             1 + exp(-loosening_rate * inner_iters)
         outer_iters += 1
     end