Connecting components by adding term to DE in ModelingToolkit

I am trying to connect components by adding function of variables from each component to one of the DEs in each component. In other words, for components 1 and 2 (and their respective variables z_1 and z_2)

\frac{d z_1}{dt} = g(x_1,y_1,z_1, \beta) +f (z_1, z_2, p)

\frac{d z_2}{dt} = g(x_2,y_2, z_2, \beta) + f(z_2, z_1, p)

where \frac{d z_1}{dt} = g(x_1,y_1,z_1, \beta) is the original DE in each component and f(z_1, z_2, p) is the connection with the other component that I would like to add.

I could presumably do this using a continuous callback in the ODEProblem, but was wondering if there is a more elegant way to do this with ModelingToolkit.

I have attempted to do this using compose, but am having issues getting the equations to combine.

using ModelingToolkit, OrdinaryDiffEq

function Lorenz(; name)
    @parameters σ ρ β
    @variables t x(t) y(t) z(t)
    D = Differential(t)

    eqs = [D(D(x)) ~ σ * (y - x),
        D(y) ~ x * (ρ - z) - y,
        D(z) ~ x * y - β * z]

    @named sys = ODESystem(eqs; name=name)
end

@named Lorenz1 = Lorenz()
@named Lorenz2 = Lorenz()

f(b1, b2, a) = a * (b1- b2)

@parameters c
connections = [D(Lorenz1.z) ~ + f(Lorenz1.z , Lorenz2.z, c),
                D(Lorenz2.z) ~ + f(Lorenz2.z , Lorenz1.z, c)]
connected = compose(ODESystem(connections, t,  name = :connected), Lorenz1, Lorenz2)
connected_simp = structural_simplify(connected)

u0 = [D(Lorenz1.x) => 2.0,
    Lorenz1.x => 1.0,
    Lorenz1.y => 0.0,
    Lorenz1.z => 0.0,
    D(Lorenz2.x) => 2.0,
    Lorenz2.x => 1.0,
    Lorenz2.y => 0.0,
    Lorenz2.z => 0.0]

p = [Lorenz1.σ => 28.0,
    Lorenz1.ρ => 10.0,
    Lorenz1.β => 8 / 3,
    Lorenz2.σ => 28.0,
    Lorenz2.ρ => 10.0,
    Lorenz2.β => 8 / 3,
    a => 0.005]


prob = ODEProblem(connected, u0, tspan, p)
sol = solve(prob, Tsit5())
using Plots;
plot(sol, idxs = (Lorenz1.x, Lorenz1.y))
julia> connected_simp = structural_simplify(connected)
ERROR: ExtraEquationsSystemException: The system is unbalanced. There are 6 highest order derivative variables and 8 equations.
More equations than variables, here are the potential extra equation(s):
 Differential(t)(Lorenz1₊z(t)) ~ Lorenz1₊x(t)*Lorenz1₊y(t) - Lorenz1₊β*Lorenz1₊z(t)
 Differential(t)(Lorenz2₊z(t)) ~ Lorenz2₊x(t)*Lorenz2₊y(t) - Lorenz2₊β*Lorenz2₊z(t)
Stacktrace:
 [1] error_reporting(state::TearingState{ODESystem}, bad_idxs::Vector{Int64}, n_highest_vars::Int64, iseqs::Bool, orig_inputs::Set{Any})
   @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/ZmLra/src/structural_transformation/utils.jl:44
 [2] check_consistency(state::TearingState{ODESystem}, orig_inputs::Set{Any})
   @ ModelingToolkit.StructuralTransformations ~/.julia/packages/ModelingToolkit/ZmLra/src/structural_transformation/utils.jl:85
 [3] _structural_simplify!(state::TearingState{ODESystem}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ ModelingToolkit.SystemStructures ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systemstructure.jl:612
 [4] _structural_simplify!
   @ ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systemstructure.jl:600 [inlined]
 [5] structural_simplify!(state::TearingState{ODESystem}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ ModelingToolkit.SystemStructures ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systemstructure.jl:566
 [6] structural_simplify(sys::ODESystem, io::Nothing; simplify::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
   @ ModelingToolkit ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systems.jl:39
 [7] structural_simplify
   @ ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systems.jl:19 [inlined]
 [8] structural_simplify(sys::ODESystem)
   @ ModelingToolkit ~/.julia/packages/ModelingToolkit/ZmLra/src/systems/systems.jl:19
 [9] top-level scope
   @ ~/connection_MWE:50

The short answer is you can’t modify equations by composing or extending, only add new equations. There used to be a nice explanation in another almost identical question, but by search skills are failing to find it.

The longer answer is that you have to prepare a place to connect systems by adding a variable whose value can be provided externally. The @connector macro was created to help with this and there is a nice tutorial here.

A simple way to do it in your case with minimal changes would be something like (untested).

function Lorenz(; name)
    ps = @parameters σ ρ β
    states = @variables t x(t) y(t) z(t) c(t)
    D = Differential(t)

    eqs = [D(D(x)) ~ σ * (y - x),
        D(y) ~ x * (ρ - z) - y,
        D(z) ~ x * y - β * z + c]

    ODESystem(eqs, t, states, ps; name)
end

and then

connections = [Lorenz1.c ~ f(Lorenz1.z , Lorenz2.z, c),
               Lorenz2.c ~ f(Lorenz2.z , Lorenz1.z, c)]

Thank you for your response! The simple way you suggest works:

using ModelingToolkit, OrdinaryDiffEq

function Lorenz(; name)
    @parameters σ ρ β
    @variables t x(t) y(t) z(t) c(t)
    D = Differential(t)

    eqs = [D(D(x)) ~ σ * (y - x),
        D(y) ~ x * (ρ - z) - y,
        D(z) ~ x * y - β * z + c]

    @named sys = ODESystem(eqs; name=name)
end

@named Lorenz1 = Lorenz() 
@named Lorenz2 = Lorenz()

f(b1, b2, a) = a * (b1- b2)
@variables t
@parameters C
D = Differential(t)
connections = [Lorenz1.c ~ f(Lorenz1.z , Lorenz2.z, C),
               Lorenz2.c ~ f(Lorenz2.z , Lorenz1.z, C)]
connected = compose(ODESystem(connections, t, name = :connected), Lorenz1, Lorenz2)
connected = structural_simplify(connected)

u0 = [D(Lorenz1.x) => 2.0,
    Lorenz1.x => 1.0,
    Lorenz1.y => 0.0,
    Lorenz1.z => 0.0,
    Lorenz1.c => 0.0, 
    D(Lorenz2.x) => 5.0,
    Lorenz2.x => 0.0,
    Lorenz2.y => 1.0,
    Lorenz2.z => 0.5,
    Lorenz2.c  => 0.0,]

p = [Lorenz1.σ => 28.0,
    Lorenz1.ρ => 10.0,
    Lorenz1.β => 8 / 3,
    Lorenz2.σ => 28.0,
    Lorenz2.ρ => 10.0,
    Lorenz2.β => 8 / 3,
    C => 0.005]

prob = ODEProblem(connected, u0, [0., 1000.], p)
sol = solve(prob, Tsit5())
using Plots;
plot(sol, idxs = (Lorenz1.x, Lorenz1.y))

The tutorial you linked is very helpful. Do I understand it correctly that using the @connector macro would do something similar to this?