diff --git a/Project.toml b/Project.toml index 0c2d34b..e5eb7a2 100644 --- a/Project.toml +++ b/Project.toml @@ -6,11 +6,15 @@ version = "0.1.0" [deps] BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" +FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" TrajectoryGamesBase = "ac1ac542-73eb-4349-ae1b-660ab3609574" [compat] BlockArrays = "0.16.43" DataStructures = "0.18.20" +FastDifferentiation = "0.4.2" +SparseArrays = "1.11.0" Symbolics = "6.19.0" TrajectoryGamesBase = "0.3.10" diff --git a/src/MCPSolver.jl b/src/MCPSolver.jl index 52be104..740f4f7 100644 --- a/src/MCPSolver.jl +++ b/src/MCPSolver.jl @@ -1,5 +1,13 @@ module MCPSolver -greet() = print("Hello World!") +using SparseArrays: SparseArrays +using FastDifferentiation: FastDifferentiation as FD +using Symbolics: Symbolics + +include("SymbolicUtils.jl") +include("sparse_utils.jl") + +include("mcp.jl") +include("solver.jl") end # module MCPSolver diff --git a/src/SymbolicUtils.jl b/src/SymbolicUtils.jl new file mode 100644 index 0000000..e831bce --- /dev/null +++ b/src/SymbolicUtils.jl @@ -0,0 +1,119 @@ +""" +Minimal abstraction on top of `Symbolics.jl` and `FastDifferentiation.jl` to make switching between the two easier. + +Taken from: ParametricMCPs.jl. +""" +module SymbolicUtils + +using Symbolics: Symbolics +using FastDifferentiation: FastDifferentiation as FD + +export SymbolicsBackend, FastDifferentiationBackend, make_variables, build_function + +struct SymbolicsBackend end +struct FastDifferentiationBackend end + +""" + make_variables(backend, name, dimension) + +Creates a vector of `dimension` where each element is a scalar symbolic variable from `backend` with the given `name`. +""" +function make_variables end + +function make_variables(::SymbolicsBackend, name::Symbol, dimension::Int) + vars = Symbolics.@variables($name[1:dimension]) |> only |> Symbolics.scalarize + + if isempty(vars) + vars = Symbolics.Num[] + end + + vars +end + +function make_variables(::FastDifferentiationBackend, name::Symbol, dimension::Int) + FD.make_variables(name, dimension) +end + +""" + build_function(backend, f_symbolic, args_symbolic...; in_place, options) + +Builds a callable function from a symbolic expression `f_symbolic` with the given `args_symbolic` as arguments. + +Depending on the `in_place` flag, the function will be built as in-place `f!(result, args...)` or out-of-place variant `restult = f(args...)`. + +`backend_options` will be forwarded to the backend specific function and differ between backends. +""" +function build_function end + +function build_function( + f_symbolic::AbstractArray{T}, + args_symbolic...; + in_place, + backend_options = (;), +) where {T<:Symbolics.Num} + f_callable, f_callable! = Symbolics.build_function( + f_symbolic, + args_symbolic...; + expression = Val{false}, + # slightly saner defaults... + (; parallel = Symbolics.ShardedForm(), backend_options...)..., + ) + in_place ? f_callable! : f_callable +end + +function build_function( + f_symbolic::AbstractArray{T}, + args_symbolic...; + in_place, + backend_options = (;), +) where {T<:FD.Node} + FD.make_function(f_symbolic, args_symbolic...; in_place, backend_options...) +end + +""" + gradient(f_symbolic, x_symbolic) + +Computes the symbolic gradient of `f_symbolic` with respect to `x_symbolic`. +""" +function gradient end + +function gradient(f_symbolic::T, x_symbolic::Vector{T}) where {T<:Symbolics.Num} + Symbolics.gradient(f_symbolic, x_symbolic) +end + +function gradient(f_symbolic::T, x_symbolic::Vector{T}) where {T<:FD.Node} + # FD does not have a gradient utility so we just flatten the jacobian here + vec(FD.jacobian([f_symbolic], x_symbolic)) +end + +""" + jacobian(f_symbolic, x_symbolic) + +Computes the symbolic Jacobian of `f_symbolic` with respect to `x_symbolic`. +""" +function jacobian end + +function jacobian(f_symbolic::Vector{T}, x_symbolic::Vector{T}) where {T<:Symbolics.Num} + Symbolics.jacobian(f_symbolic, x_symbolic) +end + +function jacobian(f_symbolic::Vector{T}, x_symbolic::Vector{T}) where {T<:FD.Node} + FD.jacobian([f_symbolic], x_symbolic) +end + +""" + sparse_jacobian(f_symbolic, x_symbolic) + +Computes the symbolic Jacobian of `f_symbolic` with respect to `x_symbolic` in a sparse format. +""" +function sparse_jacobian end + +function sparse_jacobian(f_symbolic::Vector{T}, x_symbolic::Vector{T}) where {T<:Symbolics.Num} + Symbolics.sparsejacobian(f_symbolic, x_symbolic) +end + +function sparse_jacobian(f_symbolic::Vector{T}, x_symbolic::Vector{T}) where {T<:FD.Node} + FD.sparse_jacobian(f_symbolic, x_symbolic) +end + +end diff --git a/src/mcp.jl b/src/mcp.jl new file mode 100644 index 0000000..6994471 --- /dev/null +++ b/src/mcp.jl @@ -0,0 +1,70 @@ +""" Store key elements of the primal-dual KKT system for a MCP composed of +functions G(.) and H(.) such that + 0 = G(x, y) + 0 ≤ H(x, y) ⟂ y ≥ 0. + +The primal-dual system arises when we introduce slack variable `s` and set + G(x, y) = 0 + H(x, y) - s = 0 + sᵀy - ϵ = 0 +for some ϵ > 0. Define the function `F(z; ϵ)` to return the left hand side of this +system of equations, where `z = [x; y; s]`. +""" +struct PrimalDualMCP{T1,T2} + "A callable `F(z; ϵ)` which computes the KKT error in the primal-dual system." + F::T1 + "A callable `∇F(z; ϵ)` which stores the Jacobian of the KKT error wrt z." + ∇F::T2 +end + +"Construct a PrimalDualMCP from symbolic expressions of G(.) and H(.)." +function PrimalDualMCP( + G_symbolic::Vector{T}, + H_symbolic::Vector{T}, + x_symbolic::Vector{T}, + y_symbolic::Vector{T}, + backend = SymbolicUtils.SymbolicsBackend(), + backend_options = (;) +) where {T<:Union{FD.Node,Symbolics.Num}} + # Create symbolic slack variable `s` and parameter `ϵ`. + s_symbolic = SymbolicUtils.make_variables(backend, :s, length(y_symbolic)) + ϵ_symbolic = SymbolicUtils.make_variables(backend, :ϵ, 1) + z_symbolic = [x_symbolic; y_symbolic; s_symbolic] + + F_symbolic = [ + G_symbolic; + H_symbolic - s_symbolic; + sum(s_symbolic .* y_symbolic) - ϵ_symbolic + ] + + F = let + _F = SymbolicUtils.build_function( + F_symbolic, + [z_symbolic; ϵ_symbolic]; + in_place = false, + backend_options, + ) + + (x, y, s; ϵ) -> _F([x; y; s; ϵ]) + end + + ∇F = let + ∇F_symbolic = SymbolicUtils.sparse_jacobian(F_symbolic, z_symbolic) + _∇F = SymbolicUtils.build_function( + ∇F_symbolic, + [z_symbolic; ϵ_symbolic]; + in_place = false, + backend_options, + ) + + # rows, cols, _ = SparseArrays.findnz(∇F_symbolic) + # constant_entries = get_constant_entries(∇F_symbolic, z_symbolic) + # SparseFunction(rows, cols, size(∇F_symbolic), constant_entries) do x, y, s, ϵ + # _∇F([x; y; z; ϵ]) + # end + + (x, y, s; ϵ) -> _∇F([x; y; s; ϵ]) + end + + PrimalDualMCP(F, ∇F) +end diff --git a/src/solver.jl b/src/solver.jl new file mode 100644 index 0000000..7405d19 --- /dev/null +++ b/src/solver.jl @@ -0,0 +1,20 @@ +abstract type SolverType end +struct InteriorPoint <: SolverType end + +""" Basic interior point solver, based on Nocedal & Wright, ch. 19. +Computes step directions `δz` by solving the relaxed primal-dual system, i.e. + ∇F(z; ϵ) δz = -F(z; ϵ). + +Given a step direction `δz`, performs a "fraction to the boundary" linesearch, +i.e., for `(x, s)` it chooses step size `α_s` such that + α_s = max(α ∈ [0, 1] : s + δs ≥ (1 - τ) s) +and for `y` it chooses step size `α_s` such that + α_y = max(α ∈ [0, 1] : y + δy ≥ (1 - τ) y). + +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. +""" +function solve(::InteriorPoint, mcp::MCP, x₀, y₀) + # TODO! +end diff --git a/src/sparse_utils.jl b/src/sparse_utils.jl new file mode 100644 index 0000000..2dee36a --- /dev/null +++ b/src/sparse_utils.jl @@ -0,0 +1,52 @@ +""" Sparse utils from ParametricMCPs.jl. """ + +struct SparseFunction{T1,T2} + _f::T1 + result_buffer::T2 + rows::Vector{Int} + cols::Vector{Int} + size::Tuple{Int,Int} + constant_entries::Vector{Int} + function SparseFunction(_f::T1, rows, cols, size, constant_entries = Int[]) where {T1} + length(constant_entries) <= length(rows) || + throw(ArgumentError("More constant entries than non-zero entries.")) + result_buffer = get_result_buffer(rows, cols, size) + new{T1,typeof(result_buffer)}(_f, result_buffer, rows, cols, size, constant_entries) + end +end + +(f::SparseFunction)(args...) = f._f(args...) +SparseArrays.nnz(f::SparseFunction) = length(f.rows) + +function get_result_buffer(rows::Vector{Int}, cols::Vector{Int}, size::Tuple{Int,Int}) + data = zeros(length(rows)) + SparseArrays.sparse(rows, cols, data, size...) +end + +function get_result_buffer(f::SparseFunction) + get_result_buffer(f.rows, f.cols, f.size) +end + +"Get the (sparse) linear indices of all entries that are constant in the symbolic matrix M w.r.t. symbolic vector z." +function get_constant_entries( + M_symbolic::AbstractMatrix{<:Symbolics.Num}, + z_symbolic::AbstractVector{<:Symbolics.Num}, +) + _z_syms = Symbolics.tosymbol.(z_symbolic) + findall(SparseArrays.nonzeros(M_symbolic)) do v + _vars_syms = Symbolics.tosymbol.(Symbolics.get_variables(v)) + isempty(intersect(_vars_syms, _z_syms)) + end +end + +function get_constant_entries( + M_symbolic::AbstractMatrix{<:FD.Node}, + z_symbolic::AbstractVector{<:FD.Node}, +) + _z_syms = FD.variables(z_symbolic) + # find all entries that are not a function of any of the symbols in z + findall(SparseArrays.nonzeros(M_symbolic)) do v + _vars_syms = FD.variables(v) + isempty(intersect(_vars_syms, _z_syms)) + end +end