Named state space sets and other MTK related fixes and improvements (#235)

* Informative error on nonexisting symbolic parameter

* mention init cond as parameters

* correct typo

* add a new function for named current parameters

* organize MTK tests better

* add reinit_dae = false

See SciML/ModelingToolkit.jl#3451

* remove `named_parametersd`

Needs MTK to be defined

* Named state space sets

* add tests for trajectory and named variables

* update tests to MTK v10

* some typos

* fix jacobian

Author: Datseris

CHANGELOG.md

@@ -1,3 +1,8 @@
+# v3.15
+
+- New function `named_variables(ds)` for getting the variable names of an MTK-generated dynamical system.
+- `trajectory` function now automatically makes a named state space set if the dynamical system is MTK-generated.
+
 # v3.14.0
 
 - `set_parameter!` and `current_parameter` will now throw an informative error message explicitly naming the parameter if the user tries to get/set a symbolic parameter that does not exist in the MTK-generated dynamical system.

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicalSystemsBase"
 uuid = "6e36e845-645a-534a-86f2-f5d4aa5a06b4"
 repo = "https://github.com/JuliaDynamics/DynamicalSystemsBase.jl.git"
-version = "3.14.0"
+version = "3.15.0"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -27,7 +27,7 @@ OrdinaryDiffEqTsit5 = "1.1"
 Reexport = "1"
 Roots = "1, 2"
 SciMLBase = "1.19.5, 2"
-StateSpaceSets = "2"
+StateSpaceSets = "2.5"
 Statistics = "1"
 StochasticDiffEq = "6.66.0"
 SymbolicIndexingInterface = "0.3.4"

docs/src/index.md

@@ -30,6 +30,7 @@ initial_time
 isinplace(::DynamicalSystem)
 successful_step
 referrenced_sciml_model
+named_variables
 ```
 
 ```@docs

src/core/dynamicalsystem_interface.jl

@@ -44,7 +44,7 @@ from an experimental measurement of a dynamical system with an unknown dynamic r
 
 See also the DynamicalSystems.jl tutorial online for examples making dynamical systems.
 
-## Integration with ModelingToolkit.jl
+## Integration with ModelingToolkit.jl (MTK)
 
 Dynamical systems that have been constructed from `DEProblem`s that themselves
 have been constructed from ModelingToolkit.jl keep a reference to the symbolic
@@ -93,6 +93,7 @@ unless when developing new algorithm implementations that use dynamical systems.
 - [`isinplace`](@ref)
 - [`successful_step`](@ref)
 - [`referrenced_sciml_model`](@ref)
+- [`named_variables`](@ref)
 
 ### API - alter status
 
@@ -126,7 +127,7 @@ errormsg(ds) = error("Not yet implemented for dynamical system of type $(nameof(
 
 export current_state, initial_state, current_parameters, current_parameter, initial_parameters, isinplace,
     current_time, initial_time, successful_step, isdeterministic, isdiscretetime, dynamic_rule,
-    reinit!, set_state!, set_parameter!, set_parameters!, step!, observe_state, referrenced_sciml_model
+    reinit!, set_state!, set_parameter!, set_parameters!, step!, observe_state, referrenced_sciml_model, named_variables
 
 ###########################################################################################
 # Symbolic support
@@ -148,10 +149,27 @@ referrenced_sciml_model(::Nothing) = nothing
 
 # return true if there is an actual referrenced system
 has_referrenced_model(prob::SciMLBase.DEProblem) = has_referrenced_model(referrenced_sciml_model(prob))
+has_referrenced_model(prob::DynamicalSystem) = has_referrenced_model(referrenced_sciml_model(prob))
 has_referrenced_model(::Nothing) = false
 has_referrenced_model(::SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}) = false
 has_referrenced_model(sys) = true
 
+"""
+    named_variables(ds::DynamicalSystem)
+
+If `ds` is constructed via MTK, return a vector of the variable names (as symbols).
+Otherwise return `nothing`. If `X` is a `StateSpaceSet`, you can always do
+```julia
+X = StateSpaceSet(X; names = named_variables(ds))
+```
+in downstream functions to name a set coming from `ds` (if possible).
+"""
+function named_variables(ds::DynamicalSystem)
+    mtk = referrenced_sciml_model(ds)
+    isnothing(mtk) && return nothing
+    return SymbolicIndexingInterface.getname.(SymbolicIndexingInterface.variable_symbols(mtk))
+end
+
 ###########################################################################################
 # API - obtaining information from the system
 ###########################################################################################

src/core/trajectory.jl

@@ -5,18 +5,15 @@ export trajectory
 
 Evolve `ds` for a total time of `T` and return its trajectory `X`, sampled at equal time
 intervals, and corresponding time vector. `X` is a `StateSpaceSet`.
-Optionally provide a starting state `u0` which is `current_state(ds)` by default.
-
 The returned time vector is `t = (t0+Ttr):Δt:(t0+Ttr+T)`.
 
-If time evolution diverged, or in general failed, before `T`,
+Optionally provide a starting state `u0` which is `current_state(ds)` by default.
+
+If time evolution diverged or in general failed before `T`,
 the remaining of the trajectory is set to the last valid point.
 
-`trajectory` is a very simple function provided for convenience.
-For continuous time systems, it doesn't play well with callbacks,
-use `DifferentialEquations.solve` if you want a trajectory/timeseries
-that works with callbacks, or in general you want more flexibility in the generated
-trajectory (but remember to convert the output of `solve` to a `StateSpaceSet`).
+The dimensions of `X` are automatically named if `ds` referrences an MTK model
+and if `save_idxs` remains at its default value.
 
 ## Keyword arguments
 
@@ -27,22 +24,34 @@ trajectory (but remember to convert the output of `solve` to a `StateSpaceSet`).
 * `t0 = initial_time(ds)`: Starting time.
 * `container = SVector`: Type of vector that will represent the state space points
   that will be included in the `StateSpaceSet` output. See `StateSpaceSet` for valid options.
-* `save_idxs::AbstractVector`: Which variables to output in `X`. It can be
-  any type of index that can be given to [`observe_state`](@ref).
-  Defaults to `1:dimension(ds)` (all dynamic variables).
+* `save_idxs`: Which variables to output in `X`. By default it is `nothing` (all variables).
+  It can be a vector of any type of index that can be given to [`observe_state`](@ref).
   Note: if you mix integer and symbolic indexing be sure to initialize the array
   as `Any` so that integers `1, 2, ...` are not converted to symbolic expressions.
+
+## Description
+
+`trajectory` is a very simple function provided for convenience.
+For continuous time systems, it doesn't play well with callbacks,
+use `DifferentialEquations.solve` if you want a trajectory
+that works with callbacks, or in general you want more flexibility in the generated
+trajectory (but remember to convert the output of `solve` to a `StateSpaceSet`).
 """
 function trajectory(ds::DynamicalSystem, T, u0 = current_state(ds);
         save_idxs = nothing, t0 = initial_time(ds), kwargs...
     )
     accessor = svector_access(save_idxs)
     reinit!(ds, u0; t0)
     if isdiscretetime(ds)
-        trajectory_discrete(ds, T; accessor, kwargs...)
+        X, tvec = trajectory_discrete(ds, T; accessor, kwargs...)
     else
-        trajectory_continuous(ds, T; accessor, kwargs...)
+        X, tvec = trajectory_continuous(ds, T; accessor, kwargs...)
+    end
+    # name automatically if possible
+    if isnothing(save_idxs)
+        X = StateSpaceSet(X; names = named_variables(ds))
     end
+    return X, tvec
 end
 
 function trajectory_discrete(ds, T;

test/Project.toml

@@ -10,4 +10,4 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ModelingToolkit = "9"
+ModelingToolkit = "10"

test/jacobian.jl

@@ -30,23 +30,23 @@ end
 
 @testset "MTK Jacobian" begin
     using ModelingToolkit
-    using ModelingToolkit: Num, GeneratedFunctionWrapper
+    using ModelingToolkit: GeneratedFunctionWrapper
     using DynamicalSystemsBase: SciMLBase
     @independent_variables t
     @variables u(t)[1:2]
     D = Differential(t)
 
     eqs = [D(u[1]) ~ 3.0 * u[1],
            D(u[2]) ~ -3.0 * u[2]]
-    @named sys = ODESystem(eqs, t)
-    sys = structural_simplify(sys)
+    @named sys = System(eqs, t)
+    sys = mtkcompile(sys)
 
     prob = ODEProblem(sys, [1.0, 1.0], (0.0, 1.0); jac=true)
     ds = CoupledODEs(prob)
 
     jac = jacobian(ds)
     @test jac isa GeneratedFunctionWrapper
-    @test jac([1.0, 1.0], [], 0.0) == [3 0;0 -3]
+    @test jac([1.0, 1.0], [], 0.0) == [-3 0;0 3]
 
     @testset "CoupledSDEs" begin
         # just to check if MTK @brownian does not give any problems
@@ -61,6 +61,6 @@ end
 
         jac = jacobian(sde)
         @test jac isa GeneratedFunctionWrapper
-        @test jac([1.0, 1.0], [], 0.0) == [3 0; 0 -3]
+        @test jac([1.0, 1.0], [], 0.0) == [-3 0; 0 3]
     end
 end

test/mtk_integration.jl

@@ -13,7 +13,7 @@ function fol_factory(separate = false; name)
         D(x) ~ RHS] :
           D(x) ~ (f - x) / τ
 
-    ODESystem(eqs, t; name)
+    System(eqs, t; name)
 end
 
 @named fol_1 = fol_factory()
@@ -22,97 +22,86 @@ end
 connections = [fol_1.f ~ 1.5,
     fol_2.f ~ fol_1.x]
 
-connected = compose(ODESystem(connections, t; name = :connected), fol_1, fol_2)
+connected = compose(System(connections, t; name = :connected), fol_1, fol_2)
 
-sys = structural_simplify(connected; split = false)
+sys = mtkcompile(connected; split = false)
 
 u0 = [fol_1.x => -0.5,
     fol_2.x => 1.0]
 
 p = [fol_1.τ => 2.0,
     fol_2.τ => 4.0]
 
-prob = ODEProblem(sys, u0, (0.0, 10.0), p)
+initials = vcat(u0, p)
+
+prob = ODEProblem(sys, initials, (0.0, 10.0))
 ds = CoupledODEs(prob)
 
 @testset "parameter get/set" begin
-    @test current_parameter(ds, 1) == 2.0
     @test current_parameter(ds, fol_1.τ) == 2.0
-    @test current_parameter(ds, 2) == 4.0
     @test current_parameter(ds, fol_2.τ) == 4.0
-
-    set_parameter!(ds, 1, 3.0)
-    @test current_parameter(ds, 1) == 3.0
-    @test current_parameter(ds, fol_1.τ) == 3.0
-
-    set_parameter!(ds, fol_1.τ, 2.0)
-    @test current_parameter(ds, 1) == 2.0
-    @test current_parameter(ds, fol_1.τ) == 2.0
+    set_parameter!(ds, fol_1.τ, 4.0)
+    @test current_parameter(ds, fol_1.τ) == 4.0
 end
 
 # pure parameter container
 @testset "pure parameter container" begin
     pp = deepcopy(current_parameters(ds))
-    set_parameter!(ds, fol_1.τ, 4.0, pp)
-    @test current_parameter(ds, fol_1.τ, pp) == 4.0
+    set_parameter!(ds, fol_1.τ, 2.0, pp)
+    @test current_parameter(ds, fol_1.τ, pp) == 2.0
 end
 
 @testset "states and observed variables" begin
-    @test observe_state(ds, 1) == -0.5
     @test observe_state(ds, fol_1.x) == -0.5
     @test observe_state(ds, fol_2.RHS) == -0.375
 
-    set_state!(ds, 1.5, 1)
-    @test observe_state(ds, 1) == 1.5
+    x0 = observe_state(ds, 1)
+    set_state!(ds, x0, 1)
+    @test observe_state(ds, 1) == x0
     set_state!(ds, -0.5, fol_1.x)
-    @test observe_state(ds, 1) == -0.5
+    @test observe_state(ds, fol_1.x) == -0.5
 end
 
 @testset "derivative dynamical systems" begin
     u1 = current_state(ds)
     pds = ParallelDynamicalSystem(ds, [u1, copy(u1)])
 
     set_parameter!(pds, fol_1.τ, 4.0)
-    @test current_parameter(pds, 1) == 4.0
     @test current_parameter(pds, fol_1.τ) == 4.0
     @test observe_state(pds, fol_1.x) == -0.5
     @test observe_state(pds, fol_2.RHS) == -0.375
 
     sds = StroboscopicMap(ds, 1.0)
     set_parameter!(sds, fol_1.τ, 2.0)
-    @test current_parameter(sds, 1) == 2.0
     @test current_parameter(sds, fol_1.τ) == 2.0
     @test observe_state(sds, fol_1.x) == -0.5
     @test observe_state(sds, fol_2.RHS) == -0.375
 
     prods = ProjectedDynamicalSystem(ds, [1], [0.0])
     set_parameter!(prods, fol_1.τ, 3.0)
-    @test current_parameter(prods, 1) == 3.0
     @test current_parameter(prods, fol_1.τ) == 3.0
     @test observe_state(prods, fol_1.x) == -0.5
     @test observe_state(prods, fol_2.RHS) == -0.375
 
     # notice this evolves the dynamical system!!!
-    pmap = PoincareMap(ds, (1, 0.0))
+    pmap = PoincareMap(ds, (2, 0.0))
     set_parameter!(pmap, fol_1.τ, 4.0)
-    @test current_parameter(pmap, 1) == 4.0
     @test current_parameter(pmap, fol_1.τ) == 4.0
     @test observe_state(pmap, fol_1.x) ≈ 0 atol = 1e-3 rtol = 0
 end
 
+@testset "trajectory naming" begin
+    vars = named_variables(ds)
+    @test length(vars) == 2
+    @test sort!(string.(vars)) == ["fol_1₊x", "fol_2₊x"]
+    X, tvec = trajectory(ds, 10)
+    @test X.names == vars
+end
 
 @testset "split = true" begin
-sys = structural_simplify(connected; split = true)
-
-u0 = [fol_1.x => -0.5,
-    fol_2.x => 1.0]
-
-p = [fol_1.τ => 2.0,
-    fol_2.τ => 4.0]
-
-    prob = ODEProblem(sys, u0, (0.0, 10.0), p)
+    sys = mtkcompile(connected; split = true)
+    prob = ODEProblem(sys, initials, (0.0, 10.0))
     ds = CoupledODEs(prob)
-
     @test current_parameter(ds, fol_1.τ) == 2.0
     set_parameter!(ds, fol_1.τ, 3.0)
     @test current_parameter(ds, fol_1.τ) == 3.0
@@ -123,6 +112,7 @@ end
         du[1] = p[1] * (u[2] - u[1])
         du[2] = u[1] * (28.0 - u[3]) - u[2]
         du[3] = u[1] * u[2] - (8 / 3) * u[3]
+        return nothing
     end
     u0 = [1.0; 0.0; 0.0]
     tspan = (0.0, 100.0)
@@ -136,7 +126,7 @@ end
 
     @test observe_state(ds, 1) == 1.0
 
-    @test_throws ErrorException observe_state(ds, fol_1.f)
+    @test_throws ArgumentError observe_state(ds, fol_1.f)
 end
 
 # Test that remake works also without anything initial
@@ -163,11 +153,11 @@ D = Differential(t)
     end
 end
 
-@mtkbuild roessler_model = Roessler()
+@mtkcompile roessler_model = Roessler()
 
 @testset "type: $(iip)" for iip in (true, false)
     if iip
-        prob = ODEProblem(roessler_model)
+        prob = ODEProblem(roessler_model, nothing, (0.0, Inf))
     else
         prob = ODEProblem{false}(roessler_model, nothing, (0.0, Inf); u0_constructor = x->SVector(x...))
     end
@@ -209,7 +199,7 @@ end
 end
 
 @testset "informative errors" begin
-    prob = ODEProblem(roessler_model)
+    prob = prob = ODEProblem(roessler_model, nothing, (0.0, Inf))
     ds = CoupledODEs(prob)
     @parameters XOXO = 0.5
     @test_throws "XOXO" current_parameter(ds, :XOXO)
@@ -231,10 +221,10 @@ Differential(t)(DS) ~ η2 - η3*DS - abs(DT - DS)*DS,
 η1 ~ η1_0 + r_η*t, # this symbolic variable has its own equation!
 ]
 
-sys = ODESystem(eqs, t; name = :stommel)
-sys = structural_simplify(sys; split = false)
+sys = System(eqs, t; name = :stommel)
+sys = mtkcompile(sys; split = false)
 
-prob = ODEProblem(sys)
+prob = ODEProblem(sys, nothing, (0.0, Inf))
 ds = CoupledODEs(prob)
 
 X, tvec = trajectory(ds, 10.0; Δt = 0.1, save_idxs = Any[1, 2, η1])