Add RecursiveSet and DuplicateVector, better benchmarking of Brus…

…selator (#50)

* Add `RecursiveSet`

* Put set types in subfolders

* Add `DuplicateVector`

* Make Brusselator into callable struct

benchmark/benchmark.jl

@@ -1,40 +1,33 @@
-using ADTypes
-using ADTypes: AbstractSparsityDetector
+using ADTypes: AbstractSparsityDetector, jacobian_sparsity
 using BenchmarkTools
 using SparseConnectivityTracer
-using SparseConnectivityTracer: SortedVector
+using SparseConnectivityTracer: DuplicateVector, RecursiveSet, SortedVector
 using NNlib: conv
 
 include("../test/brusselator_definition.jl")
 
 const METHODS = (
     TracerSparsityDetector(BitSet),
     TracerSparsityDetector(Set{UInt64}),
+    TracerSparsityDetector(DuplicateVector{UInt64}),
+    TracerSparsityDetector(RecursiveSet{UInt64}),
     TracerSparsityDetector(SortedVector{UInt64}),
 )
 
 function benchmark_brusselator(N::Integer, method::AbstractSparsityDetector)
-    dims = (N, N, 2)
-    A = 1.0
-    B = 1.0
-    alpha = 1.0
-    xyd = fill(1.0, N)
-    dx = 1.0
-    p = (A, B, alpha, xyd, dx, N)
-
-    u = rand(dims...)
-    du = similar(u)
-    f!(du, u) = brusselator_2d_loop(du, u, p, nothing)
-
-    return @benchmark ADTypes.jacobian_sparsity($f!, $du, $u, $method)
+    f! = Brusselator!(N)
+    x = rand(N, N, 2)
+    y = similar(x)
+
+    return @benchmark jacobian_sparsity($f!, $y, $x, $method)
 end
 
 function benchmark_conv(N, method::AbstractSparsityDetector)
     x = rand(N, N, 3, 1) # WHCN image 
     w = rand(5, 5, 3, 2)  # corresponds to Conv((5, 5), 3 => 2)
     f(x) = conv(x, w)
 
-    return @benchmark ADTypes.jacobian_sparsity($f, $x, $method)
+    return @benchmark jacobian_sparsity($f, $x, $method)
 end
 
 ## Run Brusselator benchmarks

docs/src/api.md

@@ -28,8 +28,10 @@ JacobianTracer
 HessianTracer
 ```
 
-We also define a custom alternative to sets that can deliver faster `union`:
+We also define alternative pseudo-set types that can deliver faster `union`:
 
 ```@docs
+SparseConnectivityTracer.DuplicateVector
+SparseConnectivityTracer.RecursiveSet
 SparseConnectivityTracer.SortedVector
 ```

references/show/HessianTracer_BitSet.txt

@@ -0,0 +1,3 @@
+HessianTracer{BitSet}(
+  2 => (),
+)

references/show/HessianTracer_DuplicateVector{UInt64}.txt

@@ -0,0 +1,3 @@
+HessianTracer{DuplicateVector{UInt64}}(
+  2 => (),
+)

references/show/HessianTracer_RecursiveSet{UInt64}.txt

@@ -0,0 +1,3 @@
+HessianTracer{RecursiveSet{UInt64}}(
+  2 => (),
+)

references/show/HessianTracer_Set{UInt64}.txt

@@ -0,0 +1,3 @@
+HessianTracer{Set{UInt64}}(
+  2 => (),
+)

references/show/HessianTracer_SortedVector{UInt64}.txt

@@ -0,0 +1,3 @@
+HessianTracer{SortedVector{UInt64}}(
+  2 => (),
+)

src/SparseConnectivityTracer.jl

@@ -4,7 +4,6 @@ using ADTypes: ADTypes
 import SparseArrays: sparse
 import Random: rand, AbstractRNG, SamplerType
 
-include("sortedvector.jl")
 include("tracers.jl")
 include("conversion.jl")
 include("operators.jl")
@@ -14,6 +13,10 @@ include("overload_hessian.jl")
 include("pattern.jl")
 include("adtypes.jl")
 
+include("settypes/duplicatevector.jl")
+include("settypes/recursiveset.jl")
+include("settypes/sortedvector.jl")
+
 export ConnectivityTracer, connectivity_pattern
 export JacobianTracer, jacobian_pattern
 export HessianTracer, hessian_pattern

src/pattern.jl

@@ -203,7 +203,7 @@ function hessian_pattern_to_mat(
     V = Bool[]   # values
 
     for i in keys(yt.inputs)
-        for j in yt.inputs[i]
+        for j in inputs(yt, i)
             push!(I, i)
             push!(J, j)
             push!(V, true)

src/settypes/duplicatevector.jl

@@ -0,0 +1,56 @@
+"""
+    DuplicateVector
+
+Vector that can have duplicate values, for which union is just concatenation.
+"""
+struct DuplicateVector{T<:Number}
+    data::Vector{T}
+
+    DuplicateVector{T}(data::AbstractVector{T}) where {T} = new{T}(convert(Vector{T}, data))
+    DuplicateVector{T}(x::Number) where {T} = new{T}([convert(T, x)])
+    DuplicateVector{T}() where {T} = new{T}(T[])
+end
+
+Base.eltype(::Type{DuplicateVector{T}}) where {T} = T
+
+function Base.union(dv1::DuplicateVector{T}, dv2::DuplicateVector{T}) where {T}
+    return DuplicateVector{T}(vcat(dv1.data, dv2.data))
+end
+
+Base.collect(dv::DuplicateVector) = collect(Set(dv.data))
+
+## SCT tricks
+
+function keys2set(::Type{S}, d::Dict{I}) where {I<:Integer,S<:DuplicateVector{I}}
+    return S(collect(keys(d)))
+end
+
+const EMPTY_CONNECTIVITY_TRACER_DV_U16 = ConnectivityTracer(DuplicateVector{UInt16}())
+const EMPTY_CONNECTIVITY_TRACER_DV_U32 = ConnectivityTracer(DuplicateVector{UInt32}())
+const EMPTY_CONNECTIVITY_TRACER_DV_U64 = ConnectivityTracer(DuplicateVector{UInt64}())
+
+const EMPTY_JACOBIAN_TRACER_DV_U16 = JacobianTracer(DuplicateVector{UInt16}())
+const EMPTY_JACOBIAN_TRACER_DV_U32 = JacobianTracer(DuplicateVector{UInt32}())
+const EMPTY_JACOBIAN_TRACER_DV_U64 = JacobianTracer(DuplicateVector{UInt64}())
+
+const EMPTY_HESSIAN_TRACER_DV_U16 = HessianTracer(Dict{UInt16,DuplicateVector{UInt16}}())
+const EMPTY_HESSIAN_TRACER_DV_U32 = HessianTracer(Dict{UInt32,DuplicateVector{UInt32}}())
+const EMPTY_HESSIAN_TRACER_DV_U64 = HessianTracer(Dict{UInt64,DuplicateVector{UInt64}}())
+
+function empty(::Type{ConnectivityTracer{DuplicateVector{UInt16}}})
+    return EMPTY_CONNECTIVITY_TRACER_DV_U16
+end
+function empty(::Type{ConnectivityTracer{DuplicateVector{UInt32}}})
+    return EMPTY_CONNECTIVITY_TRACER_DV_U32
+end
+function empty(::Type{ConnectivityTracer{DuplicateVector{UInt64}}})
+    return EMPTY_CONNECTIVITY_TRACER_DV_U64
+end
+
+empty(::Type{JacobianTracer{DuplicateVector{UInt16}}}) = EMPTY_JACOBIAN_TRACER_DV_U16
+empty(::Type{JacobianTracer{DuplicateVector{UInt32}}}) = EMPTY_JACOBIAN_TRACER_DV_U32
+empty(::Type{JacobianTracer{DuplicateVector{UInt64}}}) = EMPTY_JACOBIAN_TRACER_DV_U64
+
+empty(::Type{HessianTracer{DuplicateVector{UInt16},UInt16}}) = EMPTY_HESSIAN_TRACER_DV_U16
+empty(::Type{HessianTracer{DuplicateVector{UInt32},UInt32}}) = EMPTY_HESSIAN_TRACER_DV_U32
+empty(::Type{HessianTracer{DuplicateVector{UInt64},UInt64}}) = EMPTY_HESSIAN_TRACER_DV_U64

src/settypes/recursiveset.jl

@@ -0,0 +1,94 @@
+"""
+    RecursiveSet
+
+Lazy union of sets.
+"""
+struct RecursiveSet{T<:Number}
+    s::Union{Nothing,Set{T}}
+    child1::Union{Nothing,RecursiveSet{T}}
+    child2::Union{Nothing,RecursiveSet{T}}
+
+    function RecursiveSet{T}(s) where {T}
+        return new{T}(Set{T}(s), nothing, nothing)
+    end
+
+    function RecursiveSet{T}(x::Number) where {T}
+        return new{T}(Set{T}(convert(T, x)), nothing, nothing)
+    end
+
+    function RecursiveSet{T}() where {T}
+        return new{T}(Set{T}(), nothing, nothing)
+    end
+
+    function RecursiveSet{T}(rs1::RecursiveSet{T}, rs2::RecursiveSet{T}) where {T}
+        return new{T}(nothing, rs1, rs2)
+    end
+end
+
+function print_recursiveset(io::IO, rs::RecursiveSet{T}; offset) where {T}
+    if !isnothing(rs.s)
+        print(io, "RecursiveSet{$T} containing $(rs.s)")
+    else
+        print(io, "RecursiveSet{$T} with two children:")
+        print(io, "\n  ", " "^offset, "1: ")
+        print_recursiveset(io, rs.child1; offset=offset + 2)
+        print(io, "\n  ", " "^offset, "2: ")
+        print_recursiveset(io, rs.child2; offset=offset + 2)
+    end
+end
+
+function Base.show(io::IO, rs::RecursiveSet{T}) where {T}
+    return print_recursiveset(io, rs; offset=0)
+end
+
+Base.eltype(::Type{RecursiveSet{T}}) where {T} = T
+
+function Base.union(rs1::RecursiveSet{T}, rs2::RecursiveSet{T}) where {T}
+    return RecursiveSet{T}(rs1, rs2)
+end
+
+function Base.collect(rs::RecursiveSet{T}) where {T}
+    accumulator = Set{T}()
+    collect_aux!(accumulator, rs)
+    return collect(accumulator)
+end
+
+function collect_aux!(accumulator::Set{T}, rs::RecursiveSet{T})::Nothing where {T}
+    if !isnothing(rs.s)
+        union!(accumulator, rs.s::Set{T})
+    else
+        collect_aux!(accumulator, rs.child1::RecursiveSet{T})
+        collect_aux!(accumulator, rs.child2::RecursiveSet{T})
+    end
+    return nothing
+end
+
+## SCT tricks
+
+function keys2set(::Type{S}, d::Dict{I}) where {I<:Integer,S<:RecursiveSet{I}}
+    return S(keys(d))
+end
+
+const EMPTY_CONNECTIVITY_TRACER_RS_U16 = ConnectivityTracer(RecursiveSet{UInt16}())
+const EMPTY_CONNECTIVITY_TRACER_RS_U32 = ConnectivityTracer(RecursiveSet{UInt32}())
+const EMPTY_CONNECTIVITY_TRACER_RS_U64 = ConnectivityTracer(RecursiveSet{UInt64}())
+
+const EMPTY_JACOBIAN_TRACER_RS_U16 = JacobianTracer(RecursiveSet{UInt16}())
+const EMPTY_JACOBIAN_TRACER_RS_U32 = JacobianTracer(RecursiveSet{UInt32}())
+const EMPTY_JACOBIAN_TRACER_RS_U64 = JacobianTracer(RecursiveSet{UInt64}())
+
+const EMPTY_HESSIAN_TRACER_RS_U16 = HessianTracer(Dict{UInt16,RecursiveSet{UInt16}}())
+const EMPTY_HESSIAN_TRACER_RS_U32 = HessianTracer(Dict{UInt32,RecursiveSet{UInt32}}())
+const EMPTY_HESSIAN_TRACER_RS_U64 = HessianTracer(Dict{UInt64,RecursiveSet{UInt64}}())
+
+empty(::Type{ConnectivityTracer{RecursiveSet{UInt16}}}) = EMPTY_CONNECTIVITY_TRACER_RS_U16
+empty(::Type{ConnectivityTracer{RecursiveSet{UInt32}}}) = EMPTY_CONNECTIVITY_TRACER_RS_U32
+empty(::Type{ConnectivityTracer{RecursiveSet{UInt64}}}) = EMPTY_CONNECTIVITY_TRACER_RS_U64
+
+empty(::Type{JacobianTracer{RecursiveSet{UInt16}}}) = EMPTY_JACOBIAN_TRACER_RS_U16
+empty(::Type{JacobianTracer{RecursiveSet{UInt32}}}) = EMPTY_JACOBIAN_TRACER_RS_U32
+empty(::Type{JacobianTracer{RecursiveSet{UInt64}}}) = EMPTY_JACOBIAN_TRACER_RS_U64
+
+empty(::Type{HessianTracer{RecursiveSet{UInt16},UInt16}}) = EMPTY_HESSIAN_TRACER_RS_U16
+empty(::Type{HessianTracer{RecursiveSet{UInt32},UInt32}}) = EMPTY_HESSIAN_TRACER_RS_U32
+empty(::Type{HessianTracer{RecursiveSet{UInt64},UInt64}}) = EMPTY_HESSIAN_TRACER_RS_U64

src/sortedvector.jl → src/settypes/sortedvector.jl

@@ -1,23 +1,7 @@
 """
     SortedVector
 
-A wrapper for sorted vectors, designed for fast unions.
-
-# Constructor
-
-    SortedVector(data::AbstractVector; sorted=false)
-
-# Example
-
-```jldoctest
-x = SortedVector([3, 4, 2])
-x = SortedVector([1, 3, 5]; sorted=true)
-z = union(x, y)
-
-# output
-
-SortedVector([1, 2, 3, 4, 5])
-````
+Sorted vector without duplicates, designed for fast set unions with merging.
 """
 struct SortedVector{T<:Number} <: AbstractVector{T}
     data::Vector{T}
@@ -40,7 +24,6 @@ function Base.convert(::Type{SortedVector{T}}, v::Vector{T}) where {T}
     return SortedVector{T}(v; sorted=false)
 end
 
-Base.eltype(::SortedVector{T}) where {T} = T
 Base.size(v::SortedVector) = size(v.data)
 Base.getindex(v::SortedVector, i) = v.data[i]
 Base.IndexStyle(::Type{SortedVector{T}}) where {T} = IndexStyle(Vector{T})
@@ -80,3 +63,33 @@ function Base.union(v1::SortedVector{T}, v2::SortedVector{T}) where {T}
     resize!(result, result_index - 1)
     return SortedVector{T}(result; sorted=true)
 end
+
+## SCT tricks
+
+function keys2set(::Type{S}, d::Dict{I}) where {I<:Integer,S<:SortedVector{I}}
+    return S(collect(keys(d)); sorted=false)
+end
+
+const EMPTY_CONNECTIVITY_TRACER_SV_U16 = ConnectivityTracer(SortedVector{UInt16}())
+const EMPTY_CONNECTIVITY_TRACER_SV_U32 = ConnectivityTracer(SortedVector{UInt32}())
+const EMPTY_CONNECTIVITY_TRACER_SV_U64 = ConnectivityTracer(SortedVector{UInt64}())
+
+const EMPTY_JACOBIAN_TRACER_SV_U16 = JacobianTracer(SortedVector{UInt16}())
+const EMPTY_JACOBIAN_TRACER_SV_U32 = JacobianTracer(SortedVector{UInt32}())
+const EMPTY_JACOBIAN_TRACER_SV_U64 = JacobianTracer(SortedVector{UInt64}())
+
+const EMPTY_HESSIAN_TRACER_SV_U16 = HessianTracer(Dict{UInt16,SortedVector{UInt16}}())
+const EMPTY_HESSIAN_TRACER_SV_U32 = HessianTracer(Dict{UInt32,SortedVector{UInt32}}())
+const EMPTY_HESSIAN_TRACER_SV_U64 = HessianTracer(Dict{UInt64,SortedVector{UInt64}}())
+
+empty(::Type{ConnectivityTracer{SortedVector{UInt16}}}) = EMPTY_CONNECTIVITY_TRACER_SV_U16
+empty(::Type{ConnectivityTracer{SortedVector{UInt32}}}) = EMPTY_CONNECTIVITY_TRACER_SV_U32
+empty(::Type{ConnectivityTracer{SortedVector{UInt64}}}) = EMPTY_CONNECTIVITY_TRACER_SV_U64
+
+empty(::Type{JacobianTracer{SortedVector{UInt16}}}) = EMPTY_JACOBIAN_TRACER_SV_U16
+empty(::Type{JacobianTracer{SortedVector{UInt32}}}) = EMPTY_JACOBIAN_TRACER_SV_U32
+empty(::Type{JacobianTracer{SortedVector{UInt64}}}) = EMPTY_JACOBIAN_TRACER_SV_U64
+
+empty(::Type{HessianTracer{SortedVector{UInt16},UInt16}}) = EMPTY_HESSIAN_TRACER_SV_U16
+empty(::Type{HessianTracer{SortedVector{UInt32},UInt32}}) = EMPTY_HESSIAN_TRACER_SV_U32
+empty(::Type{HessianTracer{SortedVector{UInt64},UInt64}}) = EMPTY_HESSIAN_TRACER_SV_U64