Use SymbolicTracingUtils.jl

JuliaGameTheoreticPlanning · Dec 21, 2024 · 9fe8731 · 9fe8731
1 parent 0dc92a6
commit 9fe8731
Show file tree

Hide file tree

Showing 9 changed files with 60 additions and 218 deletions.
diff --git a/Project.toml b/Project.toml
@@ -4,11 +4,10 @@ authors = ["lassepe <[email protected]> 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
@@ -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
@@ -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
@@ -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
diff --git 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
@@ -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