Enhancements

Project.toml

@@ -12,6 +12,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NonconvexCore = "035190e5-69f1-488f-aaab-becca2889735"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -23,18 +25,20 @@ IterativeSolvers = "0.8, 0.9"
 LinearMaps = "3"
 MacroTools = "0.5"
 NonconvexCore = "1"
+SparseDiffTools = "1.24"
 Symbolics = "4.6"
 Zygote = "0.5, 0.6"
 julia = "1"
 
 [extras]
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NamedTupleTools = "d9ec5142-1e00-5aa0-9d6a-321866360f50"
+NonconvexIpopt = "bf347577-a06d-49ad-a669-8c0e005493b8"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [targets]
-test = ["NamedTupleTools", "NLsolve", "ReverseDiff", "SparseArrays", "StableRNGs", "Test", "Tracker"]
+test = ["NamedTupleTools", "NonconvexIpopt", "NLsolve", "ReverseDiff", "SparseArrays", "StableRNGs", "Test", "Tracker"]

src/NonconvexUtils.jl

@@ -8,13 +8,17 @@ export  ForwardDiffFunction,
         LazyJacobian,
         CustomHessianFunction,
         ImplicitFunction,
-        SymbolicFunction
+        SymbolicFunction,
+        SparseForwardDiffFunction,
+        sparsify,
+        symbolify
 
 using ChainRulesCore, AbstractDifferentiation, ForwardDiff, LinearAlgebra
-using Zygote, LinearMaps, IterativeSolvers, NonconvexCore
+using Zygote, LinearMaps, IterativeSolvers, NonconvexCore, SparseArrays
 using NonconvexCore: flatten
 using MacroTools
 using Symbolics: Symbolics
+using SparseDiffTools: SparseDiffTools
 
 include("forwarddiff_frule.jl")
 include("abstractdiff.jl")
@@ -23,5 +27,6 @@ include("trace.jl")
 include("custom.jl")
 include("implicit.jl")
 include("symbolic.jl")
+include("sparse_forwarddiff.jl")
 
 end

src/custom.jl

@@ -16,11 +16,28 @@ end
 function ChainRulesCore.frule(
     (_, Δx), f::CustomGradFunction, x::AbstractVector,
 )
-    v, ∇ = f.f(x), f.g(x)
+    v = f.f(x)
+    if f.g === nothing
+        if v isa Real
+            ∇ = zeros(eltype(v), length(x))'
+        else
+            ∇ = zeros(eltype(v), length(v), length(x))
+        end
+    else
+        ∇ = f.g(x)
+    end
     if ∇ isa AbstractVector && Δx isa AbstractVector
-        return v, ∇' * Δx
+        if !(∇ isa LazyJacobian) && issparse(∇) && nnz(∇) == 0
+            return v, zero(eltype(Δx))
+        else
+            return v, ∇' * Δx
+        end
     else
-        return v, ∇ * Δx
+        if !(∇ isa LazyJacobian) && issparse(∇) && nnz(∇) == 0
+            return v, zeros(eltype(Δx), size(∇, 1))
+        else
+            return v, ∇ * Δx
+        end
     end
 end
 @ForwardDiff_frule (f::CustomGradFunction)(x::AbstractVector{<:ForwardDiff.Dual})

src/sparse_forwarddiff.jl

@@ -0,0 +1,125 @@
+struct SparseForwardDiffFunction{F, F!, Y, J, JP, JC, JJ, JJ!, G, H, HP, HC} <: Function
+    f::F
+    f!::F!
+    y::Y
+    jac::J
+    jac_pattern::JP
+    jac_colors::JC
+    vecJ::JJ
+    vecJ!::JJ!
+    G::G
+    hess::H
+    hess_pattern::HP
+    hess_colors::HC
+end
+
+function SparseForwardDiffFunction(f, x::AbstractVector; hessian = false, jac_pattern = nothing, hess_pattern = nothing)
+    val = f(x)
+    _f = val isa Real ? x -> [f(x)] : f
+    f! = (y, x) -> begin
+        v = f(x)
+        y .= v
+        return y
+    end
+    y = val isa Real ? [val] : copy(val)
+    jac_pattern = jac_pattern === nothing ? Symbolics.jacobian_sparsity(f!, y, x) : jac_pattern
+    if nnz(jac_pattern) > 0
+        jac = float.(jac_pattern)
+        jac_colors = SparseDiffTools.matrix_colors(jac)
+    else
+        T = eltype(G)
+        jac = sparse(Int[], Int[], T[])
+        jac_colors = Int[]
+    end
+    vecJ! = (G, x) -> begin
+        _jac = SparseDiffTools.forwarddiff_color_jacobian(_f, x, colorvec = jac_colors)
+        G .= vec(Array(_jac))
+        return G
+    end
+    G = vec(Array(jac))
+    vecJ = x -> copy(vecJ!(G, x))
+    if hessian
+        hess_pattern = hess_pattern === nothing ? Symbolics.jacobian_sparsity(vecJ!, G, x) : hess_pattern
+        if nnz(hess_pattern) > 0
+            hess = float.(hess_pattern)
+            hess_colors = SparseDiffTools.matrix_colors(hess)
+        else
+            T = eltype(G)
+            hess = sparse(Int[], Int[], T[])
+            hess_colors = Int[]
+        end
+    else
+        hess = nothing
+        hess_colors = nothing
+    end
+    return SparseForwardDiffFunction(f, f!, y, jac, jac_pattern, jac_colors, vecJ, vecJ!, G, hess, hess_pattern, hess_colors)
+end
+(f::SparseForwardDiffFunction)(x) = f.f(x)
+function ChainRulesCore.rrule(f::SparseForwardDiffFunction, x::AbstractVector)
+    if f.vecJ === nothing
+        return f(x), _ -> (NoTangent(), NoTangent())
+    else
+        vecjac = SparseForwardDiffFunction(f.vecJ, f.vecJ!, f.G, f.hess, f.hess_pattern, f.hess_colors, nothing, nothing, nothing, nothing, nothing, nothing)
+        val = f(x)
+        if val isa Real
+            jac = reshape(vecjac(x), 1, length(x))
+        else
+            jac = reshape(vecjac(x), length(val), length(x))
+        end
+        return val, Δ -> begin
+            (NoTangent(), jac' * (Δ isa Real ? Δ : vec(Δ)))
+        end
+    end
+end
+function ChainRulesCore.frule((_, Δx), f::SparseForwardDiffFunction, x::AbstractVector)
+    if f.vecJ === nothing
+        val = f(x)
+        return val, zero(val)
+    else
+        vecjac = SparseForwardDiffFunction(f.vecJ, f.vecJ!, f.G, f.hess, f.hess_pattern, f.hess_colors, nothing, nothing, nothing, nothing, nothing, nothing)
+        val = f(x)
+        if val isa Real
+            jac = reshape(vecjac(x), 1, length(x))
+        else
+            jac = reshape(vecjac(x), length(val), length(x))
+        end
+        Δy = jac * (Δx isa Real ? Δx : vec(Δx))
+        return val, (val isa Real) ? only(Δy) : reshape(Δy, size(val))
+    end
+end
+@ForwardDiff_frule (f::SparseForwardDiffFunction)(x::AbstractVector{<:ForwardDiff.Dual})
+
+function sparsify(model::NonconvexCore.AbstractModel; objective = true, ineq_constraints = true, eq_constraints = true, kwargs...)
+    vmodel, v, _ = NonconvexCore.tovecmodel(model)
+    if objective
+        # Objective
+        sparse_flat_obj = SparseForwardDiffFunction(vmodel.objective, v; kwargs...)
+        obj = NonconvexCore.Objective(x -> sparse_flat_obj(flatten(x)[1]), flags = model.objective.flags)
+    else
+        obj = model.objective
+    end
+    if ineq_constraints
+        ineq = length(vmodel.ineq_constraints.fs) != 0 ? NonconvexCore.VectorOfFunctions(map(vmodel.ineq_constraints.fs) do c
+            sparse_flat_ineq = SparseForwardDiffFunction(c, v; kwargs...)
+            NonconvexCore.IneqConstraint(x -> sparse_flat_ineq(flatten(x)[1]), c.rhs, c.dim, c.flags)
+        end) : NonconvexCore.VectorOfFunctions(NonconvexCore.IneqConstraint[])
+    else
+        ineq = model.ineq_constraints
+    end
+    if eq_constraints
+        eq = length(vmodel.eq_constraints.fs) != 0 ? NonconvexCore.VectorOfFunctions(map(vmodel.eq_constraints.fs) do c
+            sparse_flat_eq = SparseForwardDiffFunction(c, v; kwargs...)
+            NonconvexCore.EqConstraint(x -> sparse_flat_eq(flatten(x)[1]), c.rhs, c.dim, c.flags)
+        end) : NonconvexCore.VectorOfFunctions(NonconvexCore.EqConstraint[])
+    else
+        eq = model.eq_constraints
+    end
+    if model isa NonconvexCore.Model
+        ModelT = NonconvexCore.Model
+    elseif model isa NonconvexCore.DictModel
+        ModelT = NonconvexCore.DictModel
+    else
+        throw("Unsupported model type.")
+    end
+    return ModelT(obj, eq, ineq, model.sd_constraints, model.box_min, model.box_max, model.init, model.integer)
+end

src/symbolic.jl

@@ -4,31 +4,64 @@ struct SymbolicFunction{F, G, H, X} <: Function
     h::H
     x::X
 end
-function SymbolicFunction(f, _x::AbstractVector; hessian = false)
+
+function SymbolicFunction(f, _x::AbstractVector; hessian = false, sparse = false, simplify = false)
     N = length(_x)
     val = f(_x)
+    _T = eltype(val)
+    T = x -> begin
+        if (eltype(x) <: Symbolics.Num) || eltype(x) === _T || !(x isa Real) && (issparse(x) && nnz(x) == 0 || isempty(x))
+            return x
+        else
+            return _T.(x)
+        end
+    end
     Symbolics.@variables x[1:N]
     if val isa Real
-        sgrad = Symbolics.jacobian([f(x)], x)
-        f_expr = Symbolics.build_function(sgrad, x)[1]
-        _g = eval(f_expr)
-        g = x -> vec(_g(x))
+        if sparse
+            sgrad = Symbolics.sparsejacobian([f(x)], x; simplify)
+            f_expr = Symbolics.build_function(sgrad, x)[1]
+            _g = eval(f_expr)
+            g = x -> vec(Matrix(T(_g(x))))
+        else
+            sgrad = Symbolics.jacobian([f(x)], x; simplify)
+            f_expr = Symbolics.build_function(sgrad, x)[1]
+            _g = eval(f_expr)
+            g = x -> vec(T(_g(x)))
+        end
         if hessian
-            shess = Symbolics.jacobian(Base.invokelatest(g, x), x)
+            if sparse
+                shess = Symbolics.sparsejacobian(Base.invokelatest(g, x), x; simplify)
+            else
+                shess = Symbolics.jacobian(Base.invokelatest(g, x), x; simplify)
+            end
             h_expr = Symbolics.build_function(shess, x)[1]
-            h = eval(h_expr)
+            _h = eval(h_expr)
+            h = x -> T(_h(x))
         else
             h = nothing
         end
     else
-        sjac = Symbolics.jacobian(f(x), x)
-        f_expr = Symbolics.build_function(sjac, x)[1]
-        g = eval(f_expr)
+        if sparse
+            sjac = Symbolics.sparsejacobian(f(x), x; simplify)
+            f_expr = Symbolics.build_function(sjac, x)[1]
+            _g = eval(f_expr)
+            g = x -> Matrix(T(_g(x)))
+        else
+            sjac = Symbolics.jacobian(f(x), x; simplify)
+            f_expr = Symbolics.build_function(sjac, x)[1]
+            _g = eval(f_expr)
+            g = x -> T(_g(x))
+        end
         if hessian
-            shess = Symbolics.jacobian(vec(Base.invokelatest(g, x)), x)
+            if sparse
+                shess = Symbolics.sparsejacobian(vec(Base.invokelatest(g, x)), x; simplify)
+            else
+                shess = Symbolics.jacobian(vec(Base.invokelatest(g, x)), x; simplify)
+            end
             h_expr = Symbolics.build_function(shess, x)[1]
             _h = eval(h_expr)
-            h = x -> reshape(_h(x), length(val), N, N)
+            h = x -> reshape(T(_h(x)), length(val), N, N)
         else
             h = nothing
         end
@@ -67,3 +100,38 @@ function ChainRulesCore.frule(
     end
 end
 @ForwardDiff_frule (f::SymbolicFunction)(x::AbstractVector{<:ForwardDiff.Dual})
+
+function symbolify(model::NonconvexCore.AbstractModel; objective = true, ineq_constraints = true, eq_constraints = true, kwargs...)
+    vmodel, v, _ = NonconvexCore.tovecmodel(model)
+    if objective
+        # Objective
+        sparse_flat_obj = SymbolicFunction(vmodel.objective, v; kwargs...)
+        obj = NonconvexCore.Objective(x -> sparse_flat_obj(flatten(x)[1]), flags = model.objective.flags)
+    else
+        obj = model.objective
+    end
+    if ineq_constraints
+        ineq = length(vmodel.ineq_constraints.fs) != 0 ? NonconvexCore.VectorOfFunctions(map(vmodel.ineq_constraints.fs) do c
+            sparse_flat_ineq = SymbolicFunction(c, v; kwargs...)
+            NonconvexCore.IneqConstraint(x -> sparse_flat_ineq(flatten(x)[1]), c.rhs, c.dim, c.flags)
+        end) : NonconvexCore.VectorOfFunctions(NonconvexCore.IneqConstraint[])
+    else
+        ineq = model.ineq_constraints
+    end
+    if eq_constraints
+        eq = length(vmodel.eq_constraints.fs) != 0 ? NonconvexCore.VectorOfFunctions(map(vmodel.eq_constraints.fs) do c
+            sparse_flat_eq = SymbolicFunction(c, v; kwargs...)
+            NonconvexCore.EqConstraint(x -> sparse_flat_eq(flatten(x)[1]), c.rhs, c.dim, c.flags)
+        end) : NonconvexCore.VectorOfFunctions(NonconvexCore.EqConstraint[])
+    else
+        eq = model.eq_constraints
+    end
+    if model isa NonconvexCore.Model
+        ModelT = NonconvexCore.Model
+    elseif model isa NonconvexCore.DictModel
+        ModelT = NonconvexCore.DictModel
+    else
+        throw("Unsupported model type.")
+    end
+    return ModelT(obj, eq, ineq, model.sd_constraints, model.box_min, model.box_max, model.init, model.integer)
+end

test/runtests.jl

@@ -1,10 +1,11 @@
 using NonconvexUtils, ForwardDiff, ReverseDiff, Tracker, Zygote
 using Test, LinearAlgebra, SparseArrays, NLsolve, IterativeSolvers
-using StableRNGs, ChainRulesCore, NonconvexCore
+using StableRNGs, ChainRulesCore, NonconvexCore, NonconvexIpopt
 
 include("forwarddiff_frule.jl")
 include("abstractdiff.jl")
 include("trace.jl")
 include("custom.jl")
 include("implicit.jl")
 include("symbolic.jl")
+include("sparse_forwarddiff.jl")

test/sparse_forwarddiff.jl

@@ -0,0 +1,39 @@
+@testset "SparseForwardDiffFunction" begin
+    @testset "Derivatives" begin
+        f = SparseForwardDiffFunction(sum, rand(3); hessian = false)
+        x = rand(3)
+        @test Zygote.gradient(f, x)[1] ≈ ForwardDiff.gradient(f, rand(3))
+
+        f = SparseForwardDiffFunction(x -> 2(x.^2) + x[1] * ones(3), rand(3); hessian = false)
+        x = rand(3)
+        @test Zygote.jacobian(f, x)[1] ≈ ForwardDiff.jacobian(f, x)
+
+        f = SparseForwardDiffFunction(sum, rand(3); hessian = true)
+        x = rand(3)
+        @test Zygote.gradient(f, x)[1] ≈ ForwardDiff.gradient(f, rand(3))
+        @test Zygote.hessian(f, x) ≈ ForwardDiff.hessian(f, rand(3))
+
+        f = SparseForwardDiffFunction(x -> sum(x)^2 + x[1], rand(3); hessian = true)
+        x = rand(3)
+        @test Zygote.hessian(f, x) ≈ ForwardDiff.hessian(f, x)
+
+        f = SparseForwardDiffFunction(x -> sum(x)^2 + x[1], rand(3); hessian = true)
+        x = rand(3)
+        @test Zygote.jacobian(x -> Zygote.gradient(f, x)[1], x)[1] ≈ ForwardDiff.hessian(f, x)
+    end
+    @testset "sparsify - first order = $first_order" for first_order in (true, false)
+        f = (x::AbstractVector) -> sqrt(x[2])
+        g = (x::AbstractVector, a, b) -> (a*x[1] + b)^3 - x[2]
+        options = IpoptOptions(first_order = first_order)
+        m = Model(f)
+        addvar!(m, [0.0, 0.0], [10.0, 10.0])
+        add_ineq_constraint!(m, x -> g(x, 2, 0))
+        add_ineq_constraint!(m, x -> g(x, -1, 1))
+
+        alg = IpoptAlg()
+        sp_model = sparsify(m)
+        r = NonconvexIpopt.optimize(sp_model, alg, [1.234, 2.345], options = options)
+        @test abs(r.minimum - sqrt(8/27)) < 1e-6
+        @test norm(r.minimizer - [1/3, 8/27]) < 1e-6    
+    end
+end