Refactor constraints

Author: vyudu

src/ModelingToolkit.jl

@@ -150,6 +150,10 @@ include("systems/imperative_affect.jl")
 include("systems/callbacks.jl")
 include("systems/problem_utils.jl")
 
+include("systems/optimization/constraints_system.jl")
+include("systems/optimization/optimizationsystem.jl")
+include("systems/optimization/modelingtoolkitize.jl")
+
 include("systems/nonlinear/nonlinearsystem.jl")
 include("systems/nonlinear/homotopy_continuation.jl")
 include("systems/diffeqs/odesystem.jl")
@@ -165,10 +169,6 @@ include("systems/discrete_system/discrete_system.jl")
 
 include("systems/jumps/jumpsystem.jl")
 
-include("systems/optimization/constraints_system.jl")
-include("systems/optimization/optimizationsystem.jl")
-include("systems/optimization/modelingtoolkitize.jl")
-
 include("systems/pde/pdesystem.jl")
 
 include("systems/sparsematrixclil.jl")

src/systems/abstractsystem.jl

@@ -983,6 +983,7 @@ for prop in [:eqs
              :structure
              :op
              :constraints
+             :constraintsystem
              :controls
              :loss
              :bcs

src/systems/diffeqs/abstractodesystem.jl

@@ -827,6 +827,12 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEProblem`")
     end
+
+    if !isnothing(get_constraintsystem(sys))
+        error("An ODESystem with constraints cannot be used to construct a regular ODEProblem. 
+              Consider a BVProblem instead.")
+    end
+
     f, u0, p = process_SciMLProblem(ODEFunction{iip, specialize}, sys, u0map, parammap;
         t = tspan !== nothing ? tspan[1] : tspan,
         check_length, warn_initialize_determined, eval_expression, eval_module, kwargs...)
@@ -866,18 +872,23 @@ Create a boundary value problem from the [`ODESystem`](@ref).
 must have either an initial guess supplied using `guesses` or a fixed initial 
 value specified using `u0map`.
 
-`constraints` are used to specify boundary conditions to the ODESystem in the
-form of equations. These values should specify values that state variables should
+Boundary value conditions are supplied to ODESystems
+in the form of a ConstraintsSystem. These equations 
+should specify values that state variables should
 take at specific points, as in `x(0.5) ~ 1`). More general constraints that 
 should hold over the entire solution, such as `x(t)^2 + y(t)^2`, should be 
-specified as one of the equations used to build the `ODESystem`. Below is an example.
+specified as one of the equations used to build the `ODESystem`.
+
+If an ODESystem without `constraints` is specified, it will be treated as an initial value problem. 
 
 ```julia
-    @parameters g
+    @parameters g t_c = 0.5
     @variables x(..) y(t) [state_priority = 10] λ(t)
     eqs = [D(D(x(t))) ~ λ * x(t)
            D(D(y)) ~ λ * y - g
            x(t)^2 + y^2 ~ 1]
+    cstr = [x(0.5) ~ 1]
+    @named cstrs = ConstraintsSystem(cstr, t)
     @mtkbuild pend = ODESystem(eqs, t)
 
     tspan = (0.0, 1.5)
@@ -889,9 +900,7 @@ specified as one of the equations used to build the `ODESystem`. Below is an exa
     bvp = SciMLBase.BVProblem{true, SciMLBase.AutoSpecialize}(pend, u0map, tspan, parammap; constraints, guesses, check_length = false)
 ```
 
-If no `constraints` are specified, the problem will be treated as an initial value problem.
-
-If the `ODESystem` has algebraic equations like `x(t)^2 + y(t)^2`, the resulting 
+If the `ODESystem` has algebraic equations, like `x(t)^2 + y(t)^2`, the resulting 
 `BVProblem` must be solved using BVDAE solvers, such as Ascher.
 """
 function SciMLBase.BVProblem(sys::AbstractODESystem, args...; kwargs...)
@@ -916,7 +925,7 @@ end
 function SciMLBase.BVProblem{iip, specialize}(sys::AbstractODESystem, u0map = [],
         tspan = get_tspan(sys),
         parammap = DiffEqBase.NullParameters();
-        constraints = nothing, guesses = Dict(),
+        guesses = Dict(),
         version = nothing, tgrad = false,
         callback = nothing,
         check_length = true,
@@ -930,21 +939,14 @@ function SciMLBase.BVProblem{iip, specialize}(sys::AbstractODESystem, u0map = []
     end
     !isnothing(callback) && error("BVP solvers do not support callbacks.")
 
-    has_alg_eqs(sys) && error("The BVProblem currently does not support ODESystems with algebraic equations.") # Remove this when the BVDAE solvers get updated, the codegen should work when it does.
+    has_alg_eqs(sys) && error("The BVProblem constructor currently does not support ODESystems with algebraic equations.") # Remove this when the BVDAE solvers get updated, the codegen should work when it does.
 
-    constraintsts = nothing
-    constraintps = nothing
     sts = unknowns(sys)
     ps = parameters(sys)
 
-    # Constraint validation
     if !isnothing(constraints)
-        constraints isa Equation || 
-            constraints isa Vector{Equation} || 
-            error("Constraints must be specified as an equation or a vector of equations.")
-
         (length(constraints) + length(u0map) > length(sts)) && 
-        error("The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) cannot exceed the total number of states.")
+        @warn "The BVProblem is overdetermined. The total number of conditions (# constraints + # fixed initial values given by u0map) exceeds the total number of states. The BVP solvers will default to doing a nonlinear least-squares optimization."
     end
 
     # ODESystems without algebraic equations should use both fixed values + guesses
@@ -957,97 +959,60 @@ function SciMLBase.BVProblem{iip, specialize}(sys::AbstractODESystem, u0map = []
     stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
     u0_idxs = has_alg_eqs(sys) ? collect(1:length(sts)) : [stidxmap[k] for (k,v) in u0map]
 
-    bc = process_constraints(sys, constraints, u0, u0_idxs, tspan, iip)
-
+    bc = generate_function_bc(sys, u0, u0_idxs, tspan, iip)
     return BVProblem{iip}(f, bc, u0, tspan, p; kwargs...)
 end
 
 get_callback(prob::BVProblem) = error("BVP solvers do not support callbacks.")
 
-# Validate that all the variables in the BVP constraints are well-formed states or parameters.
-function validate_constraint_syms(eq, constraintsts, constraintps, sts, ps, iv) 
-    for var in constraintsts
-        if length(arguments(var)) > 1
-            error("Too many arguments for variable $var.")
-        elseif isequal(arguments(var)[1], iv)
-            var ∈ sts || error("Constraint equation $eq contains a variable $var that is not a variable of the ODESystem.")
-            error("Constraint equation $eq contains a variable $var that does not have a specified argument. Such equations should be specified as algebraic equations to the ODESystem rather than a boundary constraints.")
-        else
-            operation(var)(iv) ∈ sts || error("Constraint equation $eq contains a variable $(operation(var)) that is not a variable of the ODESystem.")
-        end
-    end
-
-    for var in constraintps
-        if !iscall(var)
-            var ∈ ps || error("Constraint equation $eq contains a parameter $var that is not a parameter of the ODESystem.")
-        else
-            length(arguments(var)) > 1 && error("Too many arguments for parameter $var.")
-            operation(var) ∈ ps || error("Constraint equations contain a parameter $var that is not a parameter of the ODESystem.")
-        end
-    end
-end
-
 """
-    process_constraints(sys, constraints, u0, tspan, iip)
+    generate_function_bc(sys::ODESystem, u0, u0_idxs, tspan, iip)
 
-    Given an ODESystem with some constraints, generate the boundary condition function.
+    Given an ODESystem with constraints, generate the boundary condition function to pass to boundary value problem solvers.
 """
-function process_constraints(sys::ODESystem, constraints, u0, u0_idxs, tspan, iip)
-
+function generate_function_bc(sys::ODESystem, u0, u0_idxs, tspan, iip)
     iv = get_iv(sys)
     sts = get_unknowns(sys)
     ps = get_ps(sys)
     np = length(ps)
     ns = length(sts)
+    conssys = get_constraintsystem(sys)
+    cons = constraints(conssys)
 
     stidxmap = Dict([v => i for (i, v) in enumerate(sts)])
     pidxmap = Dict([v => i for (i, v) in enumerate(ps)])
 
     @variables sol(..)[1:ns] p[1:np]
     exprs = Any[]
 
-    constraintsts = OrderedSet()
-    constraintps = OrderedSet()
+    for st in get_unknowns(cons)
+        x = operation(st)
+        t = first(arguments(st))
+        idx = stidxmap[x(iv)]
 
-    !isnothing(constraints) && for cons in constraints
-        collect_vars!(constraintsts, constraintps, cons, iv)
-        validate_constraint_syms(cons, constraintsts, constraintps, Set(sts), Set(ps), iv)
-        expr = cons.rhs - cons.lhs
-
-        for st in constraintsts
-            x = operation(st)
-            t = arguments(st)[1]
-            idx = stidxmap[x(iv)]
-
-            expr = Symbolics.substitute(expr, Dict(x(t) => sol(t)[idx]))
-        end
+        cons = Symbolics.substitute(cons, Dict(x(t) => sol(t)[idx]))
+    end
 
-        for var in constraintps
-            if iscall(var)
-                x = operation(var)
-                t = arguments(var)[1]
-                idx = pidxmap[x]
+    for var in get_parameters(cons) 
+        if iscall(var)
+            x = operation(var)
+            t = arguments(var)[1]
+            idx = pidxmap[x]
 
-                expr = Symbolics.substitute(expr, Dict(x(t) => p[idx]))
-            else
-                idx = pidxmap[var]
-                expr = Symbolics.substitute(expr, Dict(var => p[idx]))
-            end
+            cons = Symbolics.substitute(cons, Dict(x(t) => p[idx]))
+        else
+            idx = pidxmap[var]
+            cons = Symbolics.substitute(cons, Dict(var => p[idx]))
         end
-
-        empty!(constraintsts)
-        empty!(constraintps)
-        push!(exprs, expr)
     end
 
-    init_cond_exprs = Any[]
-
+    init_conds = Any[]
     for i in u0_idxs
         expr = sol(tspan[1])[i] - u0[i]
-        push!(init_cond_exprs, expr)
+        push!(init_conds, expr)
     end
 
-    exprs = vcat(init_cond_exprs, exprs)
+    exprs = vcat(init_conds, cons)
     bcs = Symbolics.build_function(exprs, sol, p, expression = Val{false})
     if iip
         return (resid, u, p, t) -> bcs[2](resid, u, p)