Add treetransformers

Author: lkdvos

Project.toml

@@ -1,13 +1,14 @@
 name = "TensorKit"
 uuid = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 authors = ["Jutho Haegeman"]
-version = "0.13"
+version = "0.13.0"
 
 [deps]
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
@@ -32,6 +33,7 @@ LRUCache = "1.0.2"
 LinearAlgebra = "1"
 PackageExtensionCompat = "1"
 Random = "1"
+SparseArrays = "1"
 Strided = "2"
 TensorKitSectors = "0.1"
 TensorOperations = "5.1"

src/TensorKit.jl

@@ -119,6 +119,9 @@ using LinearAlgebra: norm, dot, normalize, normalize!, tr,
                      eigen, eigen!, svd, svd!,
                      isposdef, isposdef!, ishermitian,
                      Diagonal, Hermitian
+
+using SparseArrays: SparseMatrixCSC, sparse, nzrange, rowvals, nonzeros
+
 import Base.Meta
 
 using Random: Random
@@ -185,6 +188,7 @@ include("tensors/adjoint.jl")
 include("tensors/linalg.jl")
 include("tensors/vectorinterface.jl")
 include("tensors/tensoroperations.jl")
+include("tensors/treetransformers.jl")
 include("tensors/indexmanipulations.jl")
 include("tensors/truncation.jl")
 include("tensors/factorizations.jl")

src/tensors/indexmanipulations.jl

@@ -279,8 +279,8 @@ See also [`permute`](@ref), [`permute!`](@ref), [`add_braid!`](@ref), [`add_tran
                                           α::Number,
                                           β::Number,
                                           backend::AbstractBackend...) where {T,S,N₁,N₂}
-    treepermuter(f₁, f₂) = permute(f₁, f₂, p[1], p[2])
-    return add_transform!(tdst, tsrc, p, treepermuter, α, β, backend...)
+    transformer = treepermuter(tdst, tsrc, p)
+    return add_transform!(tdst, tsrc, p, transformer, α, β, backend...)
 end
 
 """
@@ -305,8 +305,8 @@ See also [`braid`](@ref), [`braid!`](@ref), [`add_permute!`](@ref), [`add_transp
     levels1 = TupleTools.getindices(levels, codomainind(tsrc))
     levels2 = TupleTools.getindices(levels, domainind(tsrc))
     # TODO: arg order for tensormaps is different than for fusiontrees
-    treebraider(f₁, f₂) = braid(f₁, f₂, levels1, levels2, p...)
-    return add_transform!(tdst, tsrc, p, treebraider, α, β, backend...)
+    transformer = treebraider(tdst, tsrc, p, (levels1, levels2))
+    return add_transform!(tdst, tsrc, p, transformer, α, β, backend...)
 end
 
 """
@@ -324,14 +324,14 @@ See also [`transpose`](@ref), [`transpose!`](@ref), [`add_permute!`](@ref), [`ad
                                             α::Number,
                                             β::Number,
                                             backend::AbstractBackend...) where {T,S,N₁,N₂}
-    treetransposer(f₁, f₂) = transpose(f₁, f₂, p[1], p[2])
-    return add_transform!(tdst, tsrc, p, treetransposer, α, β, backend...)
+    transformer = treetransposer(tdst, tsrc, p)
+    return add_transform!(tdst, tsrc, p, transformer, α, β, backend...)
 end
 
 function add_transform!(tdst::AbstractTensorMap{T,S,N₁,N₂},
                         tsrc::AbstractTensorMap,
                         (p₁, p₂)::Index2Tuple{N₁,N₂},
-                        fusiontreetransform,
+                        transformer::GenericTreeTransformer,
                         α::Number,
                         β::Number,
                         backend::AbstractBackend...) where {T,S,N₁,N₂}
@@ -341,16 +341,34 @@ function add_transform!(tdst::AbstractTensorMap{T,S,N₁,N₂},
             dest = $(codomain(tdst))←$(domain(tdst)), p₁ = $(p₁), p₂ = $(p₂)"))
     end
 
-    I = sectortype(S)
-    if p₁ == codomainind(tsrc) && p₂ == domainind(tsrc)
-        add!(tdst, tsrc, α, β)
-    elseif I === Trivial
-        _add_trivial_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
-    elseif FusionStyle(I) isa UniqueFusion
-        _add_abelian_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
-    else
-        _add_general_kernel!(tdst, tsrc, (p₁, p₂), fusiontreetransform, α, β, backend...)
+    structure_dst = transformer.structure_dst.fusiontreestructure
+    structure_src = transformer.structure_src.fusiontreestructure
+
+    rows = rowvals(transformer.matrix)
+    vals = nonzeros(transformer.matrix)
+
+    # TODO: this could be multithreaded
+    for j in axes(transformer.matrix, 2)
+        sz_dst, str_dst, offset_dst = structure_dst[j]
+        subblock_dst = StridedView(tdst.data, sz_dst, str_dst, offset_dst)
+        nzrows = nzrange(transformer.matrix, j)
+
+        # treat first entry
+        sz_src, str_src, offset_src = structure_src[rows[first(nzrows)]]
+        subblock_src = StridedView(tsrc.data, sz_src, str_src, offset_src)
+        TO.tensoradd!(subblock_dst, subblock_src, (p₁, p₂), false, α * vals[first(nzrows)],
+                      β,
+                      backend...)
+
+        # treat remaining entries
+        for i in @view(nzrows[2:end])
+            sz_src, str_src, offset_src = structure_src[rows[i]]
+            subblock_src = StridedView(tsrc.data, sz_src, str_src, offset_src)
+            TO.tensoradd!(subblock_dst, subblock_src, (p₁, p₂), false, α * vals[i], One(),
+                          backend...)
+        end
     end
+
     return tdst
 end
 
@@ -389,7 +407,7 @@ function _add_general_kernel!(tdst, tsrc, p, fusiontreetransform, α, β, backen
         tdst = scale!(tdst, β)
     end
     β′ = One()
-    if Threads.nthreads() > 1
+    if false # Threads.nthreads() > 1
         Threads.@sync for s₁ in sectors(codomain(tsrc)), s₂ in sectors(domain(tsrc))
             Threads.@spawn _add_nonabelian_sector!(tdst, tsrc, p, fusiontreetransform, s₁,
                                                    s₂, α, β′, backend...)

src/tensors/treetransformers.jl

@@ -0,0 +1,123 @@
+"""
+    TreeTransformer
+
+Supertype for structures containing the data for a tree transformation.
+"""
+abstract type TreeTransformer end
+
+function treetransformertype(Vdst, Vsrc)
+    I = sectortype(Vdst)
+    N = numind(Vdst)
+    F1 = fusiontreetype(I, numout(Vdst))
+    F2 = fusiontreetype(I, numin(Vdst))
+    F3 = fusiontreetype(I, numout(Vsrc))
+    F4 = fusiontreetype(I, numin(Vsrc))
+    return GenericTreeTransformer{sectorscalartype(I),I,N,F1,F2,F3,F4}
+end
+
+struct GenericTreeTransformer{T,I,N,F1,F2,F3,F4} <: TreeTransformer
+    matrix::SparseMatrixCSC{T,Int}
+    structure_dst::FusionBlockStructure{I,N,F1,F2}
+    structure_src::FusionBlockStructure{I,N,F3,F4}
+end
+
+function GenericTreeTransformer(transform::Function, Vsrc::HomSpace, Vdst::HomSpace)
+    structure_dst = fusionblockstructure(Vdst)
+    structure_src = fusionblockstructure(Vsrc)
+
+    ldst = length(structure_dst.fusiontreelist)
+    lsrc = length(structure_src.fusiontreelist)
+
+    rows = Int[]
+    cols = Int[]
+    vals = sectorscalartype(sectortype(Vdst))[]
+
+    for (f1, f2) in structure_src.fusiontreelist
+        row = structure_src.fusiontreeindices[(f1, f2)]
+        for ((f3, f4), coeff) in transform(f1, f2)
+            col = structure_dst.fusiontreeindices[(f3, f4)]
+            push!(rows, row)
+            push!(cols, col)
+            push!(vals, coeff)
+        end
+    end
+    matrix = sparse(rows, cols, vals, ldst, lsrc)
+
+    return GenericTreeTransformer(matrix, structure_dst, structure_src)
+end
+
+# Transpose
+# ---------
+const treetransposercache = LRU{Any,Any}(; maxsize=10^5)
+const usetreetransposercache = Ref{Bool}(true)
+
+function treetransposer(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple)
+    if usetreetransposercache[]
+        key = (space(tdst), space(tsrc), p)
+        A = treetransformertype(space(tdst), space(tsrc))
+        return _get_treetransposer(A, key)
+    else
+        return _treetransposer((space(tdst), space(tsrc), p))
+    end
+end
+@noinline function _get_treetransposer(A, key)
+    d::A = get!(treetransposercache, key) do
+        return _treetransposer(key)
+    end
+    return d
+end
+function _treetransposer((Vdst, Vsrc, p))
+    fusiontreetransform(f1, f2) = transpose(f1, f2, p...)
+    return GenericTreeTransformer(fusiontreetransform, Vsrc, Vdst)
+end
+
+# Braid
+# -----
+const treebraidercache = LRU{Any,Any}(; maxsize=10^5)
+const usetreebraidercache = Ref{Bool}(true)
+
+function treebraider(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple,
+                     l::Index2Tuple)
+    if usetreebraidercache[]
+        key = (space(tdst), space(tsrc), p, l)
+        A = treetransformertype(space(tdst), space(tsrc))
+        return _get_treebraider(A, key)
+    else
+        return _treebraider((space(tdst), space(tsrc), p, l))
+    end
+end
+@noinline function _get_treebraider(A, key)
+    d::A = get!(treebraidercache, key) do
+        return _treebraider(key)
+    end
+    return d
+end
+function _treebraider((Vdst, Vsrc, p, l))
+    fusiontreetransform(f1, f2) = braid(f1, f2, p..., l...)
+    return GenericTreeTransformer(fusiontreetransform, Vsrc, Vdst)
+end
+
+# Permute
+# -------
+const treepermutercache = LRU{Any,Any}(; maxsize=10^5)
+const usetreepermutercache = Ref{Bool}(true)
+
+function treepermuter(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple)
+    if usetreepermutercache[]
+        key = (space(tdst), space(tsrc), p)
+        A = treetransformertype(space(tdst), space(tsrc))
+        return _get_treepermuter(A, key)
+    else
+        return _treepermuter((space(tdst), space(tsrc), p))
+    end
+end
+@noinline function _get_treepermuter(A, key)
+    d::A = get!(treepermutercache, key) do
+        return _treepermuter(key)
+    end
+    return d
+end
+function _treepermuter((Vdst, Vsrc, p))
+    fusiontreetransform(f1, f2) = permute(f1, f2, p...)
+    return GenericTreeTransformer(fusiontreetransform, Vsrc, Vdst)
+end