diff --git a/lib/OptimizationBase/src/OptimizationDISparseExt.jl b/lib/OptimizationBase/src/OptimizationDISparseExt.jl index 392d82882..a72c42cbd 100644 --- a/lib/OptimizationBase/src/OptimizationDISparseExt.jl +++ b/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 diff --git a/lib/OptimizationBase/test/AD/adtests.jl b/lib/OptimizationBase/test/AD/adtests.jl index 5fbe3caef..573c048f1 100644 --- a/lib/OptimizationBase/test/AD/adtests.jl +++ b/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}( diff --git a/lib/OptimizationIpopt/test/core_tests.jl b/lib/OptimizationIpopt/test/core_tests.jl index c6fb82780..0b4f668b8 100644 --- a/lib/OptimizationIpopt/test/core_tests.jl +++ b/lib/OptimizationIpopt/test/core_tests.jl @@ -106,6 +106,43 @@ end @test sol.u ≈ [1.0, 0.0] atol = 1.0e-6 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")