Merge pull request #3652 from AayushSabharwal/as/better-metadata

refactor: change `metadata` field to be like `BasicSymbolic`

Author: AayushSabharwal

src/systems/abstractsystem.jl

@@ -2665,7 +2665,7 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
     cevs = union(get_continuous_events(basesys), get_continuous_events(sys))
     devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
-    meta = union_nothing(get_metadata(basesys), get_metadata(sys))
+    meta = merge(get_metadata(basesys), get_metadata(sys))
     syss = union(get_systems(basesys), get_systems(sys))
     args = length(ivs) == 0 ? (eqs, sts, ps) : (eqs, ivs[1], sts, ps)
     kwargs = (parameter_dependencies = dep_ps, observed = obs, continuous_events = cevs,

src/systems/connectors.jl

@@ -298,19 +298,19 @@ end
     error("Different types of connectors are in one connection statement: <$(map(nameof, ss))>")
 end
 
+abstract type IsFrame end
+
 "Return true if the system is a 3D multibody frame, otherwise return false."
 function isframe(sys)
-    (has_metadata(sys) && (md = get_metadata(sys)) isa Dict) || return false
-    get(md, :frame, false)
+    getmetadata(sys, IsFrame, false)
 end
 
+abstract type FrameOrientation end
+
 "Return orientation object of a multibody frame."
 function ori(sys)
-    @assert has_metadata(sys)
-    md = get_metadata(sys)
-    if md isa Dict && (O = get(md, :orientation, nothing)) !== nothing
-        return O
-    else
+    val = getmetadata(sys, FrameOrientation, nothing)
+    if val === nothing
         error("System $(sys.name) does not have an orientation object.")
     end
 end

src/systems/system.jl

@@ -9,6 +9,8 @@ struct Schedule{V <: BipartiteGraphs.Matching}
     dummy_sub::Dict{Any, Any}
 end
 
+const MetadataT = Base.ImmutableDict{DataType, Any}
+
 struct System <: AbstractSystem
     tag::UInt
     eqs::Vector{Equation}
@@ -38,7 +40,7 @@ struct System <: AbstractSystem
     discrete_events::Vector{SymbolicDiscreteCallback}
     connector_type::Any
     assertions::Dict{BasicSymbolic, String}
-    metadata::Any
+    metadata::MetadataT
     gui_metadata::Any # ?
     is_dde::Bool
     tstops::Vector{Any}
@@ -60,7 +62,7 @@ struct System <: AbstractSystem
             brownians, iv, observed, parameter_dependencies, var_to_name, name, description,
             defaults, guesses, systems, initialization_eqs, continuous_events, discrete_events,
             connector_type, assertions = Dict{BasicSymbolic, String}(),
-            metadata = nothing, gui_metadata = nothing,
+            metadata = MetadataT(), gui_metadata = nothing,
             is_dde = false, tstops = [], tearing_state = nothing, namespacing = true,
             complete = false, index_cache = nothing, ignored_connections = nothing,
             preface = nothing, parent = nothing, initializesystem = nothing,
@@ -119,8 +121,9 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = [];
         guesses = Dict(), systems = System[], initialization_eqs = Equation[],
         continuous_events = SymbolicContinuousCallback[], discrete_events = SymbolicDiscreteCallback[],
         connector_type = nothing, assertions = Dict{BasicSymbolic, String}(),
-        metadata = nothing, gui_metadata = nothing, is_dde = nothing, tstops = [],
-        tearing_state = nothing, ignored_connections = nothing, parent = nothing,
+        metadata = MetadataT(), gui_metadata = nothing,
+        is_dde = nothing, tstops = [], tearing_state = nothing,
+        ignored_connections = nothing, parent = nothing,
         description = "", name = nothing, discover_from_metadata = true,
         initializesystem = nothing, is_initializesystem = false, preface = [],
         checks = true)
@@ -185,6 +188,17 @@ function System(eqs::Vector{Equation}, iv, dvs, ps, brownians = [];
 
     assertions = Dict{BasicSymbolic, String}(unwrap(k) => v for (k, v) in assertions)
 
+    if isempty(metadata)
+        metadata = MetadataT()
+    elseif metadata isa MetadataT
+        metadata = metadata
+    else
+        meta = MetadataT()
+        for kvp in metadata
+            meta = Base.ImmutableDict(meta, kvp)
+        end
+        metadata = meta
+    end
     System(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)), eqs, noise_eqs, jumps, constraints,
         costs, consolidate, dvs, ps, brownians, iv, observed, parameter_dependencies,
         var_to_name, name, description, defaults, guesses, systems, initialization_eqs,
@@ -621,6 +635,17 @@ function Base.hash(sys::System, h::UInt)
     return h
 end
 
+function SymbolicUtils.getmetadata(sys::AbstractSystem, k::DataType, default)
+    meta = get_metadata(sys)
+    return get(meta, k, default)
+end
+
+function SymbolicUtils.setmetadata(sys::AbstractSystem, k::DataType, v)
+    meta = get_metadata(sys)
+    meta = Base.ImmutableDict(meta, k => v)::MetadataT
+    @set sys.metadata = meta
+end
+
 """
     $(TYPEDSIGNATURES)
 """

test/components.jl

@@ -8,6 +8,7 @@ using ModelingToolkitStandardLibrary.Electrical
 using ModelingToolkitStandardLibrary.Blocks
 using LinearAlgebra
 using ModelingToolkitStandardLibrary.Thermal
+using SymbolicUtils: getmetadata
 include("common/rc_model.jl")
 
 @testset "Basics" begin
@@ -328,8 +329,8 @@ end
 @testset "Issue#3275: Metadata retained on `complete`" begin
     @variables x(t) y(t)
     @named inner = System(D(x) ~ x, t)
-    @named outer = System(D(y) ~ y, t; systems = [inner], metadata = "test")
-    @test ModelingToolkit.get_metadata(outer) == "test"
+    @named outer = System(D(y) ~ y, t; systems = [inner], metadata = [Int => "test"])
+    @test getmetadata(outer, Int, nothing) == "test"
     sys = complete(outer)
-    @test ModelingToolkit.get_metadata(sys) == "test"
+    @test getmetadata(sys, Int, nothing) == "test"
 end

test/discrete_system.jl

@@ -5,7 +5,6 @@
 =#
 using ModelingToolkit, SymbolicIndexingInterface, Test
 using ModelingToolkit: t_nounits as t
-using ModelingToolkit: get_metadata, MTKParameters
 
 # Make sure positive shifts error
 @variables x(t)
@@ -205,11 +204,6 @@ RHS2 = RHS
 #     @test c[1] + 1 == length(sol)
 # end
 
-@variables x(t) y(t)
-testdict = Dict([:test => 1])
-@named sys = System([x(k + 1) ~ 1.0], t, [x], []; metadata = testdict)
-@test get_metadata(sys) == testdict
-
 @variables x(t) y(t) u(t)
 eqs = [u ~ 1
        x ~ x(k - 1) + u

test/nonlinearsystem.jl

@@ -1,5 +1,4 @@
 using ModelingToolkit, StaticArrays, LinearAlgebra
-using ModelingToolkit: get_metadata
 using DiffEqBase, SparseArrays
 using Test
 using NonlinearSolve
@@ -198,14 +197,6 @@ eq = [v1 ~ sin(2pi * t * h)
 @named sys = System(eq, t)
 @test length(equations(structural_simplify(sys))) == 0
 
-@variables x(t)
-@parameters a
-eqs = [0 ~ a * x]
-
-testdict = Dict([:test => 1])
-@named sys = System(eqs, [x], [a], metadata = testdict)
-@test get_metadata(sys) == testdict
-
 @testset "Remake" begin
     @parameters a=1.0 b=1.0 c=1.0
     @constants h = 1

test/odesystem.jl

@@ -799,16 +799,6 @@ let
     @test string.(independent_variables(prob.f.sys)) == ["t"]
 end
 
-@parameters C L R
-@variables q(t) p(t) F(t)
-
-eqs = [D(q) ~ -p / L - F
-       D(p) ~ q / C
-       0 ~ q / C - R * F]
-testdict = Dict([:name => "test"])
-@named sys = System(eqs, t, metadata = testdict)
-@test get_metadata(sys) == testdict
-
 @variables P(t)=NaN Q(t)=NaN
 eqs = [D(Q) ~ 1 / sin(P), D(P) ~ log(-cos(Q))]
 @named sys = System(eqs, t, [P, Q], [])
@@ -1112,16 +1102,23 @@ end
 
 # https://github.com/SciML/ModelingToolkit.jl/issues/2502
 @testset "Extend systems with a field that can be nothing" begin
-    A = Dict(:a => 1)
-    B = Dict(:b => 2)
+    A = Dict(Int => 1)
+    B = Dict(String => 2)
     @named A1 = System(Equation[], t, [], [])
     @named B1 = System(Equation[], t, [], [])
     @named A2 = System(Equation[], t, [], []; metadata = A)
     @named B2 = System(Equation[], t, [], []; metadata = B)
-    @test ModelingToolkit.get_metadata(extend(A1, B1)) == nothing
-    @test ModelingToolkit.get_metadata(extend(A1, B2)) == B
-    @test ModelingToolkit.get_metadata(extend(A2, B1)) == A
-    @test Set(ModelingToolkit.get_metadata(extend(A2, B2))) == Set(A ∪ B)
+    @test isempty(ModelingToolkit.get_metadata(extend(A1, B1)))
+    meta = ModelingToolkit.get_metadata(extend(A1, B2))
+    @test length(meta) == 1
+    @test meta[String] == 2
+    meta = ModelingToolkit.get_metadata(extend(A2, B1))
+    @test length(meta) == 1
+    @test meta[Int] == 1
+    meta = ModelingToolkit.get_metadata(extend(A2, B2))
+    @test length(meta) == 2
+    @test meta[Int] == 1
+    @test meta[String] == 2
 end
 
 # https://github.com/SciML/ModelingToolkit.jl/issues/2859

test/optimizationsystem.jl

@@ -1,7 +1,6 @@
 using ModelingToolkit, SparseArrays, Test, Optimization, OptimizationOptimJL,
       OptimizationMOI, Ipopt, AmplNLWriter, Ipopt_jll, SymbolicIndexingInterface,
       LinearAlgebra
-using ModelingToolkit: get_metadata
 
 @testset "basic" begin
     @variables x y
@@ -228,18 +227,6 @@ end
      =#
 end
 
-@testset "metadata" begin
-    @variables x
-    o1 = (x - 1)^2
-    c1 = [
-        x ~ 1
-    ]
-    testdict = Dict(["test" => 1])
-    sys1 = OptimizationSystem(o1, [x], [], name = :sys1, constraints = c1,
-        metadata = testdict)
-    @test get_metadata(sys1) == testdict
-end
-
 @testset "non-convex problem with inequalities" begin
     @variables x[1:2] [bounds = (0.0, Inf)]
     @named sys = OptimizationSystem(x[1] + x[2], [x...], [];