Merge pull request #281 from JuliaDynamics/hw/edgepairs

allow indexing of edges via pairs

Author: hexaeder

NEWS.md

@@ -7,6 +7,7 @@
 - **improved Initialization System**: Added comprehensive initialization formulas and constraints system:
   - added `@initformula` to add explicit algebraic init equations for specific variables
   - added `@initconstraint` to add additional constraints for the component initialization
+- allow access edges via Pairs, i.e. `EIndex(1=>2,:a)` references variable `:a` in edge from vertex 1 to 2. Works also with unique names of vertices like `EIndex(:a=>:b)` [#281](https://github.com/JuliaDynamics/NetworkDynamics.jl/pull/281).
 
 ## v0.9 Changelog
 ### Main changes in this release

src/symbolicindexing.jl

@@ -3,14 +3,15 @@
     idx = VIndex(comp, sub)
 
 A symbolic index for a vertex state variable.
-- `comp`: the component index, either int or a collection of ints
+- `comp`: the component index, either int, symbol or a collection
 - `sub`: the subindex, either int, symbol or a collection of those.
 
 ```
 VIndex(1, :P)      # vertex 1, variable :P
 VIndex(1:5, 1)     # first state of vertices 1 to 5
 VIndex(7, (:x,:y)) # states :x and :y of vertex 7
 VIndex(2)          # references the second vertex model
+VIndex(:a)         # references vertex with unique name :a
 ```
 
 Can be used to index into objects supporting the `SymbolicIndexingInterface`,
@@ -28,14 +29,16 @@ VIndex(ci::Union{Symbol,Int}) = VIndex(ci, nothing)
     idx = EIndex(comp, sub)
 
 A symbolic index for an edge state variable.
-- `comp`: the component index, either int or a collection of ints
+- `comp`: the component index, either int, symbol, pair or a collection
 - `sub`: the subindex, either int, symbol or a collection of those.
 
 ```
 EIndex(1, :P)      # edge 1, variable :P
 EIndex(1:5, 1)     # first state of edges 1 to 5
 EIndex(7, (:x,:y)) # states :x and :y of edge 7
 EIndex(2)          # references the second edge model
+EIndex(1=>2)       # references edge from v1 to v2
+EIndex(:a=>:b)     # references edge from (uniquely named) vertex :a to :b
 ```
 
 Can be used to index into objects supporting the `SymbolicIndexingInterface`,
@@ -53,7 +56,7 @@ EIndex(ci::Union{Symbol,Int}) = EIndex(ci, nothing)
     idx = VPIndex(comp, sub)
 
 A symbolic index into the parameter a vertex:
-- `comp`: the component index, either int or a collection of ints
+- `comp`: the component index, either int, symbol or a collection
 - `sub`: the subindex, either int, symbol or a collection of those.
 
 Can be used to index into objects supporting the `SymbolicIndexingInterface`,
@@ -69,8 +72,8 @@ end
     EPIndex{C,S} <: SymbolicStateIndex{C,S}
     idx = VEIndex(comp, sub)
 
-A symbolic index into the parameter a vertex:
-- `comp`: the component index, either int or a collection of ints
+A symbolic index into the parameter of an edge:
+- `comp`: the component index, either int, symbol, pair or a collection
 - `sub`: the subindex, either int, symbol or a collection of those.
 
 Can be used to index into objects supporting the `SymbolicIndexingInterface`,
@@ -95,18 +98,25 @@ SSI Maintainer assured that f.sys is really only used for symbolic indexig so me
 SciMLBase.__has_sys(nw::Network) = true
 Base.getproperty(nw::Network, s::Symbol) = s===:sys ? nw : getfield(nw, s)
 
-SII.symbolic_type(::Type{<:SymbolicIndex{<:Union{Symbol,Int},<:Union{Symbol,Int}}}) = SII.ScalarSymbolic()
+SII.symbolic_type(::Type{<:SymbolicIndex{<:Union{<:Pair,Symbol,Int},<:Union{Symbol,Int}}}) = SII.ScalarSymbolic()
 SII.symbolic_type(::Type{<:SymbolicIndex}) = SII.ArraySymbolic()
 
 SII.hasname(::SymbolicIndex) = false
-SII.hasname(::SymbolicIndex{<:Union{Symbol,Int},<:Union{Symbol,Int}}) = true
+SII.hasname(::SymbolicIndex{<:Union{<:Pair,Symbol,Int},<:Union{Symbol,Int}}) = true
 function SII.getname(x::SymbolicVertexIndex)
     prefix = x.compidx isa Int ? :v : Symbol()
     Symbol(prefix, Symbol(x.compidx), :₊, Symbol(x.subidx))
 end
 function SII.getname(x::SymbolicEdgeIndex)
-    prefix = x.compidx isa Int ? :e : Symbol()
-    Symbol(prefix, Symbol(x.compidx), :₊, Symbol(x.subidx))
+    if x.compidx isa Pair
+        src, dst = x.compidx
+        _src = src isa Int ? Symbol(:v, src) : Symbol(src)
+        _dst = dst isa Int ? Symbol(:v, dst) : Symbol(dst)
+        Symbol(_src, "ₜₒ", _dst, :₊, Symbol(x.subidx))
+    else
+        prefix = x.compidx isa Int ? :e : Symbol()
+        Symbol(prefix, Symbol(x.compidx), :₊, Symbol(x.subidx))
+    end
 end
 
 resolvecompidx(nw::Network, sni) = resolvecompidx(nw.im, sni)
@@ -119,20 +129,49 @@ function resolvecompidx(im::IndexManager, sni::SymbolicIndex{Symbol})
         throw(ArgumentError("Could not resolve component index for $sni, the name might not be unique?"))
     end
 end
+function resolvecompidx(im::IndexManager, sni::SymbolicEdgeIndex{<:Pair})
+    src, dst = sni.compidx
+
+    src_i = try
+        resolvecompidx(im, VIndex(src))
+    catch
+        throw(ArgumentError("Could not resolve edge source $src"))
+    end
+    dst_i = try
+        resolvecompidx(im, VIndex(dst))
+    catch
+        throw(ArgumentError("Could not resolve edge destination $dst"))
+    end
+
+    eidx = findfirst(im.edgevec) do e
+        e.src == src_i && e.dst == dst_i
+    end
+    if isnothing(eidx)
+        reverse = findfirst(im.edgevec) do e
+            e.src == dst_i && e.dst == src_i
+        end
+        err = "Invalid Index: Network does not contain edge from $(src) => $(dst)!"
+        if !isnothing(reverse)
+            err *= " Maybe you meant the reverse edge from $(dst) => $(src)?"
+        end
+        throw(ArgumentError(err))
+    end
+    return eidx
+end
 getcomp(nw::Network, sni) = getcomp(nw.im, sni)
 getcomp(im::IndexManager, sni::SymbolicEdgeIndex) = im.edgem[resolvecompidx(im, sni)]
 getcomp(im::IndexManager, sni::SymbolicVertexIndex) = im.vertexm[resolvecompidx(im, sni)]
 
 getcomprange(nw::Network, sni) = getcomprange(nw.im, sni)
 getcomprange(im::IndexManager, sni::VIndex{<:Union{Symbol,Int}}) = im.v_data[resolvecompidx(im, sni)]
-getcomprange(im::IndexManager, sni::EIndex{<:Union{Symbol,Int}}) = im.e_data[resolvecompidx(im, sni)]
+getcomprange(im::IndexManager, sni::EIndex{<:Union{<:Pair,Symbol,Int}}) = im.e_data[resolvecompidx(im, sni)]
 
 getcompoutrange(nw::Network, sni) = getcompoutrange(nw.im, sni)
 getcompoutrange(im::IndexManager, sni::VIndex{<:Union{Symbol,Int}}) = im.v_out[resolvecompidx(im, sni)]
-getcompoutrange(im::IndexManager, sni::EIndex{<:Union{Symbol,Int}}) = flatrange(im.e_out[resolvecompidx(im, sni)])
+getcompoutrange(im::IndexManager, sni::EIndex{<:Union{<:Pair,Symbol,Int}}) = flatrange(im.e_out[resolvecompidx(im, sni)])
 
 getcompprange(nw::Network, sni::SymbolicVertexIndex{<:Union{Symbol,Int}}) = nw.im.v_para[resolvecompidx(nw, sni)]
-getcompprange(nw::Network, sni::SymbolicEdgeIndex{<:Union{Symbol,Int}}) = nw.im.e_para[resolvecompidx(nw, sni)]
+getcompprange(nw::Network, sni::SymbolicEdgeIndex{<:Union{<:Pair,Symbol,Int}}) = nw.im.e_para[resolvecompidx(nw, sni)]
 
 subsym_has_idx(sym::Symbol, syms) = sym ∈ syms
 subsym_has_idx(idx::Int, syms) = 1 ≤ idx ≤ length(syms)
@@ -143,7 +182,7 @@ subsym_to_idx(idx::Int, _) = idx
 #### Iterator/Broadcast interface for ArraySymbolic types
 ####
 # TODO: not broadcasting over idx with colon is weird
-Base.broadcastable(si::SymbolicIndex{<:Union{Int,Symbol,Colon},<:Union{Int,Symbol,Colon}}) = Ref(si)
+Base.broadcastable(si::SymbolicIndex{<:Union{Int,Symbol,<:Pair,Colon},<:Union{Int,Symbol,Colon}}) = Ref(si)
 
 const _IterableComponent = SymbolicIndex{<:Union{AbstractVector,Tuple},<:Union{Int,Symbol}}
 Base.length(si::_IterableComponent) = length(si.compidx)
@@ -166,7 +205,7 @@ function Base.iterate(si::_IterableComponent, state=nothing)
     _similar(si, it[1], si.subidx), it[2]
 end
 
-const _IterableSubcomponent = SymbolicIndex{<:Union{Symbol,Int},<:Union{AbstractVector,Tuple}}
+const _IterableSubcomponent = SymbolicIndex{<:Union{<:Pair,Symbol,Int},<:Union{AbstractVector,Tuple}}
 Base.length(si::_IterableSubcomponent) = length(si.subidx)
 Base.size(si::_IterableSubcomponent) = (length(si),)
 Base.IteratorSize(si::_IterableSubcomponent) = Base.HasShape{1}()
@@ -200,9 +239,9 @@ _resolve_colon(nw::Network, sni::EIndex{Colon}) = EIndex(1:ne(nw), sni.subidx)
 _resolve_colon(nw::Network, sni::VPIndex{Colon}) = VPIndex(1:nv(nw), sni.subidx)
 _resolve_colon(nw::Network, sni::EPIndex{Colon}) = EPIndex(1:ne(nw), sni.subidx)
 _resolve_colon(nw::Network, sni::VIndex{<:Union{Symbol,Int},Colon}) = VIndex{Int, UnitRange{Int}}(sni.compidx, 1:dim(getcomp(nw,sni)))
-_resolve_colon(nw::Network, sni::EIndex{<:Union{Symbol,Int},Colon}) = EIndex{Int, UnitRange{Int}}(sni.compidx, 1:dim(getcomp(nw,sni)))
+_resolve_colon(nw::Network, sni::EIndex{<:Union{<:Pair,Symbol,Int},Colon}) = EIndex{Int, UnitRange{Int}}(sni.compidx, 1:dim(getcomp(nw,sni)))
 _resolve_colon(nw::Network, sni::VPIndex{<:Union{Symbol,Int},Colon}) = VPIndex{Int, UnitRange{Int}}(sni.compidx, 1:pdim(getcomp(nw,sni)))
-_resolve_colon(nw::Network, sni::EPIndex{<:Union{Symbol,Int},Colon}) = EPIndex{Int, UnitRange{Int}}(sni.compidx, 1:pdim(getcomp(nw,sni)))
+_resolve_colon(nw::Network, sni::EPIndex{<:Union{<:Pair,Symbol,Int},Colon}) = EPIndex{Int, UnitRange{Int}}(sni.compidx, 1:pdim(getcomp(nw,sni)))
 
 
 #### Implmentation of index provider interface
@@ -220,7 +259,7 @@ SII.all_symbols(nw::Network) = vcat(SII.all_variable_symbols(nw), SII.parameter_
 ####
 #### variable indexing
 ####
-const POTENTIAL_SCALAR_SIDX = Union{SymbolicStateIndex{<:Union{Symbol,Int},<:Union{Int,Symbol}}}
+const POTENTIAL_SCALAR_SIDX = Union{SymbolicStateIndex{<:Union{<:Pair,Symbol,Int},<:Union{Int,Symbol}}}
 function SII.is_variable(nw::Network, sni)
     if _hascolon(sni)
         SII.is_variable(nw, _resolve_colon(nw,sni))
@@ -277,8 +316,8 @@ end
 ####
 # when using an number instead of symbol only PIndex is valid
 const POTENTIAL_SCALAR_PIDX = Union{
-    SymbolicParameterIndex{<:Union{Symbol,Int},<:Union{Int,Symbol}},
-    SymbolicIndex{<:Union{Symbol,Int},Symbol}
+    SymbolicParameterIndex{<:Union{<:Pair,Symbol,Int},<:Union{Int,Symbol}},
+    SymbolicIndex{<:Union{<:Pair,Symbol,Int},Symbol}
 }
 function SII.is_parameter(nw::Network, sni)
     if _hascolon(sni)
@@ -410,7 +449,7 @@ function SII.is_observed(nw::Network, sni)
     end
 end
 _is_observed(nw::Network, _) = false
-function _is_observed(nw::Network, sni::SymbolicStateIndex{<:Union{Symbol,Int},Symbol})
+function _is_observed(nw::Network, sni::SymbolicStateIndex{<:Union{<:Pair,Symbol,Int},Symbol})
     cf = getcomp(nw, sni)
     return sni.subidx ∈ obssym_all(cf)
 end
@@ -1212,7 +1251,7 @@ Base.getindex(s::NWState, idx::ObservableExpression) = SII.getu(s, idx)(s)
 Base.getindex(s::NWParameter, idx::ObservableExpression) = SII.getp(s, idx)(s)
 
 # using getindex to access component models
-function Base.getindex(nw::Network, i::EIndex{<:Union{Symbol,Int}, Nothing})
+function Base.getindex(nw::Network, i::EIndex{<:Union{<:Pair,Symbol,Int}, Nothing})
     return getcomp(nw, i)
 end
 function Base.getindex(nw::Network, i::VIndex{<:Union{Symbol,Int}, Nothing})

test/symbolicindexing_test.jl

@@ -552,3 +552,66 @@ end
     b = @b $obsf2($(rand(dim(nw))), $(rand(pdim(nw))), NaN, $(zeros(length(idxs2)))) # 17ns 0 allocs
     @test b.allocs == 0
 end
+
+@testset "test edge indexing with Pair syntax" begin
+    # Create a simple network with named vertices for testing
+    v1 = Lib.kuramoto_second(name=:v1)
+    v2 = Lib.kuramoto_second(name=:v2)
+    v3 = Lib.kuramoto_second(name=:v3)
+    e12 = Lib.kuramoto_edge(name=:e12)
+    e23 = Lib.kuramoto_edge(name=:e23)
+    g = path_graph(3)
+    nw = Network(g, [v1, v2, v3], [e12, e23])
+
+    # Basic pair syntax tests
+    @test SII.is_observed(nw, EIndex(1=>2, :P))
+    @test SII.is_observed(nw, EIndex(:v1=>:v2, :P))
+    @test SII.is_parameter(nw, EPIndex(1=>2, :K))
+    @test SII.is_parameter(nw, EPIndex(:v1=>:v2, :K))
+
+    # Test that pairs resolve to regular indices
+    @test SII.parameter_index(nw, EPIndex(1=>2, :K)) == SII.parameter_index(nw, EPIndex(1, :K))
+    @test SII.parameter_index(nw, EPIndex(:v1=>:v2, :K)) == SII.parameter_index(nw, EPIndex(1, :K))
+    @test SII.parameter_index(nw, EPIndex(:v1=>2, :K)) == SII.parameter_index(nw, EPIndex(1, :K))
+
+    # Error cases
+    @test_throws ArgumentError SII.parameter_index(nw, EPIndex(1=>3, :K))  # no direct edge
+    @test_throws ArgumentError SII.parameter_index(nw, EPIndex(:nonexistent=>:v2, :K))  # invalid vertex
+
+    # Test with solution and state objects
+    u0 = rand(dim(nw))
+    p = rand(pdim(nw))
+    prob = ODEProblem(nw, u0, (0.0, 1.0), p)
+    sol = solve(prob, Tsit5())
+    s = NWState(nw, u0, p)
+
+    # Solution indexing
+    @test sol([0.1, 0.5], idxs=EIndex(1=>2, :P)).u ≈ sol([0.1, 0.5], idxs=EIndex(1, :P)).u
+    @test sol([0.1, 0.5], idxs=EIndex(:v1=>:v2, :P)).u ≈ sol([0.1, 0.5], idxs=EIndex(1, :P)).u
+
+    # State access
+    @test s[EIndex(1=>2, :P)] == s[EIndex(1, :P)]
+    @test s[EIndex(:v1=>:v2, :P)] == s[EIndex(1, :P)]
+    @test s[EPIndex(1=>2, :K)] == s[EPIndex(1, :K)]
+    @test s[EPIndex(:v1=>:v2, :K)] == s[EPIndex(1, :K)]
+
+    # Proxy syntax
+    @test s.e[1=>2, :P] == s[EIndex(1, :P)]
+    @test s.e[:v1=>:v2, :P] == s[EIndex(1, :P)]
+    s.p.e[1=>2, :K] = 2.71
+    @test s.p.e[1, :K] == 2.71
+    @test s.p[EPIndex(1=>2, :K)] == 2.71
+
+    # Multiple edge access
+    edges_int = [EIndex(1=>2, :P), EIndex(2=>3, :P)]
+    edges_named = [EIndex(:v1=>:v2, :P), EIndex(:v2=>:v3, :P)]
+    edges_regular = [EIndex(1, :P), EIndex(2, :P)]
+    @test s[edges_int] == s[edges_regular]
+    @test s[edges_named] == s[edges_regular]
+
+    # naming
+    @test SII.getname(EIndex(1=>2, :P)) == :v1ₜₒv2₊P
+    @test SII.getname(EIndex(:a=>:b, :P)) == :aₜₒb₊P
+    @test SII.getname(EIndex(:a=>2, :P)) == :aₜₒv2₊P
+    @test SII.getname(EIndex(1=>:b, :P)) == :v1ₜₒb₊P
+end