Merge remote-tracking branch 'vyudu/BVP-with-constraints' into BVP-with-constraints

Author: vyudu

docs/src/basics/Debugging.md

@@ -35,6 +35,39 @@ dsol = solve(dprob, Tsit5());
 Now we see that it crashed because `u1` decreased so much that it became negative and outside the domain of the `√` function.
 We could have figured that out ourselves, but it is not always so obvious for more complex models.
 
+Suppose we also want to validate that `u1 + u2 >= 2.0`. We can do this via the assertions functionality.
+
+```@example debug
+@mtkbuild sys = ODESystem(eqs, t; defaults, assertions = [(u1 + u2 >= 2.0) => "Oh no!"])
+```
+
+The assertions must be an iterable of pairs, where the first element is the symbolic condition and
+the second is a message to be logged when the condition fails. All assertions are added to the
+generated code and will cause the solver to reject steps that fail the assertions. For systems such
+as the above where the assertion is guaranteed to eventually fail, the solver will likely exit
+with a `dtmin` failure..
+
+```@example debug
+prob = ODEProblem(sys, [], (0.0, 10.0))
+sol = solve(prob, Tsit5())
+```
+
+We can use `debug_system` to log the failing assertions in each call to the RHS function.
+
+```@repl debug
+dsys = debug_system(sys; functions = []);
+dprob = ODEProblem(dsys, [], (0.0, 10.0));
+dsol = solve(dprob, Tsit5());
+```
+
+Note the logs containing the failed assertion and corresponding message. To temporarily disable
+logging in a system returned from `debug_system`, use `ModelingToolkit.ASSERTION_LOG_VARIABLE`.
+
+```@repl debug
+dprob[ModelingToolkit.ASSERTION_LOG_VARIABLE] = false;
+solve(drob, Tsit5());
+```
+
 ```@docs
 debug_system
 ```

docs/src/basics/Variable_metadata.md

@@ -54,6 +54,11 @@ using ModelingToolkit: t_nounits as t, D_nounits as D
 
 @variables i(t) [connect = Flow]
 @variables k(t) [connect = Stream]
+hasconnect(i)
+```
+
+```@example connect
+getconnect(k)
 ```
 
 ## Input or output
@@ -177,8 +182,45 @@ A variable can be marked `irreducible` to prevent it from being moved to an
 `observed` state. This forces the variable to be computed during solving so that
 it can be accessed in [callbacks](@ref events)
 
-```julia
-@variable important_value [irreducible = true]
+```@example metadata
+@variables important_value [irreducible = true]
+isirreducible(important_value)
+```
+
+## State Priority
+
+When a model is structurally simplified, the algorithm will try to ensure that the variables with higher state priority become states of the system. A variable's state priority is a number set using the `state_priority` metadata.
+
+```@example metadata
+@variables important_dof [state_priority = 10] unimportant_dof [state_priority = -2]
+state_priority(important_dof)
+```
+
+## Units
+
+Units for variables can be designated using symbolic metadata. For more information, please see the [model validation and units](@ref units) section of the docs. Note that `getunit` is not equivalent to `get_unit` - the former is a metadata getter for individual variables (and is provided so the same interface function for `unit` exists like other metadata), while the latter is used to handle more general symbolic expressions.
+
+```@example metadata
+@variables speed [unit = u"m/s"]
+hasunit(speed)
+```
+
+```@example metadata
+getunit(speed)
+```
+
+## Miscellaneous metadata
+
+User-defined metadata can be added using the `misc` metadata. This can be queried
+using the `hasmisc` and `getmisc` functions.
+
+```@example metadata
+@variables u [misc = :conserved_parameter] y [misc = [2, 4, 6]]
+hasmisc(u)
+```
+
+```@example metadata
+getmisc(y)
 ```
 
 ## Additional functions

docs/src/tutorials/initialization.md

@@ -14,8 +14,10 @@ principles of initialization of DAE systems. Take a DAE written in semi-explicit
 form:
 
 ```math
-x' = f(x,y,t)\\
-0 = g(x,y,t)
+\begin{aligned}
+    x^\prime &= f(x,y,t) \\
+    0 &= g(x,y,t)
+\end{aligned}
 ```
 
 where ``x`` are the differential variables and ``y`` are the algebraic variables.

src/ModelingToolkit.jl

@@ -247,7 +247,9 @@ export initial_state, transition, activeState, entry, ticksInState, timeInState
 export @component, @mtkmodel, @mtkbuild
 export isinput, isoutput, getbounds, hasbounds, getguess, hasguess, isdisturbance,
        istunable, getdist, hasdist,
-       tunable_parameters, isirreducible, getdescription, hasdescription
+       tunable_parameters, isirreducible, getdescription, hasdescription,
+       hasunit, getunit, hasconnect, getconnect,
+       hasmisc, getmisc, state_priority
 export ode_order_lowering, dae_order_lowering, liouville_transform
 export PDESystem
 export Differential, expand_derivatives, @derivatives

src/debugging.jl

@@ -42,3 +42,59 @@ function debug_sub(ex, funcs; kw...)
     f in funcs ? logged_fun(f, args...; kw...) :
     maketerm(typeof(ex), f, args, metadata(ex))
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+A function which returns `NaN` if `condition` fails, and `0.0` otherwise.
+"""
+function _nan_condition(condition::Bool)
+    condition ? 0.0 : NaN
+end
+
+@register_symbolic _nan_condition(condition::Bool)
+
+"""
+    $(TYPEDSIGNATURES)
+
+A function which takes a condition `expr` and returns `NaN` if it is false,
+and zero if it is true. In case the condition is false and `log == true`,
+`message` will be logged as an `@error`.
+"""
+function _debug_assertion(expr::Bool, message::String, log::Bool)
+    value = _nan_condition(expr)
+    isnan(value) || return value
+    log && @error message
+    return value
+end
+
+@register_symbolic _debug_assertion(expr::Bool, message::String, log::Bool)
+
+"""
+Boolean parameter added to models returned from `debug_system` to control logging of
+assertions.
+"""
+const ASSERTION_LOG_VARIABLE = only(@parameters __log_assertions_ₘₜₖ::Bool = false)
+
+"""
+    $(TYPEDSIGNATURES)
+
+Get a symbolic expression for all the assertions in `sys`. The expression returns `NaN`
+if any of the assertions fail, and `0.0` otherwise. If `ASSERTION_LOG_VARIABLE` is a
+parameter in the system, it will control whether the message associated with each
+assertion is logged when it fails.
+"""
+function get_assertions_expr(sys::AbstractSystem)
+    asserts = assertions(sys)
+    term = 0
+    if is_parameter(sys, ASSERTION_LOG_VARIABLE)
+        for (k, v) in asserts
+            term += _debug_assertion(k, "Assertion $k failed:\n$v", ASSERTION_LOG_VARIABLE)
+        end
+    else
+        for (k, v) in asserts
+            term += _nan_condition(k)
+        end
+    end
+    return term
+end

src/systems/abstractsystem.jl

@@ -984,6 +984,7 @@ for prop in [:eqs
              :gui_metadata
              :discrete_subsystems
              :parameter_dependencies
+             :assertions
              :solved_unknowns
              :split_idxs
              :parent
@@ -1469,6 +1470,24 @@ end
 """
 $(TYPEDSIGNATURES)
 
+Get the assertions for a system `sys` and its subsystems.
+"""
+function assertions(sys::AbstractSystem)
+    has_assertions(sys) || return Dict{BasicSymbolic, String}()
+
+    asserts = get_assertions(sys)
+    systems = get_systems(sys)
+    namespaced_asserts = mapreduce(
+        merge!, systems; init = Dict{BasicSymbolic, String}()) do subsys
+        Dict{BasicSymbolic, String}(namespace_expr(k, subsys) => v
+        for (k, v) in assertions(subsys))
+    end
+    return merge(asserts, namespaced_asserts)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
 Get the guesses for variables in the initialization system of the system `sys` and its subsystems.
 
 See also [`initialization_equations`](@ref) and [`ModelingToolkit.get_guesses`](@ref).
@@ -2249,15 +2268,20 @@ macro mtkbuild(exprs...)
     expr = exprs[1]
     named_expr = ModelingToolkit.named_expr(expr)
     name = named_expr.args[1]
-    kwargs = if length(exprs) > 1
-        NamedTuple{Tuple(ex.args[1] for ex in Base.tail(exprs))}(Tuple(ex.args[2]
-        for ex in Base.tail(exprs)))
+    kwargs = Base.tail(exprs)
+    kwargs = map(kwargs) do ex
+        @assert ex.head == :(=)
+        Expr(:kw, ex.args[1], ex.args[2])
+    end
+    if isempty(kwargs)
+        kwargs = ()
     else
-        (;)
+        kwargs = (Expr(:parameters, kwargs...),)
     end
+    call_expr = Expr(:call, structural_simplify, kwargs..., name)
     esc(quote
         $named_expr
-        $name = $structural_simplify($name; $(kwargs)...)
+        $name = $call_expr
     end)
 end
 
@@ -2279,6 +2303,13 @@ ERROR: Function /(1, sin(P(t))) output non-finite value Inf with input
   1 => 1
   sin(P(t)) => 0.0
 ```
+
+Additionally, all assertions in the system are optionally logged when they fail.
+A new parameter is also added to the system which controls whether the message associated
+with each assertion will be logged when the assertion fails. This parameter defaults to
+`true` and can be toggled by symbolic indexing with
+`ModelingToolkit.ASSERTION_LOG_VARIABLE`. For example,
+`prob.ps[ModelingToolkit.ASSERTION_LOG_VARIABLE] = false` will disable logging.
 """
 function debug_system(
         sys::AbstractSystem; functions = [log, sqrt, (^), /, inv, asin, acos], kw...)
@@ -2289,11 +2320,17 @@ function debug_system(
         error("debug_system(sys) only works on systems with no sub-systems! Consider flattening it with flatten(sys) or structural_simplify(sys) first.")
     end
     if has_eqs(sys)
-        @set! sys.eqs = debug_sub.(equations(sys), Ref(functions); kw...)
+        eqs = debug_sub.(equations(sys), Ref(functions); kw...)
+        @set! sys.eqs = eqs
+        @set! sys.ps = unique!([get_ps(sys); ASSERTION_LOG_VARIABLE])
+        @set! sys.defaults = merge(get_defaults(sys), Dict(ASSERTION_LOG_VARIABLE => true))
     end
     if has_observed(sys)
         @set! sys.observed = debug_sub.(observed(sys), Ref(functions); kw...)
     end
+    if iscomplete(sys)
+        sys = complete(sys; split = is_split(sys))
+    end
     return sys
 end
 
@@ -3032,6 +3069,11 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
         kwargs = merge(kwargs, (initialization_eqs = ieqs, guesses = guesses))
     end
 
+    if has_assertions(basesys)
+        kwargs = merge(
+            kwargs, (; assertions = merge(get_assertions(basesys), get_assertions(sys))))
+    end
+
     return T(args...; kwargs...)
 end
 

src/systems/diffeqs/abstractodesystem.jl

@@ -168,6 +168,10 @@ function generate_function(sys::AbstractODESystem, dvs = unknowns(sys),
     rhss = implicit_dae ? [_iszero(eq.lhs) ? eq.rhs : eq.rhs - eq.lhs for eq in eqs] :
            [eq.rhs for eq in eqs]
 
+    if !isempty(assertions(sys))
+        rhss[end] += unwrap(get_assertions_expr(sys))
+    end
+
     # TODO: add an optional check on the ordering of observed equations
     u = dvs
     p = reorder_parameters(sys, ps)
@@ -542,7 +546,7 @@ Create a Julia expression for an `ODEFunction` from the [`ODESystem`](@ref).
 The arguments `dvs` and `ps` are used to set the order of the dependent
 variable and parameter vectors, respectively.
 """
-struct ODEFunctionExpr{iip} end
+struct ODEFunctionExpr{iip, specialize} end
 
 struct ODEFunctionClosure{O, I} <: Function
     f_oop::O
@@ -551,7 +555,7 @@ end
 (f::ODEFunctionClosure)(u, p, t) = f.f_oop(u, p, t)
 (f::ODEFunctionClosure)(du, u, p, t) = f.f_iip(du, u, p, t)
 
-function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
+function ODEFunctionExpr{iip, specialize}(sys::AbstractODESystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
         version = nothing, tgrad = false,
         jac = false, p = nothing,
@@ -560,14 +564,12 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
         steady_state = false,
         sparsity = false,
         observedfun_exp = nothing,
-        kwargs...) where {iip}
+        kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunctionExpr`")
     end
     f_oop, f_iip = generate_function(sys, dvs, ps; expression = Val{true}, kwargs...)
 
-    dict = Dict()
-
     fsym = gensym(:f)
     _f = :($fsym = $ODEFunctionClosure($f_oop, $f_iip))
     tgradsym = gensym(:tgrad)
@@ -590,30 +592,28 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
         _jac = :($jacsym = nothing)
     end
 
+    Msym = gensym(:M)
     M = calculate_massmatrix(sys)
-
-    _M = if sparse && !(u0 === nothing || M === I)
-        SparseArrays.sparse(M)
+    if sparse && !(u0 === nothing || M === I)
+        _M = :($Msym = $(SparseArrays.sparse(M)))
     elseif u0 === nothing || M === I
-        M
+        _M = :($Msym = $M)
     else
-        ArrayInterface.restructure(u0 .* u0', M)
+        _M = :($Msym = $(ArrayInterface.restructure(u0 .* u0', M)))
     end
 
     jp_expr = sparse ? :($similar($(get_jac(sys)[]), Float64)) : :nothing
     ex = quote
-        $_f
-        $_tgrad
-        $_jac
-        M = $_M
-        ODEFunction{$iip}($fsym,
-            sys = $sys,
-            jac = $jacsym,
-            tgrad = $tgradsym,
-            mass_matrix = M,
-            jac_prototype = $jp_expr,
-            sparsity = $(sparsity ? jacobian_sparsity(sys) : nothing),
-            observed = $observedfun_exp)
+        let $_f, $_tgrad, $_jac, $_M
+            ODEFunction{$iip, $specialize}($fsym,
+                sys = $sys,
+                jac = $jacsym,
+                tgrad = $tgradsym,
+                mass_matrix = $Msym,
+                jac_prototype = $jp_expr,
+                sparsity = $(sparsity ? jacobian_sparsity(sys) : nothing),
+                observed = $observedfun_exp)
+        end
     end
     !linenumbers ? Base.remove_linenums!(ex) : ex
 end
@@ -622,6 +622,14 @@ function ODEFunctionExpr(sys::AbstractODESystem, args...; kwargs...)
     ODEFunctionExpr{true}(sys, args...; kwargs...)
 end
 
+function ODEFunctionExpr{true}(sys::AbstractODESystem, args...; kwargs...)
+    return ODEFunctionExpr{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+end
+
+function ODEFunctionExpr{false}(sys::AbstractODESystem, args...; kwargs...)
+    return ODEFunctionExpr{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
+end
+
 """
 ```julia
 DAEFunctionExpr{iip}(sys::AbstractODESystem, dvs = unknowns(sys),