diff --git a/Project.toml b/Project.toml index e27976a..664ad25 100644 --- a/Project.toml +++ b/Project.toml @@ -4,11 +4,10 @@ authors = ["lassepe and contributors"] version = "0.1.16" [deps] -FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PATHSolver = "f5f7c340-0bb3-5c69-969a-41884d311d1b" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" +SymbolicTracingUtils = "77ddf47f-b2ab-4ded-95ee-54f4fa148129" [weakdeps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" @@ -20,8 +19,7 @@ ForwardDiffExt = "ForwardDiff" [compat] ChainRulesCore = "1" -FastDifferentiation = "0.4" ForwardDiff = "0.10" PATHSolver = "1.4" -Symbolics = "4,5,6" +SymbolicTracingUtils = "0.1.1" julia = "1.10" diff --git a/src/InternalAutoDiffUtils.jl b/src/InternalAutoDiffUtils.jl index 2d4f9fd..9d0b76f 100644 --- a/src/InternalAutoDiffUtils.jl +++ b/src/InternalAutoDiffUtils.jl @@ -1,6 +1,6 @@ module InternalAutoDiffUtils -using ..ParametricMCPs: get_problem_size, get_result_buffer, get_parameter_dimension +using ..ParametricMCPs: get_problem_size, get_parameter_dimension using SparseArrays: SparseArrays using LinearAlgebra: LinearAlgebra @@ -26,13 +26,13 @@ function solve_jacobian_θ(problem, solution, θ; active_tolerance = 1e-3) end ∂f_reduce∂θ = let - ∂f∂θ = get_result_buffer(jacobian_θ!) + ∂f∂θ = jacobian_θ!.result_buffer jacobian_θ!(∂f∂θ, z_star, θ) ∂f∂θ[inactive_indices, :] end ∂f_reduced∂z_reduced = let - ∂f∂z = get_result_buffer(jacobian_z!) + ∂f∂z = jacobian_z!.result_buffer jacobian_z!(∂f∂z, z_star, θ) ∂f∂z[inactive_indices, inactive_indices] end diff --git a/src/Internals.jl b/src/Internals.jl index 34d956b..5067012 100644 --- a/src/Internals.jl +++ b/src/Internals.jl @@ -1,6 +1,7 @@ module Internals -using FastDifferentiation: FastDifferentiation as FD +using SymbolicTracingUtils: FD +using SparseArrays: SparseArrays function infer_problem_size(lower_bounds, upper_bounds) if lower_bounds isa AbstractVector @@ -31,4 +32,22 @@ function check_dimensions(args...) the_unique_dimension end +""" +Convert a Julia sparse array `M` into the \ +[COO](https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)) format required by PATH. +This implementation has been extracted from \ +[here](https://github.com/chkwon/PATHSolver.jl/blob/8e63723e51833cdbab58c39b6646f8cdf79d74a2/src/C_API.jl#L646) +""" +function _coo_from_sparse!(col, len, row, data, M) + @assert length(col) == length(len) == size(M, 1) + @assert length(row) == length(data) == SparseArrays.nnz(M) + n = length(col) + @inbounds begin + col .= @view M.colptr[1:n] + len .= diff(M.colptr) + row .= SparseArrays.rowvals(M) + data .= SparseArrays.nonzeros(M) + end +end + end diff --git a/src/ParametricMCPs.jl b/src/ParametricMCPs.jl index 9450ee5..2833839 100644 --- a/src/ParametricMCPs.jl +++ b/src/ParametricMCPs.jl @@ -2,14 +2,14 @@ module ParametricMCPs using PATHSolver: PATHSolver using SparseArrays: SparseArrays -using FastDifferentiation: FastDifferentiation as FD -using Symbolics: Symbolics +using SymbolicTracingUtils: SymbolicTracingUtils -include("Internals.jl") -include("sparse_utils.jl") +const SymbolicUtils = SymbolicTracingUtils + +# re-exporting symbolic tracing utilities for backward compatibility +export SymbolicTracingUtils -include("SymbolicUtils.jl") -export SymbolicUtils +include("Internals.jl") include("parametric_problem.jl") export ParametricMCP, get_parameter_dimension, get_problem_size diff --git a/src/SymbolicUtils.jl b/src/SymbolicUtils.jl deleted file mode 100644 index 9cac18e..0000000 --- a/src/SymbolicUtils.jl +++ /dev/null @@ -1,117 +0,0 @@ -""" -Minimal abstraction on top of `Symbolics.jl` and `FastDifferentiation.jl` to make switching between the two easier. -""" -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/parametric_problem.jl b/src/parametric_problem.jl index 39c51d7..d60b902 100644 --- a/src/parametric_problem.jl +++ b/src/parametric_problem.jl @@ -41,7 +41,7 @@ parameter vector `θ` of size `parameter_dimension` to an length `n` vector outp - `parameter_dimension`: the size of the parameter vector `θ` in `f`. Keyword arguments: -- `[backend]`: the backend (from `ParametricMCPs.SymbolicUtils`) to be used for compiling callbacks for `f` and its Jacobians needed by PATH. `SymbolicsBackend` (default) is slightly more flexible. `FastDifferentiationBackend` has reduced compilation times and reduced runtime in some cases. +- `[backend]`: the backend (from `SymbolicTracingUtils`) to be used for compiling callbacks for `f` and its Jacobians needed by PATH. `SymbolicsBackend` (default) is slightly more flexible. `FastDifferentiationBackend` has reduced compilation times and reduced runtime in some cases. - `compute_sensitivities`: whether to compile the callbacks needed for sensitivity computation. - `[problem_size]`: the number of decision variables. If not provided and `lower_bounds` or `upper_bounds` are vectors, the problem size is inferred from the length of these vectors. @@ -55,14 +55,14 @@ function ParametricMCP( lower_bounds, upper_bounds, parameter_dimension; - backend = SymbolicUtils.SymbolicsBackend(), + backend = SymbolicTracingUtils.SymbolicsBackend(), problem_size = Internals.infer_problem_size(lower_bounds, upper_bounds), kwargs..., ) problem_size = Internals.check_dimensions(lower_bounds, upper_bounds, problem_size) - z_symbolic = SymbolicUtils.make_variables(backend, :z, problem_size) - θ_symbolic = SymbolicUtils.make_variables(backend, :θ, parameter_dimension) + z_symbolic = SymbolicTracingUtils.make_variables(backend, :z, problem_size) + θ_symbolic = SymbolicTracingUtils.make_variables(backend, :θ, parameter_dimension) f_symbolic = f(z_symbolic, θ_symbolic) ParametricMCP(f_symbolic, z_symbolic, θ_symbolic, lower_bounds, upper_bounds; kwargs...) @@ -81,7 +81,7 @@ function ParametricMCP( warm_up_callbacks = true, parallel = nothing, backend_options = (;), -) where {T<:Union{FD.Node,Symbolics.Num}} +) where {T<:Union{SymbolicTracingUtils.FD.Node,SymbolicTracingUtils.Symbolics.Num}} problem_size = Internals.check_dimensions(f_symbolic, z_symbolic, lower_bounds, upper_bounds) if !isnothing(parallel) @@ -96,7 +96,7 @@ function ParametricMCP( # compile all the symbolic expressions into callable julia code f! = let # The multi-arg version of `make_function` is broken so we concatenate to a single arg here - _f! = SymbolicUtils.build_function( + _f! = SymbolicTracingUtils.build_function( f_symbolic, [z_symbolic; θ_symbolic]; in_place = true, @@ -107,31 +107,41 @@ function ParametricMCP( # same as above but for the Jacobian in z jacobian_z! = let - jacobian_z = SymbolicUtils.sparse_jacobian(f_symbolic, z_symbolic) - _jacobian_z! = SymbolicUtils.build_function( + jacobian_z = SymbolicTracingUtils.sparse_jacobian(f_symbolic, z_symbolic) + _jacobian_z! = SymbolicTracingUtils.build_function( jacobian_z, [z_symbolic; θ_symbolic]; in_place = true, backend_options, ) rows, cols, _ = SparseArrays.findnz(jacobian_z) - constant_entries = get_constant_entries(jacobian_z, z_symbolic) - SparseFunction(rows, cols, size(jacobian_z), constant_entries) do result, z, θ + constant_entries = SymbolicTracingUtils.get_constant_entries(jacobian_z, z_symbolic) + SymbolicTracingUtils.SparseFunction( + rows, + cols, + size(jacobian_z), + constant_entries, + ) do result, z, θ _jacobian_z!(result, [z; θ]) end end if compute_sensitivities jacobian_θ! = let - jacobian_θ = SymbolicUtils.sparse_jacobian(f_symbolic, θ_symbolic) - _jacobian_θ! = SymbolicUtils.build_function( + jacobian_θ = SymbolicTracingUtils.sparse_jacobian(f_symbolic, θ_symbolic) + _jacobian_θ! = SymbolicTracingUtils.build_function( jacobian_θ, [z_symbolic; θ_symbolic]; in_place = true, ) rows, cols, _ = SparseArrays.findnz(jacobian_θ) - constant_entries = get_constant_entries(jacobian_θ, θ_symbolic) - SparseFunction(rows, cols, size(jacobian_θ), constant_entries) do result, z, θ + constant_entries = SymbolicTracingUtils.get_constant_entries(jacobian_θ, θ_symbolic) + SymbolicTracingUtils.SparseFunction( + rows, + cols, + size(jacobian_θ), + constant_entries, + ) do result, z, θ _jacobian_θ!(result, [z; θ]) end end @@ -165,10 +175,10 @@ function _warm_up_callbacks(mcp::ParametricMCP) θ = zeros(get_parameter_dimension(mcp)) z = zeros(get_problem_size(mcp)) mcp.f!(z, z, θ) - jacz = ParametricMCPs.get_result_buffer(mcp.jacobian_z!) + jacz = mcp.jacobian_z!.result_buffer mcp.jacobian_z!(jacz, z, θ) if !isnothing(mcp.jacobian_θ!) - jacθ = ParametricMCPs.get_result_buffer(mcp.jacobian_θ!) + jacθ = mcp.jacobian_θ!.result_buffer mcp.jacobian_θ!(jacθ, z, θ) end nothing diff --git a/src/solver.jl b/src/solver.jl index 01b6eb2..5022a19 100644 --- a/src/solver.jl +++ b/src/solver.jl @@ -22,7 +22,7 @@ function solve( function J(n, nnz, z, col, len, row, data) jacobian_z!(jacobian_z!.result_buffer, z, θ) - _coo_from_sparse!(col, len, row, data, jacobian_z!.result_buffer) + Internals._coo_from_sparse!(col, len, row, data, jacobian_z!.result_buffer) Cint(0) end diff --git a/src/sparse_utils.jl b/src/sparse_utils.jl deleted file mode 100644 index a6700e0..0000000 --- a/src/sparse_utils.jl +++ /dev/null @@ -1,68 +0,0 @@ -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 - -""" -Convert a Julia sparse array `M` into the \ -[COO](https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)) format required by PATH. -This implementation has been extracted from \ -[here](https://github.com/chkwon/PATHSolver.jl/blob/8e63723e51833cdbab58c39b6646f8cdf79d74a2/src/C_API.jl#L646) -""" -function _coo_from_sparse!(col, len, row, data, M) - @assert length(col) == length(len) == size(M, 1) - @assert length(row) == length(data) == SparseArrays.nnz(M) - n = length(col) - @inbounds begin - col .= @view M.colptr[1:n] - len .= diff(M.colptr) - row .= SparseArrays.rowvals(M) - data .= SparseArrays.nonzeros(M) - end -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 diff --git a/test/runtests.jl b/test/runtests.jl index 15a06e5..b6265e3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,8 +7,8 @@ using FiniteDiff: FiniteDiff @testset "ParametricMCPs.jl" begin for backend in [ - ParametricMCPs.SymbolicUtils.SymbolicsBackend(), - ParametricMCPs.SymbolicUtils.FastDifferentiationBackend(), + ParametricMCPs.SymbolicTracingUtils.SymbolicsBackend(), + ParametricMCPs.SymbolicTracingUtils.FastDifferentiationBackend(), ] rng = Random.MersenneTwister(1) parameter_dimension = 2