Merge pull request #2500 from ven-k/vkb/type

Provision to enforce types in parameters, variables and structural_parameters in `@mtkmodel`

Author: ChrisRackauckas

docs/src/basics/MTKModel_Connector.md

@@ -229,12 +229,12 @@ end
 
   - `:components`: List of sub-components in the form of [[name, sub_component_name],...].
   - `: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 default values.
+  - `:structural_parameters`: Dictionary of structural parameters mapped to their metadata.
   - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. For
     parameter arrays, length is added to the metadata as `:size`.
   - `:variables`: Dictionary of symbolic variables mapped to their metadata. For
     variable arrays, length is added to the metadata as `:size`.
-  - `:kwargs`: Dictionary of keyword arguments mapped to their default values.
+  - `:kwargs`: Dictionary of keyword arguments mapped to their metadata.
   - `:independent_variable`: Independent variable, which is added while generating the Model.
   - `:equations`: List of equations (represented as strings).
 
@@ -243,13 +243,14 @@ For example, the structure of `ModelC` is:
 ```julia
 julia> ModelC.structure
 Dict{Symbol, Any} with 7 entries:
-  :components           => [[:model_a, :ModelA]]
-  :variables            => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var), :v_array=>Dict(:size=>(2, 3)))
-  :icon                 => URI("https://github.com/SciML/SciMLDocs/blob/main/docs/src/assets/logo.png")
-  :kwargs               => Dict{Symbol, Any}(:f=>:sin, :v=>:v_var, :v_array=>nothing, :model_a__k_array=>nothing, :p1=>nothing)
-  :independent_variable => t
-  :extend               => Any[[:p2, :p1], Symbol("#mtkmodel__anonymous__ModelB"), :ModelB]
-  :equations            => ["model_a.k ~ f(v)"]
+  :components            => [[:model_a, :ModelA]]
+  :variables             => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var), :v_array=>Dict(:size=>(2, 3)))
+  :icon                  => URI("https://github.com/SciML/SciMLDocs/blob/main/docs/src/assets/logo.png")
+  :kwargs                => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin), :v=>Dict{Symbol, Union{Nothing, Symbol}}(:value=>:v_var, :type=>nothing), :v_array=>Dict(:value=>nothing, :type=>nothing), :p1=>Dict(:value=>nothing))
+  :structural_parameters => Dict{Symbol, Dict}(:f=>Dict(:value=>:sin))
+  :independent_variable  => t
+  :extend                => Any[[:p2, :p1], Symbol("#mtkmodel__anonymous__ModelB"), :ModelB]
+  :equations             => ["model_a.k ~ f(v)"]
 ```
 
 ### Using conditional statements
@@ -322,11 +323,12 @@ The conditional parts are reflected in the `structure`. For `BranchOutsideTheBlo
 ```julia
 julia> BranchOutsideTheBlock.structure
 Dict{Symbol, Any} with 5 entries:
-  :components           => Any[(:if, :flag, [[:sys1, :C]], Any[])]
-  :kwargs               => Dict{Symbol, Any}(:flag=>true)
-  :independent_variable => t
-  :parameters           => Dict{Symbol, Dict{Symbol, Any}}(:a1=>Dict(:condition=>(:if, :flag, Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a1 => nothing), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a1 => Dict())]), Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a2 => nothing), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict())]))
-  :equations            => Any[(:if, :flag, ["a1 ~ 0"], ["a2 ~ 0"])]
+  :components            => Any[(:if, :flag, [[:sys1, :C]], Any[])]
+  :kwargs                => Dict{Symbol, Dict}(:flag=>Dict{Symbol, Bool}(:value=>1))
+  :structural_parameters => Dict{Symbol, Dict}(:flag=>Dict{Symbol, Bool}(:value=>1))
+  :independent_variable  => t
+  :parameters            => Dict{Symbol, Dict{Symbol, Any}}(:a1=>Dict(:condition=>(:if, :flag, Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a1 => nothing), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a1 => Dict())]), Dict{Symbol, Any}(:kwargs => Dict{Any, Any}(:a2 => nothing), :parameters => Any[Dict{Symbol, Dict{Symbol, Any}}(:a2 => Dict())]))
+  :equations             => Any[(:if, :flag, ["a1 ~ 0"], ["a2 ~ 0"])]
 ```
 
 Conditional entries are entered in the format of `(branch, condition, [case when it is true], [case when it is false])`;

src/systems/model_parsing.jl

@@ -35,8 +35,10 @@ end
 
 function _model_macro(mod, name, expr, isconnector)
     exprs = Expr(:block)
-    dict = Dict{Symbol, Any}()
-    dict[:kwargs] = Dict{Symbol, Any}()
+    dict = Dict{Symbol, Any}(
+        :kwargs => Dict{Symbol, Dict}(),
+        :structural_parameters => Dict{Symbol, Dict}()
+    )
     comps = Symbol[]
     ext = Ref{Any}(nothing)
     eqs = Expr[]
@@ -107,7 +109,8 @@ function _model_macro(mod, name, expr, isconnector)
 end
 
 function parse_variable_def!(dict, mod, arg, varclass, kwargs;
-        def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing)
+        def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
+        type::Union{Type, Nothing} = nothing)
     metatypes = [(:connection_type, VariableConnectType),
         (:description, VariableDescription),
         (:unit, VariableUnit),
@@ -125,15 +128,34 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs;
     arg isa LineNumberNode && return
     MLStyle.@match arg begin
         a::Symbol => begin
-            push!(kwargs, Expr(:kw, a, nothing))
+            if type isa Nothing
+                push!(kwargs, Expr(:kw, a, nothing))
+            else
+                push!(kwargs, Expr(:kw, Expr(:(::), a, Union{Nothing, type}), nothing))
+            end
             var = generate_var!(dict, a, varclass; indices)
-            dict[:kwargs][getname(var)] = def
+            dict[:kwargs][getname(var)] = Dict(:value => def, :type => type)
             (var, def)
         end
+        Expr(:(::), a, type) => begin
+            type = Core.eval(mod, type)
+            _type_check!(a, type)
+            parse_variable_def!(dict, mod, a, varclass, kwargs; def, type)
+        end
+        Expr(:(::), Expr(:call, a, b), type) => begin
+            type = Core.eval(mod, type)
+            def = _type_check!(def, a, type)
+            parse_variable_def!(dict, mod, a, varclass, kwargs; def, type)
+        end
         Expr(:call, a, b) => begin
-            push!(kwargs, Expr(:kw, a, nothing))
+            if type isa Nothing
+                push!(kwargs, Expr(:kw, a, nothing))
+            else
+                push!(kwargs, Expr(:kw, Expr(:(::), a, Union{Nothing, type}), nothing))
+            end
             var = generate_var!(dict, a, b, varclass; indices)
-            dict[:kwargs][getname(var)] = def
+            type !== nothing && (dict[varclass][getname(var)][:type] = type)
+            dict[:kwargs][getname(var)] = Dict(:value => def, :type => type)
             (var, def)
         end
         Expr(:(=), a, b) => begin
@@ -304,15 +326,23 @@ function parse_structural_parameters!(exprs, sps, dict, mod, body, kwargs)
     Base.remove_linenums!(body)
     for arg in body.args
         MLStyle.@match arg begin
+            Expr(:(=), Expr(:(::), a, type), b) => begin
+                type = Core.eval(mod, type)
+                b = _type_check!(Core.eval(mod, b), a, type)
+                push!(sps, a)
+                push!(kwargs, Expr(:kw, Expr(:(::), a, type), b))
+                dict[:structural_parameters][a] = dict[:kwargs][a] = Dict(
+                    :value => b, :type => type)
+            end
             Expr(:(=), a, b) => begin
                 push!(sps, a)
                 push!(kwargs, Expr(:kw, a, b))
-                dict[:kwargs][a] = b
+                dict[:structural_parameters][a] = dict[:kwargs][a] = Dict(:value => b)
             end
             a => begin
                 push!(sps, a)
                 push!(kwargs, a)
-                dict[:kwargs][a] = nothing
+                dict[:structural_parameters][a] = dict[:kwargs][a] = Dict(:value => nothing)
             end
         end
     end
@@ -336,17 +366,17 @@ function extend_args!(a, b, dict, expr, kwargs, varexpr, has_param = false)
                     end
                 end
                 push!(kwargs, Expr(:kw, x, nothing))
-                dict[:kwargs][x] = nothing
+                dict[:kwargs][x] = Dict(:value => nothing)
             end
             Expr(:kw, x) => begin
                 push!(kwargs, Expr(:kw, x, nothing))
-                dict[:kwargs][x] = nothing
+                dict[:kwargs][x] = Dict(:value => nothing)
             end
             Expr(:kw, x, y) => begin
                 b.args[i] = Expr(:kw, x, x)
                 push!(varexpr.args, :($x = $x === nothing ? $y : $x))
                 push!(kwargs, Expr(:kw, x, nothing))
-                dict[:kwargs][x] = nothing
+                dict[:kwargs][x] = Dict(:value => nothing)
             end
             Expr(:parameters, x...) => begin
                 has_param = true
@@ -851,3 +881,18 @@ function parse_conditional_model_statements(comps, dict, eqs, exprs, kwargs, mod
             $equations_blk
         end))
 end
+
+_type_check!(a, type) = return
+function _type_check!(val, a, type)
+    if val isa type
+        return val
+    else
+        try
+            return convert(type, val)
+        catch
+            (e)
+            throw(TypeError(Symbol("`@mtkmodel`"),
+                "`@structural_parameters`, while assigning to `$a`", type, typeof(val)))
+        end
+    end
+end

test/model_parsing.jl

@@ -164,7 +164,7 @@ resistor = getproperty(rc, :resistor; namespace = false)
 # Test that `C_val` passed via argument is set as default of C.
 @test getdefault(rc.capacitor.C) == C_val
 # Test that `k`'s default value is unchanged.
-@test getdefault(rc.constant.k) == RC.structure[:kwargs][:k_val]
+@test getdefault(rc.constant.k) == RC.structure[:kwargs][:k_val][:value]
 @test getdefault(rc.capacitor.v) == 0.0
 
 @test get_gui_metadata(rc.resistor).layout == Resistor.structure[:icon] ==
@@ -241,6 +241,33 @@ resistor = getproperty(rc, :resistor; namespace = false)
     @test isequal(getdefault(model.k), model.kval)
 end
 
+@testset "Type annotation" begin
+    @mtkmodel TypeModel begin
+        @structural_parameters begin
+            flag::Bool = true
+        end
+        @parameters begin
+            par0::Bool = true
+            par1::Int = 1
+            par2(t)::Int,
+            [description = "Enforced `par4` to be an Int by setting the type to the keyword-arg."]
+            par3(t)::Float64 = 1.0
+            par4(t)::Float64 = 1 # converts 1 to 1.0 of Float64 type
+        end
+    end
+
+    @named type_model = TypeModel()
+
+    @test getname.(parameters(type_model)) == [:par0, :par1, :par2, :par3, :par4]
+
+    @test_throws TypeError TypeModel(; name = :throws, flag = 1)
+    @test_throws TypeError TypeModel(; name = :throws, par0 = 1)
+    @test_throws TypeError TypeModel(; name = :throws, par1 = 1.5)
+    @test_throws TypeError TypeModel(; name = :throws, par2 = 1.5)
+    @test_throws TypeError TypeModel(; name = :throws, par3 = true)
+    @test_throws TypeError TypeModel(; name = :throws, par4 = true)
+end
+
 @testset "Defaults of subcomponents MTKModel" begin
     @mtkmodel A begin
         @parameters begin
@@ -322,7 +349,9 @@ end
     @test A.structure[:parameters] == Dict(:p => Dict())
     @test A.structure[:extend] == [[:e], :extended_e, :E]
     @test A.structure[:equations] == ["e ~ 0"]
-    @test A.structure[:kwargs] == Dict(:p => nothing, :v => nothing)
+    @test A.structure[:kwargs] ==
+          Dict{Symbol, Dict}(:p => Dict(:value => nothing, :type => nothing),
+        :v => Dict(:value => nothing, :type => nothing))
     @test A.structure[:components] == [[:cc, :C]]
 end
 
@@ -392,9 +421,9 @@ end
     @named else_in_sys = InsideTheBlock(flag = 3)
     else_in_sys = complete(else_in_sys)
 
-    @test nameof.(parameters(if_in_sys)) == [:if_parameter, :eq]
-    @test nameof.(parameters(elseif_in_sys)) == [:elseif_parameter, :eq]
-    @test nameof.(parameters(else_in_sys)) == [:else_parameter, :eq]
+    @test getname.(parameters(if_in_sys)) == [:if_parameter, :eq]
+    @test getname.(parameters(elseif_in_sys)) == [:elseif_parameter, :eq]
+    @test getname.(parameters(else_in_sys)) == [:else_parameter, :eq]
 
     @test nameof.(get_systems(if_in_sys)) == [:if_sys, :default_sys]
     @test nameof.(get_systems(elseif_in_sys)) == [:elseif_sys, :default_sys]
@@ -481,9 +510,9 @@ end
     @named ternary_out_sys = OutsideTheBlock(condition = 4)
     else_out_sys = complete(else_out_sys)
 
-    @test nameof.(parameters(if_out_sys)) == [:if_parameter, :default_parameter]
-    @test nameof.(parameters(elseif_out_sys)) == [:elseif_parameter, :default_parameter]
-    @test nameof.(parameters(else_out_sys)) == [:else_parameter, :default_parameter]
+    @test getname.(parameters(if_out_sys)) == [:if_parameter, :default_parameter]
+    @test getname.(parameters(elseif_out_sys)) == [:elseif_parameter, :default_parameter]
+    @test getname.(parameters(else_out_sys)) == [:else_parameter, :default_parameter]
 
     @test nameof.(get_systems(if_out_sys)) == [:if_sys, :default_sys]
     @test nameof.(get_systems(elseif_out_sys)) == [:elseif_sys, :default_sys]
@@ -529,8 +558,8 @@ end
     @named ternary_false = TernaryBranchingOutsideTheBlock(condition = false)
     ternary_false = complete(ternary_false)
 
-    @test nameof.(parameters(ternary_true)) == [:ternary_parameter_true]
-    @test nameof.(parameters(ternary_false)) == [:ternary_parameter_false]
+    @test getname.(parameters(ternary_true)) == [:ternary_parameter_true]
+    @test getname.(parameters(ternary_false)) == [:ternary_parameter_false]
 
     @test nameof.(get_systems(ternary_true)) == [:ternary_sys_true]
     @test nameof.(get_systems(ternary_false)) == [:ternary_sys_false]