Causal modeling in MTK?

After seeing @ChrisRackauckas slides from JuliaCon 2024, I found out that (it seems) ModelingToolkit supports causal models, through passing input and outputs around. Is this really possible? I ask because Causal.jl struggles with a model I’m trying to implement - see this issue. Any tips on porting this Causal.jl model to MTK would be appreciated! I’m not sure how to define the component models in MTK.

I whipped this together really quick and haven’t run it yet, but yeah you can define two systems, one with an input and and one with an output, and connect them. It would look something like this:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D

@mtkmodel InfectionSystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        S
        u [output = true]
    end
    @equations begin
        D(S) ~ rate_to_proportion(β*I,δt)*S
        u ~ rate_to_proportion(β*I,δt)*S
    end
end

@mtkmodel RecoverySystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        R
        u [input = true]
    end
    @equations begin
        D(R) ~ u - rate_to_proportion(γ,δt)*I
    end
end

@mtkmodel SIR begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @components begin
        infection = InfectionSystem()
        recovery = RecoverySystem()
    end
    @equations begin
        connect(infection.u, recovery.u)
    end
end

Thanks @ChrisRackauckas! The line u [output=true] throws a parse error (as it expects no gap between u and [.

What’s going on under the hood? Is it a true causal system where InfectionSystem would be run, then RecoverySystem?

    @variables begin
        R
        u, [input = true]
    end

the comma was missing.

That’s orthogonal. That’s separate compilation.

The following code throws an error when I try to instantiate the model that I can’t easily debug:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D

@inline function rate_to_proportion(r,t)
    1-exp(-r*t)
end

@mtkmodel InfectionSystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        S(t)
        I(t)
        u(t),[output=true]
    end
    @equations begin
        D(S) ~ rate_to_proportion(β*I,δt)*S
        u ~ rate_to_proportion(β*I,δt)*S
    end
end

@mtkmodel RecoverySystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        I(t)
        R(t)
        u(t),[input=true]
    end
    @equations begin
        D(R) ~ u - rate_to_proportion(γ,δt)*I
    end
end

@mtkmodel SIR begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @components begin
        infection = InfectionSystem()
        recovery = RecoverySystem()
    end
    @equations begin
        connect(infection.u, recovery.u)
    end
end

@mtkbuild sir=SIR()

Output:

ERROR: None Sym BasicSymbolic doesn't have a name
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] nameof(s::SymbolicUtils.BasicSymbolic{Real})
    @ SymbolicUtils ~/.julia/packages/SymbolicUtils/qyMYa/src/types.jl:249
  [3] nameof(n::Num)
    @ Symbolics ~/.julia/packages/Symbolics/Eas9m/src/num.jl:88
  [4] unique(f::typeof(nameof), C::Tuple{Num, Num}; seen::Nothing)
    @ Base ./set.jl:261
  [5] unique
    @ ./set.jl:247 [inlined]
  [6] connect(::Num, ::Num)
    @ Symbolics ~/.julia/packages/Symbolics/Eas9m/src/equations.jl:18
  [7] __SIR__(; name::Symbol, β::Nothing, γ::Nothing, δt::Nothing)
    @ Main ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/model_parsing.jl:620
  [8] __SIR__
    @ ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/model_parsing.jl:132 [inlined]
  [9] #_#321
    @ ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/model_parsing.jl:25 [inlined]
 [10] top-level scope
    @ ~/.julia/packages/ModelingToolkit/cAhZr/src/systems/abstractsystem.jl:1445

Using connectors is probably nicer. Here’s a version that builds:

using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D
using ModelingToolkitStandardLibrary.Blocks: RealInput, RealOutput

@inline function rate_to_proportion(r,t)
    1-exp(-r*t)
end

@mtkmodel InfectionSystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        S(t)
    end
    @components begin
        output = RealOutput()
    end
    @equations begin
        D(S) ~ rate_to_proportion(β*I,δt)*S
        output.u ~ rate_to_proportion(β*I,δt)*S
    end
end

@mtkmodel RecoverySystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        I(t)
        R(t)
    end
    @components begin
        input = RealInput()
    end
    @equations begin
        D(I) ~ rate_to_proportion(γ,δt)*I
        D(R) ~ input.u - rate_to_proportion(γ,δt)*I
    end
end

@mtkmodel SIR begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @components begin
        infection = InfectionSystem()
        recovery = RecoverySystem()
    end
    @equations begin
        connect(infection.output, recovery.input)
    end
end

@mtkbuild sir=SIR()

That doesn’t build for me; I had to add in I as a variable in InfectionSystem, but the following gives me an error.

ERROR: ExtraVariablesSystemException: The system is unbalanced. There are 6 highest order derivative variables and 5 equations.
More variables than equations, here are the potential extra variable(s):
 infection₊S(t)
 recovery₊I(t)
 recovery₊R(t)
 recovery₊input₊u(t)
Stacktrace:
  [1] error_reporting(state::TearingState{…}, bad_idxs::Vector{…}, n_highest_vars::Int64, iseqs::Bool, orig_inputs::Set{…})
    @ ModelingToolkit.StructuralTransformations C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\structural_transformation\utils.jl:50
  [2] check_consistency(state::TearingState{ODESystem}, orig_inputs::Set{Any})
    @ ModelingToolkit.StructuralTransformations C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\structural_transformation\utils.jl:85
  [3] _structural_simplify!(state::TearingState{…}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, warn_initialize_determined::Bool, dummy_derivative::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systemstructure.jl:691
  [4] structural_simplify!(state::TearingState{…}, io::Nothing; simplify::Bool, check_consistency::Bool, fully_determined::Bool, warn_initialize_determined::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systemstructure.jl:634
  [5] __structural_simplify(sys::ODESystem, io::Nothing; simplify::Bool, kwargs::@Kwargs{…})
    @ ModelingToolkit C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systems.jl:85
  [6] __structural_simplify
    @ C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systems.jl:66 [inlined]
  [7] structural_simplify(sys::ODESystem, io::Nothing; simplify::Bool, split::Bool, allow_symbolic::Bool, allow_parameter::Bool, conservative::Bool, fully_determined::Bool, kwargs::@Kwargs{})
    @ ModelingToolkit C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systems.jl:24
  [8] structural_simplify
    @ C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systems.jl:20 [inlined]
  [9] structural_simplify(sys::ODESystem)
    @ ModelingToolkit C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\systems.jl:20
 [10] top-level scope
    @ C:\Users\sdwfr\.julia\packages\ModelingToolkit\3AmNp\src\systems\abstractsystem.jl:1881
Some type information was truncated. Use `show(err)` to see complete types.
using ModelingToolkit, OrdinaryDiffEq
using ModelingToolkit: t_nounits as t, D_nounits as D
using ModelingToolkitStandardLibrary.Blocks: RealInput, RealOutput

@inline function rate_to_proportion(r,t)
    1-exp(-r*t)
end

@mtkmodel InfectionSystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        S(t)
        I(t)
    end
    @components begin
        output = RealOutput()
    end
    @equations begin
        D(S) ~ rate_to_proportion(β*I,δt)*S
        output.u ~ rate_to_proportion(β*I,δt)*S
    end
end

@mtkmodel RecoverySystem begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @variables begin
        I(t)
        R(t)
    end
    @components begin
        input = RealInput()
    end
    @equations begin
        D(I) ~ rate_to_proportion(γ,δt)*I
        D(R) ~ input.u - rate_to_proportion(γ,δt)*I
    end
end

@mtkmodel SIR begin
    @parameters begin
        β = 0.5/1000
        γ = 0.25
        δt = 0.1
    end
    @components begin
        infection = InfectionSystem()
        recovery = RecoverySystem()
    end
    @equations begin
        connect(infection.output, recovery.input)
    end
end

@mtkbuild sir=SIR()

You cannot have I in both or you would be missing a connection. This error is saying to use what I shared

…but then the InfectionSystem doesn’t compile (as it doesn’t know about I)

The models are incomplete without connections (which is known because of the inputs/outputs), but the connected model is complete and it compiles fine (the code I shared)

Sorry for not being clear - the code you shared throws this error when building the final system:

ERROR: UndefVarError: `I` not defined
Stacktrace:
 [1] __InfectionSystem__(; name::Symbol, β::Nothing, γ::Nothing, δt::Nothing, S::Nothing)

My version of MTK:

(models) pkg> status ModelingToolkit
  [961ee093] ModelingToolkit v9.26.0