Implement computed / observed states (#13)

* implement computed / observed states

* switch to testing with OrdinaryDiffEqTsit5

* bump version

* test on v1.11

.github/workflows/CI.yml

@@ -11,6 +11,7 @@ jobs:
       matrix:
         version:
           - '1.10'
+          - '1.11'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "GraphDynamics"
 uuid = "bcd5d0fe-e6b7-4ef1-9848-780c183c7f4c"
-version = "0.2.3"
+version = "0.2.4"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/GraphDynamics.jl

@@ -73,14 +73,17 @@ using SciMLBase:
     VectorContinuousCallback,
     ContinuousCallback,
     DiscreteCallback,
-    remake
+    remake,
+    ODEFunction
 
 using RecursiveArrayTools: ArrayPartition
 
 using SymbolicIndexingInterface:
     SymbolicIndexingInterface,
     setu,
-    setp
+    setp,
+    getp,
+    observed
 
 using Accessors:
     Accessors,
@@ -138,6 +141,31 @@ function get_tag end
 function get_params end
 function get_states end
 
+"""
+   computed_properties(::Subsystem{T}) :: NamedTuple{props, NTuple{N, funcs}}
+
+Signal that a subsystem has properties which can be computed on-the-fly based on it's existing properties. In the termoinology used by ModelingToolkit.jl, these are "observed states".
+
+This function takes in a `Subsystem` and returns a `NamedTuple` where each key is a property name that can be computed, and each value is a function that takes in the subsystem and returns a computed value.
+
+By default, this function returns an empty NamedTuple.
+
+Example:
+
+```julia
+struct Position end
+function GraphDynamics.computed_properties(::Subsystem{Position})
+    (;r = (;x, y) -> √(x^2 + y^2),
+      θ = (;x, y) -> atan(y, x))
+end
+
+let sys = Subsystem{Position}(states=(x=1, y=2), params=(;))
+    sys.r == √(sys.x^2 + sys.y^2)
+end
+```
+"""
+computed_properies(s::Subsystem) = ()
+
 """
      subsystem_differential(subsystem, input, t)
 
@@ -253,46 +281,44 @@ Base.size(m::ConnectionMatrix{N}) where {N} = (N, N)
 
 abstract type GraphSystem end
 
-@kwdef struct ODEGraphSystem{CM <: ConnectionMatrices, S, P, EVT, CDEP, CCEP, Ns, SNM, PNM} <: GraphSystem
+@kwdef struct ODEGraphSystem{CM <: ConnectionMatrices, S, P, EVT, CDEP, CCEP, Ns, SNM, PNM, CNM} <: GraphSystem
     connection_matrices::CM
     states_partitioned::S
     params_partitioned::P
     tstops::EVT = Float64[]
     composite_discrete_events_partitioned::CDEP = nothing
     composite_continuous_events_partitioned::CCEP = nothing
     names_partitioned::Ns
-    state_namemap::SNM = compute_namemap(names_partitioned, states_partitioned)
-    param_namemap::PNM = compute_namemap(names_partitioned, params_partitioned)
+    state_namemap::SNM = make_state_namemap(names_partitioned, states_partitioned)
+    param_namemap::PNM = make_param_namemap(names_partitioned, params_partitioned)
+    compu_namemap::CNM = make_compu_namemap(names_partitioned, states_partitioned, params_partitioned)
 end
-@kwdef struct SDEGraphSystem{CM <: ConnectionMatrices, S, P, EVT, CDEP, CCEP, Ns, SNM, PNM} <: GraphSystem
+@kwdef struct SDEGraphSystem{CM <: ConnectionMatrices, S, P, EVT, CDEP, CCEP, Ns, SNM, PNM, CNM} <: GraphSystem
     connection_matrices::CM
     states_partitioned::S
     params_partitioned::P
     tstops::EVT = Float64[]
     composite_discrete_events_partitioned::CDEP = nothing
     composite_continuous_events_partitioned::CCEP = nothing
     names_partitioned::Ns
-    state_namemap::SNM = compute_namemap(names_partitioned, states_partitioned)
-    param_namemap::PNM = compute_namemap(names_partitioned, params_partitioned)
+    state_namemap::SNM = make_state_namemap(names_partitioned, states_partitioned)
+    param_namemap::PNM = make_param_namemap(names_partitioned, params_partitioned)
+    compu_namemap::CNM = make_compu_namemap(names_partitioned, states_partitioned, params_partitioned)
 end
 
 
 #----------------------------------------------------------
 # Infrastructure for subsystems
 include("subsystems.jl")
 
-
-
 #----------------------------------------------------------
 # Problem generating API:
 include("problems.jl")
 
-
 #----------------------------------------------------------
 # Symbolically indexing the solutions of graph systems
 include("symbolic_indexing.jl")
 
-
 #----------------------------------------------------------
 # Solving graph differential equations
 include("graph_solve.jl")

src/subsystems.jl

@@ -128,6 +128,8 @@ end
 function ConstructionBase.getproperties(s::Subsystem)
     states = NamedTuple(get_states(s))
     params = NamedTuple(get_params(s))
+    # TODO: should there be observed states in here? I don't want to accidentally waste CPU cycles on them
+    # but it might be necessary? Not sure, since they can't be `setproperty`-d
     merge(states, params)
 end
 function ConstructionBase.setproperties(s::Subsystem{T, Eltype, States, Params}, patch::NamedTuple) where {T, Eltype, States, Params}
@@ -172,11 +174,8 @@ get_states(s::Subsystem) = getfield(s, :states)
 get_params(s::Subsystem) = getfield(s, :params)
 get_tag(::Subsystem{Name}) where {Name} = Name
 
-
-
 get_tag(::Type{<:Subsystem{Name}}) where {Name} = Name
 
-
 function Base.getproperty(s::Subsystem{<:Any, States, Params},
                           prop::Symbol) where {States, Params}
     states = NamedTuple(get_states(s))
@@ -186,7 +185,12 @@ function Base.getproperty(s::Subsystem{<:Any, States, Params},
     elseif prop ∈ keys(params)
         getproperty(params, prop)
     else
-        subsystem_prop_err(s, prop)
+        comp_props = computed_properies(s)
+        if prop ∈ keys(comp_props)
+            comp_props[prop](s)
+        else
+            subsystem_prop_err(s, prop)
+        end
     end
 end
 @noinline subsystem_prop_err(s::Subsystem{Name}, prop) where {Name} = error(ArgumentError(

src/symbolic_indexing.jl

@@ -15,30 +15,54 @@ struct ParamIndex #todo: this'll require some generalization to support weight p
     prop::Symbol
 end
 
-function compute_namemap(names_partitioned, states_partitioned::Tuple{Vararg{AbstractVector{<:SubsystemStates}}})
-    state_namemap = OrderedDict{Symbol, StateIndex}()
+function make_state_namemap(names_partitioned::NTuple{N, Vector{Symbol}},
+                            states_partitioned::NTuple{N, AbstractVector{<:SubsystemStates}}) where {N}
+    namemap = OrderedDict{Symbol, StateIndex}()
     for i ∈ eachindex(names_partitioned, states_partitioned)
         for j ∈ eachindex(names_partitioned[i], states_partitioned[i])
-            for (k, name) ∈ enumerate(propertynames(states_partitioned[i][j]))
+            states = states_partitioned[i][j]
+            for (k, name) ∈ enumerate(propertynames(states))
                 propname = Symbol(names_partitioned[i][j], "₊", name)
-                state_namemap[propname] = StateIndex(i, j, k)
+                namemap[propname] = StateIndex(i, j, k)
             end
         end
     end
-    state_namemap
+    namemap
 end
-function compute_namemap(names_partitioned, params_partitioned::Tuple{Vararg{AbstractVector{<:SubsystemParams}}})
-    param_namemap = OrderedDict{Symbol, ParamIndex}()
+
+function make_param_namemap(names_partitioned::NTuple{N, Vector{Symbol}},
+                            params_partitioned::NTuple{N, AbstractVector{<:SubsystemParams}}) where {N}
+    namemap = OrderedDict{Symbol, ParamIndex}()
     for i ∈ eachindex(names_partitioned, params_partitioned)
         for j ∈ eachindex(names_partitioned[i], params_partitioned[i])
-            for name ∈ propertynames(params_partitioned[i][j])
+            params = params_partitioned[i][j]
+            for name ∈ propertynames(params)
                 propname = Symbol(names_partitioned[i][j], "₊", name)
                 #TODO: this'll require some generalization to support weight params
-                param_namemap[propname] = ParamIndex(i, j, name)
+                namemap[propname] = ParamIndex(i, j, name)
+            end
+        end
+    end
+    namemap
+end
+
+
+function make_compu_namemap(names_partitioned::NTuple{N, Vector{Symbol}},
+                          states_partitioned::NTuple{N, AbstractVector{<:SubsystemStates}},
+                          params_partitioned::NTuple{N, AbstractVector{<:SubsystemParams}}) where {N}
+    namemap = OrderedDict{Symbol, ParamIndex}()
+    for i ∈ eachindex(names_partitioned, states_partitioned, params_partitioned)
+        for j ∈ eachindex(names_partitioned[i], states_partitioned[i], params_partitioned[i])
+            states = states_partitioned[i][j]
+            params = params_partitioned[i][j]
+            sys = Subsystem(states, params)
+            for name ∈ keys(computed_properies(sys))
+                propname = Symbol(names_partitioned[i][j], "₊", name)
+                namemap[propname] = ParamIndex(i, j, name)
             end
         end
     end
-    param_namemap
+    namemap
 end
 
 
@@ -118,7 +142,20 @@ function SymbolicIndexingInterface.is_time_dependent(sys::GraphSystem)
 end
 
 function SymbolicIndexingInterface.is_observed(sys::GraphSystem, sym)
-    false # TODO: support observed variables
+    haskey(sys.compu_namemap, sym)
+end
+
+function SymbolicIndexingInterface.observed(f::ODEFunction{a, b, F}, sym::Symbol) where {a, b, F<:GraphSystemFunction}
+    observed(f.f.sys, sym)
+end
+function SymbolicIndexingInterface.observed(sys::GraphSystem, sym)
+    (; tup_index, v_index, prop) = sys.compu_namemap[sym]
+    function gen(u, p, t)
+        (; params_partitioned, state_types_val) = p
+        states_partitioned = to_vec_o_states(u.x, state_types_val)
+        subsys = Subsystem(states_partitioned[tup_index][v_index], params_partitioned[tup_index][v_index])
+        getproperty(subsys, prop)
+    end
 end
 
 # function SymbolicIndexingInterface.all_solvable_symbols(sys::GraphSystem)

test/Project.toml

@@ -1,7 +1,7 @@
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 GraphDynamics = "bcd5d0fe-e6b7-4ef1-9848-780c183c7f4c"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/particle_osc_example.jl

@@ -1,4 +1,4 @@
-using GraphDynamics, OrdinaryDiffEq, Test, ForwardDiff
+using GraphDynamics, OrdinaryDiffEqTsit5, Test, ForwardDiff
 
 struct Particle end
 function GraphDynamics.subsystem_differential(sys::Subsystem{Particle}, F, t)
@@ -7,6 +7,7 @@ function GraphDynamics.subsystem_differential(sys::Subsystem{Particle}, F, t)
     dv = F/m
     SubsystemStates{Particle}((;x=dx, v=dv))
 end
+
 GraphDynamics.initialize_input(::Subsystem{Particle}) = 0.0
 
 struct Oscillator end
@@ -17,6 +18,7 @@ function GraphDynamics.subsystem_differential(sys::Subsystem{Oscillator}, F, t)
     SubsystemStates{Oscillator}((;x=dx, v=dv))
 end
 GraphDynamics.initialize_input(::Subsystem{Oscillator}) = 0.0
+GraphDynamics.computed_properies(::Subsystem{Oscillator}) = (;ω₀ = ((;m, k),) -> √(k/m))
 
 struct Spring
     k::Float64
@@ -35,6 +37,13 @@ function ((;fac)::Coulomb)(a, b)
 end
 
 
+@testset "basic" begin
+    sys = Subsystem{Oscillator}(states=(;x=-2.0, v=1.0), params=(;x₀=0.0, m=30.0, k=1.0, q=1.0))
+    @test sys.x == -2.0
+    @test sys.ω₀ == √(sys.k/sys.m)
+end
+
+
 function solve_particle_osc(;x1, x2, tspan = (0.0, 10.0), alg=Tsit5())
     # put some garbage values in here for states and params, but we'll set them to reasonable values later with the
     # u0map and param_map
@@ -89,18 +98,22 @@ end
 
 @testset "solutions" begin
     sol = solve_particle_osc(;x1=1.0, x2=-1.0)
-    @test sol[:particle1₊x][end] ≈  0.580617 rtol=1e-3
-    @test sol[:particle2₊x][end] ≈ -1.391576 rtol=1e-3
-    @test sol[:osc₊x][end]       ≈ -1.063306 rtol=1e-3
+    @test sol[:particle1₊x, end] ≈  0.580617 rtol=1e-3
+    @test sol[:particle2₊x, end] ≈ -1.391576 rtol=1e-3
+    @test sol[:osc₊x, end]       ≈ -1.063306 rtol=1e-3
+    k = GraphDynamics.getp(sol, :osc₊k)(sol)
+    m = GraphDynamics.getp(sol, :osc₊m)(sol)
+    @test sol[:osc₊ω₀, end] == √(k/m)
 end
 
 @testset "sensitivies" begin
     jac = ForwardDiff.jacobian([1.0, -1.0]) do (x1, x2)
         sol = solve_particle_osc(;x1, x2)
-        [sol[:particle1₊x][end], sol[:particle2₊x][end], sol[:osc₊x][end]]
+        [sol[:particle1₊x, end], sol[:particle2₊x, end], sol[:osc₊x, end]]
     end
     @test jac ≈ [-0.50821   -0.740152
                  -0.199444  -0.906593
                  -0.586021   0.118173] rtol=1e-3
 
 end
+