Use live parameters in sparse constraint Jacobians

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/OptimizationBase/src/OptimizationDISparseExt.jl

@@ -20,6 +20,8 @@ function instantiate_function(
         lag_h = false
     )
     adtype, soadtype = generate_sparse_adtype(adtype)
+    Tx0 = typeof(x)
+    Tp0 = typeof(p)
 
     if g == true && f.grad === nothing
         prep_grad = prepare_gradient(f.f, adtype.dense_ad, x, Constant(p))
@@ -167,17 +169,31 @@ function instantiate_function(
     cons_jac_prototype = f.cons_jac_prototype
     cons_jac_colorvec = f.cons_jac_colorvec
     if f.cons !== nothing && cons_j == true && f.cons_j === nothing
-        prep_jac = prepare_jacobian(cons_oop, adtype, x)
-        function cons_j!(J, θ)
-            jacobian!(cons_oop, J, prep_jac, adtype, θ)
-            return if size(J, 1) == 1
-                J = vec(J)
+        cons_oop_p = let f = f, num_cons = num_cons
+            function (x, p)
+                res = Vector{_cons_out_eltype(x, p)}(undef, num_cons)
+                f.cons(res, x, p)
+                return res
+            end
+        end
+        prep_jac = prepare_jacobian(cons_oop_p, adtype, x, Constant(p))
+        cons_j! = let cons_oop_p = cons_oop_p, prep_jac = prep_jac,
+                adtype = adtype, p = p, Tx0 = Tx0, Tp0 = Tp0
+            function (J, θ, p = p)
+                if _prep_valid(Tx0, θ) && _prep_valid(Tp0, p)
+                    jacobian!(cons_oop_p, J, prep_jac, adtype, θ, Constant(p))
+                else
+                    jacobian!(cons_oop_p, J, adtype, θ, Constant(p))
+                end
+                return size(J, 1) == 1 ? vec(J) : J
             end
         end
         cons_jac_prototype = prep_jac.coloring_result.A
         cons_jac_colorvec = prep_jac.coloring_result.color
     elseif cons_j === true && f.cons !== nothing
-        cons_j! = (J, θ) -> f.cons_j(J, θ, p)
+        cons_j! = let f = f, p = p
+            (J, θ, p = p) -> f.cons_j(J, θ, p)
+        end
     else
         cons_j! = nothing
     end
@@ -339,6 +355,8 @@ function instantiate_function(
         lag_h = false
     )
     adtype, soadtype = generate_sparse_adtype(adtype)
+    Tx0 = typeof(x)
+    Tp0 = typeof(p)
 
     if g == true && f.grad === nothing
         prep_grad = prepare_gradient(f.f, adtype.dense_ad, x, Constant(p))
@@ -462,17 +480,23 @@ function instantiate_function(
     cons_jac_colorvec = f.cons_jac_colorvec
     if f.cons !== nothing && cons_j == true && f.cons_j === nothing
         prep_jac = prepare_jacobian(f.cons, adtype, x, Constant(p))
-        function cons_j!(θ)
-            J = jacobian(f.cons, prep_jac, adtype, θ, Constant(p))
-            if size(J, 1) == 1
-                J = vec(J)
+        cons_j! = let f = f, prep_jac = prep_jac, adtype = adtype,
+                p = p, Tx0 = Tx0, Tp0 = Tp0
+            function (θ, p = p)
+                J = if _prep_valid(Tx0, θ) && _prep_valid(Tp0, p)
+                    jacobian(f.cons, prep_jac, adtype, θ, Constant(p))
+                else
+                    jacobian(f.cons, adtype, θ, Constant(p))
+                end
+                return size(J, 1) == 1 ? vec(J) : J
             end
-            return J
         end
         cons_jac_prototype = prep_jac.coloring_result.A
         cons_jac_colorvec = prep_jac.coloring_result.color
     elseif cons_j === true && f.cons !== nothing
-        cons_j! = (θ) -> f.cons_j(θ, p)
+        cons_j! = let f = f, p = p
+            (θ, p = p) -> f.cons_j(θ, p)
+        end
     else
         cons_j! = nothing
     end

lib/OptimizationBase/test/AD/adtests.jl

@@ -911,6 +911,48 @@ end
     @test nnz(lag_H) == 5
 end
 
+@testset "sparse constraint Jacobians use live parameters" begin
+    objective(x, p) = sum(abs2, x)
+    constraints(x, p) = [p[1] * x[1]]
+    function constraints!(res, x, p)
+        res[1] = p[1] * x[1]
+        return nothing
+    end
+    constraint_jacobian(x, p) = sparse([1], [1], [p[1]], 1, 2)
+    function constraint_jacobian!(J, x, p)
+        J[1, 1] = p[1]
+        return nothing
+    end
+
+    x = [1.0, 1.0]
+    initial_p = [2.0]
+    live_p = [3.0]
+
+    @testset "in-place = $iip, analytic = $analytic" for iip in (false, true),
+            analytic in (false, true)
+        jacobian_kwargs = analytic ?
+            (;
+                cons_j = iip ? constraint_jacobian! : constraint_jacobian,
+                cons_jac_prototype = sparse([1], [1], [1.0], 1, 2),
+            ) : (;)
+        optf = OptimizationFunction{iip}(
+            objective, AutoSparse(AutoForwardDiff());
+            cons = iip ? constraints! : constraints, jacobian_kwargs...
+        )
+        optprob = OptimizationBase.instantiate_function(
+            optf, x, optf.adtype, initial_p, 1; cons_j = true
+        )
+
+        if iip
+            J = similar(optprob.cons_jac_prototype, Float64)
+            optprob.cons_j(J, x, live_p)
+            @test vec(Array(J)) == [3.0, 0.0]
+        else
+            @test vec(Array(optprob.cons_j(x, live_p))) == [3.0, 0.0]
+        end
+    end
+end
+
 @testset "OOP" begin
     cons = (x, p) -> [x[1]^2 + x[2]^2]
     optf = OptimizationFunction{false}(

lib/OptimizationIpopt/test/core_tests.jl

@@ -71,6 +71,43 @@ end
     end
 end
 
+@testset "sparse constraint Jacobian with parameters" begin
+    objective(x, p) = sum(abs2, x)
+    function constraints!(res, x, p)
+        res[1] = p[1] * x[1]
+        return nothing
+    end
+    function constraint_jacobian!(J, x, p)
+        J[1, 1] = p[1]
+        return nothing
+    end
+
+    @testset "analytic = $analytic" for analytic in (false, true)
+        jacobian_kwargs = analytic ?
+            (;
+                cons_j = constraint_jacobian!,
+                cons_jac_prototype = sparse([1], [1], [1.0], 1, 2),
+            ) : (;)
+        f = OptimizationFunction{true}(
+            objective, AutoSparse(AutoForwardDiff());
+            cons = constraints!, jacobian_kwargs...
+        )
+        prob = OptimizationProblem(
+            f, [1.0, 1.0], [2.0]; lcons = [0.0], ucons = [0.0]
+        )
+        sol = solve(
+            prob,
+            IpoptOptimizer(;
+                hessian_approximation = "limited-memory",
+                additional_options = Dict{String, Any}("print_level" => 0)
+            )
+        )
+
+        @test SciMLBase.successful_retcode(sol)
+        @test isapprox(sol.u[1], 0.0; atol = 1.0e-8)
+    end
+end
+
 include("additional_tests.jl")
 include("advanced_features.jl")
 include("problem_types.jl")