Merge pull request #1104 from isaacsas/add_compse_dispatch

Add compose dispatch and constent handling of scoping to MTK

Author: isaacsas

src/Catalyst.jl

@@ -26,7 +26,7 @@ import Symbolics: BasicSymbolic
 using Symbolics: iscall, sorted_arguments
 using ModelingToolkit: Symbolic, value, get_unknowns, get_ps, get_iv, get_systems,
                        get_eqs, get_defaults, toparam, get_var_to_name, get_observed,
-                       getvar
+                       getvar, has_iv
 
 import ModelingToolkit: get_variables, namespace_expr, namespace_equation, get_variables!,
                         modified_unknowns!, validate, namespace_variables,

src/reaction.jl

@@ -351,20 +351,16 @@ MT.is_alg_equation(rx::Reaction) = false
 # MTK functions for extracting variables within equation type object
 MT.eqtype_supports_collect_vars(rx::Reaction) = true
 function MT.collect_vars!(unknowns, parameters, rx::Reaction, iv; depth = 0,
-        op = MT.Differential)
+        op = MT.Differential)    
+
     MT.collect_vars!(unknowns, parameters, rx.rate, iv; depth, op)
-    for sub in rx.substrates
-        MT.collect_vars!(unknowns, parameters, sub, iv; depth, op)
-    end
-    for prod in rx.products
-        MT.collect_vars!(unknowns, parameters, prod, iv; depth, op)
-    end
-    for substoich in rx.substoich
-        MT.collect_vars!(unknowns, parameters, substoich, iv; depth, op)
-    end
-    for prodstoich in rx.prodstoich
-        MT.collect_vars!(unknowns, parameters, prodstoich, iv; depth, op)
+
+    for items in (rx.substrates, rx.products, rx.substoich, rx.prodstoich)
+        for item in items
+            MT.collect_vars!(unknowns, parameters, item, iv; depth, op)
+        end
     end
+
     if hasnoisescaling(rx)
         ns = getnoisescaling(rx)
         MT.collect_vars!(unknowns, parameters, ns, iv; depth, op)

src/reactionsystem.jl

@@ -512,24 +512,8 @@ function make_ReactionSystem_internal(rxs_and_eqs::Vector, iv, us_in, ps_in;
     eqs = Equation[eq for eq in rxs_and_eqs if eq isa Equation]
 
     # Loops through all reactions, adding encountered quantities to the unknown and parameter vectors.
-    # Starts by looping through substrates + products only (so these are added to the vector first).
-    # Next, the other components of reactions (e.g. rates and stoichiometries) are added.
     for rx in rxs
-        for reactants in (rx.substrates, rx.products), spec in reactants
-            MT.isparameter(spec) ? push!(ps, spec) : push!(us, spec)
-        end
-    end
-    for rx in rxs
-        # Adds all quantities encountered in the reaction's rate.
-        findvars!(ps, us, rx.rate, ivs, vars)
-
-        # Extracts all quantities encountered within stoichiometries.
-        for stoichiometry in (rx.substoich, rx.prodstoich), sym in stoichiometry
-            (sym isa Symbolic) && findvars!(ps, us, sym, ivs, vars)
-        end
-
-        # Extract all quantities encountered in relevant `Reaction` metadata.
-        hasnoisescaling(rx) && findvars!(ps, us, getnoisescaling(rx), ivs, vars)
+        MT.collect_vars!(us, ps, rx, iv)
     end
 
     # Extracts any species, variables, and parameters that occur in (non-reaction) equations.
@@ -543,6 +527,11 @@ function make_ReactionSystem_internal(rxs_and_eqs::Vector, iv, us_in, ps_in;
         fulleqs = rxs
     end
 
+    # get variables in subsystems with scope at this level
+    for ssys in get(kwargs, :systems, [])
+        MT.collect_scoped_vars!(us, ps, ssys, iv)
+    end
+
     # Loops through all events, adding encountered quantities to the unknown and parameter vectors.
     find_event_vars!(ps, us, continuous_events, ivs, vars)
     find_event_vars!(ps, us, discrete_events, ivs, vars)
@@ -1397,6 +1386,42 @@ function MT.flatten(rs::ReactionSystem; name = nameof(rs))
         discrete_events = MT.discrete_events(rs))
 end
 
+"""
+    ModelingToolkit.compose(sys::ReactionSystem, systems::AbstractArray; name = nameof(sys))
+
+Compose the indicated [`ReactionSystem`](@ref) with one or more `AbstractSystem`s.
+
+Notes:
+- The `AbstractSystem` being added in must be an `ODESystem`, `NonlinearSystem`,
+  or `ReactionSystem` currently.
+- Returns a new `ReactionSystem` and does not modify `rs`.
+- By default, the new `ReactionSystem` will have the same name as `sys`.
+"""
+function ModelingToolkit.compose(sys::ReactionSystem, systems::AbstractArray; name = nameof(sys))
+    nsys = length(systems)
+    nsys == 0 && return sys
+    @set! sys.name = name
+    @set! sys.systems = [get_systems(sys); systems]
+    newunknowns = OrderedSet{BasicSymbolic{Real}}()
+    newparams = OrderedSet()
+    iv = has_iv(sys) ? get_iv(sys) : nothing
+    for ssys in systems
+        MT.collect_scoped_vars!(newunknowns, newparams, ssys, iv)
+    end
+
+    if !isempty(newunknowns) 
+        @set! sys.unknowns = union(get_unknowns(sys), newunknowns)
+        sort!(get_unknowns(sys), by = !isspecies)
+        @set! sys.species = filter(isspecies, get_unknowns(sys))
+    end
+
+    if !isempty(newparams)
+        @set! sys.ps = union(get_ps(sys), newparams)
+    end
+
+    return sys
+end
+
 """
     ModelingToolkit.extend(sys::AbstractSystem, rs::ReactionSystem; name::Symbol=nameof(sys))
 

test/compositional_modelling/component_based_model_creation.jl

@@ -501,3 +501,46 @@ let
     defs[rn2.X] == 30.0
     defs[rn2.Z] == 40.0
 end
+
+# test scoping in compose
+# code adapted from ModelingToolkit.jl tests
+let
+    t = default_t()
+    D = default_time_deriv()
+    @species x1(t) x2(t) 
+    @variables x3(t) x4(t) x5(t)
+    x2 = ParentScope(x2)
+    x3 = ParentScope(ParentScope(x3))
+    x4 = DelayParentScope(x4, 2)
+    x5 = GlobalScope(x5)
+    @parameters p1 p2 p3 p4 p5
+    p2 = ParentScope(p2)
+    p3 = ParentScope(ParentScope(p3))
+    p4 = DelayParentScope(p4, 2)
+    p5 = GlobalScope(p5)
+    rxs = [Reaction(p1, nothing, [x1]), Reaction(p2, [x2], nothing), 
+           D(x3) ~ p3, D(x4) ~ p4, D(x5) ~ p5]
+    @named sys1 = ReactionSystem(rxs, t)
+    @test isequal(x1, only(unknowns(sys1)))
+    @test isequal(x1, only(species(sys1)))
+    @test isequal(p1, only(parameters(sys1)))
+    @named sys2 = ReactionSystem([], t; systems = [sys1])
+    @test length(unknowns(sys2)) == 2
+    @test any(isequal(x2), unknowns(sys2))
+    @test any(isequal(x2), species(sys2))
+    @test length(parameters(sys2)) == 2
+    @test any(isequal(p2), parameters(sys2))
+    @named sys3 = ReactionSystem(Equation[], t)
+    sys3 = sys3 ∘ sys2
+    @test length(unknowns(sys3)) == 4
+    @test any(isequal(x3), unknowns(sys3))
+    @test any(isequal(x4), unknowns(sys3))
+    @test length(species(sys3)) == 2
+    @test length(parameters(sys3)) == 4
+    @test any(isequal(p3), parameters(sys3))
+    @test any(isequal(p4), parameters(sys3))
+    sys4 = complete(sys3)
+    @test length(unknowns(sys3)) == 4
+    @test length(parameters(sys4)) == 5
+    @test any(isequal(p5), parameters(sys4))
+end