Skip to content

Commit d994326

Browse files
fix: fix constraint function not wrapped in OptimizationProblem
1 parent 57ea5fc commit d994326

File tree

2 files changed

+15
-2
lines changed

2 files changed

+15
-2
lines changed

src/systems/optimization/optimizationsystem.jl

+7-2
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,12 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
399399
cons, lcons_, ucons_ = generate_function(cons_sys, checkbounds = checkbounds,
400400
linenumbers = linenumbers,
401401
expression = Val{true})
402-
cons = eval_or_rgf.(cons; eval_expression, eval_module)
402+
cons = let (cons_oop, cons_iip) = eval_or_rgf.(cons; eval_expression, eval_module)
403+
_cons(u, p) = cons_oop(u, p)
404+
_cons(resid, u, p) = cons_iip(resid, u, p)
405+
_cons(u, p::MTKParameters) = cons_oop(u, p...)
406+
_cons(resid, u, p::MTKParameters) = cons_iip(resid, u, p...)
407+
end
403408
if cons_j
404409
_cons_j = let (cons_jac_oop, cons_jac_iip) = eval_or_rgf.(
405410
generate_jacobian(cons_sys;
@@ -465,7 +470,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
465470
grad = _grad,
466471
hess = _hess,
467472
hess_prototype = hess_prototype,
468-
cons = cons[2],
473+
cons = cons,
469474
cons_j = _cons_j,
470475
cons_h = _cons_h,
471476
cons_jac_prototype = cons_jac_prototype,

test/optimizationsystem.jl

+8
Original file line numberDiff line numberDiff line change
@@ -368,3 +368,11 @@ end
368368
@test is_variable(sys, x[2])
369369
@test is_variable(sys, x[3])
370370
end
371+
372+
@testset "Constraints work with nonnumeric parameters" begin
373+
@variables x
374+
@parameters p f(::Real)
375+
@mtkbuild sys = OptimizationSystem(x^2 + f(x) * p, [x], [f, p]; constraints = [2.0 f(x) + p])
376+
prob = OptimizationProblem(sys, [x => 1.0], [p => 1.0, f => (x -> 2x)])
377+
@test abs(prob.f.cons(prob.u0, prob.p)[1]) 1.0
378+
end

0 commit comments

Comments
 (0)