Correct version of #203 (#204)

* allow using dict in set state

* improve docstring of set _parametrs, allow pair

* add tests

* oh lord yes every single test passes

* hell yeall all tests pass, all features are implemented

* add changelog and version

* make ds broadcastable

* finish up everything!

* better docstring

Author: Datseris

CHANGELOG.md

@@ -1,3 +1,19 @@
+# v3.8.0
+
+Even more amazing level of integration with ModelingToolkit.jl!
+
+- Now it is possible to partially set specific state variables by
+providing a dictionary or vector of pairs that map indices to variables in `set_state!` or `reinit!`.
+- Partial setting of states via providing a dictionary also works in `reinit!`.
+- Partial setting of states also works when initializing a `ParallelDynamicalSystem`
+  by providing a vector of dictionaries or vector of pairs.
+- It was already possible to provide a dictionary of parameter indices to values
+  in `set_parameters!`, now it was made more clear in the docstring.
+- bugfix where `reinit!` or `set_state!` of continuous time out of place parallel dynamical systems was not working correctly.
+- Dynamical systems are now broadacastable, e.g., you can do `observe_state.(ds, (a, b, c))`.
+- Fixed a bug where `set_state!(ds, val, index)` only worked for in-place systems.
+
+
 # v3.7.0
 
 Symbolic indices used in `observe_state` can now also be used in `trajectory`.

Project.toml

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

src/core/dynamicalsystem_interface.jl

@@ -105,6 +105,9 @@ unless when developing new algorithm implementations that use dynamical systems.
 """
 abstract type DynamicalSystem end
 
+# Make it broadcastable:
+Base.broadcastable(ds::DynamicalSystem) = Ref(ds)
+
 # We utilize nature of time for dispatch; continuous time dispatches to `integ`
 # and dispatches for `::DEIntegrator` are defined in `CoupledODEs` file.
 """
@@ -326,27 +329,35 @@ StateSpaceSets.dimension(ds::DynamicalSystem) = length(current_state(ds))
 # API - altering status of the system
 ###########################################################################################
 """
-    set_state!(ds::DynamicalSystem, u::AbstractArray)
+    set_state!(ds::DynamicalSystem, u::AbstractArray{Real})
 
 Set the state of `ds` to `u`, which must match dimensionality with that of `ds`.
 Also ensure that the change is notified to whatever integration protocol is used.
 """
 set_state!(ds, u) = errormsg(ds)
 
 """
-    set_state!(ds::DynamicalSystem, value::Real, index) → u
+    set_state!(ds::DynamicalSystem, value::Real, i) → u
 
-Set the `i`th variable of `ds` to `value`. The `index` can be an integer or
+Set the `i`th variable of `ds` to `value`. The index `i` can be an integer or
 a symbolic-like index for systems that reference a ModelingToolkit.jl model.
-
-Calling instead `set_state!(u, value, index, ds)` will modify the given
-state `u` and return it, leaving `ds` unaltered.
+For example:
+```julia
+i = :x # or `1` or `only(@variables(x))`
+set_state!(ds, 0.5, i)
+```
 
 **Warning:** this function should not be used with derivative dynamical systems
 such as Poincare/stroboscopic/projected dynamical systems.
+Use the method below to manipulate an array and give that to `set_state!`.
+
+
+    set_state!(u::AbstractArray, value, index, ds::DynamicalSystem)
+
+Modify the given state `u` and leave `ds` untouched.
 """
 function set_state!(ds::DynamicalSystem, value::Real, i)
-    u = current_state(ds)
+    u = Array(current_state(ds)) # ensure it works for out of place as well!
     u = set_state!(u, value, i, ds)
     set_state!(ds, u)
 end
@@ -365,6 +376,27 @@ function set_state!(u::AbstractArray, value::Real, i, ds::DynamicalSystem)
     return u
 end
 
+"""
+    set_state!(ds::DynamicalSystem, mapping::Dict)
+
+Convenience version of `set_state!` that iteratively calls `set_state!(ds, val, i)`
+for all index-value pairs `(i, val)` in `mapping`. This allows you to
+partially set only some state variables.
+"""
+function set_state!(ds::DynamicalSystem, mapping::Dict)
+    # ensure we use a mutable vector, so same code works for in-place problems
+    # (SymbolicIndexingInterface only works with mutable objects)
+    um = Array(copy(current_state(ds)))
+    set_state!(um, mapping, ds)
+    set_state!(ds, um)
+end
+function set_state!(um::Array{<:Real}, mapping::Dict, ds::DynamicalSystem)
+    for (i, value) in pairs(mapping)
+        set_state!(um, value, i, ds)
+    end
+    return um
+end
+
 """
     set_parameter!(ds::DynamicalSystem, index, value [, p])
 
@@ -403,14 +435,18 @@ end
 """
     set_parameters!(ds::DynamicalSystem, p = initial_parameters(ds))
 
-Set the parameter values in the [`current_parameters`](@ref)`(ds)` to match `p`.
-This is done as an in-place overwrite by looping over the keys of `p`.
-Hence the keys of `p` must be a subset of the keys of [`current_parameters`](@ref)`(ds)`.
+Set the parameter values in the [`current_parameters`](@ref)`(ds)` to match those in `p`.
+This is done as an in-place overwrite by looping over the keys of `p`
+hence `p` can be an arbitrary container mapping parameter indices to values
+(such as a `Vector{Real}`, `Vector{Pair}`, or `Dict`).
+
+The keys of `p` must be valid keys that can be given to [`set_parameter!`](@ref).
 """
 function set_parameters!(ds::DynamicalSystem, p = initial_parameters(ds))
     cp = current_parameters(ds)
     p === cp && return
-    for (index, value) in pairs(p)
+    iter = p isa Vector ? pairs(p) : p # allows using vector, dict, or vector{pair}.
+    for (index, value) in iter
         _set_parameter!(ds, index, value, cp)
     end
     return
@@ -442,15 +478,28 @@ SciMLBase.step!(ds::DynamicalSystem, args...) = errormsg(ds)
 
 Reset the status of `ds`, so that it is as if it has be just initialized
 with initial state `u`. Practically every function of the ecosystem that evolves
-`ds` first calls this function on it. Besides the new initial state `u`, you
+`ds` first calls this function on it. Besides the new state `u`, you
 can also configure the keywords `t0 = initial_time(ds)` and `p = current_parameters(ds)`.
 
-Note the default settings: the state and time are the initial,
-but the parameters are the current.
+    reinit!(ds::DynamicalSystem, u::Dict; kwargs...) → ds
 
-The special method `reinit!(ds, ::Nothing; kwargs...)` is also available,
-which does nothing and leaves the system as is. This is so that downstream functions
+If `u` is a `Dict` (for partially setting specific state variables in [`set_state`](@ref)),
+then the alterations in `u` are still done in the state given by the keyword
+`reference_state = copy(initial_state(ds))`.
+
+    reinit!(ds, ::Nothing; kwargs...)
+
+This method does nothing and leaves the system as is. This is so that downstream functions
 that call `reinit!` can still be used without resetting the system but rather
 continuing from its exact current state.
 """
-SciMLBase.reinit!(ds::DynamicalSystem, args...; kwargs...) = errormsg(ds)
+function SciMLBase.reinit!(ds::DynamicalSystem, mapping::Dict;
+    reference_state = copy(initial_state(ds)), kwargs...)
+    um = Array(reference_state)
+    set_state!(um, mapping, ds)
+    reinit!(ds, um; kwargs...)
+end
+
+SciMLBase.reinit!(ds::DynamicalSystem, ::Nothing; kw...) = ds
+# all extensions of `reinit!` in concrete implementaitons
+# should only implement `reinit!(ds, ::AbstractArray)` method.

src/core_systems/arbitrary_steppable.jl

@@ -32,7 +32,7 @@ the following adjustments:
 - The keyword `isdeterministic` should be set properly, as it decides whether
   downstream algorithms should error or not.
 """
-struct ArbitrarySteppable{M, F, R, ES, EP, U, P, S} <: DiscreteTimeDynamicalSystem
+struct ArbitrarySteppable{U, M, F, R, ES, EP, P, S} <: DiscreteTimeDynamicalSystem
     model::M
     step::F
     reinit::R
@@ -77,10 +77,9 @@ function SciMLBase.step!(ds::ArbitrarySteppable, n::Int = 1, stop::Bool = true)
     return ds
 end
 
-function SciMLBase.reinit!(ds::ArbitrarySteppable, u = initial_state(ds);
+function SciMLBase.reinit!(ds::ArbitrarySteppable{U}, u::U = initial_state(ds);
         p = current_parameters(ds), t0 = 0, # t0 is not used but required for downstream.
-    )
-    isnothing(u) && return
+    ) where {U}
     ds.reinit(ds.model, u, p)
     ds.t[] = 0
     return ds

src/core_systems/continuous_time_ode.jl

@@ -136,15 +136,14 @@ for f in (:initial_state, :current_parameters, :dynamic_rule,
 end
 
 SciMLBase.isinplace(::CoupledODEs{IIP}) where {IIP} = IIP
-set_state!(ds::CoupledODEs, u::AbstractArray) = set_state!(ds.integ, u)
+set_state!(ds::CoupledODEs, u::AbstractArray) = (set_state!(ds.integ, u); ds)
 
 # so that `ds` is printed
 SciMLBase.step!(ds::CoupledODEs, args...) = (step!(ds.integ, args...); ds)
 
-function SciMLBase.reinit!(ds::ContinuousTimeDynamicalSystem, u = initial_state(ds);
+function SciMLBase.reinit!(ds::ContinuousTimeDynamicalSystem, u::AbstractArray = initial_state(ds);
         p = current_parameters(ds), t0 = initial_time(ds)
     )
-    isnothing(u) && return
     set_parameters!(ds, p)
     reinit!(ds.integ, u; reset_dt = true, t0)
     return ds
@@ -174,7 +173,7 @@ end
 function set_state!(integ::DEIntegrator, u)
     if integ.u isa Array{<:Real}
         integ.u .= u
-    elseif integ.u isa Union{SVector, SMatrix}
+    elseif integ.u isa StateSpaceSets.StaticArraysCore.SArray{<:Real}
         integ.u = u
     else
         integ.u = recursivecopy(u)

src/core_systems/discrete_time_map.jl

@@ -73,7 +73,7 @@ isdeterministic(::DIM) = true
 function set_state!(ds::DeterministicIteratedMap{IIP}, u) where {IIP}
     ds.u = recursivecopy(u)
     ds.dummy = recursivecopy(u)
-    return
+    return ds
 end
 
 ###########################################################################################
@@ -115,10 +115,9 @@ end
 ###########################################################################################
 # Alterations
 ###########################################################################################
-function reinit!(ds::DIM, u = initial_state(ds);
+function reinit!(ds::DIM, u::AbstractArray = initial_state(ds);
         p = current_parameters(ds), t0 = initial_time(ds)
     )
-    isnothing(u) && return
     set_state!(ds, u)
     ds.t = t0
     set_parameters!(ds, p)

src/derived_systems/parallel_systems.jl

@@ -23,6 +23,13 @@ This struct follows the [`DynamicalSystem`](@ref) interface with the following a
 - [`current_states`](@ref) and [`initial_states`](@ref) can be used to get
   all parallel states.
 - [`reinit!`](@ref) takes in a vector of states (like `states`) for `u`.
+
+
+    ParallelDynamicalSystem(ds::DynamicalSystem, states::Vector{<:Dict})
+
+For a dynamical system referring a MTK model, one can specify states
+as a vector of dictionaries to alter the current state of `ds`
+as in [`set_state!`](@ref).
 """
 abstract type ParallelDynamicalSystem <: DynamicalSystem end
 
@@ -45,7 +52,7 @@ struct ParallelDynamicalSystemAnalytic{D, M} <: ParallelDynamicalSystem
     prob        # reference original MTK problem
 end
 
-function ParallelDynamicalSystem(ds::CoreDynamicalSystem, states)
+function ParallelDynamicalSystem(ds::CoreDynamicalSystem, states::Vector{<:AbstractArray{<:Real}})
     f, st = parallel_rule(ds, states)
     if ds isa DeterministicIteratedMap
         pds = DeterministicIteratedMap(f, st, current_parameters(ds); t0 = initial_time(ds))
@@ -60,6 +67,13 @@ function ParallelDynamicalSystem(ds::CoreDynamicalSystem, states)
     return ParallelDynamicalSystemAnalytic{typeof(pds), M}(pds, dynamic_rule(ds), prob)
 end
 
+function ParallelDynamicalSystem(ds::CoreDynamicalSystem, mappings::Vector{<:Dict})
+    # convert to vector of arrays:
+    u = Array(current_state(ds))
+    states = [set_state!(copy(u), mapping, ds) for mapping in mappings]
+    return ParallelDynamicalSystem(ds, states)
+end
+
 # Out of place: everywhere the same
 function parallel_rule(ds::CoreDynamicalSystem{false}, states)
     f = dynamic_rule(ds)
@@ -74,7 +88,7 @@ function parallel_rule(ds::CoreDynamicalSystem{false}, states)
     return parallel_f, st
 end
 
-# In place: for where `Vector{Vector}` is possible
+# In place: for where `AbstractVector{Vector}` is possible
 function parallel_rule(ds::DeterministicIteratedMap{true}, states)
     f = dynamic_rule(ds)
     S = typeof(correct_state(Val{true}(), first(states)))
@@ -88,7 +102,7 @@ function parallel_rule(ds::DeterministicIteratedMap{true}, states)
     return parallel_f, st
 end
 
-# In place, uses matrix with each column a parallel state
+# In place, but ODEs use matrix with each column a parallel state
 function parallel_rule(ds::CoupledODEs{true}, states)
     st = Matrix(hcat(states...))
     f = dynamic_rule(ds)
@@ -136,6 +150,7 @@ initial_states(pdsa::ParallelDynamicalSystemAnalytic) = initial_state(pdsa.ds)
 function set_state!(pdsa::ParallelDynamicalSystemAnalytic, u::AbstractArray, i::Int = 1)
     current_states(pdsa)[i] = u
     set_state!(pdsa.ds, current_states(pdsa))
+    return pdsa
 end
 
 # States IO for matrix state
@@ -208,6 +223,6 @@ function set_state!(pdtds::PDTDS, u, i::Int = 1)
     end
 end
 
-function SciMLBase.reinit!(pdtds::PDTDS, states = initial_states(pdtds); kwargs...)
+function SciMLBase.reinit!(pdtds::PDTDS, states::AbstractVector = initial_states(pdtds); kwargs...)
     for (ds, s) in zip(pdtds.systems, states); reinit!(ds, s; kwargs...); end
 end

src/derived_systems/poincare/hyperplane.jl

@@ -1,7 +1,5 @@
 export PlaneCrossing, poincaresos
 
-# TODO: Remove keyword `save_idxs`!
-
 ####################################################################################################
 # Hyperplane definitions
 ####################################################################################################

src/derived_systems/poincare/poincaremap.jl

@@ -6,6 +6,8 @@ export PoincareMap, current_crossing_time
 
 const ROOTS_ALG = Roots.A42()
 
+# TODO: Remove keyword `save_idxs`!
+
 ###########################################################################################
 # Type definition
 ###########################################################################################
@@ -164,7 +166,7 @@ function SciMLBase.step!(pmap::PoincareMap)
 end
 SciMLBase.step!(pmap::PoincareMap, n::Int, s = true) = (for _ ∈ 1:n; step!(pmap); end; pmap)
 
-function SciMLBase.reinit!(pmap::PoincareMap, u0 = initial_state(pmap);
+function SciMLBase.reinit!(pmap::PoincareMap, u0::AbstractArray = initial_state(pmap);
         t0 = initial_time(pmap), p = current_parameters(pmap)
     )
     isnothing(u0) && return pmap
@@ -184,7 +186,7 @@ end
 function set_state!(pmap::PoincareMap, u)
     set_state!(pmap.ds, u)
     step!(pmap) # always step once to reach the PSOS
-    return
+    return pmap
 end
 
 function _recreate_state_from_poincare_plane(pmap::PoincareMap, u0)

src/derived_systems/projected_system.jl

@@ -113,7 +113,7 @@ function SciMLBase.step!(prods::ProjectedDynamicalSystem, args...)
 	return prods
 end
 
-function SciMLBase.reinit!(prods::ProjectedDynamicalSystem{P, D, <:AbstractVector}, y = initial_state(prods); kwargs...) where {P, D}
+function SciMLBase.reinit!(prods::ProjectedDynamicalSystem{P, D, <:AbstractVector}, y::AbstractArray = initial_state(prods); kwargs...) where {P, D}
     isnothing(y) && return(prods)
     u = prods.u
     u[prods.projection] .= y
@@ -122,7 +122,7 @@ function SciMLBase.reinit!(prods::ProjectedDynamicalSystem{P, D, <:AbstractVecto
     return prods
 end
 
-function SciMLBase.reinit!(prods::ProjectedDynamicalSystem{P, D, <:Function}, y = initial_state(prods); kwargs...) where {P, D}
+function SciMLBase.reinit!(prods::ProjectedDynamicalSystem{P, D, <:Function}, y::AbstractArray = initial_state(prods); kwargs...) where {P, D}
     isnothing(y) || reinit!(prods.ds, prods.complete_state(y); kwargs...)
     return prods
 end