Add TranslateVarArray for simplified indexing

src/SparseVariables.jl

@@ -10,6 +10,8 @@ include("dictionaries.jl")
 include("indexedarray.jl")
 include("tables.jl")
 
+include("translate.jl")
+
 export SparseArray
 export IndexedVarArray
 export insertvar!

src/translate.jl

@@ -0,0 +1,156 @@
+"""
+    TranslateVarArray{V,N,T}
+
+    Structure for holding an optimization variable with a sparse structure with extra indexing
+    Translate from abstract types to type stable id for performance in dictionaries
+"""
+struct TranslateVarArray{V<:AbstractVariableRef,N,T} <: AbstractSparseArray{V,N}
+    f::Function
+    data::Dictionary{T,V}
+    index_names::NamedTuple
+    index_cache::Vector{Dictionary}
+end
+
+"""
+    translate
+    Function to translate from type instance to type stable (e.g. Int or String) for performance in dictionaries.
+    Extend this for types (e.g. Node or Link types) to improve performance
+"""
+function translate(x)
+    return x
+end
+function translate(::Type{T}) where {T}
+    return T
+end
+
+_data(sa::TranslateVarArray) = sa.data
+
+"""
+    insertvar!(var::IndexedVarArray{V,N,T}, index...)
+
+Insert a new variable with the given index only after checking if keys are valid and not already defined.
+"""
+function insertvar!(var::TranslateVarArray{V,N,T}, index...) where {V,N,T}
+    return insertvar!(var, SafeInsert(), index...)
+end
+function insertvar!(
+    var::TranslateVarArray{V,N,T},
+    ::SafeInsert = SafeInsert(),
+    index...,
+) where {V,N,T}
+    tindex = translate.(index)
+    !valid_index(var, index) && throw(BoundsError(var, index))# "Not a valid index for $(var.name): $index"g
+    already_defined(var, tindex) && error("$index already defined for array")
+    var[tindex] = var.f(tindex...)
+    clear_cache!(var)
+    return var[tindex]
+end
+
+# Extension for standard JuMP macros
+function Containers.container(
+    f::Function,
+    indices,
+    D::Type{TranslateVarArray},
+    names,
+)
+    iva_names = NamedTuple{tuple(names...)}(indices.prod.iterators)
+    T = Tuple{translate.(eltype.(indices.prod.iterators))...}
+    N = length(names)
+    V = first(Base.return_types(f))
+    return TranslateVarArray{V,N,T}(
+        f,
+        Dictionary{T,V}(),
+        iva_names,
+        Vector{Dictionary}(undef, 2^N),
+    )
+end
+
+@generated function _getindex(
+    sa::TranslateVarArray{T,N},
+    tpl::Tuple,
+) where {T,N}
+    lookup = true
+    slice = true
+    for t in fieldtypes(tpl)
+        if !isfixed(t)
+            lookup = false
+            if !iscolon(t)
+                slice = false
+            end
+        end
+    end
+
+    if lookup
+        return :(get(_data(sa), translate.(tpl), zero(T)))
+    elseif !slice
+        return :(retval = select(_data(sa), translate.(tpl));
+        length(retval) > 0 ? retval : zero(T))
+    else    # Return selection or zero if empty to avoid reduction of empty iterate
+        return :(retval = _select_var(sa, translate.(tpl));
+        length(retval) > 0 ? retval : zero(T))
+    end
+end
+
+function Base.firstindex(sa::TranslateVarArray, d)
+    return first(sort(sa.index_names[d]))
+end
+function Base.lastindex(sa::TranslateVarArray, d)
+    return last(sort(sa.index_names[d]))
+end
+
+function build_cache!(cache, pat, sa::TranslateVarArray{V,N,T}) where {V,N,T}
+    if isempty(cache)
+        for v in keys(sa)
+            vred = _active(v, translate.(pat))
+            nv = get!(cache, vred, T[])
+            push!(nv, v)
+        end
+    end
+    return cache
+end
+
+function _select_cached(sa::TranslateVarArray{V,N,T}, pat) where {V,N,T}
+    # TODO: Benchmark to find good cutoff-value for caching
+    # TODO: Return same type for type stability
+    tpat = translate.(pat)
+    length(_data(sa)) < 100 && return _select_gen(keys(_data(sa)), tpat)
+    cache =
+        _getcache(sa, tpat)::Dictionary{_decode_nonslices(sa, tpat),Vector{T}}
+    build_cache!(cache, tpat, sa)
+    vals = _dropslices_gen(tpat)
+    return get!(cache, vals, T[])
+end
+
+function _getcache(sa::TranslateVarArray{V,N,T}, pat::P) where {V,N,T,P}
+    t = _get_cache_index(pat)
+    if isassigned(sa.index_cache, t)
+        return sa.index_cache[t]
+    else
+        sa.index_cache[t] = Dictionary{_decode_nonslices(sa, t),Vector{T}}()
+    end
+    return sa.index_cache[t]
+end
+
+"""
+    _decode_nonslices(::IndexedVarArray{V,N,T}, ::P)
+
+Reconstruct types of a pattern from the array types and the pattern type
+"""
+@generated function _decode_nonslices(
+    ::TranslateVarArray{V,N,T},
+    ::P,
+) where {V,N,T,P}
+    fts = fieldtypes(T)
+    fts2 = fieldtypes(P)
+    t = Tuple{
+        (translate(fts[i]) for (i, v) in enumerate(fts2) if v != Colon)...,
+    }
+    return :($t)
+end
+
+function _decode_nonslices(::TranslateVarArray{V,N,T}, v::Integer) where {V,N,T}
+    fts = fieldtypes(T)
+    return Tuple{
+        (fts[i] for (i, c) in enumerate(last(bitstring(v), N)) if c == '1')...,
+    }
+end

test/runtests.jl

@@ -324,3 +324,50 @@ end
     @test sum(x) == sum(x[:, :])
     @test typeof(sum(x)) <: GenericAffExpr{Float64,MockVariableRef}
 end
+
+@testset "TranslateVarArray" begin
+    # Demo custom types
+    abstract type Node end
+    struct Source <: Node
+        id::Int
+    end
+    struct Sink <: Node
+        id::Int
+    end
+
+    # Translation to type-stable index
+    SV.translate(n::Node) = n.id
+    SV.translate(::Type{<:Node}) = Int
+
+    # Test variable construction
+    m = Model()
+    # length < 100 -> filtering (no cache)
+    @variable(
+        m,
+        x[i = Source.(1:10), j = Sink.(1:5)],
+        container = SV.TranslateVarArray
+    )
+    for i in Source.(1:10), j in Sink.(1:5)
+        insertvar!(x, i, j)
+    end
+    @test length(x) == 50
+    @test typeof(x) == SV.TranslateVarArray{VariableRef,2,Tuple{Int,Int}}
+    @test length(x[:, 1]) == 10
+    @test length(x[1, :]) == 5
+    @test length(x[1:2, :]) == 10
+
+    # length > 100 -> cache
+    @variable(
+        m,
+        y[i = Source.(1:10), j = Sink.(1:15)],
+        container = SV.TranslateVarArray
+    )
+    for i in Source.(1:10), j in Sink.(1:15)
+        insertvar!(y, i, j)
+    end
+    @test length(y) == 150
+    @test typeof(y) == SV.TranslateVarArray{VariableRef,2,Tuple{Int,Int}}
+    @test length(y[:, 1]) == 10
+    @test length(y[1, :]) == 15
+    @test length(y[1:2, :]) == 30
+end