Reduce cache footprint by decoupling degeneracy-dependent data (#387)

Author: lkdvos

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
 using LRUCache
 using OhMyThreads
 using ScopedValues
@@ -200,6 +201,23 @@ include("fusiontrees/fusiontrees.jl")
 #-------------------------------------------
 include("spaces/vectorspaces.jl")
 
+# ElementarySpace types
+include("spaces/cartesianspace.jl")
+include("spaces/complexspace.jl")
+include("spaces/generalspace.jl")
+include("spaces/gradedspace.jl")
+include("spaces/planarspace.jl")
+
+# CompositeSpace types
+include("spaces/productspace.jl")
+include("spaces/deligne.jl")
+
+# HomSpace
+include("spaces/homspace.jl")
+
+# Derived information
+include("spaces/structure.jl")
+
 # Multithreading settings
 #-------------------------
 const TRANSFORMER_THREADS = Ref(1)

src/auxiliary/dicts.jl

@@ -263,3 +263,25 @@ 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, H <: Function, E <: Function}
+    val::T
+    hashf::H
+    eqf::E
+end
+
+Hashed(val, hashf = Base.hash, eqf = Base.isequal) =
+    Hashed{typeof(val), typeof(hashf), typeof(eqf)}(val, hashf, eqf)
+
+Base.parent(h::Hashed) = h.val
+Base.hash(h::Hashed, seed::UInt) = h.hashf(parent(h), seed)
+# Note: requires the equality functions to be equal to avoid asymmetric results
+Base.isequal(h1::Hashed{<:Any, <:Any, E}, h2::Hashed{<:Any, <:Any, E}) where {E} =
+    h1.eqf(parent(h1), parent(h2))

src/fusiontrees/braiding_manipulations.jl

@@ -299,44 +299,44 @@ Base.@assume_effects :foldable function _fsdicttype(::Type{T}) where {I, N₁, N
     return Pair{FusionTreeBlock{I, N₁, N₂, Tuple{F₁, F₂}}, Matrix{E}}
 end
 
-@cached function fsbraid(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: FSPBraidKey{I, N₁, N₂}}
-    ((f₁, f₂), (p1, p2), (l1, l2)) = key
-    p = linearizepermutation(p1, p2, length(f₁), length(f₂))
-    levels = (l1..., reverse(l2)...)
-    (f, f0), coeff1 = repartition((f₁, f₂), N₁ + N₂)
-    f′, coeff2 = braid(f, p, levels)
-    (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁)
-    return (f₁′, f₂′) => coeff1 * coeff2 * coeff3
-end
-@cached function fsbraid(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: FSBBraidKey{I, N₁, N₂}}
-    src, (p1, p2), (l1, l2) = key
+@cached function fsbraid(key::K)::_fsdicttype(K) where {I, N₁, N₂, K <: Union{FSPBraidKey{I, N₁, N₂}, FSBBraidKey{I, N₁, N₂}}}
+    if K <: FSPBraidKey
+        ((f₁, f₂), (p1, p2), (l1, l2)) = key
+        p = linearizepermutation(p1, p2, length(f₁), length(f₂))
+        levels = (l1..., reverse(l2)...)
+        (f, f0), coeff1 = repartition((f₁, f₂), N₁ + N₂)
+        f′, coeff2 = braid(f, p, levels)
+        (f₁′, f₂′), coeff3 = repartition((f′, f0), N₁)
+        return (f₁′, f₂′) => coeff1 * coeff2 * coeff3
 
-    p = linearizepermutation(p1, p2, numout(src), numin(src))
-    levels = (l1..., reverse(l2)...)
+    else
+        src, (p1, p2), (l1, l2) = key
 
-    dst, U = repartition(src, numind(src))
+        p = linearizepermutation(p1, p2, numout(src), numin(src))
+        levels = (l1..., reverse(l2)...)
 
-    for s in permutation2swaps(p)
-        inv = levels[s] > levels[s + 1]
-        dst, U_tmp = artin_braid(dst, s; inv)
-        U = U_tmp * U
-        l = levels[s]
-        levels = TupleTools.setindex(levels, levels[s + 1], s)
-        levels = TupleTools.setindex(levels, l, s + 1)
-    end
+        dst, U = repartition(src, numind(src))
 
-    if N₂ == 0
-        return dst => U
-    else
-        dst, U_tmp = repartition(dst, N₁)
-        U = U_tmp * U
-        return dst => U
+        for s in permutation2swaps(p)
+            inv = levels[s] > levels[s + 1]
+            dst, U_tmp = artin_braid(dst, s; inv)
+            U = U_tmp * U
+            l = levels[s]
+            levels = TupleTools.setindex(levels, levels[s + 1], s)
+            levels = TupleTools.setindex(levels, l, s + 1)
+        end
+
+        if N₂ == 0
+            return dst => U
+        else
+            dst, U_tmp = repartition(dst, N₁)
+            U = U_tmp * U
+            return dst => U
+        end
     end
 end
 
-CacheStyle(::typeof(fsbraid), k::FSPBraidKey{I}) where {I} =
-    FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache()
-CacheStyle(::typeof(fsbraid), k::FSBBraidKey{I}) where {I} =
+CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} =
     FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache()
 
 """

src/spaces/gradedspace.jl

@@ -191,6 +191,21 @@ function Base.:(==)(V₁::GradedSpace, V₂::GradedSpace)
     return sectortype(V₁) == sectortype(V₂) && (V₁.dims == V₂.dims) && V₁.dual == V₂.dual
 end
 
+function sectorhash(V::GradedSpace{I, NTuple{N, Int}}, h::UInt) where {I, N}
+    return hash(iszero.(V.dims), hash(isdual(V), h))
+end
+function sectorequal(V₁::GradedSpace{I, D}, V₂::GradedSpace{I, D}) where {I, N, D <: NTuple{N, Int}}
+    return isdual(V₁) == isdual(V₂) && all(zip(V₁.dims, V₂.dims)) do (d₁, d₂)
+        return iszero(d₁) == iszero(d₂)
+    end
+end
+function sectorhash(V::GradedSpace{I, <:SectorDict}, h::UInt) where {I}
+    return hash(keys(V.dims), hash(isdual(V), h))
+end
+function sectorequal(V₁::GradedSpace{I, D}, V₂::GradedSpace{I, D}) where {I, D <: SectorDict}
+    return isdual(V₁) == isdual(V₂) && keys(V₁.dims) == keys(V₂.dims)
+end
+
 Base.summary(io::IO, V::GradedSpace) = print(io, type_repr(typeof(V)))
 
 function Base.show(io::IO, V::GradedSpace)