Merge pull request #192 from JuliaDynamics/hw/obsexp

add ObservableExpressions to create derived quantities on the go

Author: hexaeder

Project.toml

@@ -29,17 +29,19 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 StyledStrings = "f489334b-da3d-4c2e-b8f0-e476e12c162b"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
-Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [extensions]
 CUDAExt = ["CUDA", "Adapt"]
-MTKExt = "ModelingToolkit"
+MTKExt = ["ModelingToolkit", "Symbolics"]
+SymbolicsExt = ["Symbolics", "MacroTools"]
 
 [compat]
 Adapt = "4.0.4"
@@ -53,6 +55,7 @@ Graphs = "1"
 InteractiveUtils = "1"
 KernelAbstractions = "0.9.18"
 LinearAlgebra = "1"
+MacroTools = "0.5.15"
 Mixers = "0.1.2"
 ModelingToolkit = "9"
 NNlib = "0.9.13"

docs/src/API.md

@@ -82,6 +82,7 @@ VIndex
 EIndex
 VPIndex
 EPIndex
+@obsex
 ```
 
 ### Index generators

docs/src/symbolic_indexing.md

@@ -114,3 +114,20 @@ Observables can be accessed like any other state, for example, the flows in the
 ```@example si
 plot(sol; idxs=eidxs(nw, :, :flow))
 ```
+
+## Derived `ObservableExpressions` using `@obsex`
+
+Sometimes it is usefull to plot or observe some simple derived quantity.For that,
+one can used the [`@obsex`](@ref) macro, to define simple derived quantities.
+
+For example, we can directly plot the storage difference with respect to storage of node 1.
+
+```@example si
+plot(sol; idxs=@obsex(vidxs(nw,:,:storage) .- VIndex(1,:storage)))
+```
+
+Other examples are the calculation of magnitude and argument of complex values which are modeld in real and imaginary part.
+```
+@obsex mag = sqrt(VIndex(1, :u_r)^2 + VIndex(2, :u_i)^2)
+```
+

ext/MTKExt.jl

@@ -4,7 +4,7 @@ using ModelingToolkit: Symbolic, iscall, operation, arguments, build_function
 using ModelingToolkit: ModelingToolkit, Equation, ODESystem, Differential
 using ModelingToolkit: full_equations, get_variables, structural_simplify, getname, unwrap
 using ModelingToolkit: full_parameters, unknowns, independent_variables, observed, defaults
-using ModelingToolkit.Symbolics: Symbolics, fixpoint_sub, substitute
+using Symbolics: Symbolics, fixpoint_sub, substitute
 using RecursiveArrayTools: RecursiveArrayTools
 using ArgCheck: @argcheck
 using LinearAlgebra: Diagonal, I

ext/SymbolicsExt.jl

@@ -0,0 +1,62 @@
+module SymbolicsExt
+
+using Symbolics: Symbolics, @variables, Num
+using MacroTools: postwalk, @capture
+using NetworkDynamics: NetworkDynamics, ObservableExpression, SymbolicVertexIndex, SymbolicEdgeIndex, SII
+
+function NetworkDynamics.generate_observable_expression(ex::Expr)
+    if @capture(ex, capname_ = capcontent_)
+        name = capname
+        content = capcontent
+    else
+        name = nothing
+        content = ex
+    end
+    # manually hygen on "mapping", otherwise we cannot escape x
+    mapping_sym = gensym("mapping")
+    symbolic_expr = postwalk(content) do x
+        if x ∈ (:+, :-, :^, :/, :.+, :.-, :.^, :./, Symbol(":")) || x isa QuoteNode
+            x
+        else
+            :(NetworkDynamics.collect_symbol!($mapping_sym, $x))
+        end
+    end
+    quote
+        $(esc(mapping_sym)) = Dict()
+        expr = $(esc(symbolic_expr))
+        NetworkDynamics.ObservableExpression($(esc(mapping_sym)), expr, $(Meta.quot(name)))
+    end
+end
+
+function NetworkDynamics.ObservableExpression(mapping, expr::Vector, names)
+    if names isa Symbol
+        names = [Symbol(names, NetworkDynamics.subscript(i)) for i in 1:length(expr)]
+    elseif isnothing(names)
+        names = Iterators.repeated(nothing, length(expr))
+    end
+    [NetworkDynamics.ObservableExpression(mapping, e, n) for (e, n) in zip(expr, names)]
+end
+
+function NetworkDynamics.ObservableExpression(mapping, expr, name)
+    nwidx = collect(keys(mapping))
+    syms = collect(values(mapping))
+    f = Symbolics.build_function(expr, syms; expression=Val(false))
+    NetworkDynamics.ObservableExpression(nwidx, f, expr, name)
+end
+
+NetworkDynamics.collect_symbol!(_, x) = x
+function NetworkDynamics.collect_symbol!(mapping, x::Union{AbstractVector, Tuple})
+    map(el -> NetworkDynamics.collect_symbol!(mapping, el), x)
+end
+function NetworkDynamics.collect_symbol!(mapping, nwindex::Union{SymbolicVertexIndex,SymbolicEdgeIndex})
+    if haskey(mapping, nwindex)
+        return mapping[nwindex]
+    else
+        name = SII.getname(nwindex)
+        sym = only(@variables $name)
+        mapping[nwindex] = sym
+        return sym
+    end
+end
+
+end # module

src/NetworkDynamics.jl

@@ -61,6 +61,7 @@ include("external_inputs.jl")
 export VIndex, EIndex, VPIndex, EPIndex, NWState, NWParameter, uflat, pflat
 export vidxs, eidxs, vpidxs, epidxs
 export save_parameters!
+export @obsex
 include("symbolicindexing.jl")
 
 export has_metadata, get_metadata, set_metadata!

src/symbolicindexing.jl

@@ -402,6 +402,10 @@ function observed_symbols(nw::Network)
 end
 
 function SII.observed(nw::Network, snis)
+    if (snis isa AbstractVector || snis isa Tuple) && any(sni -> sni isa ObservableExpression, snis)
+        throw(ArgumentError("Cannot mix normal symbolic indices with @obsex currently!"))
+    end
+
     _snis = _expand_and_collect(nw, snis)
     isscalar = _snis isa SymbolicIndex
     if isscalar
@@ -1079,3 +1083,80 @@ extract_nw(p::IndexingProxy) = extract_nw(p.s)
 function extract_nw(::Nothing)
     throw(ArgumentError("Needs system context to generate matching indices. Pass Network, sol, prob, ..."))
 end
+
+####
+#### Observable Expressions
+####
+struct ObservableExpression{VT,F,Ex,N}
+    inputs::Vector{VT}
+    f::F
+    ex::Ex
+    name::N
+end
+function Base.show(io::IO, mime::MIME"text/plain", obsex::ObservableExpression)
+    print(io, "ObservableExpression(")
+    isnothing(obsex.name) || print(io, obsex.name, " = ")
+    show(io, mime, obsex.ex)
+    print(io, ")")
+end
+
+"""
+    @obsex([name =] expression)
+
+Define observable expressions, which are simple combinations of knonw
+states/parameters/observables. `@obsex(...)` returns an `ObservableExpression`
+which can be used as an symbolic index. This is mainly intended for quick
+plotting or export of common "derived" variables, such as the argument of a
+2-component complex state. For example:
+
+    sol(t; idxs=@obsex(arg = atan(VIndex(1,:u_i), VIndex(1,:u_r))]
+    sol(t; idxs=@obsex(δrel = VIndex(1,:δ) - VIndex(2,:δ)))
+
+"""
+macro obsex(ex)
+    generate_observable_expression(ex)
+end
+
+function generate_observable_expression(::Any)
+    error("@obsex can only be used when Symbolics.jl is loaded.")
+end
+
+# define function stub to overload in SymbolicsExt
+function collect_symbol! end
+
+function SII.is_observed(nw::Network, obsex::ObservableExpression)
+    true
+end
+
+SII.symbolic_type(::Type{<:ObservableExpression}) = SII.ScalarSymbolic()
+SII.hasname(::ObservableExpression) = true
+function SII.getname(obsex::ObservableExpression)
+    if obsex.name isa Symbol
+        return obsex.name
+    else
+        io = IOBuffer()
+        show(io, MIME"text/plain"(), obsex.ex)
+        str = String(take!(io))
+        return Symbol(replace(str, " "=>""))
+    end
+end
+
+function SII.observed(nw::Network, obsex::ObservableExpression)
+    inputf = SII.observed(nw, obsex.inputs)
+    (u, p, t) -> begin
+        input = inputf(u, p, t)
+        obsex.f(input)
+    end
+end
+function SII.observed(nw::Network, obsexs::AbstractVector{<:ObservableExpression})
+    inputfs = map(obsex -> SII.observed(nw, obsex.inputs), obsexs)
+    (u, p, t) -> begin
+        map(obsexs, inputfs) do obsex, inputf
+            input = inputf(u, p, t)
+            obsex.f(input)
+        end
+    end
+end
+
+Base.getindex(s::NWState, idx::ObservableExpression) = SII.getu(s, idx)(s)
+Base.getindex(s::NWParameter, idx::ObservableExpression) = SII.getp(s, idx)(s)

test/Project.toml

@@ -41,6 +41,7 @@ Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 

test/symbolicindexing_test.jl

@@ -3,6 +3,7 @@ using Graphs
 using OrdinaryDiffEqTsit5
 using Chairmarks
 using Test
+using Symbolics
 import SymbolicIndexingInterface as SII
 using NetworkDynamics: VIndex, EIndex, VPIndex, EPIndex, _resolve_colon
 
@@ -488,3 +489,24 @@ end
     @test s[EIndex(:e2, :src₊δin)] == s[VIndex(:v2, :δ)]
     @test s[EIndex(:e2, :dst₊δin)] == s[VIndex(:v3, :δ)]
 end
+
+@testset "test observed expressions" begin
+    v1 = Lib.kuramoto_second(name=:v1, vidx=1, insym=[:Pin])
+    v2 = Lib.kuramoto_second(name=:v2, vidx=2, insym=[:Pin])
+    v3 = Lib.kuramoto_second(name=:v3, vidx=3, insym=[:Pin])
+    e1 = Lib.kuramoto_edge(name=:e1, src=1, dst=2, insym=[:δin])
+    e2 = Lib.kuramoto_edge(name=:e2, src=2, dst=3, insym=[:δin])
+    nw = Network([v1,v2,v3], [e1,e2])
+    s = NWState(nw, rand(dim(nw)), rand(pdim(nw)))
+
+    obsex = @obsex(VIndex(1,:δ) + VIndex(2,:δ))
+    @test s[obsex] == s[VIndex(1,:δ)] + s[VIndex(2,:δ)]
+    @test s[@obsex VIndex(1,:δ) - EIndex(:e1, :src₊δin)] == 0
+
+    @test SII.getname(@obsex(VIndex(1,:δ) + VIndex(2,:δ))) == Symbol("v1₊δ+v2₊δ")
+    @test SII.getname(@obsex(δ²=VIndex(1,:δ)^2)) == :δ²
+
+    @test s[@obsex(vidxs(s, :, :δ) .- VIndex(1, :δ))] == [0, s[VIndex(2,:δ)] - s[VIndex(1,:δ)], s[VIndex(3,:δ)] - s[VIndex(1,:δ)]]
+
+    obsex = @obsex(δ_rel = vidxs(s, :, :δ) .- VIndex(1, :δ))
+end