test: add more comprehensive tests for FMIComponent

AayushSabharwal · AayushSabharwal · commit c8b42eb04477 · 2025-01-26T19:34:08.000+05:30
diff --git a/test/extensions/fmi.jl b/test/extensions/fmi.jl
@@ -81,83 +81,229 @@ const FMU_DIR = joinpath(@__DIR__, "fmus")
     end
 end
 
-@testset "IO Model" begin
-    @testset "v2, ME" begin
-        fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :ME)
-        @named adder = MTK.FMIComponent(Val(2); fmu, type = :ME)
-        @test MTK.isinput(adder.a)
-        @test MTK.isinput(adder.b)
-        @test MTK.isoutput(adder.out)
-        @test !MTK.isinput(adder.c) && !MTK.isoutput(adder.c)
+@mtkmodel SimpleAdder begin
+    @variables begin
+        a(t)
+        b(t)
+        c(t)
+        out(t)
+        out2(t)
+    end
+    @parameters begin
+        value = 1.0
+    end
+    @equations begin
+        out ~ a + b + value
+        D(c) ~ out
+        out2 ~ 2c
+    end
+end
+
+@mtkmodel StateSpace begin
+    @variables begin
+        x(t)
+        y(t)
+        u(t)
+    end
+    @parameters begin
+        A = 1.0
+        B = 1.0
+        C = 1.0
+        _D = 1.0
+    end
+    @equations begin
+        D(x) ~ A * x + B * u
+        y ~ C * x + _D * u
+    end
+end
 
+@testset "IO Model" begin
+    function build_simple_adder(adder)
         @variables a(t) b(t) c(t) [guess = 1.0]
         @mtkbuild sys = ODESystem(
             [adder.a ~ a, adder.b ~ b, D(a) ~ t,
                 D(b) ~ adder.out + adder.c, c^2 ~ adder.out + adder.value],
             t;
             systems = [adder])
-
         # c will be solved for by initialization
         # this tests that initialization also works with FMUs
-        prob = ODEProblem(sys, [sys.adder.c => 1.0, sys.a => 1.0, sys.b => 1.0], (0.0, 1.0))
-        sol = solve(prob, Rodas5P(autodiff = false))
+        prob = ODEProblem(sys, [sys.adder.c => 2.0, sys.a => 1.0, sys.b => 1.0],
+            (0.0, 1.0), [sys.adder.value => 2.0])
+        return sys, prob
+    end
+
+    @named adder = SimpleAdder()
+    truesys, trueprob = build_simple_adder(adder)
+    truesol = solve(trueprob, abstol = 1e-8, reltol = 1e-8)
+    @test SciMLBase.successful_retcode(truesol)
+
+    @testset "v2, ME" begin
+        fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :ME)
+        @named adder = MTK.FMIComponent(Val(2); fmu, type = :ME)
+        @test MTK.isinput(adder.a)
+        @test MTK.isinput(adder.b)
+        @test MTK.isoutput(adder.out)
+        @test MTK.isoutput(adder.out2)
+        @test !MTK.isinput(adder.c) && !MTK.isoutput(adder.c)
+
+        sys, prob = build_simple_adder(adder)
+        sol = solve(prob, Rodas5P(autodiff = false), abstol = 1e-8, reltol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
+
+        @test truesol(sol.t;
+            idxs = [truesys.a, truesys.b, truesys.c, truesys.adder.c]).u≈sol[[
+            sys.a, sys.b, sys.c, sys.adder.c]] rtol=1e-7
     end
     @testset "v2, CS" begin
         fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :CS)
         @named adder = MTK.FMIComponent(
-            Val(2); fmu, type = :CS, communication_step_size = 0.001)
+            Val(2); fmu, type = :CS, communication_step_size = 1e-3,
+            reinitializealg = BrownFullBasicInit())
         @test MTK.isinput(adder.a)
         @test MTK.isinput(adder.b)
         @test MTK.isoutput(adder.out)
+        @test MTK.isoutput(adder.out2)
         @test !MTK.isinput(adder.c) && !MTK.isoutput(adder.c)
-        @variables a(t) b(t) c(t) [guess = 1.0]
-        @mtkbuild sys = ODESystem(
-            [adder.a ~ a, adder.b ~ b, D(a) ~ t,
-                D(b) ~ adder.out + adder.c, c^2 ~ adder.out + adder.value],
-            t;
-            systems = [adder])
 
-        # c will be solved for by initialization
-        # this tests that initialization also works with FMUs
-        prob = ODEProblem(sys, [sys.adder.c => 1.0, sys.a => 1.0, sys.b => 1.0],
-            (0.0, 1.0); use_scc = false)
-        sol = solve(prob, Rodas5P(autodiff = false))
+        sys, prob = build_simple_adder(adder)
+        sol = solve(prob, Rodas5P(autodiff = false), abstol = 1e-8, reltol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
+
+        @test truesol(sol.t; idxs = [truesys.a, truesys.b, truesys.c]).u≈sol.u rtol=1e-2
+        # sys.adder.c is a discrete variable
+        @test truesol(sol.t; idxs = truesys.adder.c)≈sol(sol.t; idxs = sys.adder.c) rtol=1e-3
     end
 
+    function build_sspace_model(sspace)
+        @variables u(t)=1.0 x(t)=1.0 y(t) [guess = 1.0]
+        @mtkbuild sys = ODESystem(
+            [sspace.u ~ u, D(u) ~ t, D(x) ~ sspace.x + sspace.y, y^2 ~ sspace.y + sspace.x], t;
+            systems = [sspace]
+        )
+
+        prob = ODEProblem(
+            sys, [sys.sspace.x => 1.0], (0.0, 1.0), [sys.sspace.A => 2.0]; use_scc = false)
+        return sys, prob
+    end
+
+    @named sspace = StateSpace()
+    truesys, trueprob = build_sspace_model(sspace)
+    truesol = solve(trueprob, abstol = 1e-8, reltol = 1e-8)
+    @test SciMLBase.successful_retcode(truesol)
+
     @testset "v3, ME" begin
         fmu = loadFMU(joinpath(FMU_DIR, "StateSpace.fmu"); type = :ME)
         @named sspace = MTK.FMIComponent(Val(3); fmu, type = :ME)
         @test MTK.isinput(sspace.u)
         @test MTK.isoutput(sspace.y)
         @test !MTK.isinput(sspace.x) && !MTK.isoutput(sspace.x)
-        @variables u(t)=1.0 x(t)=1.0 y(t) [guess = 1.0]
-        @mtkbuild sys = ODESystem(
-            [sspace.u ~ u, D(u) ~ t, D(x) ~ sspace.x + sspace.y, y^2 ~ sspace.y + x], t;
-            systems = [sspace]
-        )
 
-        prob = ODEProblem(sys, [sys.sspace.x => 1.0], (0.0, 1.0); use_scc = false)
-        sol = solve(prob, Rodas5P(autodiff = false))
+        sys, prob = build_sspace_model(sspace)
+        sol = solve(prob, Rodas5P(autodiff = false); abstol = 1e-8, reltol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
+
+        @test truesol(sol.t;
+            idxs = [truesys.u, truesys.x, truesys.y, truesys.sspace.x]).u≈sol[[
+            sys.u, sys.x, sys.y, sys.sspace.x]] rtol=1e-7
     end
 
     @testset "v3, CS" begin
         fmu = loadFMU(joinpath(FMU_DIR, "StateSpace.fmu"); type = :CS)
         @named sspace = MTK.FMIComponent(
-            Val(3); fmu, communication_step_size = 1e-3, type = :CS)
+            Val(3); fmu, communication_step_size = 1e-4, type = :CS,
+            reinitializealg = BrownFullBasicInit())
         @test MTK.isinput(sspace.u)
         @test MTK.isoutput(sspace.y)
         @test !MTK.isinput(sspace.x) && !MTK.isoutput(sspace.x)
-        @variables u(t)=1.0 x(t)=1.0 y(t) [guess = 1.0]
+
+        sys, prob = build_sspace_model(sspace)
+        sol = solve(prob, Rodas5P(autodiff = false); abstol = 1e-8, reltol = 1e-8)
+        @test SciMLBase.successful_retcode(sol)
+
+        @test truesol(
+            sol.t; idxs = [truesys.u, truesys.x, truesys.y]).u≈sol[[sys.u, sys.x, sys.y]] rtol=1e-2
+        @test truesol(sol.t; idxs = truesys.sspace.x).u≈sol(sol.t; idxs = sys.sspace.x).u rtol=1e-2
+    end
+end
+
+@testset "FMUs in a loop" begin
+    function build_looped_adders(adder1, adder2)
+        @variables x(t) = 1
         @mtkbuild sys = ODESystem(
-            [sspace.u ~ u, D(u) ~ t, D(x) ~ sspace.x + sspace.y, y^2 ~ sspace.y + x], t;
-            systems = [sspace]
-        )
+            [D(x) ~ x, adder1.a ~ adder2.out2,
+                adder2.a ~ adder1.out2, adder1.b ~ 1.0, adder2.b ~ 2.0],
+            t;
+            systems = [adder1, adder2])
+        prob = ODEProblem(
+            sys, [adder1.c => 1.0, adder2.c => 1.0, adder1.a => 2.0], (0.0, 1.0))
+        return sys, prob
+    end
+    @named adder1 = SimpleAdder()
+    @named adder2 = SimpleAdder()
+    truesys, trueprob = build_looped_adders(adder1, adder2)
+    truesol = solve(trueprob, Tsit5(), reltol = 1e-8)
+    @test SciMLBase.successful_retcode(truesol)
+
+    @testset "v2, ME" begin
+        fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :ME)
+        @named adder1 = MTK.FMIComponent(Val(2); fmu, type = :ME)
+        @named adder2 = MTK.FMIComponent(Val(2); fmu, type = :ME)
+        sys, prob = build_looped_adders(adder1, adder2)
+        sol = solve(prob, Rodas5P(autodiff = false); reltol = 1e-8)
+        @test SciMLBase.successful_retcode(sol)
+        @test truesol(sol.t;
+            idxs = [truesys.adder1.c, truesys.adder2.c]).u≈sol(
+            sol.t; idxs = [sys.adder1.c, sys.adder2.c]).u rtol=1e-7
+    end
+    @testset "v2, CS" begin
+        fmu = loadFMU(joinpath(FMU_DIR, "SimpleAdder.fmu"); type = :CS)
+        @named adder1 = MTK.FMIComponent(
+            Val(2); fmu, type = :CS, communication_step_size = 1e-3)
+        @named adder2 = MTK.FMIComponent(
+            Val(2); fmu, type = :CS, communication_step_size = 1e-3)
+        sys, prob = build_looped_adders(adder1, adder2)
+        sol = solve(prob, Tsit5(); reltol = 1e-8)
+        @test truesol(sol.t;
+            idxs = [truesys.adder1.c, truesys.adder2.c]).u≈sol(
+            sol.t; idxs = [sys.adder1.c, sys.adder2.c]).u rtol=1e-3
+    end
+
+    function build_looped_sspace(sspace1, sspace2)
+        @variables x(t) = 1
+        @mtkbuild sys = ODESystem([D(x) ~ x, sspace1.u ~ sspace2.x, sspace2.u ~ sspace1.y],
+            t; systems = [sspace1, sspace2])
+        prob = ODEProblem(sys, [sspace1.x => 1.0, sspace2.x => 1.0], (0.0, 1.0))
+        return sys, prob
+    end
+    @named sspace1 = StateSpace()
+    @named sspace2 = StateSpace()
+    truesys, trueprob = build_looped_sspace(sspace1, sspace2)
+    truesol = solve(trueprob, Rodas5P(), reltol = 1e-8)
+    @test SciMLBase.successful_retcode(truesol)
+
+    @testset "v3, ME" begin
+        fmu = loadFMU(joinpath(FMU_DIR, "StateSpace.fmu"); type = :ME)
+        @named sspace1 = MTK.FMIComponent(Val(3); fmu, type = :ME)
+        @named sspace2 = MTK.FMIComponent(Val(3); fmu, type = :ME)
+        sys, prob = build_looped_sspace(sspace1, sspace2)
+        sol = solve(prob, Rodas5P(autodiff = false); reltol = 1e-8)
+        @test SciMLBase.successful_retcode(sol)
+        @test truesol(sol.t;
+            idxs = [truesys.sspace1.x, truesys.sspace2.x]).u≈sol(
+            sol.t; idxs = [sys.sspace1.x, sys.sspace2.x]).u rtol=1e-7
+    end
 
-        prob = ODEProblem(sys, [sys.sspace.x => 1.0], (0.0, 1.0); use_scc = false)
-        sol = solve(prob, Rodas5P(autodiff = false))
+    @testset "v3, CS" begin
+        fmu = loadFMU(joinpath(FMU_DIR, "StateSpace.fmu"); type = :CS)
+        @named sspace1 = MTK.FMIComponent(
+            Val(3); fmu, type = :CS, communication_step_size = 1e-4)
+        @named sspace2 = MTK.FMIComponent(
+            Val(3); fmu, type = :CS, communication_step_size = 1e-4)
+        sys, prob = build_looped_sspace(sspace1, sspace2)
+        sol = solve(prob, Rodas5P(autodiff = false); reltol = 1e-8)
         @test SciMLBase.successful_retcode(sol)
+        @test truesol(sol.t;
+            idxs = [truesys.sspace1.x, truesys.sspace2.x]).u≈sol(
+            sol.t; idxs = [sys.sspace1.x, sys.sspace2.x]).u rtol=1e-2
     end
 end
