Merge pull request #2437 from ven-k/vkb/at-defaults

feat: add `@defaults` to `@mtkmodel`

Author: YingboMa

docs/src/basics/MTKModel_Connector.md

@@ -25,6 +25,7 @@ equations.
 
   - `@components`: for listing sub-components of the system
   - `@constants`: for declaring constants
+  - `@defaults`: for passing `defaults` to ODESystem
   - `@equations`: for the list of equations
   - `@extend`: for extending a base system and unpacking its unknowns
   - `@icon` : for embedding the model icon
@@ -66,6 +67,7 @@ end
     @variables begin
         v(t) = v_var
         v_array(t)[1:2, 1:3]
+        v_for_defaults(t)
     end
     @extend ModelB(; p1)
     @components begin
@@ -79,6 +81,9 @@ end
     @equations begin
         model_a.k ~ f(v)
     end
+    @defaults begin
+        v_for_defaults => 2.0
+    end
 end
 ```
 
@@ -172,6 +177,11 @@ getdefault(model_c3.model_a.k_array[2])
 
   - List all the equations here
 
+#### `@defaults` begin block
+
+  - Default values can be passed as pairs.
+  - This is equivalent to passing `defaults` argument to `ODESystem`.
+
 #### A begin block
 
   - Any other Julia operations can be included with dedicated begin blocks.
@@ -239,6 +249,7 @@ end
 
   - `:components`: The list of sub-components in the form of [[name, sub_component_name],...].
   - `:constants`: Dictionary of constants mapped to its metadata.
+  - `:defaults`: Dictionary of variables and default values specified in the `@defaults`.
   - `:extend`: The list of extended unknowns, name given to the base system, and name of the base system.
   - `:structural_parameters`: Dictionary of structural parameters mapped to their metadata.
   - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. For
@@ -253,15 +264,16 @@ For example, the structure of `ModelC` is:
 
 ```julia
 julia> ModelC.structure
-Dict{Symbol, Any} with 9 entries:
+Dict{Symbol, Any} with 10 entries:
   :components            => Any[Union{Expr, Symbol}[:model_a, :ModelA], Union{Expr, Symbol}[:model_array_a, :ModelA, :(1:N)], Union{Expr, Symbol}[:model_array_b, :ModelA, :(1:N)]]
-  :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var, :type=>Real), :v_array=>Dict(:type=>Real, :size=>(2, 3)))
+  :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var, :type=>Real), :v_array=>Dict(:type=>Real, :size=>(2, 3)), :v_for_defaults=>Dict(:type=>Real))
   :icon                  => URI("https://github.com/SciML/SciMLDocs/blob/main/docs/src/assets/logo.png")
-  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :N=>Dict(:value=>2), :v=>Dict{Symbol, Any}(:value=>:v_var, :type=>Real), :v_array=>Dict{Symbol, Union{Nothing, UnionAll}}(:value=>nothing, :type=>AbstractArray{Real}), :p1=>Dict(:value=>nothing))
+  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :N=>Dict(:value=>2), :v=>Dict{Symbol, Any}(:value=>:v_var, :type=>Real), :v_array=>Dict{Symbol, Union{Nothing, UnionAll}}(:value=>nothing, :type=>AbstractArray{Real}), :v_for_defaults=>Dict{Symbol, Union{Nothing, DataType}}(:value=>nothing, :type=>Real), :p1=>Dict(:value=>nothing))
   :structural_parameters => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :N=>Dict(:value=>2))
   :independent_variable  => t
   :constants             => Dict{Symbol, Dict}(:c=>Dict{Symbol, Any}(:value=>1, :type=>Int64, :description=>"Example constant."))
   :extend                => Any[[:p2, :p1], Symbol("#mtkmodel__anonymous__ModelB"), :ModelB]
+  :defaults              => Dict{Symbol, Any}(:v_for_defaults=>2.0)
   :equations             => Any["model_a.k ~ f(v)"]
 ```
 
@@ -327,6 +339,9 @@ used inside the if-elseif-else statements.
             a2 ~ 0
         end
     end
+    @defaults begin
+        a1 => 10
+    end
 end
 ```
 

src/systems/model_parsing.jl

@@ -39,6 +39,7 @@ function _model_macro(mod, name, expr, isconnector)
     exprs = Expr(:block)
     dict = Dict{Symbol, Any}(
         :constants => Dict{Symbol, Dict}(),
+        :defaults => Dict{Symbol, Any}(),
         :kwargs => Dict{Symbol, Dict}(),
         :structural_parameters => Dict{Symbol, Dict}()
     )
@@ -54,6 +55,7 @@ function _model_macro(mod, name, expr, isconnector)
     push!(exprs.args, :(parameters = []))
     push!(exprs.args, :(systems = ODESystem[]))
     push!(exprs.args, :(equations = Equation[]))
+    push!(exprs.args, :(defaults = Dict{Num, Union{Number, Symbol, Function}}()))
 
     Base.remove_linenums!(expr)
     for arg in expr.args
@@ -99,8 +101,11 @@ function _model_macro(mod, name, expr, isconnector)
     gui_metadata = isassigned(icon) > 0 ? GUIMetadata(GlobalRef(mod, name), icon[]) :
                    GUIMetadata(GlobalRef(mod, name))
 
+    @inline pop_structure_dict!.(
+        Ref(dict), [:constants, :defaults, :kwargs, :structural_parameters])
+
     sys = :($ODESystem($Equation[equations...], $iv, variables, parameters;
-        name, systems, gui_metadata = $gui_metadata))
+        name, systems, gui_metadata = $gui_metadata, defaults))
 
     if ext[] === nothing
         push!(exprs.args, :(var"#___sys___" = $sys))
@@ -122,6 +127,8 @@ function _model_macro(mod, name, expr, isconnector)
     :($name = $Model($f, $dict, $isconnector))
 end
 
+pop_structure_dict!(dict, key) = length(dict[key]) == 0 && pop!(dict, key)
+
 function update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
         varclass, where_types)
     if indices isa Nothing
@@ -355,6 +362,8 @@ function parse_model!(exprs, comps, ext, eqs, icon, vs, ps, sps,
     elseif mname == Symbol("@icon")
         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)
     else
         error("$mname is not handled.")
     end
@@ -400,6 +409,28 @@ function parse_constants!(exprs, dict, body, mod)
     end
 end
 
+push_additional_defaults!(dict, a, b::Number) = dict[:defaults][a] = b
+push_additional_defaults!(dict, a, b::QuoteNode) = dict[:defaults][a] = b.value
+function push_additional_defaults!(dict, a, b::Expr)
+    dict[:defaults][a] = readable_code(b)
+end
+
+function parse_system_defaults!(exprs, defaults_body, dict)
+    for default_arg in defaults_body.args[end].args
+        # for arg in default_arg.args
+        MLStyle.@match default_arg begin
+            # For cases like `p => 1` and `p => f()`. In both cases the definitions of
+            # `a`, here `p` and when `b` is a function, here `f` are available while
+            # defining the model
+            Expr(:call, :(=>), a, b) => begin
+                push!(exprs, :(defaults[$a] = $b))
+                push_additional_defaults!(dict, a, b)
+            end
+            _ => error("Invalid `defaults` entry $default_arg $(typeof(a)) $(typeof(b))")
+        end
+    end
+end
+
 function parse_structural_parameters!(exprs, sps, dict, mod, body, kwargs)
     Base.remove_linenums!(body)
     for arg in body.args

test/model_parsing.jl

@@ -1,7 +1,8 @@
 using ModelingToolkit, Test
-using ModelingToolkit: get_gui_metadata, get_systems, get_connector_type,
-                       get_ps, getdefault, getname, scalarize, symtype,
-                       VariableDescription, RegularConnector
+using ModelingToolkit: get_connector_type, get_defaults, get_gui_metadata,
+                       get_systems, get_ps, getdefault, getname, scalarize, symtype,
+                       VariableDescription,
+                       RegularConnector
 using URIs: URI
 using Distributions
 using DynamicQuantities, OrdinaryDiffEq
@@ -219,17 +220,26 @@ end
             j(t) = jval, [description = "j(t)"]
             k = kval, [description = "k"]
             l(t)[1:2, 1:3] = 2, [description = "l is more than 1D"]
+            n # test defaults with Number input
+            n2 # test defaults with Function input
         end
         @structural_parameters begin
             m = 1
             func
         end
+        begin
+            g() = 5
+        end
+        @defaults begin
+            n => 1.0
+            n2 => g()
+        end
     end
 
     kval = 5
     @named model = MockModel(; b2 = [1, 3], kval, cval = 1, func = identity)
 
-    @test lastindex(parameters(model)) == 29
+    @test lastindex(parameters(model)) == 31
 
     @test all(getdescription.([model.e2...]) .== "e2")
     @test all(getdescription.([model.h2...]) .== "h2(t)")
@@ -256,6 +266,10 @@ end
     @test all(getdefault.(scalarize(model.l)) .== 2)
     @test isequal(getdefault(model.j), model.jval)
     @test isequal(getdefault(model.k), model.kval)
+    @test get_defaults(model)[model.n] == 1.0
+    @test get_defaults(model)[model.n2] == 5
+
+    @test MockModel.structure[:defaults] == Dict(:n => 1.0, :n2 => "g()")
 end
 
 @testset "Type annotation" begin

test/runtests.jl

@@ -1,5 +1,6 @@
 using SafeTestsets, Pkg, Test
 
+#=
 const GROUP = get(ENV, "GROUP", "All")
 
 function activate_extensions_env()
@@ -91,3 +92,7 @@ end
         @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl")
     end
 end
+
+=#
+
+@safetestset "Model Parsing Test" include("model_parsing.jl")