create batches based on hash for faster comparison

Author: hexaeder

src/callbacks.jl

@@ -235,7 +235,7 @@ function wrap_component_callbacks(nw)
     end
     # group the callbacks such that they are in groups which are "batchequal"
     # batchequal groups can be wrapped into a single callback
-    idx_per_type = _find_identical(callbacks, 1:length(components))
+    idx_per_type = find_identical(callbacks; equality=_batchequal)
     batches = []
     for typeidx in idx_per_type
         batchcomps = components[typeidx]
@@ -251,21 +251,22 @@ function wrap_component_callbacks(nw)
     end
     return batches
 end
-function batchequal(a::ContinousComponentCallback, b::ContinousComponentCallback)
-    batchequal(a.condition, b.condition) || return false
-    batchequal(a.kwargs, b.kwargs)       || return false
+_batchequal(a, b) = false
+function _batchequal(a::ContinousComponentCallback, b::ContinousComponentCallback)
+    _batchequal(a.condition, b.condition) || return false
+    _batchequal(a.kwargs, b.kwargs)       || return false
     return true
 end
-function batchequal(a::VectorContinousComponentCallback, b::VectorContinousComponentCallback)
-    batchequal(a.condition, b.condition) || return false
-    batchequal(a.kwargs, b.kwargs)       || return false
+function _batchequal(a::VectorContinousComponentCallback, b::VectorContinousComponentCallback)
+    _batchequal(a.condition, b.condition) || return false
+    _batchequal(a.kwargs, b.kwargs)       || return false
     a.len == b.len                       || return false
     return true
 end
-function batchequal(a::T, b::T) where {T <: Union{ComponentCondition, ComponentAffect}}
+function _batchequal(a::T, b::T) where {T <: Union{ComponentCondition, ComponentAffect}}
     typeof(a) == typeof(b)
 end
-function batchequal(a::NamedTuple, b::NamedTuple)
+function _batchequal(a::NamedTuple, b::NamedTuple)
     length(a) == length(b) || return false
     for (k, v) in a
         haskey(b, k) || return false

src/component_functions.jl

@@ -579,19 +579,6 @@ dispatchT(::T) where {T<:ComponentModel} = dispatchT(T)
 dispatchT(T::Type{<:VertexModel}) = VertexModel
 dispatchT(T::Type{<:EdgeModel}) = EdgeModel
 
-# TODO: introduce batchequal hash for faster batching of component models
-batchequal(a, b) = false
-function batchequal(a::ComponentModel, b::ComponentModel)
-    compf(a) === compf(b)  || return false
-    compg(a) === compg(b)  || return false
-    fftype(a) == fftype(b) || return false
-    dim(a)    == dim(b)    || return false
-    outdim(a) == outdim(b) || return false
-    pdim(a)   == pdim(b)   || return false
-    extdim(a) == extdim(b) || return false
-    return true
-end
-
 """
     _construct_comp(::Type{T}, kwargs) where {T}
 

src/construction.jl

@@ -128,10 +128,18 @@ function Network(g::AbstractGraph,
 
         # batch identical edge and vertex model
         @timeit_debug "batch identical vertexes" begin
-            vidxs = _find_identical(vertexm, 1:nv(g))
+            vidxs = if all_same_v
+                [collect(1:nv(g))]
+            else
+                _find_identical_components(_vertexm)
+            end
         end
         @timeit_debug "batch identical edges" begin
-            eidxs = _find_identical(edgem, 1:ne(g))
+            eidxs = if all_same_e
+                [collect(1:ne(g))]
+            else
+                _find_identical_components(_edgem)
+            end
         end
 
         # create vertex batches and initialize with index manager
@@ -198,6 +206,26 @@ function Network(g::AbstractGraph,
     return nw
 end
 
+function _find_identical_components(models)
+    # identical components are based on identical _component_hash
+    # those can have different metadata but are considered identical when in comes to batching
+    hashs = _component_hash.(models)
+    find_identical(hashs)
+end
+# hash condition: components with same hash will end up in the same batch
+function _component_hash(c::ComponentModel)
+    hash((
+        typeof(c), # same type
+        compf(c),  # same f-function
+        compg(c),  # same g-function
+        fftype(c), # same feedforward type
+        dim(c),    # same state dimension
+        outdim(c), # same output dimension
+        pdim(c),   # same parameter dimension
+        extdim(c), # same external input dimension
+    ))
+end
+
 function Network(vertexfs, edgefs; kwargs...)
     vertexfs = vertexfs isa VertexModel ? [vertexfs] : vertexfs
     edgefs   = edgefs isa EdgeModel     ? [edgefs]   : edgefs
@@ -324,28 +352,6 @@ function batch_by_idxs(v::AbstractVector, batches::Vector{Vector{Int}})
     [v[batch] for batch in batches]
 end
 
-_find_identical(v::ComponentModel, indices) = [collect(indices)]
-function _find_identical(v::Vector{T}, indices) where {T<:ComponentModel}
-    idxs_per_type = Vector{Int}[]
-    unique_comp = T[]
-    for i in eachindex(v)
-        found = false
-        for j in eachindex(unique_comp)
-            if batchequal(v[i], unique_comp[j])
-                found = true
-                push!(idxs_per_type[j], indices[i])
-                break
-            end
-        end
-        if !found
-            push!(unique_comp, v[i])
-            push!(idxs_per_type, [indices[i]])
-        end
-    end
-    @assert length(unique_comp) == length(idxs_per_type)
-    return idxs_per_type
-end
-
 function construct_mass_matrix(im; type=nothing)
     vertexd = filter(pairs(im.vertexm)) do (_, c)
         hasproperty(c,:mass_matrix) && c.mass_matrix != LinearAlgebra.I

src/utils.jl

@@ -140,3 +140,31 @@ abstract type SymbolicParameterIndex{C,S} <: SymbolicIndex{C,S} end
 
 flatten_sym(v::NamedTuple) = reduce(vcat, values(v))
 flatten_sym(v::AbstractVector{Symbol}) = v
+
+"""
+    find_identical(v::Vector;; equality)
+
+Find identical elements in a vector `v` using the `equality` function.
+Returns a vector of vectors where each vector contains the indices of identical elements.
+"""
+function find_identical(v::Vector{T}; equality=isequal) where {T}
+    indices = eachindex(v)
+    idxs_per_type = Vector{Int}[]
+    unique_comp = T[]
+    for i in eachindex(v)
+        found = false
+        for j in eachindex(unique_comp)
+            if equality(v[i], unique_comp[j])
+                found = true
+                push!(idxs_per_type[j], indices[i])
+                break
+            end
+        end
+        if !found
+            push!(unique_comp, v[i])
+            push!(idxs_per_type, [indices[i]])
+        end
+    end
+    @assert length(unique_comp) == length(idxs_per_type)
+    return idxs_per_type
+end

test/utils_test.jl

@@ -37,23 +37,23 @@ using NetworkDynamics
     end
 
     @testset "find_identical" begin
-        using NetworkDynamics: _find_identical
+        using NetworkDynamics: _find_identical_components
 
         v1 = Lib.kuramoto_second()
-        @test _find_identical(v1, 1:10) == [collect(1:10)]
+        @test _find_identical_components([v1 for _ in 1:10]) == [collect(1:10)]
         v2 = Lib.diffusion_vertex()
         v3 = Lib.kuramoto_first()
 
         # v2 and v3 are equal when it comes to the function!!
         vs = [v1,v2,v3,v2,v2,v1,v1,v3]
 
-        @test _find_identical(vs, eachindex(vs)) == [[1,6,7],[2,4,5],[3,8]]
+        @test _find_identical_components(vs) == [[1,6,7],[2,4,5],[3,8]]
 
         es = [Lib.diffusion_edge(),
               Lib.diffusion_edge_closure(),
               Lib.diffusion_edge_closure(),
               Lib.diffusion_edge_fid()]
-        @test _find_identical(es, eachindex(es)) == [[1], [2], [3], [4]]
+        @test _find_identical_components(es) == [[1], [2], [3], [4]]
     end
 
     @testset "algin strings" begin