Specialise GradedSpace functions based on storage type

src/TensorKit.jl

@@ -44,7 +44,7 @@ export infimum, supremum, isisomorphic, ismonomorphic, isepimorphic
 export sectortype, sectors, hassector
 export unit, rightunit, leftunit, allunits, isunit, otimes, deligneproduct, timereversed
 export Nsymbol, Fsymbol, Rsymbol, Bsymbol, frobenius_schur_phase, frobenius_schur_indicator, twist, fusiontensor
-export sectorscalartype, fusionscalartype, braidingscalartype
+export sectorscalartype, fusionscalartype, braidingscalartype, dimscalartype
 
 # Export methods for fusion trees
 export fusiontrees, braid, permute, transpose

src/auxiliary/dicts.jl

@@ -89,14 +89,14 @@ end
 Base.empty(::SortedVectorDict, ::Type{K}, ::Type{V}) where {K, V} = SortedVectorDict{K, V}()
 Base.empty!(d::SortedVectorDict) = (empty!(d.keys); empty!(d.values); return d)
 
-# _searchsortedfirst(v::Vector, k) = searchsortedfirst(v, k)
-function _searchsortedfirst(v::Vector, k)
-    i = 1
-    @inbounds while i <= length(v) && isless(v[i], k)
-        i += 1
-    end
-    return i
-end
+_searchsortedfirst(v::Vector, k) = searchsortedfirst(v, k)
+# function _searchsortedfirst(v::Vector, k)
+#     i = 1
+#     @inbounds while i <= length(v) && isless(v[i], k)
+#         i += 1
+#     end
+#     return i
+# end
 
 function Base.delete!(d::SortedVectorDict{K}, k) where {K}
     key = convert(K, k)
@@ -186,6 +186,55 @@ function Base.:(==)(d1::SortedVectorDict, d2::SortedVectorDict)
     return true
 end
 
+# merge over two SORTED vector pairs representing keys and values
+# - combine(v1,v2): value for a key present in both operands
+# - only1(v1) / only2(v2): value for a key present in only one operand;
+# pass `nothing` to drop such keys entirely (e.g. for an intersection)
+# zero results are dropped (either from `combine` or `only1`/`only2`), matching how GradedSpace never stores an explicit zero dimension
+# k1 and k2 originate from GradedSpace.dims.keys, which are guaranteed to be sorted
+function _sortedmerge(k1::Vector{I}, v1::Vector{Int}, k2::Vector{I}, v2::Vector{Int}, combine, only1, only2) where {I}
+    n1, n2 = length(k1), length(k2)
+    ks, vs = Vector{I}(), Vector{Int}()
+    sizehint!(ks, n1 + n2)
+    sizehint!(vs, n1 + n2)
+    i, j = 1, 1
+    @inbounds while i <= n1 && j <= n2
+        if k1[i] == k2[j]
+            d = combine(v1[i], v2[j])
+            if !iszero(d)
+                push!(ks, k1[i])
+                push!(vs, d)
+            end
+            i += 1
+            j += 1
+        elseif k1[i] < k2[j]
+            _mergeonly!(ks, vs, k1[i], v1[i], only1)
+            i += 1
+        else
+            _mergeonly!(ks, vs, k2[j], v2[j], only2)
+            j += 1
+        end
+    end
+    @inbounds while i <= n1
+        _mergeonly!(ks, vs, k1[i], v1[i], only1)
+        i += 1
+    end
+    @inbounds while j <= n2
+        _mergeonly!(ks, vs, k2[j], v2[j], only2)
+        j += 1
+    end
+    return ks, vs
+end
+@inline _mergeonly!(ks, vs, k, v, ::Nothing) = nothing
+@inline function _mergeonly!(ks, vs, k, v, f)
+    d = f(v)
+    if !iszero(d)
+        push!(ks, k)
+        push!(vs, d)
+    end
+    return nothing
+end
+
 """
     Hashed(value, hashfunction = Base.hash, isequal = Base.isequal)
 

src/factorizations/factorizations.jl

@@ -6,7 +6,7 @@ module Factorizations
 export copy_oftype, factorisation_scalartype, one!, truncspace
 
 using ..TensorKit
-using ..TensorKit: AdjointTensorMap, SectorDict, SectorVector,
+using ..TensorKit: AdjointTensorMap, SectorDict, SectorVector, findindex,
     blocktype, foreachblock, one!,
     similar_diagonal, similarstoragetype
 

src/factorizations/truncation.jl

@@ -37,6 +37,31 @@ _blocklength(ax::Base.OneTo, ind::AbstractVector{Bool}) = count(ind)
 function truncate_space(V::ElementarySpace, inds)
     return spacetype(V)(c => _blocklength(dim(V, c), ind) for (c, ind) in pairs(inds))
 end
+function truncate_space(V::GradedSpace{I, NTuple{N, Int}}, inds) where {I <: Sector, N}
+    vals = values(I)
+    dualV = isdual(V)
+    newdims = zeros(Int, N)
+    for (c, ind) in pairs(inds)
+        n_read = findindex(vals, dualV ? dual(c) : c) # dual-adjusted index for reading V.dims
+        n_write = findindex(vals, c) # output is never dual, so c is fine as-is
+        newdims[n_write] = _blocklength(V.dims[n_read], ind) # dim(c) = dim(dual(c))
+    end
+    return typeof(V)(NTuple{N, Int}(newdims), false)
+end
+function truncate_space(V::GradedSpace{I, <:SectorDict}, inds) where {I <: Sector}
+    dualV = isdual(V)
+    ks, vs = Vector{I}(), Vector{Int}() # accumulate and sort once at the end
+    for (c, ind) in pairs(inds)
+        d = get(V.dims, dualV ? dual(c) : c, 0)
+        len = _blocklength(d, ind)
+        if !iszero(len)
+            push!(ks, c)
+            push!(vs, len)
+        end
+    end
+    perm = sortperm(ks)
+    return typeof(V)(SectorDict{I, Int}(ks[perm], vs[perm]), false)
+end
 
 function truncate_domain!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, inds)
     for (c, b) in blocks(tdst)

src/spaces/gradedspace.jl

@@ -30,17 +30,17 @@ end
 sectortype(::Type{<:GradedSpace{I}}) where {I <: Sector} = I
 
 function GradedSpace{I, NTuple{N, Int}}(dims; dual::Bool = false) where {I, N}
-    d = ntuple(n -> 0, N)
-    isset = ntuple(n -> false, N)
+    d = zeros(Int, N)
+    isset = falses(N)
     for (c, dc) in dims
         k = convert(I, c)
         i = findindex(values(I), k)
-        k = dc < 0 && throw(ArgumentError(lazy"Sector $k has negative dimension $dc"))
+        dc < 0 && throw(ArgumentError(lazy"Sector $k has negative dimension $dc"))
         isset[i] && throw(ArgumentError(lazy"Sector $c appears multiple times"))
-        isset = TupleTools.setindex(isset, true, i)
-        d = TupleTools.setindex(d, dc, i)
+        isset[i] = true
+        d[i] = dc
     end
-    return GradedSpace{I, NTuple{N, Int}}(d, dual)
+    return GradedSpace{I, NTuple{N, Int}}(NTuple{N, Int}(d), dual)
 end
 function GradedSpace{I, NTuple{N, Int}}(dims::Pair; dual::Bool = false) where {I, N}
     return GradedSpace{I, NTuple{N, Int}}((dims,); dual = dual)
@@ -89,9 +89,20 @@ GradedSpace(g::AbstractDict; dual::Bool = false) = GradedSpace(g...; dual = dual
 field(::Type{<:GradedSpace}) = ℂ
 InnerProductStyle(::Type{<:GradedSpace}) = EuclideanInnerProduct()
 
-function dim(V::GradedSpace)
-    init = 0 * dim(first(allunits(sectortype(V))))
-    return sum(c -> dim(c) * dim(V, c), sectors(V); init = init)
+function dim(V::GradedSpace{I, <:AbstractDict}) where {I <: Sector}
+    init = zero(dimscalartype(I))
+    return sum(((c, d),) -> dim(c) * d, V.dims; init)
+end
+function dim(V::GradedSpace{I, NTuple{N, Int}}) where {I <: Sector, N}
+    init = zero(dimscalartype(I))
+    D = init
+    vals = values(I)
+    @inbounds for n in 1:N
+        d = V.dims[n]
+        iszero(d) && continue
+        D += dim(vals[n]) * d # dim(c) = dim(dual(c))
+    end
+    return D
 end
 function dim(V::GradedSpace{I, <:AbstractDict}, c::I) where {I <: Sector}
     return get(V.dims, isdual(V) ? dual(c) : c, 0)
@@ -126,54 +137,101 @@ function unitspace(S::Type{<:GradedSpace{I}}) where {I <: Sector}
 end
 zerospace(S::Type{<:GradedSpace}) = S()
 
-# TODO: the following methods can probably be implemented more efficiently for
-# `FiniteGradedSpace`, but we don't expect them to be used often in hot loops, so
-# these generic definitions (which are still quite efficient) are good for now.
-function ⊕(V₁::GradedSpace{I}, V₂::GradedSpace{I}) where {I <: Sector}
+function ⊕(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector}
+    dual1 = isdual(V₁)
+    dual1 == isdual(V₂) || throw(SpaceMismatch("Direct sum of a vector space and a dual space does not exist"))
+    ks, vs = _sortedmerge(V₁.dims.keys, V₁.dims.values, V₂.dims.keys, V₂.dims.values, +, identity, identity)
+    return typeof(V₁)(SectorDict{I, Int}(ks, vs), dual1)
+end
+function ⊕(V₁::GradedSpace{I, <:Tuple}, V₂::GradedSpace{I, <:Tuple}) where {I <: Sector}
     dual1 = isdual(V₁)
     dual1 == isdual(V₂) ||
         throw(SpaceMismatch("Direct sum of a vector space and a dual space does not exist"))
-    dims = SectorDict{I, Int}()
-    for c in union(sectors(V₁), sectors(V₂))
-        cout = ifelse(dual1, dual(c), c)
-        dims[cout] = dim(V₁, c) + dim(V₂, c)
+    newdims = map(+, V₁.dims, V₂.dims)
+    return typeof(V₁)(newdims, dual1)
+end
+function ⊖(V::GradedSpace{I, <:Tuple}, W::GradedSpace{I, <:Tuple}) where {I <: Sector}
+    dualV = isdual(V)
+    V ≿ W && dualV == isdual(W) || throw(SpaceMismatch("$(W) is not a subspace of $(V)"))
+    newdims = map(-, V.dims, W.dims)
+    return typeof(V)(newdims, dualV)
+end
+function ⊖(V::GradedSpace{I, <:SectorDict}, W::GradedSpace{I, <:SectorDict}) where {I <: Sector}
+    dualV = isdual(V)
+    V ≿ W && dualV == isdual(W) || throw(SpaceMismatch("$(W) is not a subspace of $(V)"))
+    ks, vs = _sortedmerge(V.dims.keys, V.dims.values, W.dims.keys, W.dims.values, -, identity, nothing)
+    return typeof(V)(SectorDict{I, Int}(ks, vs), dualV)
+end
+
+function fuse(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector}
+    dual1, dual2 = isdual(V₁), isdual(V₂)
+    acc = Dict{I, Int}() # SectorDict `get` within the double for loop accumulates O(N^2) ` findindex` calls -> sort afterwards
+    k1, k2 = V₁.dims.keys, V₂.dims.keys
+    v1, v2 = V₁.dims.values, V₂.dims.values
+    @inbounds for na in eachindex(k1)
+        a₀, da = k1[na], v1[na]
+        a = dual1 ? dual(a₀) : a₀
+        for nb in eachindex(k2)
+            b₀, db = k2[nb], v2[nb]
+            b = dual2 ? dual(b₀) : b₀
+            dab = da * db
+            for c in a ⊗ b
+                acc[c] = get(acc, c, 0) + Nsymbol(a, b, c) * dab
+            end
+        end
     end
-    return typeof(V₁)(dims; dual = dual1)
-end
-function ⊖(V::GradedSpace{I}, W::GradedSpace{I}) where {I <: Sector}
-    dual = isdual(V)
-    V ≿ W && dual == isdual(W) ||
-        throw(SpaceMismatch("$(W) is not a subspace of $(V)"))
-    return typeof(V)(c => dim(V, c) - dim(W, c) for c in sectors(V); dual)
-end
-
-function fuse(V₁::GradedSpace{I}, V₂::GradedSpace{I}) where {I <: Sector}
-    dims = SectorDict{I, Int}()
-    for a in sectors(V₁), b in sectors(V₂)
-        for c in a ⊗ b
-            dims[c] = get(dims, c, 0) + Nsymbol(a, b, c) * dim(V₁, a) * dim(V₂, b)
+    ks0 = collect(keys(acc))
+    vs0 = collect(values(acc))
+    perm = sortperm(ks0)
+    return typeof(V₁)(SectorDict{I, Int}(ks0[perm], vs0[perm]), false)
+end
+function fuse(V₁::GradedSpace{I, NTuple{N, Int}}, V₂::GradedSpace{I, NTuple{N, Int}}) where {I <: Sector, N}
+    vals = values(I)
+    dual1, dual2 = isdual(V₁), isdual(V₂)
+    newdims = zeros(Int, N)
+    @inbounds for na in 1:N
+        da = V₁.dims[na]
+        iszero(da) && continue
+        a₀ = vals[na] # avoid call to sectors(V₁)
+        a = dual1 ? dual(a₀) : a₀
+        for nb in 1:N
+            db = V₂.dims[nb]
+            iszero(db) && continue
+            b₀ = vals[nb] # idem for V₂
+            b = dual2 ? dual(b₀) : b₀
+            dab = da * db
+            for c in a ⊗ b
+                nc = findindex(vals, c)
+                newdims[nc] += Nsymbol(a, b, c) * dab
+            end
         end
     end
-    return typeof(V₁)(dims)
+    return typeof(V₁)(NTuple{N, Int}(newdims), false)
 end
 
-function infimum(V₁::GradedSpace{I}, V₂::GradedSpace{I}) where {I <: Sector}
+function infimum(V₁::GradedSpace{I, <:Tuple}, V₂::GradedSpace{I, <:Tuple}) where {I <: Sector}
+    Visdual = isdual(V₁)
+    Visdual == isdual(V₂) || throw(SpaceMismatch("Infimum of space and dual space does not exist"))
+    newdims = map(min, V₁.dims, V₂.dims)
+    return typeof(V₁)(newdims, Visdual)
+end
+function infimum(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector}
     Visdual = isdual(V₁)
-    Visdual == isdual(V₂) ||
-        throw(SpaceMismatch("Infimum of space and dual space does not exist"))
-    return typeof(V₁)(
-        (Visdual ? dual(c) : c) => min(dim(V₁, c), dim(V₂, c))
-            for c in intersect(sectors(V₁), sectors(V₂)); dual = Visdual
-    )
-end
-function supremum(V₁::GradedSpace{I}, V₂::GradedSpace{I}) where {I <: Sector}
+    Visdual == isdual(V₂) || throw(SpaceMismatch("Infimum of space and dual space does not exist"))
+    ks, vs = _sortedmerge(V₁.dims.keys, V₁.dims.values, V₂.dims.keys, V₂.dims.values, min, nothing, nothing)
+    return typeof(V₁)(SectorDict{I, Int}(ks, vs), Visdual)
+end
+function supremum(V₁::GradedSpace{I, <:Tuple}, V₂::GradedSpace{I, <:Tuple}) where {I <: Sector}
+    Visdual = isdual(V₁)
+    Visdual == isdual(V₂) || throw(SpaceMismatch("Supremum of space and dual space does not exist"))
+    newdims = map(max, V₁.dims, V₂.dims)
+    return typeof(V₁)(newdims, Visdual)
+end
+function supremum(V₁::GradedSpace{I, <:SectorDict}, V₂::GradedSpace{I, <:SectorDict}) where {I <: Sector}
     Visdual = isdual(V₁)
-    Visdual == isdual(V₂) ||
-        throw(SpaceMismatch("Supremum of space and dual space does not exist"))
-    return typeof(V₁)(
-        (Visdual ? dual(c) : c) => max(dim(V₁, c), dim(V₂, c))
-            for c in union(sectors(V₁), sectors(V₂)); dual = Visdual
-    )
+    Visdual == isdual(V₂) || throw(SpaceMismatch("Supremum of space and dual space does not exist"))
+    ks, vs = _sortedmerge(V₁.dims.keys, V₁.dims.values, V₂.dims.keys, V₂.dims.values, max, identity, identity)
+    return typeof(V₁)(SectorDict{I, Int}(ks, vs), Visdual)
 end
 
 hassector(V::GradedSpace{I}, s::I) where {I <: Sector} = dim(V, s) != 0