feat: At and costs in @mtkmodel

Author: vyudu

src/systems/diffeqs/odesystem.jl

@@ -247,7 +247,7 @@ end
 function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         controls = Num[],
         observed = Equation[],
-        constraintsystem = nothing,
+        constraints = Any[],
         costs = Num[],
         consolidate = nothing,
         systems = ODESystem[],
@@ -276,11 +276,30 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     @assert all(control -> any(isequal.(control, ps)), controls) "All controls must also be parameters."
+
+    constraintsystem = nothing
+    if !isempty(constraints)
+        @show constraints
+        constraintsystem = process_constraint_system(constraints, dvs, ps, iv)
+        for p in parameters(constraintsystem)
+            !in(p, Set(ps)) && push!(ps, p)
+        end
+    end
+
+    if !isempty(costs)
+        coststs, costps = process_costs(costs, dvs, ps, iv)
+        for p in costps
+            !in(p, Set(ps)) && push!(ps, p)
+        end
+    end
+    costs = wrap.(costs)
+
     iv′ = value(iv)
     ps′ = value.(ps)
     ctrl′ = value.(controls)
     dvs′ = value.(dvs)
     dvs′ = filter(x -> !isdelay(x, iv), dvs′)
+
     parameter_dependencies, ps′ = process_parameter_dependencies(
         parameter_dependencies, ps′)
     if !(isempty(default_u0) && isempty(default_p))
@@ -350,7 +369,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         metadata, gui_metadata, is_dde, tstops, checks = checks)
 end
 
-function ODESystem(eqs, iv; constraints = Equation[], costs = Num[], kwargs...)
+function ODESystem(eqs, iv; kwargs...)
     diffvars, allunknowns, ps, eqs = process_equations(eqs, iv)
 
     for eq in get(kwargs, :parameter_dependencies, Equation[])
@@ -382,29 +401,7 @@ function ODESystem(eqs, iv; constraints = Equation[], costs = Num[], kwargs...)
     end
     algevars = setdiff(allunknowns, diffvars)
 
-    consvars = OrderedSet()
-    constraintsystem = nothing
-    if !isempty(constraints)
-        constraintsystem = process_constraint_system(constraints, allunknowns, new_ps, iv)
-        for st in get_unknowns(constraintsystem)
-            iscall(st) ?
-            !in(operation(st)(iv), allunknowns) && push!(consvars, st) :
-            !in(st, allunknowns) && push!(consvars, st)
-        end
-        for p in parameters(constraintsystem)
-            !in(p, new_ps) && push!(new_ps, p)
-        end
-    end
-
-    if !isempty(costs)
-        coststs, costps = process_costs(costs, allunknowns, new_ps, iv)
-        for p in costps
-            !in(p, new_ps) && push!(new_ps, p)
-        end
-    end
-    costs = wrap.(costs)
-
-    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars, consvars))),
+    return ODESystem(eqs, iv, collect(Iterators.flatten((diffvars, algevars))),
         collect(new_ps); constraintsystem, costs, kwargs...)
 end
 
@@ -760,7 +757,7 @@ end
 Build the constraint system for the ODESystem.
 """
 function process_constraint_system(
-        constraints::Vector{Equation}, sts, ps, iv; consname = :cons)
+        constraints::Vector, sts, ps, iv; consname = :cons)
     isempty(constraints) && return nothing
 
     constraintsts = OrderedSet()
@@ -800,7 +797,7 @@ Return the set of additional parameters found in the system, e.g. in x(p) ~ 3 th
 parameter of the system.
 """
 function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
-    sts = sysvars
+    sts = Set(sysvars)
 
     for var in auxvars
         if !iscall(var)
@@ -810,6 +807,7 @@ function validate_vars_and_find_ps!(auxvars, auxps, sysvars, iv)
             throw(ArgumentError("Too many arguments for variable $var."))
         elseif length(arguments(var)) == 1
             arg = only(arguments(var))
+            @show sts
             operation(var)(iv) ∈ sts ||
                 throw(ArgumentError("Variable $var is not a variable of the ODESystem. Called variables must be variables of the ODESystem."))
 

src/systems/model_parsing.jl

@@ -65,6 +65,8 @@ function _model_macro(mod, fullname::Union{Expr, Symbol}, expr, isconnector)
     ps, sps, vs, = [], [], []
     c_evts = []
     d_evts = []
+    cons = []
+    costs = []
     kwargs = OrderedCollections.OrderedSet()
     where_types = Union{Symbol, Expr}[]
 
@@ -80,7 +82,7 @@ function _model_macro(mod, fullname::Union{Expr, Symbol}, expr, isconnector)
     for arg in expr.args
         if arg.head == :macrocall
             parse_model!(exprs.args, comps, ext, eqs, icon, vs, ps,
-                sps, c_evts, d_evts, dict, mod, arg, kwargs, where_types)
+                sps, c_evts, d_evts, cons, costs, dict, mod, arg, kwargs, where_types)
         elseif arg.head == :block
             push!(exprs.args, arg)
         elseif arg.head == :if
@@ -117,16 +119,19 @@ function _model_macro(mod, fullname::Union{Expr, Symbol}, expr, isconnector)
     push!(exprs.args, :(push!(systems, $(comps...))))
     push!(exprs.args, :(push!(variables, $(vs...))))
 
+
     gui_metadata = isassigned(icon) > 0 ? GUIMetadata(GlobalRef(mod, name), icon[]) :
                    GUIMetadata(GlobalRef(mod, name))
 
+    consolidate = get(dict, :consolidate, nothing)
     description = get(dict, :description, "")
 
     @inline pop_structure_dict!.(
         Ref(dict), [:constants, :defaults, :kwargs, :structural_parameters])
 
     sys = :($type($(flatten_equations)(equations), $iv, variables, parameters;
-        name, description = $description, systems, gui_metadata = $gui_metadata, defaults))
+        name, description = $description, systems, gui_metadata = $gui_metadata, defaults,
+        costs = [$(costs...)], constraints = [$(cons...)], consolidate = $consolidate))
 
     if length(ext) == 0
         push!(exprs.args, :(var"#___sys___" = $sys))
@@ -610,9 +615,10 @@ function get_var(mod::Module, b)
 end
 
 function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps, c_evts, d_evts,
-        dict, mod, arg, kwargs, where_types)
+        cons, costs, dict, mod, arg, kwargs, where_types)
     mname = arg.args[1]
     body = arg.args[end]
+    @show dict
     if mname == Symbol("@description")
         parse_description!(body, dict)
     elseif mname == Symbol("@components")
@@ -637,7 +643,13 @@ function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps, c_evts, d_evts,
         isassigned(icon) && error("This model has more than one icon.")
         parse_icon!(body, dict, icon, mod)
     elseif mname == Symbol("@defaults")
-        parse_system_defaults!(exprs, arg, dict)
+        parse_system_defaults!(exprs, dict, body)
+    elseif mname == Symbol("@constraints")
+        parse_costs!(cons, dict, body)
+    elseif mname == Symbol("@costs")
+        parse_constraints!(costs, dict, body)
+    elseif mname == Symbol("@consolidate")
+        parse_consolidate!(body, dict)
     else
         error("$mname is not handled.")
     end
@@ -1149,6 +1161,33 @@ function parse_discrete_events!(d_evts, dict, body)
     end
 end
 
+function parse_constraints!(cons, dict, body) 
+    dict[:constraints] = [] 
+    Base.remove_linenums!(body)
+    for arg in body.args
+        push!(cons, arg)
+        push!(dict[:constraints], readable_code.(cons)...)
+    end
+end
+
+function parse_costs!(costs, dict, body) 
+    @show dict
+    dict[:costs] = [] 
+    Base.remove_linenums!(body)
+    for arg in body.args
+        push!(costs, arg)
+        push!(dict[:costs], readable_code.(costs)...)
+    end
+end
+
+function parse_consolidate!(body, dict)
+    if !(occursin("->", string(body)) || occursin("=", string(body)))
+        error("Consolidate must be a function definition.")
+    else
+        dict[:consolidate] = body
+    end
+end
+
 function parse_icon!(body::String, dict, icon, mod)
     icon_dir = get(ENV, "MTK_ICONS_DIR", joinpath(DEPOT_PATH[1], "mtk_icons"))
     dict[:icon] = icon[] = if isfile(body)

src/variables.jl

@@ -612,3 +612,42 @@ getunshifted(x::Symbolic) = Symbolics.getmetadata(x, VariableUnshifted, nothing)
 
 getshift(x::Num) = getshift(unwrap(x))
 getshift(x::Symbolic) = Symbolics.getmetadata(x, VariableShift, 0)
+
+###################
+### Evaluate at ###
+###################
+struct At <: Symbolics.Operator 
+    t::Union{Symbolic, Number}
+end
+
+function (A::At)(x::Symbolic)
+    if symbolic_type(x) == NotSymbolic() || !iscall(x)
+        if x isa Symbolics.CallWithMetadata
+            return x(A.t)
+        else
+            return x
+        end
+    end
+
+    if iscall(x) && operation(x) == getindex
+        arr = arguments(x)[1]
+        term(getindex, A(arr), arguments(x)[2:end]...)
+    elseif operation(x) isa Differential
+        x = default_toterm(x)
+        A(x)
+    else
+        length(arguments(x)) !== 1 && error("Variable $x has too many arguments. At can only be applied to one-argument variables.")
+        (symbolic_type(only(arguments(x))) !== ScalarSymbolic()) && return x
+        return operation(x)(A.t)
+    end
+end
+
+function (A::At)(x::Union{Num, Symbolics.Arr})
+    wrap(A(unwrap(x)))
+end
+SymbolicUtils.isbinop(::At) = false
+
+Base.nameof(::At) = :At
+Base.show(io::IO, A::At) = print(io, "At(", A.t, ")")
+Base.:(==)(A1::At, A2::At) = isequal(A1.t, A2.t)
+Base.hash(A::At, u::UInt) = hash(A.t, u)

test/model_parsing.jl

@@ -1026,3 +1026,35 @@ end
     @named sys = Float2Bool()
     @test typeof(sys) == DiscreteSystem
 end
+
+@testset "Constraints, costs, consolidate" begin
+    @mtkmodel Example begin
+        @variables begin
+            x(t)
+            y(t)
+        end
+        @equations begin
+            x ~ y
+        end
+        @constraints begin
+            At(0.3)(x) ~ 3
+            y ≲ 4
+        end
+        @costs begin
+            x + y
+            At(1)(y)^2
+        end
+        @consolidate f(u) = u[1]^2 + log(u[2])
+    end
+
+    @named ex = Example()
+    ex = complete(ex)
+
+    costs = ModelingToolkit.get_costs(ex) 
+    constrs = ModelingToolkit.get_constraints(ModelingToolkit.get_constraintsystem(ex))
+    @test isequal(costs[1], ex.x + ex.y)
+    @test isequal(costs[2], At(1)(ex.y)^2)
+    @test isequal(constrs[1], -3 + At(0.3)(ex.x) ~ 0)
+    @test isequal(constrs[2], -4 + ex.y ≲ 0)
+    @test ModelingToolkit.get_consolidate(ex)([1, 2]) ≈ 1 + log(2)
+end

test/variable_utils.jl

@@ -158,3 +158,29 @@ end
     @test !isinitial(c)
     @test !isinitial(x)
 end
+
+@testset "At" begin
+    @independent_variables u
+    @variables x(t) v(..) w(t)[1:3]
+    @parameters y z(u, t) r[1:3]
+
+    @test At(1)(x) isa Num
+    @test isequal(At(1)(y), y)
+    @test_throws ErrorException At(1)(z)
+    @test isequal(At(1)(v), v(1))
+    @test isequal(At(1)(v(t)), v(1))
+    @test isequal(At(1)(v(2)), v(2))
+
+    arr = At(1)(w)
+    var = At(1)(w[1])
+    @test arr isa Symbolics.Arr 
+    @test var isa Num
+    
+    @test isequal(At(1)(r), r)
+    @test isequal(At(1)(r[2]), r[2])
+
+    _x = ModelingToolkit.unwrap(x)
+    @test At(1)(_x) isa Symbolics.BasicSymbolic
+    @test only(arguments(At(1)(_x))) == 1
+    @test At(1)(D(x)) isa Num
+end