Symboli indices in trajectory (#200)

* allow arbitrary observe indices in continuous time trajectory

* allow using correct time in observe state

* allow using any observe index in trajectory

* bump version

Author: Datseris

CHANGELOG.md

@@ -1,3 +1,9 @@
+# v3.7.0
+
+Symbolic indices used in `observe_state` can now also be used in `trajectory`.
+
+It is possible to specify time to observe a state at now.
+
 # v3.6.0
 
 New convenience functions `state_name, parameter_name` for obtaining a `name::String`
@@ -15,7 +21,7 @@ The referenced MTK model corresponding to the dynamical system can be obtained w
 
 See also the online overarching tutorial for an example.
 
-At the moment this is only possible for `CoupledODEs` and its derivative systems, as these are the only systems that can be made by a `DEProblem`, however it should be easy to allow e.g., `DeterministicIterated` to be created by a `DiscreteProblem`.
+At the moment this is only possible for `CoupledODEs` and its derivative systems, as these are the only systems that can be made by a `DEProblem`, however it should be easy to allow e.g., `DeterministicIteratedMap` to be created by a `DiscreteProblem`.
 
 Also:
 

Project.toml

@@ -1,7 +1,7 @@
 name = "DynamicalSystemsBase"
 uuid = "6e36e845-645a-534a-86f2-f5d4aa5a06b4"
 repo = "https://github.com/JuliaDynamics/DynamicalSystemsBase.jl.git"
-version = "3.6.0"
+version = "3.7.0"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

src/core/dynamicalsystem_interface.jl

@@ -172,35 +172,36 @@ See also [`initial_state`](@ref), [`observe_state`](@ref).
 current_state(ds::DynamicalSystem) = ds.u
 
 """
-    observe_state(ds::DynamicalSystem, i [,u = current_state(ds)]) → x::Real
+    observe_state(ds::DynamicalSystem, i, u = current_state(ds)) → x::Real
 
 Return the state `u` of `ds` _observed_ at "index" `i`. Possibilities are:
 
 - `i::Int` returns the `i`-th dynamic variable.
-- `i::Function` returns `f(current_state(ds))`, which is asserted to be a real number.
+- `i::Function` returns `f(current_state(ds))`.
 - `i::SymbolLike` returns the value of the corresponding symbolic variable.
    This is valid only for dynamical systems referrencing a ModelingToolkit.jl model
    which also has `i` as one of its listed variables (either uknowns or observed).
    Here `i` can be anything can be anything
    that could index the solution object `sol = ModelingToolkit.solve(...)`,
    such as a `Num` or `Symbol` instance with the name of the symbolic variable.
+   In this case, a last fourth optional positional argument `t` defaults to
+   `current_time(ds)` and is the time to observe the state at.
 
 For [`ProjectedDynamicalSystem`](@ref), this function assumes that the
 state of the system is the full state space state, not the projected one
 (this makes the most sense for allowing MTK-based indexing).
 
 Use [`state_name`](@ref) for an accompanying name.
 """
-function observe_state(ds::DynamicalSystem, index, u::AbstractArray = current_state(ds))
+function observe_state(ds::DynamicalSystem, index, u::AbstractArray = current_state(ds), t = current_time(ds))
     if index isa Function
         return index(u)
-    elseif index isa Int
+    elseif index isa Integer
         return u[index]
     elseif has_referrenced_model(ds)
         prob = referrenced_sciml_prob(ds)
         ugetter = SymbolicIndexingInterface.observed(prob, index)
         p = current_parameters(ds)
-        t = current_time(ds)
         return ugetter(u, p, t)
     else
         throw(ArgumentError("Invalid index to observe state, or if symbolic index, the "*

src/core/trajectory.jl

@@ -19,7 +19,9 @@ to the last valid point.
   have access to unicode, the keyword `Dt` can be used instead.
 * `Ttr = 0`: Transient time to evolve the initial state before starting saving states.
 * `t0 = initial_time(ds)`: Starting time.
-* `save_idxs::AbstractVector{Int}`: Which variables to output in `X` (by default all).
+* `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).
 """
 function trajectory(ds::DynamicalSystem, T, u0 = initial_state(ds);
         save_idxs = nothing, t0 = initial_time(ds), kwargs...
@@ -43,10 +45,10 @@ function trajectory_discrete(ds, T;
     X = isnothing(accessor) ? dimension(ds) : length(accessor)
     data = Vector{SVector{X, ET}}(undef, L)
     Ttr ≠ 0 && step!(ds, Ttr)
-    data[1] = obtain_access(current_state(ds), accessor)
+    data[1] = obtain_state(current_state(ds), accessor)
     for i in 2:L
         step!(ds, Δt)
-        data[i] = SVector{X, ET}(obtain_access(current_state(ds), accessor))
+        data[i] = SVector{X, ET}(obtain_state(ds, current_state(ds), accessor))
         if !successful_step(ds)
             # Diverged trajectory; set final state to remaining set
             # and exit iteration early
@@ -72,7 +74,7 @@ function trajectory_continuous(ds, T; Dt = 0.1, Δt = Dt, Ttr = 0.0, accessor =
             step!(ds)
             successful_step(ds) || break
         end
-        sol[i] = SVector{X, ET}(obtain_access(ds(t), accessor))
+        sol[i] = SVector{X, ET}(obtain_state(ds, t, accessor))
         if !successful_step(ds)
             # Diverged trajectory; set final state to remaining set
             # and exit iteration early
@@ -86,9 +88,21 @@ function trajectory_continuous(ds, T; Dt = 0.1, Δt = Dt, Ttr = 0.0, accessor =
     return StateSpaceSet(sol), tvec
 end
 
-# Util functions for `trajectory`
+# Util functions for `trajectory` to make indexing static vectors
+# as efficient as possible
 svector_access(::Nothing) = nothing
-svector_access(x::AbstractArray) = SVector{length(x), Int}(x...)
+svector_access(x::AbstractVector{Int}) = SVector{length(x), Int}(x...)
 svector_access(x::Int) = SVector{1, Int}(x)
-obtain_access(u, ::Nothing) = u
-obtain_access(u, i::SVector) = u[i]
+svector_access(x::AbstractVector) = x
+
+# this is a special wrapper around `observe_state` to make indexing
+# static arrays more performant without unecessary allocations
+obtain_state(u, ::Nothing) = u
+obtain_state(u, i::SVector) = u[i]
+function obtain_state(ds::DynamicalSystem, t::Real, i::Union{Nothing, SVector})
+    return obtain_state(ds(t), i)
+end
+function obtain_state(ds::DynamicalSystem, t::Real, idxs::AbstractVector)
+    u = ds(t)
+    return [observe_state(ds, i, u, t) for i in idxs]
+end