test: add type inference tests

Author: SebastianM-C

Project.toml

@@ -119,6 +119,7 @@ DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -137,4 +138,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg"]
+test = ["AmplNLWriter", "BenchmarkTools", "ControlSystemsBase", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET"]

test/mtkparameters.jl

@@ -3,6 +3,7 @@ using ModelingToolkit: t_nounits as t, D_nounits as D, MTKParameters
 using SymbolicIndexingInterface
 using SciMLStructures: SciMLStructures, canonicalize, Tunable, Discrete, Constants
 using ForwardDiff
+using JET
 
 @parameters a b c d::Integer e[1:3] f[1:3, 1:3]::Int g::Vector{AbstractFloat} h::String
 @named sys = ODESystem(
@@ -121,3 +122,74 @@ ps = MTKParameters(sys, [p => 1.0, q => 2.0, r => 3.0])
 newps = remake_buffer(sys, ps, Dict(p => 1.0f0))
 @test newps.tunable[1] isa Vector{Float32}
 @test newps.tunable[1] == [1.0f0, 2.0f0, 3.0f0]
+
+# JET tests
+
+# scalar parameters only
+function level1()
+    @parameters p1=0.5 [tunable = true] p2 = 1 [tunable=true] p3 = 3 [tunable = false] p4=3 [tunable = true] y0=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p2 * x * y
+           D(y) ~ -p3 * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+# scalar and vector parameters
+function level2()
+    @parameters p1=0.5 [tunable = true] (p23[1:2]=[1, 3.0]) [tunable = true] p4=3 [tunable = false] y0=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p23[1] * x * y
+           D(y) ~ -p23[2] * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+# scalar and vector parameters with different scalar types
+function level3()
+    @parameters p1=0.5 [tunable = true] (p23[1:2]=[1, 3.0]) [tunable = true] p4::Int=3 [tunable = true] y0::Int=1
+    @variables x(t)=2 y(t)=y0
+    D = Differential(t)
+
+    eqs = [D(x) ~ p1 * x - p23[1] * x * y
+           D(y) ~ -p23[2] * y + p4 * x * y]
+
+    sys = structural_simplify(complete(ODESystem(
+        eqs, t, tspan = (0, 3.0), name = :sys, parameter_dependencies = [y0 => 2p4])))
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(sys)
+end
+
+@testset "level$i" for (i, prob) in enumerate([level1(), level2(), level3()])
+    ps = prob.p
+    @testset "Type stability of $portion" for portion in [Tunable(), Discrete(), Constants()]
+        @test_call canonicalize(portion, ps)
+        # @inferred canonicalize(portion, ps)
+        broken =
+            (i ∈ [2,3] && portion == Tunable())
+
+        # broken because the size of a vector of vectors can't be determined at compile time
+        @test_opt broken=broken target_modules = (ModelingToolkit,) canonicalize(
+            portion, ps)
+
+        buffer, repack, alias = canonicalize(portion, ps)
+
+        @test_call SciMLStructures.replace(portion, ps, ones(length(buffer)))
+        @inferred SciMLStructures.replace(portion, ps, ones(length(buffer)))
+        @test_opt target_modules=(ModelingToolkit,) SciMLStructures.replace(
+            portion, ps, ones(length(buffer)))
+
+        @test_call target_modules = (ModelingToolkit,) SciMLStructures.replace!(
+            portion, ps, ones(length(buffer)))
+        @inferred SciMLStructures.replace!(portion, ps, ones(length(buffer)))
+        @test_opt target_modules=(ModelingToolkit,) SciMLStructures.replace!(
+            portion, ps, ones(length(buffer)))
+    end
+end