refactor: add new *_symbols methods, corresponding singletons, update tests

Author: AayushSabharwal

docs/Project.toml

@@ -1,6 +1,8 @@
 [deps]
-SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
 [compat]
 Documenter = "0.27"

docs/pages.jl

@@ -2,6 +2,7 @@
 
 pages = [
     "Home" => "index.md",
-    "API" => "api.md",
     "Tutorial" => "tutorial.md",
+    "Usage" => "usage.md",
+    "API" => "api.md",
 ]

docs/src/api.md

@@ -14,8 +14,10 @@ is_observed
 observed
 is_time_dependent
 constant_structure
-all_solvable_symbols
+all_variable_symbols
 all_symbols
+solvedvariables
+allvariables
 parameter_values
 getp
 setp

docs/usage.md

@@ -0,0 +1,131 @@
+# Using the SymbolicIndexingInterface
+
+This tutorial will cover ways to use the interface for types that implement it.
+Consider the following example:
+
+```@example Usage
+using ModelingToolkit, OrdinaryDiffEq, SymbolicIndexingInterface
+
+@parameters σ ρ β
+@variables t x(t) y(t) z(t) w(t)
+D = Differential(t)
+
+eqs = [D(D(x)) ~ σ * (y - x),
+    D(y) ~ x * (ρ - z) - y,
+    D(z) ~ x * y - β * z,
+    w ~ x + y + z]
+
+@named sys = ODESystem(eqs)
+sys = structural_simplify(sys)
+```
+
+The system has 4 state variables, 3 parameters and one observed variable:
+```@example Usage
+observed(sys)
+```
+
+Solving the system,
+```@example Usage
+u0 = [D(x) => 2.0,
+    x => 1.0,
+    y => 0.0,
+    z => 0.0]
+
+p = [σ => 28.0,
+    ρ => 10.0,
+    β => 8 / 3]
+
+tspan = (0.0, 100.0)
+prob = ODEProblem(sys, u0, tspan, p, jac = true)
+sol = solve(prob, Tsit5())
+```
+
+We can obtain the timeseries of any time-dependent variable using `getindex`
+```@example Usage
+sol[x]
+```
+
+This also works for arrays or tuples of variables, including observed quantities and
+independent variables:
+```@example Usage
+sol[[x, y]]
+```
+
+```@example Usage
+sol[(t, w)]
+```
+
+If necessary, `Symbol`s can be used to refer to variables. This is only valid for
+symbolic variables for which [`hasname`](@ref) returns `true`. The `Symbol` used must
+match the one returned by [`getname`](@ref) for the variable.
+```@example Usage
+hasname(x)
+```
+
+```@example Usage
+getname(x)
+```
+
+```@example Usage
+sol[(:x, :w)]
+```
+
+Note how when indexing with an array, the returned type is a `Vector{Array{Float64}}`,
+and when using a `Tuple`, the returned type is `Vector{Tuple{Float64, Float64}}`.
+To obtain the value of all state variables, we can use the shorthand:
+```@example Usage
+sol[solvedvariables] # equivalent to sol[variable_symbols(sol)]
+```
+
+This does not include the observed variable `w`. To include observed variables in the
+output, the following shorthand is used:
+```@example Usage
+sol[allvariables] # equivalent to sol[all_variable_symbols(sol)]
+```
+
+Parameters cannot be obtained using this syntax, and instead require using [`getp`](@ref) and [`setp`](@ref).
+
+```@example Usage
+σ_getter = getp(sys, σ)
+σ_getter(sol) # can also pass `prob`
+```
+
+Note that this also supports arrays/tuples of parameter symbols:
+
+```@example Usage
+σ_ρ_getter = getp(sys, (σ, ρ))
+σ_ρ_getter(sol) 
+```
+
+Now suppose the system has to be solved with a different value of the parameter `β`.
+
+```@example Usage
+β_setter = setp(sys, β)
+β_setter(prob, 3)
+```
+
+The updated parameter values can be checked using [`parameter_values`](@ref).
+
+```@example Usage
+parameter_values(prob)
+```
+
+Solving the new system, note that the parameter getter functions still work on the new
+solution object.
+
+```@example Usage
+sol2 = solve(prob, Tsit5())
+σ_getter(sol)
+```
+
+```@example Usage
+σ_ρ_getter(sol)
+```
+
+To set the entire parameter vector at once, [`parameter_values`](@ref) can be used
+(note the usage of broadcasted assignment).
+
+```@example Usage
+parameter_values(prob) .= [29.0, 11.0, 2.5]
+parameter_values(prob)
+```

src/interface.jl

@@ -33,6 +33,10 @@ variable_index(sys, sym, i) = variable_index(symbolic_container(sys), sym, i)
 Return a vector of the symbolic variables being solved for in the system `sys`. If
 `constant_structure(sys) == false` this accepts an additional parameter indicating
 the current time index. The returned vector should not be mutated.
+
+For types that implement `Base.getindex` with symbolic indices using this interface,
+The shorthand `sys[solvedvariables]` can be used as shorthand for
+`sys[variable_symbols(sys)]`. See: [`solvedvariables`](@ref).
 """
 variable_symbols(sys) = variable_symbols(symbolic_container(sys))
 variable_symbols(sys, i) = variable_symbols(symbolic_container(sys), i)
@@ -112,12 +116,16 @@ number of variables or parameters over time.
 constant_structure(sys) = constant_structure(symbolic_container(sys))
 
 """
-    all_solvable_symbols(sys)
+    all_variables(sys)
+
+Return a vector of pairs, where the first element of each pair is a symbolic variable
+and the second is its initial value. This includes observed quantities.
 
-Return an array of all symbols in the system that can be solved for. This includes
-observed variables, but not parameters or independent variables.
+For types that implement `Base.getindex` with symbolic indices using this interface,
+The shorthand `sys[allvariables]` can be used as shorthand for
+`sys[all_variable_symbols(sys)]`. See: [`allvariables`](@ref).
 """
-all_solvable_symbols(sys) = all_solvable_symbols(symbolic_container(sys))
+all_variable_symbols(sys) = all_variable_symbols(symbolic_container(sys))
 
 """
     all_symbols(sys)
@@ -126,3 +134,27 @@ Return an array of all symbols in the system. This includes parameters and indep
 variables.
 """
 all_symbols(sys) = all_symbols(symbolic_container(sys))
+
+struct SolvedVariables end
+
+"""
+    const solvedvariables = SolvedVariables()
+
+This singleton is used as a shortcut to allow indexing all solution variables
+(excluding observed quantities). It has a [`symbolic_type`](@ref) of
+[`ScalarSymbolic`](@ref). See: [`variable_symbols`](@ref).
+"""
+const solvedvariables = SolvedVariables()
+symbolic_type(::Type{AllVariables}) = ScalarSymbolic()
+
+struct AllVariables end
+
+"""
+    const allvariables = AllVariables()
+
+This singleton is used as a shortcut to allow indexing all solution variables
+(including observed quantities). It has a [`symbolic_type`](@ref) of
+[`ScalarSymbolic`](@ref). See [`all_variable_symbols`](@ref).
+"""
+const allvariables = AllVariables()
+symbolic_type(::Type{AllVariables}) = ScalarSymbolic()

src/symbol_cache.jl

@@ -71,7 +71,7 @@ function is_time_dependent(sc::SymbolCache)
     end
 end
 constant_structure(::SymbolCache) = true
-all_solvable_symbols(sc::SymbolCache) = variable_symbols(sc)
+all_variable_symbols(sc::SymbolCache) = variable_symbols(sc)
 all_symbols(sc::SymbolCache) = vcat(variable_symbols(sc), parameter_symbols(sc), independent_variable_symbols(sc))
 
 function Base.copy(sc::SymbolCache)

test/example_test.jl

@@ -53,7 +53,7 @@ function SymbolicIndexingInterface.observed(sys::SystemMockup, sym, states = not
 end
 SymbolicIndexingInterface.is_time_dependent(sys::SystemMockup) = isequal(sys.indepvar, :t)
 SymbolicIndexingInterface.constant_structure(sys::SystemMockup) = sys.static
-SymbolicIndexingInterface.all_solvable_symbols(sys::SystemMockup) = sys.vars
+SymbolicIndexingInterface.all_variable_symbols(sys::SystemMockup) = sys.vars
 function SymbolicIndexingInterface.all_symbols(sys::SystemMockup)
     vcat(sys.vars, sys.params, independent_variable_symbols(sys))
 end
@@ -83,7 +83,7 @@ sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], :t)
 @test variable_symbols(sys) == [:x, :y, :z]
 @test parameter_symbols(sys) == [:a, :b, :c]
 @test independent_variable_symbols(sys) == [:t]
-@test all_solvable_symbols(sys) == [:x, :y, :z]
+@test all_variable_symbols(sys) == [:x, :y, :z]
 @test sort(all_symbols(sys)) == [:a, :b, :c, :t, :x, :y, :z]
 
 sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], nothing)
@@ -99,7 +99,7 @@ sys = SystemMockup(true, [:x, :y, :z], [:a, :b, :c], nothing)
 @test variable_symbols(sys) == [:x, :y, :z]
 @test parameter_symbols(sys) == [:a, :b, :c]
 @test independent_variable_symbols(sys) == []
-@test all_solvable_symbols(sys) == [:x, :y, :z]
+@test all_variable_symbols(sys) == [:x, :y, :z]
 @test sort(all_symbols(sys)) == [:a, :b, :c, :x, :y, :z]
 
 sys = SystemMockup(false, [:x, :y, :z], [:a, :b, :c], :t)
@@ -123,5 +123,5 @@ end
 @test variable_symbols(sys, 3) == [:x, :y, :z]
 @test parameter_symbols(sys) == [:a, :b, :c]
 @test independent_variable_symbols(sys) == [:t]
-@test all_solvable_symbols(sys) == [:x, :y, :z]
+@test all_variable_symbols(sys) == [:x, :y, :z]
 @test sort(all_symbols(sys)) == [:a, :b, :c, :t, :x, :y, :z]

test/fallback_test.jl

@@ -22,3 +22,5 @@ all_syms = [:x, :y, :z, :a, :b, :t]
 @test variable_symbols(sys) == variable_symbols(sc)
 @test parameter_symbols(sys) == parameter_symbols(sc)
 @test independent_variable_symbols(sys) == independent_variable_symbols(sc)
+@test all_variable_symbols(sys) == variable_symbols(sc)
+@test all_symbols(sys) == all_symbols(sc)

test/symbol_cache_test.jl

@@ -14,7 +14,7 @@ sc = SymbolCache([:x, :y, :z], [:a, :b], [:t])
 @test variable_symbols(sc) == [:x, :y, :z]
 @test parameter_symbols(sc) == [:a, :b]
 @test independent_variable_symbols(sc) == [:t]
-@test all_solvable_symbols(sc) == [:x, :y, :z]
+@test all_variable_symbols(sc) == [:x, :y, :z]
 @test sort(all_symbols(sc)) == [:a, :b, :t, :x, :y, :z]
 
 sc = SymbolCache([:x, :y], [:a, :b])
@@ -33,15 +33,15 @@ sc = SymbolCache()
 @test all(.!is_independent_variable.((sc,), [:x, :y, :a, :b, :t]))
 @test independent_variable_symbols(sc) == []
 @test !is_time_dependent(sc)
-@test all_solvable_symbols(sc) == []
+@test all_variable_symbols(sc) == []
 @test all_symbols(sc) == []
 
 sc = SymbolCache(nothing, nothing, :t)
 @test all(.!is_independent_variable.((sc,), [:x, :y, :a, :b]))
 @test is_independent_variable(sc, :t)
 @test independent_variable_symbols(sc) == [:t]
 @test is_time_dependent(sc)
-@test all_solvable_symbols(sc) == []
+@test all_variable_symbols(sc) == []
 @test all_symbols(sc) == [:t]
 
 sc2 = copy(sc)