Add structural_simplify support for OptimizationSystem

src/systems/abstractsystem.jl

@@ -1031,15 +1031,17 @@ This will convert all `inputs` to parameters and allow them to be unconnected, i
 simplification will allow models where `n_states = n_equations - n_inputs`.
 """
 function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
-                             simplify_constants = true, kwargs...)
+                             simplify_constants = true, check_consistency = true, kwargs...)
     sys = expand_connections(sys)
     sys isa DiscreteSystem && return sys
     state = TearingState(sys)
     has_io = io !== nothing
     has_io && markio!(state, io...)
     state, input_idxs = inputs_to_parameters!(state, io)
     sys, ag = alias_elimination!(state; kwargs...)
-    check_consistency(state, ag)
+    if check_consistency
+        ModelingToolkit.check_consistency(state, ag)
+    end
     sys = dummy_derivative(sys, state, ag; simplify)
     fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
     @set! sys.observed = topsort_equations(observed(sys), fullstates)

src/systems/optimization/optimizationsystem.jl

@@ -543,3 +543,32 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
         end
     end
 end
+
+function structural_simplify(sys::OptimizationSystem; kwargs...)
+    sys = flatten(sys)
+    cons = constraints(sys)
+    econs = Equation[]
+    icons = similar(cons, 0)
+    for e in cons
+        if e isa Equation
+            push!(econs, e)
+        else
+            push!(icons, e)
+        end
+    end
+    nlsys = NonlinearSystem(econs, states(sys), parameters(sys); name = :___tmp_nlsystem)
+    snlsys = structural_simplify(nlsys; check_consistency = false, kwargs...)
+    obs = observed(snlsys)
+    subs = Dict(eq.lhs => eq.rhs for eq in observed(snlsys))
+    seqs = equations(snlsys)
+    sizehint!(icons, length(icons) + length(seqs))
+    for eq in seqs
+        push!(icons, substitute(eq, subs))
+    end
+    newsts = setdiff(states(sys), keys(subs))
+    @set! sys.constraints = icons
+    @set! sys.observed = [observed(sys); obs]
+    @set! sys.op = substitute(equations(sys), subs)
+    @set! sys.states = newsts
+    return sys
+end

test/optimizationsystem.jl

@@ -69,25 +69,27 @@ end
 end
 
 @testset "equality constraint" begin
-    @variables x y
+    @variables x y z
     @parameters a b
-    loss = (a - x)^2 + b * (y - x^2)^2
-    cons = [1.0 ~ x^2 + y^2]
-    @named sys = OptimizationSystem(loss, [x, y], [a, b], constraints = cons)
-    prob = OptimizationProblem(sys, [x => 0.0, y => 0.0], [a => 1.0, b => 1.0],
+    loss = (a - x)^2 + b * z^2
+    cons = [1.0 ~ x^2 + y^2
+            z ~ y - x^2]
+    @named sys = OptimizationSystem(loss, [x, y, z], [a, b], constraints = cons)
+    sys = structural_simplify(sys)
+    prob = OptimizationProblem(sys, [x => 0.0, y => 0.0, z => 0.0], [a => 1.0, b => 1.0],
                                grad = true, hess = true)
     sol = solve(prob, IPNewton())
     @test sol.minimum < 1.0
-    @test sol.u≈[0.808, 0.589] atol=1e-3
-    @test sol[x]^2 + sol[y]^2 ≈ 1.0
+    @test sol.u≈[0.808, -0.064] atol=1e-3
+    @test_broken sol[x]^2 + sol[y]^2 ≈ 1.0
     sol = solve(prob, Ipopt.Optimizer(); print_level = 0)
     @test sol.minimum < 1.0
-    @test sol.u≈[0.808, 0.589] atol=1e-3
-    @test sol[x]^2 + sol[y]^2 ≈ 1.0
+    @test sol.u≈[0.808, -0.064] atol=1e-3
+    @test_broken sol[x]^2 + sol[y]^2 ≈ 1.0
     sol = solve(prob, AmplNLWriter.Optimizer(Ipopt_jll.amplexe))
     @test sol.minimum < 1.0
-    @test sol.u≈[0.808, 0.589] atol=1e-3
-    @test sol[x]^2 + sol[y]^2 ≈ 1.0
+    @test sol.u≈[0.808, -0.064] atol=1e-3
+    @test_broken sol[x]^2 + sol[y]^2 ≈ 1.0
 end
 
 @testset "rosenbrock" begin