diff --git a/Project.toml b/Project.toml index b6747a7..91d5ee3 100644 --- a/Project.toml +++ b/Project.toml @@ -7,12 +7,11 @@ version = "0.1.4" BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" +SymbolicTracingUtils = "77ddf47f-b2ab-4ded-95ee-54f4fa148129" TrajectoryGamesBase = "ac1ac542-73eb-4349-ae1b-660ab3609574" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" @@ -20,12 +19,11 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" BlockArrays = "0.16.43, 1" ChainRulesCore = "1.25.0" DataStructures = "0.18.20" -FastDifferentiation = "0.4.2" FiniteDiff = "2.26.2" ForwardDiff = "0.10.38" LinearAlgebra = "1.11.0" SparseArrays = "1.11.0" -Symbolics = "6.19.0" +SymbolicTracingUtils = "0.1.1" TrajectoryGamesBase = "0.3.10" Zygote = "0.6.73" julia = "1.11" diff --git a/src/AutoDiff.jl b/src/AutoDiff.jl index ba43933..49ec381 100644 --- a/src/AutoDiff.jl +++ b/src/AutoDiff.jl @@ -13,6 +13,7 @@ using ..MixedComplementarityProblems: MixedComplementarityProblems using ChainRulesCore: ChainRulesCore using ForwardDiff: ForwardDiff using LinearAlgebra: LinearAlgebra +using SymbolicTracingUtils: SymbolicTracingUtils function _solve_jacobian_θ(mcp::MixedComplementarityProblems.PrimalDualMCP, solution, θ) !isnothing(mcp.∇F_θ!) || throw( diff --git a/src/MixedComplementarityProblems.jl b/src/MixedComplementarityProblems.jl index d23eaf6..0125f35 100644 --- a/src/MixedComplementarityProblems.jl +++ b/src/MixedComplementarityProblems.jl @@ -1,14 +1,11 @@ module MixedComplementarityProblems using SparseArrays: SparseArrays -using FastDifferentiation: FastDifferentiation as FD -using Symbolics: Symbolics using LinearAlgebra: LinearAlgebra, I, norm, eigvals using BlockArrays: blocks, blocksizes using TrajectoryGamesBase: to_blockvector +using SymbolicTracingUtils: SymbolicTracingUtils as SymbolicTracingUtils -include("SymbolicUtils.jl") -include("sparse_utils.jl") include("mcp.jl") include("solver.jl") include("game.jl") diff --git a/src/SymbolicUtils.jl b/src/SymbolicUtils.jl deleted file mode 100644 index eac2eb7..0000000 --- a/src/SymbolicUtils.jl +++ /dev/null @@ -1,122 +0,0 @@ -""" -Minimal abstraction on top of `Symbolics.jl` and `FastDifferentiation.jl` to make switching between the two easier. - -Taken from: https://github.com/JuliaGameTheoreticPlanning/ParametricMCPs.jl/blob/main/src/SymbolicUtils.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/game.jl b/src/game.jl index cb8b934..71dff34 100644 --- a/src/game.jl +++ b/src/game.jl @@ -43,13 +43,21 @@ function ParametricGame(; # Define primal and dual variables for the game, and game parameters.. # Note that BlockArrays can handle blocks of zero size. - backend = SymbolicUtils.SymbolicsBackend() - x = SymbolicUtils.make_variables(backend, :x, sum(dims.x)) |> to_blockvector(dims.x) - λ = SymbolicUtils.make_variables(backend, :λ, sum(dims.λ)) |> to_blockvector(dims.λ) - μ = SymbolicUtils.make_variables(backend, :μ, sum(dims.μ)) |> to_blockvector(dims.μ) - λ̃ = SymbolicUtils.make_variables(backend, :λ̃, dims.λ̃) - μ̃ = SymbolicUtils.make_variables(backend, :μ̃, dims.μ̃) - θ = SymbolicUtils.make_variables(backend, :θ, sum(dims.θ)) |> to_blockvector(dims.θ) + backend = SymbolicTracingUtils.SymbolicsBackend() + x = + SymbolicTracingUtils.make_variables(backend, :x, sum(dims.x)) |> + to_blockvector(dims.x) + λ = + SymbolicTracingUtils.make_variables(backend, :λ, sum(dims.λ)) |> + to_blockvector(dims.λ) + μ = + SymbolicTracingUtils.make_variables(backend, :μ, sum(dims.μ)) |> + to_blockvector(dims.μ) + λ̃ = SymbolicTracingUtils.make_variables(backend, :λ̃, dims.λ̃) + μ̃ = SymbolicTracingUtils.make_variables(backend, :μ̃, dims.μ̃) + θ = + SymbolicTracingUtils.make_variables(backend, :θ, sum(dims.θ)) |> + to_blockvector(dims.θ) # Build symbolic expressions for objectives and constraints. fs = map(problems, blocks(θ)) do p, θi @@ -70,7 +78,7 @@ function ParametricGame(; L = f - (isnothing(g) ? 0 : sum(λi .* g)) - (isnothing(h) ? 0 : sum(μi .* h)) - (isnothing(g̃) ? 0 : sum(λ̃ .* g̃)) - (isnothing(h̃) ? 0 : sum(μ̃ .* h̃)) - SymbolicUtils.gradient(L, xi) + SymbolicTracingUtils.gradient(L, xi) end # Build MCP representation. diff --git a/src/mcp.jl b/src/mcp.jl index 6ed7ef2..73f9280 100644 --- a/src/mcp.jl +++ b/src/mcp.jl @@ -31,12 +31,12 @@ function PrimalDualMCP( constrained_dimension, parameter_dimension, compute_sensitivities = false, - backend = SymbolicUtils.SymbolicsBackend(), + backend = SymbolicTracingUtils.SymbolicsBackend(), backend_options = (;), ) - x_symbolic = SymbolicUtils.make_variables(backend, :x, unconstrained_dimension) - y_symbolic = SymbolicUtils.make_variables(backend, :y, constrained_dimension) - θ_symbolic = SymbolicUtils.make_variables(backend, :θ, parameter_dimension) + x_symbolic = SymbolicTracingUtils.make_variables(backend, :x, unconstrained_dimension) + y_symbolic = SymbolicTracingUtils.make_variables(backend, :y, constrained_dimension) + θ_symbolic = SymbolicTracingUtils.make_variables(backend, :θ, parameter_dimension) G_symbolic = G(x_symbolic, y_symbolic; θ = θ_symbolic) H_symbolic = H(x_symbolic, y_symbolic; θ = θ_symbolic) @@ -60,17 +60,17 @@ function PrimalDualMCP( θ_symbolic::Vector{T}; compute_sensitivities = false, backend_options = (;), -) where {T<:Union{FD.Node,Symbolics.Num}} +) where {T<:Union{SymbolicTracingUtils.FD.Node,SymbolicTracingUtils.Symbolics.Num}} # Create symbolic slack variable `s` and parameter `ϵ`. - if T == FD.Node - backend = SymbolicUtils.FastDifferentiationBackend() + if T == SymbolicTracingUtils.FD.Node + backend = SymbolicTracingUtils.FastDifferentiationBackend() else - @assert T === Symbolics.Num - backend = SymbolicUtils.SymbolicsBackend() + @assert T === SymbolicTracingUtils.Symbolics.Num + backend = SymbolicTracingUtils.SymbolicsBackend() end - s_symbolic = SymbolicUtils.make_variables(backend, :s, length(y_symbolic)) - ϵ_symbolic = only(SymbolicUtils.make_variables(backend, :ϵ, 1)) + s_symbolic = SymbolicTracingUtils.make_variables(backend, :s, length(y_symbolic)) + ϵ_symbolic = only(SymbolicTracingUtils.make_variables(backend, :ϵ, 1)) z_symbolic = [x_symbolic; y_symbolic; s_symbolic] F_symbolic = [ @@ -80,7 +80,7 @@ function PrimalDualMCP( ] F! = let - _F! = SymbolicUtils.build_function( + _F! = SymbolicTracingUtils.build_function( F_symbolic, x_symbolic, y_symbolic, @@ -95,8 +95,8 @@ function PrimalDualMCP( end ∇F_z! = let - ∇F_symbolic = SymbolicUtils.sparse_jacobian(F_symbolic, z_symbolic) - _∇F! = SymbolicUtils.build_function( + ∇F_symbolic = SymbolicTracingUtils.sparse_jacobian(F_symbolic, z_symbolic) + _∇F! = SymbolicTracingUtils.build_function( ∇F_symbolic, x_symbolic, y_symbolic, @@ -108,8 +108,9 @@ function PrimalDualMCP( ) rows, cols, _ = SparseArrays.findnz(∇F_symbolic) - constant_entries = get_constant_entries(∇F_symbolic, z_symbolic) - SparseFunction( + constant_entries = + SymbolicTracingUtils.get_constant_entries(∇F_symbolic, z_symbolic) + SymbolicTracingUtils.SparseFunction( (result, x, y, s; θ, ϵ) -> _∇F!(result, x, y, s, θ, ϵ), rows, cols, @@ -121,8 +122,8 @@ function PrimalDualMCP( ∇F_θ! = !compute_sensitivities ? nothing : let - ∇F_symbolic = SymbolicUtils.sparse_jacobian(F_symbolic, θ_symbolic) - _∇F! = SymbolicUtils.build_function( + ∇F_symbolic = SymbolicTracingUtils.sparse_jacobian(F_symbolic, θ_symbolic) + _∇F! = SymbolicTracingUtils.build_function( ∇F_symbolic, x_symbolic, y_symbolic, @@ -134,8 +135,9 @@ function PrimalDualMCP( ) rows, cols, _ = SparseArrays.findnz(∇F_symbolic) - constant_entries = get_constant_entries(∇F_symbolic, θ_symbolic) - SparseFunction( + constant_entries = + SymbolicTracingUtils.get_constant_entries(∇F_symbolic, θ_symbolic) + SymbolicTracingUtils.SparseFunction( (result, x, y, s; θ, ϵ) -> _∇F!(result, x, y, s, θ, ϵ), rows, cols, @@ -156,11 +158,11 @@ function PrimalDualMCP( upper_bounds::Vector; parameter_dimension, compute_sensitivities = false, - backend = SymbolicUtils.SymbolicsBackend(), + backend = SymbolicTracingUtils.SymbolicsBackend(), backend_options = (;), ) - z_symbolic = SymbolicUtils.make_variables(backend, :z, length(lower_bounds)) - θ_symbolic = SymbolicUtils.make_variables(backend, :θ, parameter_dimension) + z_symbolic = SymbolicTracingUtils.make_variables(backend, :z, length(lower_bounds)) + θ_symbolic = SymbolicTracingUtils.make_variables(backend, :θ, parameter_dimension) K_symbolic = K(z_symbolic; θ = θ_symbolic) PrimalDualMCP( @@ -185,7 +187,7 @@ function PrimalDualMCP( upper_bounds::Vector; compute_sensitivities = false, backend_options = (;), -) where {T<:Union{FD.Node,Symbolics.Num}} +) where {T<:Union{SymbolicTracingUtils.FD.Node,SymbolicTracingUtils.Symbolics.Num}} @assert all(isinf.(upper_bounds)) && all(isinf.(lower_bounds) .|| lower_bounds .== 0) unconstrained_indices = findall(isinf, lower_bounds) diff --git a/src/sparse_utils.jl b/src/sparse_utils.jl deleted file mode 100644 index 7fcd066..0000000 --- a/src/sparse_utils.jl +++ /dev/null @@ -1,63 +0,0 @@ -""" -Utility for encoding functions which return sparse matrices. - -Code from: https://github.com/JuliaGameTheoreticPlanning/ParametricMCPs.jl/blob/main/src/sparse_utils.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...; kwargs...) = f._f(args...; kwargs...) -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 = [zs.node_value for zs in 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 = [vs.node_value for vs in FD.variables(v)] - isempty(intersect(_vars_syms, _z_syms)) - end -end