Merge pull request #2492 from SciML/optproboop

OptimizationSystem: Generate oop constraint derivatives

Author: ChrisRackauckas

src/systems/optimization/optimizationsystem.jl

@@ -364,14 +364,32 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
             linenumbers = linenumbers,
             expression = Val{false})
         if cons_j
-            _cons_j = generate_jacobian(cons_sys; expression = Val{false},
-                sparse = cons_sparse)[2]
+            _cons_j = let (cons_jac_oop, cons_jac_iip) = generate_jacobian(cons_sys;
+                    checkbounds = checkbounds,
+                    linenumbers = linenumbers,
+                    parallel = parallel, expression = Val{false},
+                    sparse = cons_sparse)
+                _cons_j(u, p) = cons_jac_oop(u, p)
+                _cons_j(J, u, p) = (cons_jac_iip(J, u, p); J)
+                _cons_j(u, p::MTKParameters) = cons_jac_oop(u, p...)
+                _cons_j(J, u, p::MTKParameters) = (cons_jac_iip(J, u, p...); J)
+                _cons_j
+            end
         else
             _cons_j = nothing
         end
         if cons_h
-            _cons_h = generate_hessian(cons_sys; expression = Val{false},
-                sparse = cons_sparse)[2]
+            _cons_h = let (cons_hess_oop, cons_hess_iip) = generate_hessian(
+                    cons_sys, checkbounds = checkbounds,
+                    linenumbers = linenumbers,
+                    sparse = cons_sparse, parallel = parallel,
+                    expression = Val{false})
+                _cons_h(u, p) = cons_hess_oop(u, p)
+                _cons_h(J, u, p) = (cons_hess_iip(J, u, p); J)
+                _cons_h(u, p::MTKParameters) = cons_hess_oop(u, p...)
+                _cons_h(J, u, p::MTKParameters) = (cons_hess_iip(J, u, p...); J)
+                _cons_h
+            end
         else
             _cons_h = nothing
         end

test/optimizationsystem.jl

@@ -312,3 +312,31 @@ end
     @test all(prob.f.cons_expr[i].lhs isa Symbolics.Symbolic
     for i in 1:length(prob.f.cons_expr))
 end
+
+@testset "Derivatives, iip and oop" begin
+    @variables x y
+    @parameters a b
+    loss = (a - x)^2 + b * (y - x^2)^2
+    cons2 = [x^2 + y^2 ~ 0, y * sin(x) - x ~ 0]
+    sys = complete(OptimizationSystem(
+        loss, [x, y], [a, b], name = :sys2, constraints = cons2))
+    prob = OptimizationProblem(sys, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0],
+        grad = true, hess = true, cons_j = true, cons_h = true)
+
+    G1 = Array{Float64}(undef, 2)
+    H1 = Array{Float64}(undef, 2, 2)
+    J = Array{Float64}(undef, 2, 2)
+    H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
+
+    prob.f.grad(G1, [1.0, 1.0], [1.0, 100.0])
+    @test prob.f.grad([1.0, 1.0], [1.0, 100.0]) == G1
+
+    prob.f.hess(H1, [1.0, 1.0], [1.0, 100.0])
+    @test prob.f.hess([1.0, 1.0], [1.0, 100.0]) == H1
+
+    prob.f.cons_j(J, [1.0, 1.0], [1.0, 100.0])
+    @test prob.f.cons_j([1.0, 1.0], [1.0, 100.0]) == J
+
+    prob.f.cons_h(H3, [1.0, 1.0], [1.0, 100.0])
+    @test prob.f.cons_h([1.0, 1.0], [1.0, 100.0]) == H3
+end