Merge pull request #973 from SciML/ys/connect

YingboMa · web-flow · commit f12f472f630f · 2021-04-17T12:12:28.000-04:00
Add connectors
diff --git a/examples/electrical_components.jl b/examples/electrical_components.jl
@@ -3,11 +3,23 @@ using ModelingToolkit, OrdinaryDiffEq
 
 # Basic electric components
 @parameters t
-function Pin(;name)
+@connector function Pin(;name)
     @variables v(t) i(t)
     ODESystem(Equation[], t, [v, i], [], name=name, defaults=[v=>1.0, i=>1.0])
 end
 
+function ModelingToolkit.connect(::Type{Pin}, ps...)
+    eqs = [
+           0 ~ sum(p->p.i, ps) # KCL
+          ]
+    # KVL
+    for i in 1:length(ps)-1
+        push!(eqs, ps[i].v ~ ps[i+1].v)
+    end
+
+    return eqs
+end
+
 function Ground(;name)
     @named g = Pin()
     eqs = [g.v ~ 0]
@@ -70,15 +82,3 @@ function Inductor(; name, L = 1.0)
           ]
     ODESystem(eqs, t, [v, i], [L], systems=[p, n], defaults=Dict(L => val), name=name)
 end
-
-function connect_pins(ps...)
-    eqs = [
-           0 ~ sum(p->p.i, ps) # KCL
-          ]
-    # KVL
-    for i in 1:length(ps)-1
-        push!(eqs, ps[i].v ~ ps[i+1].v)
-    end
-
-    return eqs
-end
diff --git a/examples/rc_model.jl b/examples/rc_model.jl
@@ -9,9 +9,9 @@ V = 1.0
 @named ground = Ground()
 
 rc_eqs = [
-          connect_pins(source.p, resistor.p)
-          connect_pins(resistor.n, capacitor.p)
-          connect_pins(capacitor.n, source.n, ground.g)
+          connect(source.p, resistor.p)
+          connect(resistor.n, capacitor.p)
+          connect(capacitor.n, source.n, ground.g)
          ]
 
 @named rc_model = ODESystem(rc_eqs, t, systems=[resistor, capacitor, source, ground])
diff --git a/examples/serial_inductor.jl b/examples/serial_inductor.jl
@@ -7,10 +7,10 @@ include("electrical_components.jl")
 @named ground = Ground()
 
 eqs = [
-       connect_pins(source.p, resistor.p)
-       connect_pins(resistor.n, inductor1.p)
-       connect_pins(inductor1.n, inductor2.p)
-       connect_pins(source.n, inductor2.n, ground.g)
+       connect(source.p, resistor.p)
+       connect(resistor.n, inductor1.p)
+       connect(inductor1.n, inductor2.p)
+       connect(source.n, inductor2.n, ground.g)
       ]
 
 @named ll_model = ODESystem(eqs, t, systems=[source, resistor, inductor1, inductor2, ground])
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
@@ -149,7 +149,7 @@ export SteadyStateProblem, SteadyStateProblemExpr
 export JumpProblem, DiscreteProblem
 export NonlinearSystem, OptimizationSystem
 export ControlSystem
-export alias_elimination, flatten
+export alias_elimination, flatten, connect, @connector
 export ode_order_lowering, liouville_transform
 export runge_kutta_discretize
 export PDESystem
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
@@ -157,6 +157,7 @@ for prop in [
              :domain
              :depvars
              :indvars
+             :connection_type
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
@@ -579,3 +580,64 @@ function check_eqs_u0(eqs, dvs, u0)
     end
     return nothing
 end
+
+function with_connection_type(expr)
+    @assert expr isa Expr && (expr.head == :function || (expr.head == :(=) &&
+                                       expr.args[1] isa Expr &&
+                                       expr.args[1].head == :call))
+
+    sig = expr.args[1]
+    body = expr.args[2]
+
+    fname = sig.args[1]
+    args = sig.args[2:end]
+
+    quote
+        struct $fname
+            $(gensym()) -> 1 # this removes the default constructor
+        end
+        function $fname($(args...))
+            function f()
+                $body
+            end
+            res = f()
+            $isdefined(res, :connection_type) ? $Setfield.@set!(res.connection_type = $fname) : res
+        end
+    end
+end
+
+macro connector(expr)
+    esc(with_connection_type(expr))
+end
+
+promote_connect_rule(::Type{T}, ::Type{S}) where {T, S} = Union{}
+promote_connect_rule(::Type{T}, ::Type{T}) where {T} = T
+promote_connect_type(t1::Type, t2::Type, ts::Type...) = promote_connect_rule(promote_connect_rule(t1, t2), ts...)
+@inline function promote_connect_type(::Type{T}, ::Type{S}) where {T,S}
+    promote_connect_result(
+        T,
+        S,
+        promote_connect_rule(T,S),
+        promote_connect_rule(S,T)
+    )
+end
+
+promote_connect_result(::Type, ::Type, ::Type{T}, ::Type{Union{}}) where {T} = T
+promote_connect_result(::Type, ::Type, ::Type{Union{}}, ::Type{S}) where {S} = S
+promote_connect_result(::Type, ::Type, ::Type{T}, ::Type{T}) where {T} = T
+function promote_connect_result(::Type{T}, ::Type{S}, ::Type{P1}, ::Type{P2}) where {T,S,P1,P2}
+    throw(ArgumentError("connection promotion for $T and $S resulted in $P1 and $P2. " *
+                        "Define promotion only in one direction."))
+end
+
+throw_connector_promotion(T, S) = throw(ArgumentError("Don't know how to connect systems of type $S and $T"))
+promote_connect_result(::Type{T},::Type{S},::Type{Union{}},::Type{Union{}}) where {T,S} = throw_connector_promotion(T,S)
+
+promote_connect_type(::Type{T}, ::Type{T}) where {T} = T
+function promote_connect_type(T, S)
+    error("Don't know how to connect systems of type $S and $T")
+end
+
+function connect(syss...)
+    connect(promote_connect_type(map(get_connection_type, syss)...), syss...)
+end
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -69,6 +69,10 @@ struct ODESystem <: AbstractODESystem
     structure: structural information of the system
     """
     structure::Any
+    """
+    type: type of the system
+    """
+    connection_type::Any
 end
 
 function ODESystem(
@@ -79,6 +83,7 @@ function ODESystem(
                    default_u0=Dict(),
                    default_p=Dict(),
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
+                   connection_type=nothing,
                   )
     iv′ = value(iv)
     dvs′ = value.(dvs)
@@ -98,7 +103,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, nothing)
+    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, nothing, connection_type)
 end
 
 iv_from_nested_derivative(x::Term) = operation(x) isa Differential ? iv_from_nested_derivative(arguments(x)[1]) : arguments(x)[1]
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
@@ -71,6 +71,10 @@ struct SDESystem <: AbstractODESystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
+    """
+    type: type of the system
+    """
+    connection_type::Any
 end
 
 function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
@@ -79,7 +83,9 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
                    default_u0=Dict(),
                    default_p=Dict(),
                    defaults=_merge(Dict(default_u0), Dict(default_p)),
-                   name = gensym(:SDESystem))
+                   name = gensym(:SDESystem),
+                   connection_type=nothing,
+                   )
     iv′ = value(iv)
     dvs′ = value.(dvs)
     ps′ = value.(ps)
@@ -94,7 +100,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
     Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
     Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
-    SDESystem(deqs, neqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults)
+    SDESystem(deqs, neqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
 end
 
 function generate_diffusion_function(sys::SDESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
diff --git a/src/systems/jumps/jumpsystem.jl b/src/systems/jumps/jumpsystem.jl
@@ -47,6 +47,10 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
+    """
+    type: type of the system
+    """
+    connection_type::Any
 end
 
 function JumpSystem(eqs, iv, states, ps;
@@ -55,7 +59,9 @@ function JumpSystem(eqs, iv, states, ps;
                     default_u0=Dict(),
                     default_p=Dict(),
                     defaults=_merge(Dict(default_u0), Dict(default_p)),
-                    name = gensym(:JumpSystem))
+                    name = gensym(:JumpSystem),
+                    connection_type=nothing,
+                    )
 
     ap = ArrayPartition(MassActionJump[], ConstantRateJump[], VariableRateJump[])
     for eq in eqs
@@ -75,7 +81,7 @@ function JumpSystem(eqs, iv, states, ps;
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
 
-    JumpSystem{typeof(ap)}(ap, value(iv), value.(states), value.(ps), observed, name, systems, defaults)
+    JumpSystem{typeof(ap)}(ap, value(iv), value.(states), value.(ps), observed, name, systems, defaults, connection_type)
 end
 
 function generate_rate_function(js, rate)
diff --git a/src/systems/nonlinear/nonlinearsystem.jl b/src/systems/nonlinear/nonlinearsystem.jl
@@ -43,21 +43,27 @@ struct NonlinearSystem <: AbstractSystem
     structure: structural information of the system
     """
     structure::Any
+    """
+    type: type of the system
+    """
+    connection_type::Any
 end
 
 function NonlinearSystem(eqs, states, ps;
-                         observed = [],
-                         name = gensym(:NonlinearSystem),
+                         observed=[],
+                         name=gensym(:NonlinearSystem),
                          default_u0=Dict(),
                          default_p=Dict(),
                          defaults=_merge(Dict(default_u0), Dict(default_p)),
-                         systems = NonlinearSystem[])
+                         systems=NonlinearSystem[],
+                         connection_type=nothing,
+                         )
     if !(isempty(default_u0) && isempty(default_p))
         Base.depwarn("`default_u0` and `default_p` are deprecated. Use `defaults` instead.", :NonlinearSystem, force=true)
     end
     defaults = todict(defaults)
     defaults = Dict(value(k) => value(v) for (k, v) in pairs(defaults))
-    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, defaults, nothing)
+    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, defaults, nothing, connection_type)
 end
 
 function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)
diff --git a/src/systems/pde/pdesystem.jl b/src/systems/pde/pdesystem.jl
@@ -52,8 +52,16 @@ struct PDESystem <: ModelingToolkit.AbstractSystem
     parameters are not supplied in `ODEProblem`.
     """
     defaults::Dict
-    @add_kwonly function PDESystem(eqs, bcs, domain, indvars, depvars, ps = SciMLBase.NullParameters(), defaults = Dict())
-        new(eqs, bcs, domain, indvars, depvars, ps, defaults)
+    """
+    type: type of the system
+    """
+    connection_type::Any
+    @add_kwonly function PDESystem(eqs, bcs, domain, indvars, depvars,
+                                   ps=SciMLBase.NullParameters();
+                                   defaults=Dict(),
+                                   connection_type=nothing,
+                                  )
+        new(eqs, bcs, domain, indvars, depvars, ps, defaults, connection_type)
     end
 end
 
diff --git a/src/systems/reaction/reactionsystem.jl b/src/systems/reaction/reactionsystem.jl
@@ -60,10 +60,16 @@ struct Reaction{S, T <: Number}
     `true` if `rate` represents the full reaction rate law.
     """
     only_use_rate::Bool
+    """
+    type: type of the system
+    """
+    connection_type::Any
 end
 
 function Reaction(rate, subs, prods, substoich, prodstoich;
-                  netstoich=nothing, only_use_rate=false, kwargs...)
+                  netstoich=nothing, only_use_rate=false,
+                  connection_type=nothing,
+                  kwargs...)
 
     (isnothing(prods)&&isnothing(subs)) && error("A reaction requires a non-nothing substrate or product vector.")
     (isnothing(prodstoich)&&isnothing(substoich)) && error("Both substrate and product stochiometry inputs cannot be nothing.")
@@ -80,7 +86,7 @@ function Reaction(rate, subs, prods, substoich, prodstoich;
     subs = value.(subs)
     prods = value.(prods)
     ns = isnothing(netstoich) ? get_netstoich(subs, prods, substoich, prodstoich) : netstoich
-    Reaction(value(rate), subs, prods, substoich, prodstoich, ns, only_use_rate)
+    Reaction(value(rate), subs, prods, substoich, prodstoich, ns, only_use_rate, connection_type)
 end
 
 
diff --git a/test/connectors.jl b/test/connectors.jl
@@ -0,0 +1,31 @@
+using Test, ModelingToolkit
+
+@parameters t
+
+@connector function Foo(;name)
+    @variables x(t)
+    ODESystem(Equation[], t, [x], [], defaults=Dict(x=>1.0))
+end
+
+@connector function Goo(;name)
+    @variables x(t)
+    @parameters p
+    ODESystem(Equation[], t, [x], [p], defaults=Dict(x=>1.0, p=>1.0))
+end
+
+ModelingToolkit.connect(::Type{<:Foo}, sys1, sys2) = [sys1.x ~ sys2.x]
+@named f1 = Foo()
+@named f2 = Foo()
+@named g = Goo()
+
+@test isequal(connect(f1, f2), [f1.x ~ f2.x])
+@test_throws ArgumentError connect(f1, g)
+
+# Note that since there're overloadings, these tests are not re-runable.
+ModelingToolkit.promote_connect_rule(::Type{<:Foo}, ::Type{<:Goo}) = Foo
+@test isequal(connect(f1, g), [f1.x ~ g.x])
+ModelingToolkit.promote_connect_rule(::Type{<:Goo}, ::Type{<:Foo}) = Foo
+@test isequal(connect(f1, g), [f1.x ~ g.x])
+# test conflict
+ModelingToolkit.promote_connect_rule(::Type{<:Goo}, ::Type{<:Foo}) = Goo
+@test_throws ArgumentError connect(f1, g)
diff --git a/test/runtests.jl b/test/runtests.jl
