Vectorize fusiontree manipulations (QuantumKitHub#261)

Co-authored-by: Olivier Gauthé <olivier.gauthe.2011+github@polytechnique.org>
Co-authored-by: Jutho Haegeman <jutho.haegeman@ugent.be>

Author: 3 people

Project.toml

@@ -54,7 +54,7 @@ Random = "1"
 SafeTestsets = "0.1"
 ScopedValues = "1.3.0"
 Strided = "2"
-TensorKitSectors = "0.3.5"
+TensorKitSectors = "0.3.6"
 TensorOperations = "5.1"
 Test = "1"
 TestExtras = "0.2,0.3"
@@ -86,4 +86,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = ["ArgParse", "Adapt", "Aqua", "AllocCheck", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "LinearAlgebra", "SafeTestsets", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote", "Mooncake", "JET"]
+test = ["ArgParse", "Adapt", "Aqua", "AllocCheck", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "JET", "LinearAlgebra", "SafeTestsets", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote", "Mooncake"]

docs/make.jl

@@ -1,6 +1,16 @@
+# if examples is not the current active environment, switch to it
+if Base.active_project() != joinpath(@__DIR__, "Project.toml")
+    using Pkg
+    Pkg.activate(@__DIR__)
+    Pkg.develop(PackageSpec(; path = joinpath(@__DIR__, "..")))
+    Pkg.resolve()
+    Pkg.instantiate()
+end
+
 using Documenter
 using Random
 using TensorKit
+using TensorKit: FusionTreePair, Index2Tuple
 using TensorKit.TensorKitSectors
 using TensorKit.MatrixAlgebraKit
 using DocumenterInterLinks
@@ -40,6 +50,9 @@ mathengine = MathJax3(
     )
 )
 
+# docstrings don't need `using TensorKit`
+DocMeta.setdocmeta!(TensorKit, :DocTestSetup, :(using TensorKit); recursive = true)
+
 makedocs(;
     modules = [TensorKit, TensorKitSectors],
     sitename = "TensorKit.jl",

docs/src/lib/fusiontrees.md

@@ -13,8 +13,7 @@ FusionTree
 ## Methods for defining and generating fusion trees
 
 ```@docs
-fusiontrees(uncoupled::NTuple{N,I}, coupled::I,
-                     isdual::NTuple{N,Bool}) where {N,I<:Sector}
+fusiontrees(uncoupled::NTuple{N, I}, coupled::I, isdual::NTuple{N, Bool}) where {N, I<:Sector}
 ```
 
 ## Methods for manipulating fusion trees
@@ -23,31 +22,29 @@ For manipulating single fusion trees, the following internal methods are defined
 ```@docs
 insertat
 split
+TensorKit.join
 merge
 elementary_trace
-planar_trace(f::FusionTree{I,N}, q1::IndexTuple{N₃}, q2::IndexTuple{N₃}) where {I<:Sector,N,N₃}
+planar_trace(f::FusionTree, q::Index2Tuple)
 artin_braid
-braid(f::FusionTree{I,N}, levels::NTuple{N,Int}, p::NTuple{N,Int}) where {I<:Sector,N}
-permute(f::FusionTree{I,N}, p::NTuple{N,Int}) where {I<:Sector,N}
+braid(f::FusionTree{I, N}, p::IndexTuple{N}, levels::IndexTuple{N}) where {I, N}
+permute(f::FusionTree{I, N}, p::IndexTuple{N}) where {I, N}
 ```
 
 These can be composed to implement elementary manipulations of fusion-splitting tree pairs, according to the following methods
 
-```julia
-# TODO: add documentation for the following methods
+```@docs
 TensorKit.bendright
 TensorKit.bendleft
 TensorKit.foldright
 TensorKit.foldleft
-TensorKit.cycleclockwise
-TensorKit.cycleanticlockwise
 ```
 
 Finally, these are used to define large manipulations of fusion-splitting tree pairs, which are then used in the index manipulation of `AbstractTensorMap` objects.
 The following methods defined on fusion splitting tree pairs have an associated definition for tensors.
 ```@docs
-repartition(::FusionTree{I,N₁}, ::FusionTree{I,N₂}, ::Int) where {I<:Sector,N₁,N₂}
-transpose(::FusionTree{I}, ::FusionTree{I}, ::IndexTuple{N₁}, ::IndexTuple{N₂}) where {I<:Sector,N₁,N₂}
-braid(::FusionTree{I}, ::FusionTree{I}, ::IndexTuple, ::IndexTuple, ::IndexTuple{N₁}, ::IndexTuple{N₂}) where {I<:Sector,N₁,N₂}
-permute(::FusionTree{I}, ::FusionTree{I}, ::IndexTuple{N₁}, ::IndexTuple{N₂}) where {I<:Sector,N₁,N₂}
+repartition(src::Union{FusionTreePair, FusionTreeBlock}, N::Int)
+Base.transpose(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple)
+braid(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, levels::Index2Tuple)
+permute(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple)
 ```

docs/src/man/fusiontrees.md

@@ -123,8 +123,8 @@ using OhMyThreads
 using ScopedValues
 
 using TensorKitSectors
-import TensorKitSectors: dim, BraidingStyle, FusionStyle, ⊠, ⊗
-import TensorKitSectors: dual, type_repr
+import TensorKitSectors: dim, BraidingStyle, FusionStyle, ⊠, ⊗, ×
+import TensorKitSectors: dual, type_repr, fusiontensor
 import TensorKitSectors: twist
 
 using Base: @boundscheck, @propagate_inbounds, @constprop,
@@ -215,6 +215,19 @@ function set_num_transformer_threads(n::Int)
     return TRANSFORMER_THREADS[] = n
 end
 
+const TREEMANIPULATION_THREADS = Ref(1)
+
+get_num_manipulation_threads() = TREEMANIPULATION_THREADS[]
+
+function set_num_manipulation_threads(n::Int)
+    N = Base.Threads.nthreads()
+    if n > N
+        n = N
+        Strided._set_num_threads_warn(n)
+    end
+    return TREEMANIPULATION_THREADS[] = n
+end
+
 # Definitions and methods for tensors
 #-------------------------------------
 # general definitions

src/TensorKit.jl

@@ -27,6 +27,20 @@ function permutation2swaps(perm)
     return swaps
 end
 
+# one-argument version: check whether `p` is a cyclic permutation (of `1:length(p)`)
+function iscyclicpermutation(p)
+    N = length(p)
+    @inbounds for i in 1:N
+        p[mod1(i + 1, N)] == mod1(p[i] + 1, N) || return false
+    end
+    return true
+end
+# two-argument version: check whether `v1` is a cyclic permutation of `v2`
+function iscyclicpermutation(v1, v2)
+    length(v1) == length(v2) || return false
+    return iscyclicpermutation(indexin(v1, v2))
+end
+
 _kron(A, B, C, D...) = _kron(_kron(A, B), C, D...)
 function _kron(A, B)
     sA = size(A)