Skip to content

Commit

Permalink
feat: support ternary operators inside the @equations
Browse files Browse the repository at this point in the history 
- and ternary conditional statements at the top level of the block
- adds relevant tests
  • Loading branch information
@ven-k
ven-k committed Nov 19, 2023
1 parent 597bd0a commit 8b75e01
Show file tree
Hide file tree
Showing 2 changed files with 141 additions and 67 deletions.
19 changes: 14 additions & 5 deletions src/systems/model_parsing.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -559,9 +559,14 @@ function handle_y_vars(y, dict, mod, varclass, kwargs)
end

function handle_if_x_equations!(condition, dict, ifexpr, x)
push!(ifexpr.args, condition, :(push!(equations, $(x.args...))))
# push!(dict[:equations], [:if, readable_code(condition), readable_code.(x.args)])
readable_code.(x.args)
if Meta.isexpr(x, :block)
push!(ifexpr.args, condition, :(push!(equations, $(x.args...))))
return readable_code.(x.args)
else
push!(ifexpr.args, condition, :(push!(equations, $x)))
return readable_code(x)
end
# push!(dict[:equations], [:if, readable_code(condition), readable_code.(x.args)])
end

function handle_if_y_equations!(ifexpr, y, dict)
Expand All @@ -574,8 +579,12 @@ function handle_if_y_equations!(ifexpr, y, dict)
push!(ifexpr.args, elseifexpr)
(eq_entry...,)
else
push!(ifexpr.args, :(push!(equations, $(y.args...))))
readable_code.(y.args)
if Meta.isexpr(y, :block)
push!(ifexpr.args, :(push!(equations, $(y.args...))))
else
push!(ifexpr.args, :(push!(equations, $(y))))
end
readable_code.(y.args)
end
end

Expand Down
189 changes: 127 additions & 62 deletions test/model_parsing.jl
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
using ModelingToolkit, Test
using ModelingToolkit: get_gui_metadata,
get_ps, getdefault, getname, scalarize,
VariableDescription, RegularConnector
using ModelingToolkit: get_gui_metadata, get_systems, get_connector_type,
get_ps, getdefault, getname, scalarize, VariableDescription, RegularConnector
using URIs: URI
using Distributions
using Unitful
Expand Down Expand Up @@ -303,7 +302,7 @@ end
end

@named aa = A()
@test aa.connector_type == RegularConnector()
@test get_connector_type(aa) == RegularConnector()

@test A.isconnector == true

Expand All @@ -314,7 +313,7 @@ end
@test A.structure[:components] == [[:cc, :C]]
end

@testset "Conditional components, equations and parameters" begin
@testset "Conditional statements inside the blocks" begin
@mtkmodel C begin end

# Conditional statements inside @components, @equations
Expand All @@ -326,21 +325,21 @@ end
@parameters begin
eq = flag == 1 ? 1 : 0
if flag == 1
a1
if_parameter
elseif flag == 2
a2
elseif_parameter
else
a3
else_parameter
end
end
@components begin
sys0 = C()
default_sys = C()
if flag == 1
sys1 = C()
if_sys = C()
elseif flag == 2
sys2 = C()
elseif_sys = C()
else
sys3 = C()
else_sys = C()
end
end
@equations begin
Expand All @@ -352,94 +351,160 @@ end
else
eq ~ 3
end
flag == 1 ? eq ~ 4 : eq ~ 5
end
end

@named in_sys_1 = InsideTheBlock()
in_sys_1 = complete(in_sys_1)
@named in_sys_2 = InsideTheBlock(flag = 2)
in_sys_2 = complete(in_sys_2)
@named in_sys_3 = InsideTheBlock(flag = 3)
in_sys_3 = complete(in_sys_3)

@test nameof.(parameters(in_sys_1)) == [:a1, :eq]
@test nameof.(parameters(in_sys_2)) == [:a2, :eq]
@test nameof.(parameters(in_sys_3)) == [:a3, :eq]

@test nameof.(in_sys_1.systems) == [:sys1, :sys0]
@test nameof.(in_sys_2.systems) == [:sys2, :sys0]
@test nameof.(in_sys_3.systems) == [:sys3, :sys0]

@test all([in_sys_1.eq ~ 0, in_sys_1.eq ~ 1] .∈ [equations(in_sys_1)])
@test all([in_sys_2.eq ~ 0, in_sys_2.eq ~ 2] .∈ [equations(in_sys_2)])
@test all([in_sys_3.eq ~ 0, in_sys_3.eq ~ 3] .∈ [equations(in_sys_3)])
@named if_in_sys = InsideTheBlock()
if_in_sys = complete(if_in_sys)
@named elseif_in_sys = InsideTheBlock(flag = 2)
elseif_in_sys = complete(elseif_in_sys)
@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 nameof.(get_systems(if_in_sys)) == [:if_sys, :default_sys]
@test nameof.(get_systems(elseif_in_sys)) == [:elseif_sys, :default_sys]
@test nameof.(get_systems(else_in_sys)) == [:else_sys, :default_sys]

@test all([
if_in_sys.eq ~ 0,
if_in_sys.eq ~ 1,
if_in_sys.eq ~ 4,
] .∈ [equations(if_in_sys)])
@test all([
elseif_in_sys.eq ~ 0,
elseif_in_sys.eq ~ 2,
elseif_in_sys.eq ~ 5,
] .∈ [equations(elseif_in_sys)])
@test all([
else_in_sys.eq ~ 0,
else_in_sys.eq ~ 3,
else_in_sys.eq ~ 5,
] .∈ [equations(else_in_sys)])

@test getdefault(if_in_sys.eq) == 1
@test getdefault(elseif_in_sys.eq) == 0
end

@test getdefault(in_sys_1.eq) == 1
@test getdefault(in_sys_2.eq) == 0
@testset "Conditional statements outside the blocks" begin
@mtkmodel C begin end

# Branching statement outside the begin blocks
@mtkmodel OutsideTheBlock begin
@structural_parameters begin
condition = 2
condition = 0
end

@parameters begin
a0
default_parameter
end
@components begin
sys0 = C()
default_sys = C()
end
@equations begin
a0 ~ 0
default_parameter ~ 0
end

if condition == 1
@parameters begin
a1
if_parameter
end
@equations begin
a1 ~ 0
if_parameter ~ 0
end
@components begin
sys1 = C()
if_sys = C()
end
elseif condition == 2
@parameters begin
a2
elseif_parameter
end
@equations begin
a2 ~ 0
elseif_parameter ~ 0
end
@components begin
sys2 = C()
elseif_sys = C()
end
else
@parameters begin
a3
else_parameter
end
@equations begin
a3 ~ 0
else_parameter ~ 0
end
@components begin
sys3 = C()
else_sys = C()
end
end
end

@named out_sys_1 = OutsideTheBlock(condition = 1)
out_sys_1 = complete(out_sys_1)
@named out_sys_2 = OutsideTheBlock(condition = 2)
out_sys_2 = complete(out_sys_2)
@named out_sys_3 = OutsideTheBlock(condition = 10)
out_sys_3 = complete(out_sys_3)

@test nameof.(out_sys_1.systems) == [:sys1, :sys0]
@test nameof.(out_sys_2.systems) == [:sys2, :sys0]
@test nameof.(out_sys_3.systems) == [:sys3, :sys0]

@test Equation[out_sys_1.a1 ~ 0
out_sys_1.a0 ~ 0] == equations(out_sys_1)
@test Equation[out_sys_2.a2 ~ 0
out_sys_1.a0 ~ 0] == equations(out_sys_2)
@test Equation[out_sys_3.a3 ~ 0
out_sys_1.a0 ~ 0] == equations(out_sys_3)
@named if_out_sys = OutsideTheBlock(condition = 1)
if_out_sys = complete(if_out_sys)
@named elseif_out_sys = OutsideTheBlock(condition = 2)
elseif_out_sys = complete(elseif_out_sys)
@named else_out_sys = OutsideTheBlock(condition = 10)
else_out_sys = complete(else_out_sys)
@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 nameof.(get_systems(if_out_sys)) == [:if_sys, :default_sys]
@test nameof.(get_systems(elseif_out_sys)) == [:elseif_sys, :default_sys]
@test nameof.(get_systems(else_out_sys)) == [:else_sys, :default_sys]

@test Equation[if_out_sys.if_parameter ~ 0
if_out_sys.default_parameter ~ 0] == equations(if_out_sys)
@test Equation[elseif_out_sys.elseif_parameter ~ 0
elseif_out_sys.default_parameter ~ 0] == equations(elseif_out_sys)
@test Equation[else_out_sys.else_parameter ~ 0
else_out_sys.default_parameter ~ 0] == equations(else_out_sys)

@mtkmodel TernaryBranchingOutsideTheBlock begin
@structural_parameters begin
condition = true
end
condition ? begin
@parameters begin
ternary_parameter_true
end
@equations begin
ternary_parameter_true ~ 0
end
@components begin
ternary_sys_true = C()
end
end : begin
@parameters begin
ternary_parameter_false
end
@equations begin
ternary_parameter_false ~ 0
end
@components begin
ternary_sys_false = C()
end
end
end

@named ternary_true = TernaryBranchingOutsideTheBlock()
ternary_true = complete(ternary_true)

@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 nameof.(get_systems(ternary_true)) == [:ternary_sys_true]
@test nameof.(get_systems(ternary_false)) == [:ternary_sys_false]

@test Equation[ternary_true.ternary_parameter_true ~ 0] == equations(ternary_true)
@test Equation[ternary_false.ternary_parameter_false ~ 0] == equations(ternary_false)
end

0 comments on commit 8b75e01

Please sign in to comment.