Reduce cache footprint by decoupling degeneracy-dependent data

Project.toml

@@ -1,9 +1,10 @@
 name = "TensorKit"
 uuid = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
-authors = ["Jutho Haegeman, Lukas Devos"]
 version = "0.16.3"
+authors = ["Jutho Haegeman, Lukas Devos"]
 
 [deps]
+Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
@@ -41,6 +42,7 @@ CUDA = "5.9"
 ChainRulesCore = "1"
 ChainRulesTestUtils = "1"
 Combinatorics = "1"
+Dictionaries = "0.4"
 FiniteDifferences = "0.12"
 GPUArrays = "11.3.1"
 JET = "0.9, 0.10, 0.11"

docs/src/lib/tensors.md

@@ -97,7 +97,6 @@ In `TensorMap` instances, all data is gathered in a single `AbstractVector`, whi
 
 To obtain information about the structure of the data, you can use:
 ```@docs
-fusionblockstructure(::AbstractTensorMap)
 dim(::AbstractTensorMap)
 blocksectors(::AbstractTensorMap)
 hasblock(::AbstractTensorMap, ::Sector)

ext/TensorKitMooncakeExt/utility.jl

@@ -62,7 +62,8 @@ end
 Mooncake.tangent_type(::Type{<:VectorSpace}) = Mooncake.NoTangent
 Mooncake.tangent_type(::Type{<:HomSpace}) = Mooncake.NoTangent
 
-@zero_derivative DefaultCtx Tuple{typeof(TensorKit.fusionblockstructure), Any}
+@zero_derivative DefaultCtx Tuple{typeof(TensorKit.sectorstructure), Any}
+@zero_derivative DefaultCtx Tuple{typeof(TensorKit.degeneracystructure), Any}
 
 @zero_derivative DefaultCtx Tuple{typeof(TensorKit.select), HomSpace, Index2Tuple}
 @zero_derivative DefaultCtx Tuple{typeof(TensorKit.flip), HomSpace, Any}

src/TensorKit.jl

@@ -118,6 +118,7 @@ const TO = TensorOperations
 
 using MatrixAlgebraKit
 
+using Dictionaries: Dictionaries, Dictionary, Indices, gettoken, gettokenvalue, set!
 using LRUCache
 using OhMyThreads
 using ScopedValues
@@ -218,6 +219,7 @@ end
 # Definitions and methods for tensors
 #-------------------------------------
 # general definitions
+include("tensors/tensorstructure.jl")
 include("tensors/abstracttensor.jl")
 include("tensors/backends.jl")
 include("tensors/blockiterator.jl")

src/auxiliary/dicts.jl

@@ -263,3 +263,24 @@ function Base.:(==)(d1::SortedVectorDict, d2::SortedVectorDict)
     end
     return true
 end
+
+"""
+    Hashed(value, hashfunction = Base.hash, isequal = Base.isequal)
+
+Wrapper struct to alter the `hash` and `isequal` implementations of a given value.
+This is useful in the contexts of dictionaries, where you either want to customize the hashfunction,
+or consider various values as equal with a different notion of equality.
+"""
+struct Hashed{T, Hash, Eq}
+    val::T
+end
+
+Hashed(val, hash = Base.hash, eq = Base.isequal) = Hashed{typeof(val), hash, eq}(val)
+
+Base.parent(h::Hashed) = h.val
+
+# hash overload
+Base.hash(h::Hashed{T, Hash}, seed::UInt) where {T, Hash} = Hash(parent(h), seed)
+
+# isequal overload
+Base.isequal(h1::H, h2::H) where {Eq, H <: Hashed{<:Any, <:Any, Eq}} = Eq(parent(h1), parent(h2))

src/spaces/homspace.jl

@@ -41,8 +41,7 @@ spacetype(::Type{<:HomSpace{S}}) where {S} = S
 
 const TensorSpace{S <: ElementarySpace} = Union{S, ProductSpace{S}}
 const TensorMapSpace{S <: ElementarySpace, N₁, N₂} = HomSpace{
-    S, ProductSpace{S, N₁},
-    ProductSpace{S, N₂},
+    S, ProductSpace{S, N₁}, ProductSpace{S, N₂},
 }
 
 numout(::Type{TensorMapSpace{S, N₁, N₂}}) where {S, N₁, N₂} = N₁
@@ -62,49 +61,23 @@ end
 →(dom::VectorSpace, codom::VectorSpace) = ←(codom, dom)
 
 function Base.show(io::IO, W::HomSpace)
-    if length(W.codomain) == 1
-        print(io, W.codomain[1])
-    else
-        print(io, W.codomain)
-    end
-    print(io, " ← ")
-    return if length(W.domain) == 1
-        print(io, W.domain[1])
-    else
-        print(io, W.domain)
-    end
+    return print(
+        io,
+        numout(W) == 1 ? codomain(W)[1] : codomain(W),
+        " ← ",
+        numin(W) == 1 ? domain(W)[1] : domain(W)
+    )
 end
 
 """
-    blocksectors(W::HomSpace)
+    blocksectors(W::HomSpace) -> Indices{I}
 
-Return an iterator over the different unique coupled sector labels, i.e. the intersection
-of the different fusion outputs that can be obtained by fusing the sectors present in the
-domain, as well as from the codomain.
+Return an `Indices` of all coupled sectors for `W`. The result is cached based on the
+sector structure of `W` (ignoring degeneracy dimensions).
 
-See also [`hasblock`](@ref).
+See also [`hasblock`](@ref), [`blockstructure`](@ref).
 """
-function blocksectors(W::HomSpace)
-    sectortype(W) === Trivial &&
-        return OneOrNoneIterator(dim(domain(W)) != 0 && dim(codomain(W)) != 0, Trivial())
-
-    codom = codomain(W)
-    dom = domain(W)
-    N₁ = length(codom)
-    N₂ = length(dom)
-    I = sectortype(W)
-    if N₁ == N₂ == 0
-        return allunits(I)
-    elseif N₁ == 0
-        return filter!(isunit, collect(blocksectors(dom))) # module space cannot end in empty space
-    elseif N₂ == 0
-        return filter!(isunit, collect(blocksectors(codom)))
-    elseif N₂ <= N₁
-        return filter!(c -> hasblock(codom, c), collect(blocksectors(dom)))
-    else
-        return filter!(c -> hasblock(dom, c), collect(blocksectors(codom)))
-    end
-end
+blocksectors(W::HomSpace) = sectorstructure(W).blocksectors
 
 """
     hasblock(W::HomSpace, c::Sector)
@@ -116,27 +89,27 @@ See also [`blocksectors`](@ref).
 hasblock(W::HomSpace, c::Sector) = hasblock(codomain(W), c) && hasblock(domain(W), c)
 
 """
-    dim(W::HomSpace)
+    dim(W::HomSpace) -> Int
 
 Return the total dimension of a `HomSpace`, i.e. the number of linearly independent
 morphisms that can be constructed within this space.
 """
-function dim(W::HomSpace)
-    d = 0
-    for c in blocksectors(W)
-        d += blockdim(codomain(W), c) * blockdim(domain(W), c)
-    end
-    return d
-end
+dim(W::HomSpace) = degeneracystructure(W).totaldim
 
 dims(W::HomSpace) = (dims(codomain(W))..., dims(domain(W))...)
 
 """
-    fusiontrees(W::HomSpace)
+    fusiontrees(W::HomSpace) -> Indices{Tuple{F₁,F₂}}
 
-Return the fusiontrees corresponding to all valid fusion channels of a given `HomSpace`.
+Return an `Indices` of all valid fusion tree pairs `(f₁, f₂)` for `W`, providing a
+bijection to sequential integer positions via `gettoken`/`gettokenvalue`. The result is
+cached based on the sector structure of `W` (ignoring degeneracy dimensions), so
+`HomSpace`s that share the same sectors, dualities, and index count will reuse the same
+object.
+
+See also [`sectorstructure`](@ref), [`subblockstructure`](@ref).
 """
-fusiontrees(W::HomSpace) = fusionblockstructure(W).fusiontreelist
+fusiontrees(W::HomSpace) = sectorstructure(W).fusiontrees
 
 # Operations on HomSpaces
 # -----------------------
@@ -290,125 +263,3 @@ function removeunit(P::HomSpace, ::Val{i}) where {i}
         return codomain(P) ← removeunit(domain(P), Val(i - numout(P)))
     end
 end
-
-# Block and fusion tree ranges: structure information for building tensors
-#--------------------------------------------------------------------------
-
-# sizes, strides, offset
-const StridedStructure{N} = Tuple{NTuple{N, Int}, NTuple{N, Int}, Int}
-
-struct FusionBlockStructure{I, N, F₁, F₂}
-    totaldim::Int
-    blockstructure::SectorDict{I, Tuple{Tuple{Int, Int}, UnitRange{Int}}}
-    fusiontreelist::Vector{Tuple{F₁, F₂}}
-    fusiontreestructure::Vector{StridedStructure{N}}
-    fusiontreeindices::FusionTreeDict{Tuple{F₁, F₂}, Int}
-end
-
-function fusionblockstructuretype(W::HomSpace)
-    N₁ = length(codomain(W))
-    N₂ = length(domain(W))
-    N = N₁ + N₂
-    I = sectortype(W)
-    F₁ = fusiontreetype(I, N₁)
-    F₂ = fusiontreetype(I, N₂)
-    return FusionBlockStructure{I, N, F₁, F₂}
-end
-
-@cached function fusionblockstructure(W::HomSpace)::fusionblockstructuretype(W)
-    codom = codomain(W)
-    dom = domain(W)
-    N₁ = length(codom)
-    N₂ = length(dom)
-    I = sectortype(W)
-    F₁ = fusiontreetype(I, N₁)
-    F₂ = fusiontreetype(I, N₂)
-
-    # output structure
-    blockstructure = SectorDict{I, Tuple{Tuple{Int, Int}, UnitRange{Int}}}() # size, range
-    fusiontreelist = Vector{Tuple{F₁, F₂}}()
-    fusiontreestructure = Vector{Tuple{NTuple{N₁ + N₂, Int}, NTuple{N₁ + N₂, Int}, Int}}() # size, strides, offset
-
-    # temporary data structures
-    splittingtrees = Vector{F₁}()
-    splittingstructure = Vector{Tuple{Int, Int}}()
-
-    # main computational routine
-    blockoffset = 0
-    for c in blocksectors(W)
-        empty!(splittingtrees)
-        empty!(splittingstructure)
-
-        offset₁ = 0
-        for f₁ in fusiontrees(codom, c)
-            push!(splittingtrees, f₁)
-            d₁ = dim(codom, f₁.uncoupled)
-            push!(splittingstructure, (offset₁, d₁))
-            offset₁ += d₁
-        end
-        blockdim₁ = offset₁
-        strides = (1, blockdim₁)
-
-        offset₂ = 0
-        for f₂ in fusiontrees(dom, c)
-            s₂ = f₂.uncoupled
-            d₂ = dim(dom, s₂)
-            for (f₁, (offset₁, d₁)) in zip(splittingtrees, splittingstructure)
-                push!(fusiontreelist, (f₁, f₂))
-                totaloffset = blockoffset + offset₂ * blockdim₁ + offset₁
-                subsz = (dims(codom, f₁.uncoupled)..., dims(dom, f₂.uncoupled)...)
-                @assert !any(isequal(0), subsz)
-                substr = _subblock_strides(subsz, (d₁, d₂), strides)
-                push!(fusiontreestructure, (subsz, substr, totaloffset))
-            end
-            offset₂ += d₂
-        end
-        blockdim₂ = offset₂
-        blocksize = (blockdim₁, blockdim₂)
-        blocklength = blockdim₁ * blockdim₂
-        blockrange = (blockoffset + 1):(blockoffset + blocklength)
-        blockoffset = last(blockrange)
-        blockstructure[c] = (blocksize, blockrange)
-    end
-
-    fusiontreeindices = sizehint!(
-        FusionTreeDict{Tuple{F₁, F₂}, Int}(), length(fusiontreelist)
-    )
-    for (i, f₁₂) in enumerate(fusiontreelist)
-        fusiontreeindices[f₁₂] = i
-    end
-    totaldim = blockoffset
-    structure = FusionBlockStructure(
-        totaldim, blockstructure, fusiontreelist, fusiontreestructure, fusiontreeindices
-    )
-    return structure
-end
-
-function _subblock_strides(subsz, sz, str)
-    sz_simplify = Strided.StridedViews._simplifydims(sz, str)
-    strides = Strided.StridedViews._computereshapestrides(subsz, sz_simplify...)
-    isnothing(strides) &&
-        throw(ArgumentError("unexpected error in computing subblock strides"))
-    return strides
-end
-
-function CacheStyle(::typeof(fusionblockstructure), W::HomSpace)
-    return GlobalLRUCache()
-end
-
-# Diagonal ranges
-#----------------
-# TODO: is this something we want to cache?
-function diagonalblockstructure(W::HomSpace)
-    ((numin(W) == numout(W) == 1) && domain(W) == codomain(W)) ||
-        throw(SpaceMismatch("Diagonal only support on V←V with a single space V"))
-    structure = SectorDict{sectortype(W), UnitRange{Int}}() # range
-    offset = 0
-    dom = domain(W)[1]
-    for c in blocksectors(W)
-        d = dim(dom, c)
-        structure[c] = offset .+ (1:d)
-        offset += d
-    end
-    return structure
-end

src/spaces/productspace.jl

@@ -96,6 +96,9 @@ sectors(P::ProductSpace) = _sectors(P, sectortype(P))
 function _sectors(P::ProductSpace{<:ElementarySpace, N}, ::Type{Trivial}) where {N}
     return OneOrNoneIterator(dim(P) != 0, ntuple(n -> Trivial(), N))
 end
+function _sectors(P::ProductSpace{<:ElementarySpace, 0}, ::Type{I}) where {I <: Sector}
+    return Iterators.map(u -> (u,), allunits(I))
+end
 function _sectors(P::ProductSpace{<:ElementarySpace, N}, ::Type{<:Sector}) where {N}
     return product(map(sectors, P.spaces)...)
 end
@@ -147,23 +150,11 @@ that make up the `ProductSpace` instance.
 function blocksectors(P::ProductSpace{S, N}) where {S, N}
     I = sectortype(S)
     if I == Trivial
-        return OneOrNoneIterator(dim(P) != 0, Trivial())
+        return dim(P) != 0 ? Indices([Trivial()]) : Indices(Trivial[])
     end
-    bs = Vector{I}()
-    if N == 0
-        append!(bs, allunits(I))
-    elseif N == 1
-        for s in sectors(P)
-            push!(bs, first(s))
-        end
-    else
-        for s in sectors(P)
-            for c in ⊗(s...)
-                if !(c in bs)
-                    push!(bs, c)
-                end
-            end
-        end
+    bs = Indices{I}()
+    for s in sectors(P), c in ⊗(s...)
+        set!(bs, c)
     end
     return sort!(bs)
 end