diff --git a/src/TensorKitSectors.jl b/src/TensorKitSectors.jl index 69a5b54c..8346ee70 100644 --- a/src/TensorKitSectors.jl +++ b/src/TensorKitSectors.jl @@ -58,7 +58,7 @@ end # imports # ------- -using Base: SizeUnknown, HasLength, IsInfinite +using Base: SizeUnknown, HasLength, HasShape, IsInfinite using Base: HasEltype, EltypeUnknown using Base.Iterators: product, filter using Base: @assume_effects diff --git a/src/auxiliary.jl b/src/auxiliary.jl index 6ed53605..078c8740 100644 --- a/src/auxiliary.jl +++ b/src/auxiliary.jl @@ -19,16 +19,31 @@ function _kron_promote(A₁, B₁, sz₁, sz₂) end # Manhattan based distance enumeration: I is supposed to be one-based index -# TODO: is there any way to make this faster? # forward mapping from multidimensional to single Manhattan index @inline function num_manhattan_points(d::Int, sz::Dims{1}) return Int(sz[1] > d) end + +# when sz[k] > d for every k, then x_k < sz[k] - 1 is never violated by any split of d units among N axes (one x_k is maximally d, and sz[k] - 1 ≥ d covers that) +# so the bound drops -> calculate ways to write d as an ordered sum of N nonnegative integers, +# which is the number of compositions of d into N parts +@inline function _compositions(d::Int, ::Val{N}) where {N} + try + return binomial(d + N - 1, N - 1) + catch e + e isa OverflowError || rethrow() + return nothing # fallback to recursion + end +end @inline function num_manhattan_points(d::Int, sz::Dims{N}) where {N} d == 0 && return 1 + if N ≥ 3 && all(>(d), sz) + c = _compositions(d, Val(N)) + isnothing(c) || return c + end num = 0 - for i in 1:min(sz[1], d + 1) + for i in 1:min(sz[1], d + 1) # for N = 2 the loop is already linear num += num_manhattan_points(d - i + 1, Base.tail(sz)) end return num @@ -79,7 +94,7 @@ function manhattan_to_multidimensional_index(index::Int, sz::Dims{N}) where {N} index == 1 && return ntuple(one, Val(N)) offset = 1 d = 1 - while true + while true # reason for the boundscheck currentlayer = num_manhattan_points(d, sz) if index <= offset + currentlayer break @@ -92,3 +107,44 @@ function manhattan_to_multidimensional_index(index::Int, sz::Dims{N}) where {N} index -= offset return invertlocaloffset(d, index - 1, sz) end + +# tabulation of the mapping to avoid repeated calls to manhattan_to_multidimensional_index +# and to_manhattan_index +struct MTable{N} + multi::Vector{NTuple{N, Int}} # Manhattan index -> multi-index + lin::Array{Int, N} # multi-index -> Manhattan index +end + +function build_table(sz::Dims{N}) where {N} + n = prod(sz) + multi = Vector{NTuple{N, Int}}(undef, n) + k = 0 + for J in CartesianIndices(sz) + k += 1 + multi[k] = Tuple(J) + end + sort!(multi; by = J -> (sum(J), J)) # sort first by distance, matches isless on ProductSector + lin = Array{Int, N}(undef, sz) + for i in 1:n + lin[CartesianIndex(multi[i])] = i + end + return MTable{N}(multi, lin) +end + +const TABLES = IdDict{DataType, Any}() +const TABLE_LOCK = ReentrantLock() # dictionaries aren't thread-safe for concurrent mutation +# within TensorKit no problem: OhMyThreads allows concurrent calls into e.g. sectors(V) + +@inline function mtable(::Type{I}, sz::Dims{N}) where {I, N} + t = get(TABLES, I, nothing) + if t === nothing + t = @lock TABLE_LOCK get!(() -> build_table(sz), TABLES, I) + end + return t::MTable{N} # re-establish concrete typing for downstream +end + +# refuse to tabulate absurdly large label sets +# GradedSpace NTuple storage is only used for finite sectors +# not sure what a reasonable limit is, might be edited to smaller value +# depending on findings of NTuple vs Dict performance tests +const MAXTABLE = 2^20 diff --git a/src/product.jl b/src/product.jl index b0ed5554..9f290114 100644 --- a/src/product.jl +++ b/src/product.jl @@ -39,12 +39,29 @@ end function _size(::SectorValues{I}) where {I <: ProductSector} return map(s -> _length(values(s)), _sectors(I)) end -function Base.getindex(P::SectorValues{I}, i::Int) where {I <: ProductSector} - inds = manhattan_to_multidimensional_index(i, _size(P)) +# finite <=> the product iterator has a shape, which is exactly the condition +# under which Vect[I] uses NTuple storage (pre potential cutoff) +@inline function tabulate(P::SectorValues, sz::Dims) + return Base.IteratorSize(P) isa Union{HasLength, HasShape} && prod(sz) ≤ MAXTABLE +end +Base.@propagate_inbounds function Base.getindex(P::SectorValues{I}, i::Int) where {I <: ProductSector} + sz = _size(P) + @boundscheck checksectorindex(P, i) + inds = if tabulate(P, sz) + @inbounds mtable(I, sz).multi[i] + else + manhattan_to_multidimensional_index(i, sz) + end return I(getindex.(values.(_sectors(I)), inds)) end function findindex(P::SectorValues{I}, c::I) where {I <: ProductSector} - return to_manhattan_index(findindex.(values.(_sectors(I)), Tuple(c)), _size(P)) + J = findindex.(values.(_sectors(I)), Tuple(c)) + sz = _size(P) + return if tabulate(P, sz) + @inbounds mtable(I, sz).lin[CartesianIndex(J)] + else + to_manhattan_index(J, sz) + end end function Base.iterate(P::SectorValues{I}, i = 1) where {I <: ProductSector} diff --git a/src/sectors.jl b/src/sectors.jl index e658ba69..d2dcede4 100644 --- a/src/sectors.jl +++ b/src/sectors.jl @@ -74,16 +74,17 @@ Base.IteratorEltype(::Type{<:SectorValues}) = HasEltype() Base.eltype(::Type{SectorValues{I}}) where {I <: Sector} = I Base.values(::Type{I}) where {I <: Sector} = SectorValues{I}() -Base.@propagate_inbounds function Base.getindex( - v::SectorValues{I}, i::Int - ) where {I <: Sector} - @boundscheck begin - if Base.IteratorSize(v) === HasLength() - 1 ≤ i ≤ length(v) || throw(BoundsError(v, i)) - else - 1 ≤ i || throw(BoundsError(v, i)) - end +@inline function checksectorindex(v::SectorValues, i::Integer) + if Base.IteratorSize(v) isa Union{HasLength, HasShape} + 1 ≤ i ≤ length(v) || throw(BoundsError(v, i)) + else + 1 ≤ i || throw(BoundsError(v, i)) end + return nothing +end + +Base.@propagate_inbounds function Base.getindex(v::SectorValues{I}, i::Int) where {I <: Sector} + @boundscheck checksectorindex(v, i) for (j, c) in enumerate(v) j == i && return c end