[Awaiting Review] Add metadata field to systems (#1768)

Co-authored-by: Yingbo Ma <[email protected]>

Author: xtalax

docs/src/basics/AbstractSystem.md

@@ -60,6 +60,7 @@ Optionally, a system could have:
 - `get_defaults(sys)`: A `Dict` that maps variables into their default values.
 - `independent_variables(sys)`: The independent variables of a system.
 - `get_noiseeqs(sys)`: Noise equations of the current-level system.
+- `get_metadata(sys)`: Any metadata about the system or its origin to be used by downstream packages.
 
 Note that if you know a system is an `AbstractTimeDependentSystem` you could use `get_iv` to get the 
 unique independent variable directly, rather than using `independent_variables(sys)[1]`, which is clunky and may cause problems if `sys` is an `AbstractMultivariateSystem` because there may be more than one independent variable. `AbstractTimeIndependentSystem`s do not have a method `get_iv`, and `independent_variables(sys)` will return a size-zero result for such. For an `AbstractMultivariateSystem`, `get_ivs` is equivalent.

src/systems/abstractsystem.jl

@@ -192,7 +192,8 @@ for prop in [:eqs
              :preface
              :torn_matching
              :tearing_state
-             :substitutions]
+             :substitutions
+             :metadata]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
     @eval begin

src/systems/diffeqs/odesystem.jl

@@ -115,11 +115,16 @@ struct ODESystem <: AbstractODESystem
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
 
     function ODESystem(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad,
                        jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
                        torn_matching, connector_type, connections, preface, cevents,
-                       devents, tearing_state = nothing, substitutions = nothing;
+                       devents, tearing_state = nothing, substitutions = nothing,
+                       metadata = nothing;
                        checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
@@ -133,7 +138,7 @@ struct ODESystem <: AbstractODESystem
         new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
             connector_type, connections, preface, cevents, devents, tearing_state,
-            substitutions)
+            substitutions, metadata)
     end
 end
 
@@ -149,7 +154,8 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
                    preface = nothing,
                    continuous_events = nothing,
                    discrete_events = nothing,
-                   checks = true)
+                   checks = true,
+                   metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = scalarize(deqs)
@@ -186,7 +192,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     ODESystem(deqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac,
               ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing,
               connector_type, nothing, preface, cont_callbacks, disc_callbacks,
-              checks = checks)
+              metadata, checks = checks)
 end
 
 function ODESystem(eqs, iv = nothing; kwargs...)

src/systems/diffeqs/sdesystem.jl

@@ -96,10 +96,14 @@ struct SDESystem <: AbstractODESystem
     true at the end of an integration step.
     """
     discrete_events::Vector{SymbolicDiscreteCallback}
-
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
     function SDESystem(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac,
                        ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-                       cevents, devents; checks::Union{Bool, Int} = true)
+                       cevents, devents, metadata = nothing;
+                       checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
             check_parameters(ps, iv)
@@ -110,7 +114,8 @@ struct SDESystem <: AbstractODESystem
             all_dimensionless([dvs; ps; iv]) || check_units(deqs, neqs)
         end
         new(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac,
-            Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents)
+            Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents,
+            metadata)
     end
 end
 
@@ -125,7 +130,8 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
                    connector_type = nothing,
                    checks = true,
                    continuous_events = nothing,
-                   discrete_events = nothing)
+                   discrete_events = nothing,
+                   metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = scalarize(deqs)
@@ -160,7 +166,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
 
     SDESystem(deqs, neqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac,
               ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-              cont_callbacks, disc_callbacks; checks = checks)
+              cont_callbacks, disc_callbacks, metadata; checks = checks)
 end
 
 function SDESystem(sys::ODESystem, neqs; kwargs...)

src/systems/discrete_system/discrete_system.jl

@@ -67,10 +67,15 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
 
     function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name,
                             systems, defaults, preface, connector_type,
-                            tearing_state = nothing, substitutions = nothing;
+                            tearing_state = nothing, substitutions = nothing,
+                            metadata = nothing;
                             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
@@ -80,7 +85,7 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
             all_dimensionless([dvs; ps; iv; ctrls]) || check_units(discreteEqs)
         end
         new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, defaults,
-            preface, connector_type, tearing_state, substitutions)
+            preface, connector_type, tearing_state, substitutions, metadata)
     end
 end
 
@@ -99,6 +104,7 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
                         defaults = _merge(Dict(default_u0), Dict(default_p)),
                         preface = nothing,
                         connector_type = nothing,
+                        metadata = nothing,
                         kwargs...)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
@@ -125,7 +131,7 @@ function DiscreteSystem(eqs::AbstractVector{<:Equation}, iv, dvs, ps;
         throw(ArgumentError("System names must be unique."))
     end
     DiscreteSystem(eqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, name, systems,
-                   defaults, preface, connector_type, kwargs...)
+                   defaults, preface, connector_type, metadata, kwargs...)
 end
 
 function DiscreteSystem(eqs, iv = nothing; kwargs...)

src/systems/jumps/jumpsystem.jl

@@ -83,9 +83,13 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
     state value or parameter.*
     """
     discrete_events::Vector{SymbolicDiscreteCallback}
-
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
     function JumpSystem{U}(ap::U, iv, states, ps, var_to_name, observed, name, systems,
-                           defaults, connector_type, devents;
+                           defaults, connector_type, devents,
+                           metadata = nothing;
                            checks::Union{Bool, Int} = true) where {U <: ArrayPartition}
         if checks == true || (checks & CheckComponents) > 0
             check_variables(states, iv)
@@ -95,7 +99,7 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
             all_dimensionless([states; ps; iv]) || check_units(ap, iv)
         end
         new{U}(ap, iv, states, ps, var_to_name, observed, name, systems, defaults,
-               connector_type, devents)
+               connector_type, devents, metadata)
     end
 end
 
@@ -110,6 +114,7 @@ function JumpSystem(eqs, iv, states, ps;
                     checks = true,
                     continuous_events = nothing,
                     discrete_events = nothing,
+                    metadata = nothing,
                     kwargs...)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
@@ -147,7 +152,8 @@ function JumpSystem(eqs, iv, states, ps;
     disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
 
     JumpSystem{typeof(ap)}(ap, value(iv), states, ps, var_to_name, observed, name, systems,
-                           defaults, connector_type, disc_callbacks; checks = checks)
+                           defaults, connector_type, disc_callbacks, metadata,
+                           checks = checks)
 end
 
 function generate_rate_function(js::JumpSystem, rate)

src/systems/nonlinear/nonlinearsystem.jl

@@ -62,15 +62,20 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
     substitutions: substitutions generated by tearing.
     """
     substitutions::Any
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
 
     function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems,
                              defaults, connector_type, connections, tearing_state = nothing,
-                             substitutions = nothing; checks::Union{Bool, Int} = true)
+                             substitutions = nothing, metadata = nothing;
+                             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckUnits) > 0
             all_dimensionless([states; ps]) || check_units(eqs)
         end
         new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults,
-            connector_type, connections, tearing_state, substitutions)
+            connector_type, connections, tearing_state, substitutions, metadata)
     end
 end
 
@@ -84,7 +89,8 @@ function NonlinearSystem(eqs, states, ps;
                          connector_type = nothing,
                          continuous_events = nothing, # this argument is only required for ODESystems, but is added here for the constructor to accept it without error
                          discrete_events = nothing,   # this argument is only required for ODESystems, but is added here for the constructor to accept it without error
-                         checks = true)
+                         checks = true,
+                         metadata = nothing)
     continuous_events === nothing || isempty(continuous_events) ||
         throw(ArgumentError("NonlinearSystem does not accept `continuous_events`, you provided $continuous_events"))
     discrete_events === nothing || isempty(discrete_events) ||
@@ -119,7 +125,7 @@ function NonlinearSystem(eqs, states, ps;
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults,
-                    connector_type, nothing, checks = checks)
+                    connector_type, nothing, metadata, checks = checks)
 end
 
 function calculate_jacobian(sys::NonlinearSystem; sparse = false, simplify = false)

src/systems/optimization/optimizationsystem.jl

@@ -40,16 +40,20 @@ struct OptimizationSystem <: AbstractTimeIndependentSystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
     function OptimizationSystem(op, states, ps, var_to_name, observed,
-                                constraints, name, systems, defaults;
+                                constraints, name, systems, defaults, metadata = nothing;
                                 checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckUnits) > 0
             check_units(op)
             check_units(observed)
             all_dimensionless([states; ps]) || check_units(constraints)
         end
         new(op, states, ps, var_to_name, observed,
-            constraints, name, systems, defaults)
+            constraints, name, systems, defaults, metadata)
     end
 end
 
@@ -61,7 +65,8 @@ function OptimizationSystem(op, states, ps;
                             defaults = _merge(Dict(default_u0), Dict(default_p)),
                             name = nothing,
                             systems = OptimizationSystem[],
-                            checks = true)
+                            checks = true,
+                            metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     if !(isempty(default_u0) && isempty(default_p))
@@ -83,7 +88,7 @@ function OptimizationSystem(op, states, ps;
     OptimizationSystem(value(op), states, ps, var_to_name,
                        observed,
                        constraints,
-                       name, systems, defaults; checks = checks)
+                       name, systems, defaults, metadata; checks = checks)
 end
 
 function calculate_gradient(sys::OptimizationSystem)

src/systems/pde/pdesystem.jl

@@ -64,18 +64,24 @@ struct PDESystem <: ModelingToolkit.AbstractMultivariateSystem
     name: the name of the system
     """
     name::Symbol
+    """
+    metadata: metadata for the system, to be used by downstream packages.
+    """
+    metadata::Any
     @add_kwonly function PDESystem(eqs, bcs, domain, ivs, dvs,
                                    ps = SciMLBase.NullParameters();
                                    defaults = Dict(),
                                    systems = [],
                                    connector_type = nothing,
+                                   metadata = nothing,
                                    checks::Union{Bool, Int} = true,
                                    name)
         if checks == true || (checks & CheckUnits) > 0
             all_dimensionless([dvs; ivs; ps]) || check_units(eqs)
         end
         eqs = eqs isa Vector ? eqs : [eqs]
-        new(eqs, bcs, domain, ivs, dvs, ps, defaults, connector_type, systems, name)
+        new(eqs, bcs, domain, ivs, dvs, ps, defaults, connector_type, systems, name,
+            metadata)
     end
 end