From f7e3a42b531dc94be459f55fb5ed0a5320f3de43 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 11:42:07 -0400 Subject: [PATCH 01/13] Precompute planar trace kernels in a `TraceTransformer` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `planartrace!` now precomputes the subblock positions and recoupling coefficients in a cached `TraceTransformer`, and executes them with a sector-independent kernel, analogous to the index manipulations. This replaces the per-tree loop that built ad-hoc single-pair `FusionTreeBlock`s. - `planar_trace` is defined for `Union{FusionTreePair, FusionTreeBlock}`, with unique-fusion pairs returning `dst => coeff` and blocks `dst => U`, and throws when the traced sectors are not dual to each other - the transformer skips such (sub)blocks upfront and groups the data by destination, such that every destination subblock is written once and `β` is handled inside the kernel - blocks with several destination trees are traced (and permuted) into a buffer, recoupled with a single matmul and added into the destination - `FusionTreeBlock`s hash their uncoupled sectors and duality flags Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 1 + docs/src/lib/fusiontrees.md | 1 + src/fusiontrees/duality_manipulations.jl | 115 ++++++++++++++--------- src/fusiontrees/fusiontrees.jl | 9 +- src/planar/planaroperations.jl | 83 +++++++++++++--- src/tensors/treetransformers.jl | 100 ++++++++++++++++++++ test/symmetries/doubletree.jl | 24 ++--- 7 files changed, 258 insertions(+), 75 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index b91e0e68b..130351ddb 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,6 +30,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Require TensorKitSectors v0.3.10 - `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) - The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) +- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s (internal) ### Deprecated diff --git a/docs/src/lib/fusiontrees.md b/docs/src/lib/fusiontrees.md index 57033eca6..759eeec12 100644 --- a/docs/src/lib/fusiontrees.md +++ b/docs/src/lib/fusiontrees.md @@ -48,4 +48,5 @@ 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) +planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, p::Index2Tuple, q::Index2Tuple) ``` diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index e399f87ab..d6324f9d4 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -577,58 +577,85 @@ CacheStyle(::typeof(fstranspose), k::FSBTransposeKey{I}) where {I} = # -> composite manipulations that depend on the duality (rigidity) and pivotal structure # -> planar manipulations that do not require braiding, everything is in Fsymbol (A/Bsymbol) -function planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) +# leg positions of a planar trace of a tree pair with `N₁` outgoing and `N₂` incoming legs, after +# repartitioning all legs to the outgoing side: the traced pairs `q′`, the remaining legs `r` in +# increasing order and the permutation `p′` that is applied after removing the traced legs +function _planar_trace_indices(::Val{N₁}, ::Val{N₂}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) where {N₁, N₂} length(q₁) == length(q₂) || throw(ArgumentError(lazy"trace index tuples q₁ and q₂ must have equal length, got $(length(q₁)) and $(length(q₂))")) - I = sectortype(f₁) N = length(p₁) + length(p₂) + 2 * length(q₁) - length(f₁) + length(f₂) == N || - throw(ArgumentError(lazy"fusion tree pair has $(length(f₁) + length(f₂)) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) - if isempty(q₁) - return transpose((f₁, f₂), (p₁, p₂)) - end - - linearindex = ( - ntuple(identity, Val(length(f₁)))..., - reverse(length(f₁) .+ ntuple(identity, Val(length(f₂))))..., - ) - + N₁ + N₂ == N || + throw(ArgumentError(lazy"fusion tree pair has $(N₁ + N₂) indices, but permutation expects $N = $(length(p₁)) + $(length(p₂)) + 2×$(length(q₁))")) + linearindex = (ntuple(identity, Val(N₁))..., reverse(N₁ .+ ntuple(identity, Val(N₂)))...) q₁′ = TupleTools.getindices(linearindex, q₁) q₂′ = TupleTools.getindices(linearindex, q₂) - p₁′, p₂′ = let q′ = (q₁′..., q₂′...) - ( - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)), - map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)), - ) - end + q′ = (q₁′..., q₂′...) + p₁′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₁)) + p₂′ = map(l -> l - count(l .> q′), TupleTools.getindices(linearindex, p₂)) + r = TupleTools.sort(TupleTools.getindices(linearindex, (p₁..., p₂...))) + return (p₁′, p₂′), (q₁′, q₂′), r +end - T = fusionscalartype(I) - F₁ = fusiontreetype(I, length(p₁)) - F₂ = fusiontreetype(I, length(p₂)) - newtrees = FusionTreeDict{Tuple{F₁, F₂}, T}() - if FusionStyle(I) isa UniqueFusion - (f₁′, f₂′), coeff′ = repartition((f₁, f₂), N) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - (f12′′′, coeff′′′) = transpose((f₁′′, f₂′), (p₁′, p₂′)) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end +# the sectors of the traced legs have to be dual to each other for a nonzero trace, where the +# sector of an incoming leg is the dual of the corresponding uncoupled sector of the fusion tree +function trace_sectors_match(src::Union{FusionTreePair, FusionTreeBlock}, (q₁, q₂)::Index2Tuple) + f₁, f₂ = src isa FusionTreePair ? src : first(fusiontrees(src)) + N₁ = length(f₁) + legsector(i) = i <= N₁ ? f₁.uncoupled[i] : dual(f₂.uncoupled[i - N₁]) + return all(map((i, j) -> legsector(i) == dual(legsector(j)), q₁, q₂)) +end + +""" + planar_trace((f₁, f₂)::FusionTreePair, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> (f₁′, f₂′) => coeff + planar_trace(src::FusionTreeBlock, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) -> dst => U + +Perform a planar trace of the legs `q₁[i]` with `q₂[i]` of a fusion tree pair, followed by a +planar permutation `(p₁, p₂)` of the remaining legs. The sectors of the traced legs have to be +dual to each other. + +For `FusionStyle(I) isa UniqueFusion`, the input is a single fusion tree pair, which is mapped +to a single output tree pair with a coefficient. For generic `FusionStyle`s, the input and +output consist of `FusionTreeBlock`s, where `src` and `dst` contain all tree pairs of the +traced and remaining legs respectively, and `U[j, i]` is the coefficient of the `j`th tree pair +of `dst` in the trace of the `i`th tree pair of `src`. +""" +function planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, (p₁, p₂)::Index2Tuple, (q₁, q₂)::Index2Tuple) + src isa FusionTreePair && @assert FusionStyle(src) === UniqueFusion() + p′, q′, r = _planar_trace_indices(Val(numout(src)), Val(numin(src)), (p₁, p₂), (q₁, q₂)) + trace_sectors_match(src, (q₁, q₂)) || + throw(ArgumentError(lazy"sectors of the traced legs $q₁ and $q₂ are not dual to each other")) + isempty(q₁) && return transpose(src, (p₁, p₂)) + + # bend all legs to the outgoing side, where the trace only acts on the splitting trees + src′, U₁ = repartition(src, numind(src)) + traced, U₂ = _planar_trace_splitting(src′, q′, r) + # for multiple fusion channels, the remaining legs need not couple to the unit + if traced isa FusionTreeBlock && isempty(traced) + I = sectortype(src) + F = fusiontreetype(I, length(p₁), length(p₂)) + return FusionTreeBlock(F[]) => zeros(fusionscalartype(I), 0, length(src)) else - # TODO: this is a bit of a hack to fix the traces for now - src = FusionTreeBlock([(f₁, f₂)]) - dst, U = repartition(src, N) - for ((f₁′, f₂′), coeff′) in zip(fusiontrees(dst), U) - for (f₁′′, coeff′′) in planar_trace(f₁′, (q₁′, q₂′)) - src′ = FusionTreeBlock([(f₁′′, f₂′)]) - dst′, U′ = transpose(src′, (p₁′, p₂′)) - for (f12′′′, coeff′′′) in zip(fusiontrees(dst′), U′) - coeff = coeff′ * coeff′′ * coeff′′′ - iszero(coeff) || (newtrees[f12′′′] = get(newtrees, f12′′′, zero(coeff)) + coeff) - end - end - end + dst, U₃ = transpose(traced, p′) + return dst => U₃ * U₂ * U₁ end - return newtrees +end + +function _planar_trace_splitting((f, f₀)::FusionTreePair, q′, r) + f′, coeff = only(planar_trace(f, q′)) + return (f′, f₀) => coeff +end +function _planar_trace_splitting(src::FusionTreeBlock, q′, r) + I = sectortype(src) + f₀, = first(fusiontrees(src)) + uncoupled = TupleTools.getindices(f₀.uncoupled, r) + isdual = TupleTools.getindices(f₀.isdual, r) + dst = FusionTreeBlock{I}((uncoupled, ()), (isdual, ())) + U = zeros(fusionscalartype(I), length(dst), length(src)) + positions = Dict(f => i for (i, (f, _)) in enumerate(fusiontrees(dst))) + for (j, (f, _)) in enumerate(fusiontrees(src)), (f′, coeff) in planar_trace(f, q′) + U[positions[f′], j] += coeff + end + return dst => U end """ diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index 40a2961b6..a852bbe8d 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -231,9 +231,12 @@ function Base.hash(f::FusionTree{I}, h::UInt) where {I} end return h end -# Hashing the first tree pair covers the block's defining data `(uncoupled, isdual)`, while -# also disambiguating the ad-hoc single-pair blocks constructed in `planar_trace` -Base.hash(b::FusionTreeBlock, h::UInt) = isempty(b) ? h : hash(first(fusiontrees(b)), h) +# a block is defined by the uncoupled sectors and duality flags shared by all its tree pairs +function Base.hash(b::FusionTreeBlock, h::UInt) + isempty(b) && return h + f₁, f₂ = first(fusiontrees(b)) + return hash(f₁.uncoupled, hash(f₁.isdual, hash(f₂.uncoupled, hash(f₂.isdual, h)))) +end function Base.:(==)(f₁::FusionTree{I, N}, f₂::FusionTree{I, N}) where {I <: Sector, N} f₁.coupled == f₂.coupled || return false diff --git a/src/planar/planaroperations.jl b/src/planar/planaroperations.jl index 7c24e1399..5cae4ea8f 100644 --- a/src/planar/planaroperations.jl +++ b/src/planar/planaroperations.jl @@ -92,24 +92,79 @@ function planartrace!( end @timeit_debug GLOBAL_TIMER "planartrace!" begin - if iszero(β) - fill!(C, β) - elseif !isone(β) - rmul!(C, β) + transformer = treetracer(storagetype(C), space(C), space(A), (p₁, p₂), (q₁, q₂)) + dst, src = _transform_subblocks(C, A, transformer) + ntasks = get_num_transformer_threads(A) + trace_transform_kernel!(dst, src, (p₁, p₂), (q₁, q₂), transformer, α, β, backend, allocator, ntasks) + end + return C +end + +# kernel for computing `tdst = β * tdst + α * tr(tsrc)`, with the planar trace encoded in `transformer` +function trace_transform_kernel!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + transformer::TraceTransformer, α, β, backend, allocator, ntasks::Int + ) + bufsize = buffersize(transformer) + if bufsize == 0 # every block has a single destination subblock + taskforeach(transformer.data, ntasks) do (U, inds_dst, inds_src) + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, nothing, α, β, backend, allocator) + end + else + # see `add_transform_kernel!` for the workspace strategy + cp = TO.allocator_checkpoint!(allocator) + @timeit_debug GLOBAL_TIMER "alloc: buffers" buffers = [ + TO.tensoralloc(storagetype(dst), bufsize, Val(true), allocator) + for _ in 1:min(length(transformer.data), ntasks) + ] + taskforeach(transformer.data, buffers) do (U, inds_dst, inds_src), buffer + _trace_transform_block!(dst, src, p, q, U, inds_dst, inds_src, buffer, α, β, backend, allocator) + end + foreach(Base.Fix2(TO.tensorfree!, allocator), buffers) + TO.allocator_reset!(allocator, cp) + end + return nothing +end + +function _trace_transform_block!( + dst::TransformSubblocks, src::TransformSubblocks, p::Index2Tuple, q::Index2Tuple, + U, inds_dst, inds_src, buffer, α, β, backend, allocator + ) + if size(U, 1) == 1 # single destination: trace every source directly into it + idst = only(inds_dst) + if isempty(inds_src) + @inbounds iszero(β) ? zerovector!(dst[idst]) : scale!(dst[idst], β) + return nothing end - β′ = One() - for (f₁, f₂) in fusiontrees(A) - for ((f₁′, f₂′), coeff) in planar_trace((f₁, f₂), (p₁, p₂), (q₁, q₂)) - @timeit_debug GLOBAL_TIMER "dense: trace" TO.tensortrace!( - C[f₁′, f₂′], - A[f₁, f₂], (p₁, p₂), (q₁, q₂), false, - α * coeff, β′, - backend, allocator - ) + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds begin + TO.tensortrace!(dst[idst], src[first(inds_src)], p, q, false, α * U[1, 1], β, backend, allocator) + for i in 2:length(inds_src) + TO.tensortrace!(dst[idst], src[inds_src[i]], p, q, false, α * U[1, i], One(), backend, allocator) end end + else # trace → recoupling matmul → unpack + rows, cols = size(U) + sz = size(@inbounds(dst[first(inds_dst)])) + ptriv = (ntuple(identity, length(sz)), ()) + blocksize = prod(sz) + buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) + buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) + + # the permutation is applied while tracing, such that unpacking is a plain addition + @timeit_debug GLOBAL_TIMER "dense: trace" @inbounds for (i, isrc) in enumerate(inds_src) + TO.tensortrace!( + sreshape(view(buffer_src, :, i), sz), src[isrc], p, q, false, + One(), Zero(), backend, allocator + ) + end + @timeit_debug GLOBAL_TIMER "dense: recouple mul!" _recouple!(buffer, buffer_dst, buffer_src, U) + @timeit_debug GLOBAL_TIMER "dense: unpack" @inbounds for (j, idst) in enumerate(inds_dst) + TO.tensoradd!( + dst[idst], sreshape(view(buffer_dst, :, j), sz), ptriv, false, α, β, backend, allocator + ) + end end - return C + return nothing end # insert default backend diff --git a/src/tensors/treetransformers.jl b/src/tensors/treetransformers.jl index b274c6f8c..73d910775 100644 --- a/src/tensors/treetransformers.jl +++ b/src/tensors/treetransformers.jl @@ -202,4 +202,104 @@ end return TreeTransformer(A, fusiontreetransform, p, Vdst, Vsrc, conjsrc) end + +""" + TraceTransformer{T, N₁, N₂} + +Planar trace of a tensor with `N₂` indices into a tensor with `N₁` indices, stored as +`(U, inds_dst, inds_src)` where `U[j, i]` is the coefficient with which the trace of source +subblock `inds_src[i]` contributes to destination subblock `inds_dst[j]`. The data is grouped by +destination, such that every destination subblock appears in exactly one entry, where those +without contributions have no source subblocks. The subblock structures of the destination and source +spaces are kept alongside, such that the [`StridedSubblocks`](@ref) of both tensors can be +created without further lookups. +""" +struct TraceTransformer{T, N₁, N₂} + data::Vector{GenericTransformerData{T}} + structure_dst::Vector{StridedStructure{N₁}} + structure_src::Vector{StridedStructure{N₂}} +end + +function TraceTransformer(::Type{A}, Vdst::HomSpace, Vsrc::HomSpace, p::Index2Tuple, q::Index2Tuple) where {A} + t₀ = Base.time() + src_trees, dst_trees = fusiontrees(Vsrc), fusiontrees(Vdst) + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + + # destination positions => (columns of the recoupling matrix, source positions) + groups = Dict{Vector{Int}, Tuple{Vector{Vector{T}}, Vector{Int}}}() + function addcolumn!(inds_dst, u, isrc) + any(!iszero, u) || return nothing + us, inds_src = get!(() -> (Vector{T}[], Int[]), groups, inds_dst) + push!(us, u) + push!(inds_src, isrc) + return nothing + end + @timeit_debug GLOBAL_TIMER "symmetry: trace coefficients" if FusionStyle(sectortype(Vdst)) isa UniqueFusion + for (isrc, f) in enumerate(src_trees) + trace_sectors_match(f, q) || continue + f′, coeff = planar_trace(f, p, q) + _, (_, idst) = gettoken(dst_trees, f′) + addcolumn!([idst], [convert(T, coeff)], isrc) + end + else + for fs_src in fusionblocks(Vsrc) + trace_sectors_match(fs_src, q) || continue + fs_dst, U = planar_trace(fs_src, p, q) + inds_dst = map(fusiontrees(fs_dst)) do f + _, (_, idst) = gettoken(dst_trees, f) + return idst + end + for (i, f) in enumerate(fusiontrees(fs_src)) + _, (_, isrc) = gettoken(src_trees, f) + addcolumn!(inds_dst, convert(Vector{T}, U[:, i]), isrc) + end + end + end + data = GenericTransformerData{T}[(reduce(hcat, us), inds_dst, inds_src) for (inds_dst, (us, inds_src)) in groups] + # destination subblocks without contributions are only scaled + covered = falses(length(dst_trees)) + foreach(inds_dst -> covered[inds_dst] .= true, keys(groups)) + for idst in findall(!, covered) + push!(data, (zeros(T, 1, 0), [idst], Int[])) + end + + structure_dst = degeneracystructure(Vdst).subblockstructure + structure_src = degeneracystructure(Vsrc).subblockstructure + # sort by (approximate) weight to facilitate multi-threading strategies + weights = map(data) do (U, inds_dst, _) + return length(U) * prod(structure_dst[first(inds_dst)][1]) + end + Base.permute!(data, sortperm(weights; rev = true)) + transformer = TraceTransformer(data, structure_dst, structure_src) + + Δt = Base.time() - t₀ + @debug("TraceTransformer for $Vsrc to $Vdst via $p and $q", nblocks = length(data), Δt) + + return transformer +end + +""" + buffersize(transformer::TraceTransformer) -> Int + +Compute the workspace size required to trace, recouple and unpack the largest block with +multiple destination subblocks, i.e. `prod(sz_dst) * (rows + cols)` where `(rows, cols) = size(U)`. +""" +function buffersize(transformer::TraceTransformer) + structure_dst = transformer.structure_dst + return maximum(transformer.data; init = 0) do (U, inds_dst, _) + return size(U, 1) == 1 ? 0 : prod(structure_dst[first(inds_dst)][1]) * sum(size(U)) + end +end + +function tracetransformertype(::Type{A}, Vdst, Vsrc) where {A} + T = recoupling_scalartype(A, sectorscalartype(sectortype(Vdst))) + return TraceTransformer{T, numind(Vdst), numind(Vsrc)} +end + +@cached function treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + )::tracetransformertype(A, Vdst, Vsrc) where {TA} + return TraceTransformer(A, Vdst, Vsrc, p, q) +end + # default cachestyle is GlobalLRUCache diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index 180e19bdc..b4f1f8290 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -249,28 +249,24 @@ using TensorKitSectors end @testset "Double fusion tree: planar trace" begin + p1 = ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,)) + p2, q2 = ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) if FusionStyle(I) isa UniqueFusion f1, f1 = src - dst, U = transpose((f1, f1), ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip((dst,), (U,)) + dst, U = transpose((f1, f1), p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = @constinferred TK.planar_trace(dst, p2, q2) + d2 = Dict(dst2 => U2 * U) else f1, f1 = first(fusiontrees(src)) src′ = FusionTreeBlock{I}((f1.uncoupled, f1.uncoupled), (f1.isdual, f1.isdual)) - dst, U = transpose(src′, ((N + 1, 1:N..., ((2N):-1:(N + 3))...), (N + 2,))) - d1 = zip(fusiontrees(dst), U[:, 1]) + dst, U = transpose(src′, p1) + @test TK.trace_sectors_match(dst, q2) + dst2, U2 = TK.planar_trace(dst, p2, q2) + d2 = Dict(zip(fusiontrees(dst2), U2 * U[:, 1])) end f1front, = TK.split(f1, N - 1) - T = sectorscalartype(I) - d2 = Dict{typeof((f1front, f1front)), T}() - for ((f1′, f2′), coeff′) in d1 - for ((f1′′, f2′′), coeff′′) in TK.planar_trace( - (f1′, f2′), ((2:N...,), (1, ((2N):-1:(N + 3))...)), ((N + 1,), (N + 2,)) - ) - coeff = coeff′ * coeff′′ - d2[(f1′′, f2′′)] = get(d2, (f1′′, f2′′), zero(coeff)) + coeff - end - end for ((f1_, f2_), coeff) in d2 if (f1_, f2_) == (f1front, f1front) @test coeff ≈ dim(f1.coupled) / dim(f1front.coupled) From 57b59bc2b7609edfb8c186664adb0ac7021b10d3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 12:06:16 -0400 Subject: [PATCH 02/13] Add a GPU kernel for planar traces GPU tensors pack the `TraceTransformer` into `isbits` descriptors, cached per device storagetype, and trace with one thread per destination element, fusing the sum over the traced entries and the recoupling into a single kernel, analogous to the generic tree transformers. Other tensor types, such as adjoints, are first converted to `TensorMap`s: tracing subblock views with cuTENSOR fails for data offsets that are not aligned (e.g. `Float64` data at odd offsets). `_linear_index` is now written as an explicit loop, since `sum` of an empty tuple, as occurs for a full trace, does not compile for GPUs. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 +- ext/TensorKitGPUArraysExt.jl | 167 ++++++++++++++++++++++++++++++++++- test/cuda/planar.jl | 18 ++++ 3 files changed, 182 insertions(+), 5 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 130351ddb..74720d59a 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -30,7 +30,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Require TensorKitSectors v0.3.10 - `@cached` can be used from other modules, such as package extensions, which then own the cache of the methods they add; these caches are shown with their module in `global_cache_info` after registration in `__init__` (internal) - The `TreeTransformer`s used in index manipulations are now constructed and cached per storagetype of the destination tensor, which is also a dispatch point for storage-specific transformers with their own cache (internal) -- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s (internal) +- `planartrace!` precomputes and caches a `TraceTransformer` with the subblock positions and recoupling coefficients, and executes it with a kernel shared by all sectortypes, analogous to the index manipulations; `planar_trace` is now also defined for `FusionTreeBlock`s, and GPU tensors use a dedicated kernel (internal) ### Deprecated diff --git a/ext/TensorKitGPUArraysExt.jl b/ext/TensorKitGPUArraysExt.jl index 5e013865d..610df92ff 100644 --- a/ext/TensorKitGPUArraysExt.jl +++ b/ext/TensorKitGPUArraysExt.jl @@ -12,8 +12,9 @@ using TensorKit.TensorOperations: linearize, DefaultAllocator using TensorKit.Factorizations using TensorKit.Factorizations: AbstractAlgorithm using TensorKit: SectorDict, tensormaptype, scalar, similarstoragetype, AdjointTensorMap, scalartype, project_symmetric_and_check -using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple -import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel! +using TensorKit: UniqueTreeTransformer, GenericTreeTransformer, TraceTransformer, @cached, TensorMapSpace, Index2Tuple, IndexTuple +using TensorKit: SubblockIterator +import TensorKit: randisometry, rand, randn, fill_braidingsubblock!, add_transform_kernel!, trace_transform_kernel! function TensorKit.fill_braidingsubblock!(data::TD, val) where {T, TD <: Union{<:AnyGPUMatrix{T}, <:StridedViews.StridedView{T, 4, <:AnyGPUArray{T}}}} # COV_EXCL_START @@ -271,6 +272,75 @@ function DeviceGenericTreeTransformer( ) end +""" + TraceTransformerBlock{N} + +Descriptor for a destination block of a `TraceTransformer`, indexing into the flat +`coeffs`/`structs_dst`/`structs_src` vectors of a `DeviceTraceTransformer`. +""" +struct TraceTransformerBlock{N} + sz::NTuple{N, Int} + dense_strides::NTuple{N, Int} + rows::Int + cols::Int + u_offset::Int + dst_offset::Int + src_offset::Int +end + +# (permuted strides of the remaining legs, sizes, dense strides and summed strides of the traced +# pairs, offset): the traced entries of a source subblock are found through a single stride per pair +const TraceSourceStructure{N, M} = Tuple{NTuple{N, Int}, NTuple{M, Int}, NTuple{M, Int}, NTuple{M, Int}, Int} + +struct DeviceTraceTransformer{VB <: AbstractVector{<:TraceTransformerBlock}, VO <: AbstractVector{Int}, VC <: AbstractVector{<:Number}, VD <: AbstractVector{<:Tuple{<:Tuple{Vararg{Int}}, Int}}, VS <: AbstractVector{<:TraceSourceStructure}} + blocks::VB + work_offsets::VO + nwork::Int + coeffs::VC # every `U`, concatenated in column-major order + structs_dst::VD + structs_src::VS +end + +function DeviceTraceTransformer(transformer::TraceTransformer{T, N₁, N₂}, p, (q₁, q₂)) where {T, N₁, N₂} + lp = linearize(p) + M = length(q₁) + blocks = TraceTransformerBlock{N₁}[] + coeffs = T[] + structs_dst = TreeStructure{N₁}[] + structs_src = TraceSourceStructure{N₁, M}[] + + (; structure_dst, structure_src) = transformer + for (U, inds_dst, inds_src) in transformer.data + # all destination trees in a block share the same subblock size + size_dst = first(structure_dst[first(inds_dst)]) + push!( + blocks, TraceTransformerBlock{N₁}( + size_dst, _dense_strides(size_dst), size(U, 1), size(U, 2), + length(coeffs), length(structs_dst), length(structs_src) + ) + ) + append!(coeffs, U) + for idst in inds_dst + _, strides_dst, offset_dst = structure_dst[idst] + push!(structs_dst, (strides_dst, offset_dst)) + end + for isrc in inds_src + size_src, strides_src, offset_src = structure_src[isrc] + size_traced = TupleTools.getindices(size_src, q₁) + strides_traced = TupleTools.getindices(strides_src, q₁) .+ TupleTools.getindices(strides_src, q₂) + push!( + structs_src, ( + TupleTools.getindices(strides_src, lp), size_traced, + _dense_strides(size_traced), strides_traced, offset_src, + ) + ) + end + end + + work_offsets, nwork = _work_offsets(blk.rows * prod(blk.sz) for blk in blocks) + return DeviceTraceTransformer(blocks, work_offsets, nwork, coeffs, structs_dst, structs_src) +end + # Cache the device transformer directly, constructing the cached host version only on a miss. @cached function TensorKit.treebraider( A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, @@ -287,6 +357,13 @@ end return device_transformer(A, host, linearize(p)) end +@cached function TensorKit.treetracer( + A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, q::Index2Tuple + ) where {TA <: AnyGPUVector} + host = TensorKit.treetracer(Vector{scalartype(A)}, Vdst, Vsrc, p, q) + return device_transformer(A, host, p, q) +end + const DEFAULT_DEVICECACHE_SIZE = 1000 function __init__() @@ -295,12 +372,14 @@ function __init__() TensorKit.GLOBAL_CACHES, ( Symbol("TensorKitGPUArraysExt.GLOBAL_TREEBRAIDER_CACHE") => GLOBAL_TREEBRAIDER_CACHE, Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRANSPOSER_CACHE") => GLOBAL_TREETRANSPOSER_CACHE, + Symbol("TensorKitGPUArraysExt.GLOBAL_TREETRACER_CACHE") => GLOBAL_TREETRACER_CACHE, ) ) # device-caches should probably have smaller sizes: resize!(GLOBAL_TREEBRAIDER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) resize!(GLOBAL_TREETRANSPOSER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) + resize!(GLOBAL_TREETRACER_CACHE; maxsize = DEFAULT_DEVICECACHE_SIZE) return nothing end @@ -332,6 +411,18 @@ function device_transformer(A, transformer::GenericTreeTransformer, p) end end +function device_transformer(A, transformer::TraceTransformer, p, q) + packed = DeviceTraceTransformer(transformer, p, q) + return GPUArrays.@uncached begin + proto = A(undef, 0) + upload(x) = copy!(similar(proto, eltype(x), size(x)), x) + DeviceTraceTransformer( + upload(packed.blocks), upload(packed.work_offsets), packed.nwork, + upload(packed.coeffs), upload(packed.structs_dst), upload(packed.structs_src) + ) + end +end + # COV_EXCL_START # kernels are not reachable by coverage @@ -360,7 +451,12 @@ end # finds the overall linear index in the output and input arrays corresponding to the **sublock** # coordinates currently being worked on @inline function _linear_index(coords::NTuple{N, Int}, strides::NTuple{N, Int}, offset) where {N} - return offset + sum(ntuple(n -> coords[n] * strides[n], Val(N))) + 1 + # an explicit loop, since `sum` of an empty tuple (e.g. a full trace) does not compile for GPUs + i = offset + 1 + for n in 1:N + i += coords[n] * strides[n] + end + return i end # One thread per destination element in `data_dst`. `op` is `identity` or `conj`, and is @@ -426,6 +522,44 @@ end end end +# COV_EXCL_START +# One thread per destination element, as for the generic transformers, where every source +# additionally contributes the sum over its traced entries. +@kernel function batched_trace_kernel!( + data_dst, data_src, blocks, work_offsets, coeffs, structs_dst, structs_src, α, β, nwork + ) + # w = @index(Global, Linear) - 1 does not work due to JLArrays bug + w = @index(Global, Linear) + w -= 1 + if w < nwork + b = _searchblock(work_offsets, w) + blk = @inbounds blocks[b] + local_w = w - (@inbounds work_offsets[b]) + blocksize = prod(blk.sz) + i = local_w ÷ blocksize # 0-based destination tree + coords = _coordinates(local_w % blocksize, blk.sz, blk.dense_strides) + + st_dst, offs_dst = @inbounds structs_dst[blk.dst_offset + i + 1] + i_dst = _linear_index(coords, st_dst, offs_dst) + + acc = zero(promote_type(eltype(data_src), eltype(coeffs))) + @inbounds for j in 1:blk.cols + coeff = coeffs[blk.u_offset + 1 + i + (j - 1) * blk.rows] + iszero(coeff) && continue + pst_src, sz_tr, dense_tr, st_tr, offs_src = structs_src[blk.src_offset + j] + i_src = _linear_index(coords, pst_src, offs_src) + tr = zero(eltype(data_src)) + for t in 0:(prod(sz_tr) - 1) + tr += data_src[i_src + _linear_index(_coordinates(t, sz_tr, dense_tr), st_tr, -1)] + end + acc += coeff * tr + end + # destinations without contributions are overwritten for `β == 0` + @inbounds data_dst[i_dst] = iszero(β) ? α * acc : α * acc + β * data_dst[i_dst] + end +end +# COV_EXCL_STOP + function _launch_unique!(data_dst, data_src, op, transformer, α, β, ::Val{N}) where {N} nwork = transformer.nwork nwork == 0 && return nothing @@ -450,7 +584,7 @@ end # we don't want/need the `StridedSubblocks` here because that information is already in the transformer TensorKit._transform_subblocks( tdst::TensorMap, tsrc::TensorMap, - transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer} + transformer::Union{DeviceUniqueTreeTransformer, DeviceGenericTreeTransformer, DeviceTraceTransformer} ) = (tdst.data, tsrc.data) function TensorKit.add_transform_kernel!( @@ -476,4 +610,29 @@ function TensorKit.add_transform_kernel!( return nothing end +function TensorKit.trace_transform_kernel!( + dst::AnyGPUVector, src::AnyGPUVector, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + device.nwork == 0 && return nothing + batched_trace_kernel!(get_backend(dst))( + dst, src, device.blocks, device.work_offsets, device.coeffs, + device.structs_dst, device.structs_src, α, β, device.nwork; ndrange = device.nwork + ) + return nothing +end + +# other tensor types are first converted to `TensorMap`s: tracing subblock by subblock with +# `tensortrace!` is not an option, since cuTENSOR requires aligned data for the subblock views +function TensorKit.trace_transform_kernel!( + dst::SubblockIterator, src::SubblockIterator, p::Index2Tuple, q::Index2Tuple, + device::DeviceTraceTransformer, α, β, backend, allocator, ntasks::Int + ) + _tensormap(t) = t isa TensorMap ? t : copy!(similar(t, scalartype(t), space(t)), t) + tdst, tsrc = _tensormap(dst.t), _tensormap(src.t) + trace_transform_kernel!(tdst.data, tsrc.data, p, q, device, α, β, backend, allocator, ntasks) + tdst === dst.t || copy!(dst.t, tdst) + return nothing +end + end diff --git a/test/cuda/planar.jl b/test/cuda/planar.jl index 2d992023b..8cec768db 100644 --- a/test/cuda/planar.jl +++ b/test/cuda/planar.jl @@ -74,6 +74,24 @@ for V in spacelist end end +for V in spacelist + I = sectortype(first(V)) + Istr = type_repr(I) + @timedtestset "planartrace + CUDA with symmetry: $Istr" begin + # (space, p, q) for a trace of neighbouring legs and a nested trace + cases = ( + (V[1] ⊗ V[2] ⊗ V[3] ← V[4] ⊗ V[3], ((1, 2), (4,)), ((3,), (5,))), + (V[1] ⊗ V[2] ⊗ V[3] ⊗ V[4] ← V[3] ⊗ V[4], ((1,), (2,)), ((3, 4), (5, 6))), + ) + for (W, p, q) in cases, (α, β) in ((0.7, 0.3), (1.0, 0.0)) + A = randn(ComplexF64, W) + C = randn(ComplexF64, TensorKit.select(W, p)) + Cgpu = planartrace!(adapt(CuArray, C), adapt(CuArray, A), p, q, α, β) + @test adapt(Vector{ComplexF64}, Cgpu) ≈ planartrace!(C, A, p, q, α, β) + end + end +end + @testset "planar methods" verbose = true begin @testset "planaradd" begin A = CUDA.randn(ℂ^2 ⊗ ℂ^3 ← ℂ^6 ⊗ ℂ^5 ⊗ ℂ^4) From c8a4b65dca90ab009617d29341df751b2e3e1d27 Mon Sep 17 00:00:00 2001 From: leburgel Date: Wed, 14 Oct 2026 12:59:05 +0200 Subject: [PATCH 03/13] Keep the intermediate blocks of `fsbraid` in production order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bends and Artin braids of a fusion tree block used to construct their destination block in canonical order, together with a dictionary from trees to rows, at every step. The intermediate blocks of `fsbraid` now keep their trees in the order in which they are produced (`TreeRows`, `canonical = Val(false)` for `bendleft`, `bendright` and `artin_braid`, `repartition_noncanonical`), and only the result is brought into canonical order (`canonicalize`). Small blocks find rows by a linear search. In `artin_braid`, the channels are obtained by filtering `a ⊗ d` instead of intersecting two fusion products, and the symbols that do not depend on the channel are computed once. Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 63 +++++++++++++---------- src/fusiontrees/duality_manipulations.jl | 41 +++++++++------ src/fusiontrees/fusiontrees.jl | 49 ++++++++++++++++++ 3 files changed, 110 insertions(+), 43 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 368e326a7..f1d5e6891 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -78,7 +78,16 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} return f′ => coeff end -function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where {I, N} +# the inner lines `c′` of `a ⊗ d → c′ → c′ ⊗ b → e`, i.e. `intersect(a ⊗ d, e ⊗ conj(b))`, +# without constructing the second fusion product for sectors with a simple unit +function _artin_channels(a::I, b::I, d::I, e::I) where {I <: Sector} + UnitStyle(I) isa SimpleUnit || return intersect(a ⊗ d, e ⊗ conj(b)) + return Iterators.filter(c′ -> !iszero(Nsymbol(c′, b, e)), a ⊗ d) +end + +function artin_braid( + src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, canonical::Val = Val(true) + ) where {I, N} 1 <= i < N || throw(ArgumentError("Cannot swap outputs i=$i and i+1 out of only $N outputs")) isempty(src) && return src => zeros(sectorscalartype(I), 0, 0) @@ -91,12 +100,11 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where isdual = src.isdual[1] isdual′ = TupleTools.setindex(isdual, isdual[i], i + 1) isdual′ = TupleTools.setindex(isdual′, isdual[i + 1], i) - dst = FusionTreeBlock{I}((uncoupled′, ()), (isdual′, ()); sizehint = length(src)) + rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical) oneT = one(sectorscalartype(I)) - indexmap = treeindex_map(dst) - U = zeros(sectorscalartype(I), length(dst), length(src)) + U = zeros(sectorscalartype(I), max(length(rows.trees), length(src)), length(src)) # see `bendright` if isunit(a) || isunit(b) # braiding with trivial sector: simple and always possible for (col, (f, f₂)) in enumerate(fusiontrees(src)) @@ -113,10 +121,10 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where vertices′ = TupleTools.setindex(vertices′, vertices[i - 1], i) end f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = oneT end - return dst => U + return FusionTreeBlock(rows.trees) => U end BraidingStyle(I) isa NoBraiding && @@ -132,7 +140,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where if FusionStyle(I) isa MultiplicityFreeFusion R = oftype(oneT, (inv ? conj(Rsymbol(b, a, c)) : Rsymbol(a, b, c))) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = R else # GenericFusion μ = vertices[1] @@ -142,7 +150,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where iszero(R) && continue vertices′ = TupleTools.setindex(vertices, ν, 1) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = R end end @@ -155,29 +163,30 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where c = inner_extended[i] e = inner_extended[i + 1] if FusionStyle(I) isa MultiplicityFreeFusion - for c′ in intersect(a ⊗ d, e ⊗ conj(b)) + R = inv ? Rsymbol(d, c, e) : Rsymbol(c, d, e) + for c′ in _artin_channels(a, b, d, e) coeff = if inv - conj(Rsymbol(d, c, e) * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) + conj(R * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) else - Rsymbol(c, d, e) * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) + R * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) end iszero(coeff) && continue inner′ = TupleTools.setindex(inner, c′, i - 1) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = coeff end else # GenericFusion - for c′ in intersect(a ⊗ d, e ⊗ conj(b)) - Rmat1 = inv ? Rsymbol(d, c, e)' : Rsymbol(c, d, e) + Rmat1 = inv ? Rsymbol(d, c, e)' : Rsymbol(c, d, e) + μ = vertices[i - 1] + ν = vertices[i] + for c′ in _artin_channels(a, b, d, e) Rmat2 = inv ? Rsymbol(d, a, c′)' : Rsymbol(a, d, c′) Fmat = Fsymbol(d, a, b, e, c′, c) - μ = vertices[i - 1] - ν = vertices[i] - for σ in 1:Nsymbol(a, d, c′) - for λ in 1:Nsymbol(c′, b, e) + for σ in axes(Rmat2, 1) + for λ in axes(Fmat, 2) coeff = zero(oneT) - for ρ in 1:Nsymbol(d, c, e), κ in 1:Nsymbol(d, a, c′) + for ρ in axes(Rmat1, 2), κ in axes(Rmat2, 2) coeff += Rmat1[ν, ρ] * conj(Fmat[κ, λ, μ, ρ]) * conj(Rmat2[σ, κ]) end @@ -186,7 +195,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where vertices′ = TupleTools.setindex(vertices′, λ, i) inner′ = TupleTools.setindex(inner, c′, i - 1) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′, vertices′) - row = indexmap[treeindex_data((f′, f₂))] + row = row!(rows, (f′, f₂)) @inbounds U[row, col] = coeff end end @@ -194,7 +203,7 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where end end - return dst => U + return FusionTreeBlock(rows.trees) => U end # braid fusion tree @@ -316,13 +325,15 @@ end p = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) - dst, U = repartition(src, numind(src)) + # the intermediate blocks keep their trees in the order in which they are produced, + # and only the result is brought into canonical order + dst, U = repartition_noncanonical(src, Val(numind(src))) braid_style = BraidingStyle(I) for s in permutation2swaps(p) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] - dst, U_tmp = artin_braid(dst, s; inv) + dst, U_tmp = artin_braid(dst, s; inv, canonical = Val(false)) U = U_tmp * U l = levels[s] levels = TupleTools.setindex(levels, levels[s + 1], s) @@ -330,11 +341,11 @@ end end if N₂ == 0 - return dst => U + return canonicalize(dst, U) else - dst, U_tmp = repartition(dst::FusionTreeBlock{I}, N₁) + dst, U_tmp = repartition_noncanonical(dst::FusionTreeBlock{I}, Val(N₁)) U = U_tmp * U - return dst => U + return canonicalize(dst, U) end end end diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index d6324f9d4..0277532a2 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -66,7 +66,7 @@ function bendright((f₁, f₂)::FusionTreePair) return (f₁′, f₂′) => coeff end -function bendright(src::FusionTreeBlock) +function bendright(src::FusionTreeBlock, canonical::Val = Val(true)) I = sectortype(src) N₁ = numout(src) N₂ = numin(src) @@ -80,9 +80,10 @@ function bendright(src::FusionTreeBlock) (src.isdual[2]..., !(src.isdual[1][N₁])), ) - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) - U = zeros(fusionscalartype(I), length(dst), length(src)) + rows = TreeRows(I, uncoupled_dst, isdual_dst, length(src), canonical) + # the canonical destination block has more trees than `src` if `src` holds only some of the + # trees with its uncoupled sectors (as in `planar_trace`) + U = zeros(fusionscalartype(I), max(length(rows.trees), length(src)), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) @@ -92,7 +93,7 @@ function bendright(src::FusionTreeBlock) end if FusionStyle(I) isa MultiplicityFreeFusion coeff = coeff₀ * Bsymbol(a, b, c) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = coeff else Bmat = Bsymbol(a, b, c) @@ -106,12 +107,12 @@ function bendright(src::FusionTreeBlock) iszero(coeff) && continue vertices₂ = N₂ > 0 ? (f₂.vertices..., ν) : () f₂′ = FusionTree(uncoupled₂, coupled₂, isdual₂, inner₂, vertices₂) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @doc """ @@ -145,7 +146,7 @@ end # !! note that this is more or less a copy of bendright through # (f1, f2) => conj(coeff) for ((f2, f1), coeff) in bendleft(src) -function bendleft(src::FusionTreeBlock) +function bendleft(src::FusionTreeBlock, canonical::Val = Val(true)) I = sectortype(src) N₁ = numout(src) N₂ = numin(src) @@ -159,9 +160,10 @@ function bendleft(src::FusionTreeBlock) TupleTools.front(src.isdual[2]), ) - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) - U = zeros(fusionscalartype(I), length(dst), length(src)) + rows = TreeRows(I, uncoupled_dst, isdual_dst, length(src), canonical) + # the canonical destination block has more trees than `src` if `src` holds only some of the + # trees with its uncoupled sectors (as in `planar_trace`) + U = zeros(fusionscalartype(I), max(length(rows.trees), length(src)), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₂′, f₁′) = _bendright_treepair((f₂, f₁)) @@ -171,7 +173,7 @@ function bendleft(src::FusionTreeBlock) end if FusionStyle(I) isa MultiplicityFreeFusion coeff = coeff₀ * Bsymbol(a, b, c) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = conj(coeff) else Bmat = Bsymbol(a, b, c) @@ -185,12 +187,12 @@ function bendleft(src::FusionTreeBlock) iszero(coeff) && continue vertices₁ = N₁ > 0 ? (f₁.vertices..., ν) : () f₁′ = FusionTree(uncoupled₁, coupled₁, isdual₁, inner₁, vertices₁) - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] = conj(coeff) end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @doc """ @@ -478,7 +480,7 @@ return dst, U # it, so that it is visible to the generator's world age. Otherwise a precompile workload that # first triggers `repartition` fails with `UndefVarError: _repartition_body` (at runtime the # first call happens at a later world age, which masks the ordering issue). -function _repartition_body(N) +function _repartition_body(N, canonical = true) if N == 0 ex = quote T = fusionscalartype(sectortype(src)) @@ -491,13 +493,14 @@ function _repartition_body(N) end else f = N < 0 ? bendleft : bendright + args = canonical ? () : (Val(false),) ex_rep = Expr(:block) for _ in 1:(abs(N) - 1) - push!(ex_rep.args, :((dst, Utmp) = $f(dst))) + push!(ex_rep.args, :((dst, Utmp) = $f(dst, $(args...)))) push!(ex_rep.args, :(U = Utmp * U)) end ex = quote - dst, U = $f(src) + dst, U = $f(src, $(args...)) $ex_rep return dst => U end @@ -507,6 +510,10 @@ end @generated function repartition(src::Union{FusionTreePair, FusionTreeBlock}, ::Val{N}) where {N} return _repartition_body(numout(src) - N) end +# the same for a block, with the trees of the result in the order in which they are produced +@generated function repartition_noncanonical(src::FusionTreeBlock, ::Val{N}) where {N} + return _repartition_body(numout(src) - N, false) +end """ transpose((f₁, f₂)::FusionTreePair{I}, p::Index2Tuple{N₁, N₂}) where {I, N₁, N₂} diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index a852bbe8d..1d698d059 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -117,6 +117,8 @@ const FusionTreePair{I, N₁, N₂} = Tuple{FusionTree{I, N₁}, FusionTree{I, N Collection of fusion-splitting tree pairs that share the same uncoupled sectors. Mostly internal structure to speed up non-`UniqueFusion` fusiontree manipulation computations. +Braiding a block requires it to be complete, i.e. to contain all the fusion-splitting tree pairs +with its uncoupled sectors and duality flags, as `FusionTreeBlock{I}(uncoupled, isdual)` does. """ struct FusionTreeBlock{I, N₁, N₂, F <: FusionTreePair{I, N₁, N₂}} trees::Vector{F} @@ -276,6 +278,53 @@ function treeindex_map(fs::FusionTreeBlock) return fusiontreedict(I)(treeindex_data(f) => ind for (ind, f) in enumerate(fusiontrees(fs))) end +# The trees of the destination block of a block manipulation, and the row of each of them in the +# transformation matrix. With `Val(true)`, these are the trees of the destination block in their +# canonical order; with `Val(false)`, the trees are added in the order in which they are produced, +# which avoids constructing the destination block. The latter is used for the intermediate blocks +# of a composite manipulation, of which only the result is brought into canonical order. +# With `Val(false)`, the source block has to be complete: the destination block then has as many +# trees (`n`), and an incomplete source block raises an error if its destination block has more. +struct TreeRows{F, D} + trees::Vector{F} + index::Union{D, Nothing} # `nothing`: few trees, the rows are found by a linear search + n::Int +end +function TreeRows(::Type{I}, uncoupled, isdual, n::Int, ::Val{canonical}) where {I, canonical} + trees = if canonical + fusiontrees(FusionTreeBlock{I}(uncoupled, isdual; sizehint = n)) + else + sizehint!(Vector{fusiontreetype(I, length(uncoupled[1]), length(uncoupled[2]))}(), n) + end + K = Base.promote_op(treeindex_data, eltype(trees)) + index = n > 32 ? Dict{K, Int}(treeindex_data(f) => i for (i, f) in enumerate(trees)) : nothing + return TreeRows{eltype(trees), Dict{K, Int}}(trees, index, n) +end +function row!(rows::TreeRows, f) + key = treeindex_data(f) + i = if isnothing(rows.index) + findfirst(f′ -> treeindex_data(f′) == key, rows.trees) + else + get(rows.index, key, nothing) + end + isnothing(i) || return i + length(rows.trees) < rows.n || + throw(ArgumentError("braiding requires a complete `FusionTreeBlock`")) + push!(rows.trees, f) + isnothing(rows.index) || (rows.index[key] = length(rows.trees)) + return length(rows.trees) +end + +# bring a block whose trees are not in canonical order, together with the transformation matrix +# whose rows correspond to these trees, into canonical order +function canonicalize(src::FusionTreeBlock{I}, U::AbstractMatrix) where {I} + length(src) == 1 && return src => U + rows = TreeRows(I, src.uncoupled, src.isdual, length(src), Val(true)) + U′ = zeros(eltype(U), length(rows.trees), size(U, 2)) + U′[map(f -> row!(rows, f), fusiontrees(src)), :] = U + return FusionTreeBlock(rows.trees) => U′ +end + # Facilitate getting correct fusion tree types Base.@assume_effects :foldable function fusiontreetype(::Type{I}, N::Int) where {I <: Sector} From dc0469dbd2c249bd1624f25af8ca0b7554caef2b Mon Sep 17 00:00:00 2001 From: leburgel Date: Wed, 14 Oct 2026 12:59:05 +0200 Subject: [PATCH 04/13] Recycle tree vectors and matrices between the braids in `fsbraid` The Artin braids keep the size of a block, so `fsbraid` passes the tree vector and the matrices of the previous step on to `artin_braid` and multiplies in place. Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 18 +++++++++++++----- src/fusiontrees/fusiontrees.jl | 10 +++++++--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index f1d5e6891..af51e5dd1 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -86,7 +86,8 @@ function _artin_channels(a::I, b::I, d::I, e::I) where {I <: Sector} end function artin_braid( - src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, canonical::Val = Val(true) + src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, canonical::Val = Val(true), + buffers = (nothing, nothing) # a tree vector and a matrix that may be overwritten ) where {I, N} 1 <= i < N || throw(ArgumentError("Cannot swap outputs i=$i and i+1 out of only $N outputs")) isempty(src) && return src => zeros(sectorscalartype(I), 0, 0) @@ -100,11 +101,12 @@ function artin_braid( isdual = src.isdual[1] isdual′ = TupleTools.setindex(isdual, isdual[i], i + 1) isdual′ = TupleTools.setindex(isdual′, isdual[i + 1], i) - rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical) + rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical, buffers[1]) oneT = one(sectorscalartype(I)) - U = zeros(sectorscalartype(I), max(length(rows.trees), length(src)), length(src)) # see `bendright` + U = buffers[2] isa Matrix{typeof(oneT)} ? fill!(buffers[2], zero(oneT)) : + zeros(sectorscalartype(I), max(length(rows.trees), length(src)), length(src)) # see `bendright` if isunit(a) || isunit(b) # braiding with trivial sector: simple and always possible for (col, (f, f₂)) in enumerate(fusiontrees(src)) @@ -329,12 +331,18 @@ end # and only the result is brought into canonical order dst, U = repartition_noncanonical(src, Val(numind(src))) + # the swaps keep the size of the block, so that the tree vector and the matrices of the + # previous step are recycled (except for the trees of `src`, which are part of the key) + buffers = (nothing, nothing) + U_spare = nothing braid_style = BraidingStyle(I) for s in permutation2swaps(p) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] - dst, U_tmp = artin_braid(dst, s; inv, canonical = Val(false)) - U = U_tmp * U + dst′, U_tmp = artin_braid(dst, s; inv, canonical = Val(false), buffers) + buffers = (dst.trees === src.trees ? nothing : dst.trees, U_tmp) + dst = dst′ + U_spare, U = U, U_spare isa typeof(U_tmp) ? mul!(U_spare, U_tmp, U) : U_tmp * U l = levels[s] levels = TupleTools.setindex(levels, levels[s + 1], s) levels = TupleTools.setindex(levels, l, s + 1) diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index 1d698d059..013b930a4 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -282,7 +282,8 @@ end # transformation matrix. With `Val(true)`, these are the trees of the destination block in their # canonical order; with `Val(false)`, the trees are added in the order in which they are produced, # which avoids constructing the destination block. The latter is used for the intermediate blocks -# of a composite manipulation, of which only the result is brought into canonical order. +# of a composite manipulation, of which only the result is brought into canonical order; the tree +# vector of an intermediate block that is no longer needed can then be passed on as `buffer`. # With `Val(false)`, the source block has to be complete: the destination block then has as many # trees (`n`), and an incomplete source block raises an error if its destination block has more. struct TreeRows{F, D} @@ -290,11 +291,14 @@ struct TreeRows{F, D} index::Union{D, Nothing} # `nothing`: few trees, the rows are found by a linear search n::Int end -function TreeRows(::Type{I}, uncoupled, isdual, n::Int, ::Val{canonical}) where {I, canonical} +function TreeRows( + ::Type{I}, uncoupled, isdual, n::Int, ::Val{canonical}, buffer = nothing + ) where {I, canonical} + F = fusiontreetype(I, length(uncoupled[1]), length(uncoupled[2])) trees = if canonical fusiontrees(FusionTreeBlock{I}(uncoupled, isdual; sizehint = n)) else - sizehint!(Vector{fusiontreetype(I, length(uncoupled[1]), length(uncoupled[2]))}(), n) + buffer isa Vector{F} ? empty!(buffer) : sizehint!(Vector{F}(), n) end K = Base.promote_op(treeindex_data, eltype(trees)) index = n > 32 ? Dict{K, Int}(treeindex_data(f) => i for (i, f) in enumerate(trees)) : nothing From f19849bb9b7c3c84e4c97aebe80dafadc62c95af Mon Sep 17 00:00:00 2001 From: leburgel Date: Wed, 14 Oct 2026 12:59:06 +0200 Subject: [PATCH 05/13] Braid single-tree blocks as a product of scalars A block with a single tree stays a single tree under every step of `fsbraid`, so its braid is computed with the per-tree manipulations, without blocks or matrices. Their bodies move into private `_artin_braid_single` and `_bendright_single`, which take the single coefficient of each symbol with `only`; the public methods keep their `UniqueFusion` assertion. Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 30 +++++++++++++++++++---- src/fusiontrees/duality_manipulations.jl | 21 +++++++++++++--- 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index af51e5dd1..73ad1b6da 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -16,8 +16,13 @@ tree with non-zero coefficient, namely `f` with coefficient `1`. This keyword ha if `BraidingStyle(sectortype(f)) isa SymmetricBraiding`. """ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} - 1 <= i < N || throw(ArgumentError(lazy"Cannot swap outputs i=$i and i+1 out of only $N outputs")) @assert FusionStyle(I) === UniqueFusion() + return _artin_braid_single(f, i, inv) +end +# also for a tree that is the only one with its uncoupled sectors, whose braid is again a single +# tree (`only` throws otherwise) +function _artin_braid_single(f::FusionTree{I, N}, i, inv::Bool) where {I, N} + 1 <= i < N || throw(ArgumentError(lazy"Cannot swap outputs i=$i and i+1 out of only $N outputs")) uncoupled = f.uncoupled a, b = uncoupled[i], uncoupled[i + 1] @@ -53,7 +58,7 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} if i == 1 c = N > 2 ? inner[1] : coupled′ - R = oftype(oneT, (inv ? conj(Rsymbol(b, a, c)) : Rsymbol(a, b, c))) + R = oftype(oneT, (inv ? conj(only(Rsymbol(b, a, c))) : only(Rsymbol(a, b, c)))) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner, vertices) return f′ => R end @@ -64,13 +69,13 @@ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} a = inner_extended[i - 1] c = inner_extended[i] e = inner_extended[i + 1] - c′ = only(a ⊗ d) + c′ = FusionStyle(I) isa UniqueFusion ? only(a ⊗ d) : only(_artin_channels(a, b, d, e)) coeff = oftype( oneT, if inv - conj(Rsymbol(d, c, e) * Fsymbol(d, a, b, e, c′, c)) * Rsymbol(d, a, c′) + conj(only(Rsymbol(d, c, e)) * only(Fsymbol(d, a, b, e, c′, c))) * only(Rsymbol(d, a, c′)) else - Rsymbol(c, d, e) * conj(Fsymbol(d, a, b, e, c′, c) * Rsymbol(a, d, c′)) + only(Rsymbol(c, d, e)) * conj(only(Fsymbol(d, a, b, e, c′, c)) * only(Rsymbol(a, d, c′))) end ) inner′ = TupleTools.setindex(inner, c′, i - 1) @@ -327,6 +332,21 @@ end p = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) + if length(src) == 1 # then so is every intermediate block: the braid is a product of scalars + (f, f₀), coeff = _bend_single(_bendleft_single, only(fusiontrees(src)), Val(numin(src))) + for s in permutation2swaps(p) + _check_levels(BraidingStyle(I), levels, s) + inv = levels[s] > levels[s + 1] + f, c = _artin_braid_single(f, s, inv) + coeff *= c + l = levels[s] + levels = TupleTools.setindex(levels, levels[s + 1], s) + levels = TupleTools.setindex(levels, l, s + 1) + end + (f₁′, f₂′), c = _bend_single(_bendright_single, (f, f₀), Val(N₂)) + return FusionTreeBlock([(f₁′, f₂′)]) => fill(coeff * c, 1, 1) + end + # the intermediate blocks keep their trees in the order in which they are produced, # and only the result is brought into canonical order dst, U = repartition_noncanonical(src, Val(numind(src))) diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 0277532a2..ca49a3a47 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -54,15 +54,18 @@ function _bendright_treepair((f₁, f₂)::FusionTreePair) end function bendright((f₁, f₂)::FusionTreePair) - I = sectortype((f₁, f₂)) + @assert FusionStyle((f₁, f₂)) === UniqueFusion() + return _bendright_single((f₁, f₂)) +end +# also for a final fusion vertex without multiplicity (`only` throws otherwise) +function _bendright_single((f₁, f₂)::FusionTreePair) N₁ = numout((f₁, f₂)) - @assert FusionStyle(I) === UniqueFusion() (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) # compute the coefficient coeff₀ = sqrtdim(c) * invsqrtdim(a) f₁.isdual[N₁] && (coeff₀ *= conj(frobenius_schur_phase(dual(b)))) - coeff = coeff₀ * Bsymbol(a, b, c) + coeff = coeff₀ * only(Bsymbol(a, b, c)) return (f₁′, f₂′) => coeff end @@ -143,6 +146,18 @@ function bendleft((f₁, f₂)::FusionTreePair) (f₂′, f₁′), coeff = bendright((f₂, f₁)) return (f₁′, f₂′) => conj(coeff) end +function _bendleft_single((f₁, f₂)::FusionTreePair) + (f₂′, f₁′), coeff = _bendright_single((f₂, f₁)) + return (f₁′, f₂′) => conj(coeff) +end + +# `n` bends of a tree pair that stays a single tree pair, and the product of their coefficients +_bend_single(bend, f, ::Val{0}) = f => one(sectorscalartype(sectortype(f))) +function _bend_single(bend, f, ::Val{n}) where {n} + f′, c′ = bend(f) + f″, c″ = _bend_single(bend, f′, Val(n - 1)) + return f″ => c′ * c″ +end # !! note that this is more or less a copy of bendright through # (f1, f2) => conj(coeff) for ((f2, f1), coeff) in bendleft(src) From a56726a00c42706fda104a622b2f97ea890dac5e Mon Sep 17 00:00:00 2001 From: leburgel Date: Wed, 14 Oct 2026 14:47:21 +0200 Subject: [PATCH 06/13] Test braiding a block with a single tree; changelog Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 1 + test/symmetries/doubletree.jl | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 74720d59a..459a557ff 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -44,6 +44,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. - Constructing a `FusionTreeBlock` only takes the fusion product of the shorter side, and `fusionblocks` groups the cached fusion trees of the space instead of constructing a block for every combination of sectors (internal) +- Braiding fusion tree blocks no longer constructs and indexes the canonical intermediate blocks of every step: these keep their trees in the order in which they are produced, their buffers are recycled, and blocks with a single tree are braided as a product of scalars (internal) ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index b4f1f8290..4c5987451 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -235,6 +235,31 @@ using TensorKitSectors end end + FusionStyle(I) isa UniqueFusion || !(BraidingStyle(I) isa HasBraiding) || + @testset "Double fusion tree: braiding a block with a single tree" begin + a, b = random_fusion(I, Val(2)) + cs = filter(c -> Nsymbol(a, b, c) == 1, collect(a ⊗ b)) + if !isempty(cs) + src1 = FusionTreeBlock{I}(((a, b), (rand(cs),)), ((rand(Bool), rand(Bool)), (rand(Bool),))) + @test length(src1) == 1 + for n in 0:3 + p = (randperm(3)...,) + p1, p2 = p[1:n], p[(n + 1):3] + ip = invperm(p) + ilevels = TupleTools.getindices((1, 2, 3), p) + dst, U = @constinferred TensorKit.braid(src1, (p1, p2), ((1, 2), (3,))) + @test length(dst) == 1 && size(U) == (1, 1) + dst′, U′ = @constinferred TensorKit.braid(dst, (ip[1:2], ip[3:3]), (ilevels[1:n], ilevels[(n + 1):3])) + @test dst′ == src1 + @test _isone(U * U′) + if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) + A1 = fusiontensor(only(fusiontrees(src1))) + @test permutedims(A1, (p1..., p2...)) ≈ only(U) * fusiontensor(only(fusiontrees(dst))) + end + end + end + end + UnitStyle(I) isa SimpleUnit && @testset "Double fusion tree: blocks of a HomSpace" begin V = Vect[I](Dict(s => 1 for s in (out..., incoming))) W = V ⊗ V' ← V ⊗ V' # never empty: the unit couples both sides From 2cbcffcaed1d890b0b55e3c98f47b2b4bb46a9a1 Mon Sep 17 00:00:00 2001 From: leburgel Date: Tue, 20 Oct 2026 11:30:33 +0200 Subject: [PATCH 07/13] Assert the block type in the block branch of `fsbraid` for JET Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 73ad1b6da..74a5dfeec 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -328,6 +328,7 @@ end return (f₁′, f₂′) => coeff1 * coeff2 * coeff3 else src, (p1, p2), (l1, l2) = key + src = src::FusionTreeBlock{I} p = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) From d3d0a8c7a2be6b09271246f648cb8192f04c255c Mon Sep 17 00:00:00 2001 From: leburgel Date: Fri, 23 Oct 2026 08:27:21 +0200 Subject: [PATCH 08/13] Fill a given matrix in `artin_braid!`, and choose the single braid channel by dispatch Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 36 +++++++++++++++-------- 1 file changed, 23 insertions(+), 13 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 74a5dfeec..6025eb695 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -69,7 +69,7 @@ function _artin_braid_single(f::FusionTree{I, N}, i, inv::Bool) where {I, N} a = inner_extended[i - 1] c = inner_extended[i] e = inner_extended[i + 1] - c′ = FusionStyle(I) isa UniqueFusion ? only(a ⊗ d) : only(_artin_channels(a, b, d, e)) + c′ = _artin_channel(a, b, d, e) coeff = oftype( oneT, if inv @@ -89,13 +89,24 @@ function _artin_channels(a::I, b::I, d::I, e::I) where {I <: Sector} UnitStyle(I) isa SimpleUnit || return intersect(a ⊗ d, e ⊗ conj(b)) return Iterators.filter(c′ -> !iszero(Nsymbol(c′, b, e)), a ⊗ d) end - -function artin_braid( - src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, canonical::Val = Val(true), - buffers = (nothing, nothing) # a tree vector and a matrix that may be overwritten +# the only such line, for a fusion tree that braids into a single one +_artin_channel(a, b, d, e) = _artin_channel(FusionStyle(a), a, b, d, e) +_artin_channel(::UniqueFusion, a, b, d, e) = only(a ⊗ d) +_artin_channel(::FusionStyle, a, b, d, e) = first(_artin_channels(a, b, d, e)) + +function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where {I, N} + n = length(src) # braiding a complete block keeps its size + return artin_braid!(Matrix{sectorscalartype(I)}(undef, n, n), src, i; inv) # zeroed by `artin_braid!` +end +# `U` is overwritten with the braid matrix of `src`, a complete block; the trees of the result are +# in canonical order or in the order in which they are produced, which recycles `trees` if given +function artin_braid!( + U::AbstractMatrix, src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, + canonical::Val = Val(true), trees = nothing ) where {I, N} 1 <= i < N || throw(ArgumentError("Cannot swap outputs i=$i and i+1 out of only $N outputs")) - isempty(src) && return src => zeros(sectorscalartype(I), 0, 0) + size(U) == (length(src), length(src)) || throw(DimensionMismatch(lazy"braid matrix of size $(size(U)) for a block of $(length(src)) trees")) + isempty(src) && return src => U uncoupled = src.uncoupled[1] a, b = uncoupled[i], uncoupled[i + 1] @@ -106,12 +117,10 @@ function artin_braid( isdual = src.isdual[1] isdual′ = TupleTools.setindex(isdual, isdual[i], i + 1) isdual′ = TupleTools.setindex(isdual′, isdual[i + 1], i) - rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical, buffers[1]) + rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical, trees) oneT = one(sectorscalartype(I)) - - U = buffers[2] isa Matrix{typeof(oneT)} ? fill!(buffers[2], zero(oneT)) : - zeros(sectorscalartype(I), max(length(rows.trees), length(src)), length(src)) # see `bendright` + fill!(U, zero(oneT)) if isunit(a) || isunit(b) # braiding with trivial sector: simple and always possible for (col, (f, f₂)) in enumerate(fusiontrees(src)) @@ -354,14 +363,15 @@ end # the swaps keep the size of the block, so that the tree vector and the matrices of the # previous step are recycled (except for the trees of `src`, which are part of the key) - buffers = (nothing, nothing) + trees = nothing + U_tmp = Matrix{sectorscalartype(I)}(undef, length(dst), length(dst)) U_spare = nothing braid_style = BraidingStyle(I) for s in permutation2swaps(p) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] - dst′, U_tmp = artin_braid(dst, s; inv, canonical = Val(false), buffers) - buffers = (dst.trees === src.trees ? nothing : dst.trees, U_tmp) + dst′, _ = artin_braid!(U_tmp, dst, s; inv, canonical = Val(false), trees) + trees = dst.trees === src.trees ? nothing : dst.trees dst = dst′ U_spare, U = U, U_spare isa typeof(U_tmp) ? mul!(U_spare, U_tmp, U) : U_tmp * U l = levels[s] From 0161d093fac28a0521e07bcf73e5b1dbf1c4f345 Mon Sep 17 00:00:00 2001 From: leburgel Date: Thu, 29 Oct 2026 06:51:13 +0100 Subject: [PATCH 09/13] Give `TreeRows` an inner constructor, so that Aqua finds no unbound type parameter Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/fusiontrees.jl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index 013b930a4..be1c889a1 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -290,6 +290,8 @@ struct TreeRows{F, D} trees::Vector{F} index::Union{D, Nothing} # `nothing`: few trees, the rows are found by a linear search n::Int + # no default outer constructor, of which `D` would be unbound if `index === nothing` + TreeRows{F, D}(trees::Vector{F}, index, n::Int) where {F, D} = new{F, D}(trees, index, n) end function TreeRows( ::Type{I}, uncoupled, isdual, n::Int, ::Val{canonical}, buffer = nothing From 077c459d81fceb93bf0af88ed04f84c2b82c58a8 Mon Sep 17 00:00:00 2001 From: leburgel Date: Thu, 29 Oct 2026 14:12:20 +0100 Subject: [PATCH 10/13] Braid single trees through `artin_braid` itself, take the only braid channel with `only`, and filter the braid channels for every unit style Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 6025eb695..2f50953ba 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -14,14 +14,11 @@ applying `artin_braid(f′, i; inv = true)` to all the outputs `f′` of `artin_braid(f, i; inv = false)` and collecting the results should yield a single fusion tree with non-zero coefficient, namely `f` with coefficient `1`. This keyword has no effect if `BraidingStyle(sectortype(f)) isa SymmetricBraiding`. + +For sectors without `UniqueFusion`, the braid of `f` has to be a single fusion tree, which is the +case if `f` is the only fusion tree with its uncoupled and coupled sectors; this throws otherwise. """ function artin_braid(f::FusionTree{I, N}, i; inv::Bool = false) where {I, N} - @assert FusionStyle(I) === UniqueFusion() - return _artin_braid_single(f, i, inv) -end -# also for a tree that is the only one with its uncoupled sectors, whose braid is again a single -# tree (`only` throws otherwise) -function _artin_braid_single(f::FusionTree{I, N}, i, inv::Bool) where {I, N} 1 <= i < N || throw(ArgumentError(lazy"Cannot swap outputs i=$i and i+1 out of only $N outputs")) uncoupled = f.uncoupled @@ -84,15 +81,12 @@ function _artin_braid_single(f::FusionTree{I, N}, i, inv::Bool) where {I, N} end # the inner lines `c′` of `a ⊗ d → c′ → c′ ⊗ b → e`, i.e. `intersect(a ⊗ d, e ⊗ conj(b))`, -# without constructing the second fusion product for sectors with a simple unit -function _artin_channels(a::I, b::I, d::I, e::I) where {I <: Sector} - UnitStyle(I) isa SimpleUnit || return intersect(a ⊗ d, e ⊗ conj(b)) - return Iterators.filter(c′ -> !iszero(Nsymbol(c′, b, e)), a ⊗ d) -end -# the only such line, for a fusion tree that braids into a single one +# without constructing the second fusion product +_artin_channels(a, b, d, e) = Iterators.filter(c′ -> !iszero(Nsymbol(c′, b, e)), a ⊗ d) +# the only such line, for a fusion tree that braids into a single one (`only` throws otherwise) _artin_channel(a, b, d, e) = _artin_channel(FusionStyle(a), a, b, d, e) _artin_channel(::UniqueFusion, a, b, d, e) = only(a ⊗ d) -_artin_channel(::FusionStyle, a, b, d, e) = first(_artin_channels(a, b, d, e)) +_artin_channel(::FusionStyle, a, b, d, e) = only(_artin_channels(a, b, d, e)) function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where {I, N} n = length(src) # braiding a complete block keeps its size @@ -347,7 +341,7 @@ end for s in permutation2swaps(p) _check_levels(BraidingStyle(I), levels, s) inv = levels[s] > levels[s + 1] - f, c = _artin_braid_single(f, s, inv) + f, c = artin_braid(f, s; inv) coeff *= c l = levels[s] levels = TupleTools.setindex(levels, levels[s + 1], s) From 5f60bb7877db14f94c69fb6f7533e38a5b066b34 Mon Sep 17 00:00:00 2001 From: leburgel Date: Thu, 29 Oct 2026 14:12:20 +0100 Subject: [PATCH 11/13] Set the inner line and the first vertex of a braided tree outside the loops over the multiplicities Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 2f50953ba..df7380323 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -193,7 +193,9 @@ function artin_braid!( for c′ in _artin_channels(a, b, d, e) Rmat2 = inv ? Rsymbol(d, a, c′)' : Rsymbol(a, d, c′) Fmat = Fsymbol(d, a, b, e, c′, c) + inner′ = TupleTools.setindex(inner, c′, i - 1) for σ in axes(Rmat2, 1) + vertices_σ = TupleTools.setindex(vertices, σ, i - 1) for λ in axes(Fmat, 2) coeff = zero(oneT) for ρ in axes(Rmat1, 2), κ in axes(Rmat2, 2) @@ -201,9 +203,7 @@ function artin_braid!( conj(Rmat2[σ, κ]) end iszero(coeff) && continue - vertices′ = TupleTools.setindex(vertices, σ, i - 1) - vertices′ = TupleTools.setindex(vertices′, λ, i) - inner′ = TupleTools.setindex(inner, c′, i - 1) + vertices′ = TupleTools.setindex(vertices_σ, λ, i) f′ = FusionTree{I}(uncoupled′, coupled′, isdual′, inner′, vertices′) row = row!(rows, (f′, f₂)) @inbounds U[row, col] = coeff From 941044b56726afbfba2bbc8feb5f60145f9ec22f Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 17:20:52 -0400 Subject: [PATCH 12/13] Braid a single tree pair through the tree manipulations themselves `bendright` and `bendleft` of a tree pair no longer assert `UniqueFusion`, but rely on `only` to throw if the bent vertex has a multiplicity, which replaces their `_single` variants. The shortcut of `braid` for a symmetric braiding is restricted to `UniqueFusion`, where the `Fsymbol`s are trivial, such that a block with a single tree is braided as its tree pair. Co-Authored-By: Claude Opus 5.5 --- src/fusiontrees/braiding_manipulations.jl | 42 ++++++++++------------- src/fusiontrees/duality_manipulations.jl | 21 ++---------- 2 files changed, 20 insertions(+), 43 deletions(-) diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index df7380323..70f634118 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -235,9 +235,8 @@ braid(f::FusionTree{I, N}, p::IndexTuple{N}, levels::IndexTuple{N}) where {I, N} braid(f, (p, ()), (levels, ())) function braid(f::FusionTree{I, N}, (p, _)::Index2Tuple{N, 0}, (levels, _)::Index2Tuple{N, 0}) where {I, N} TupleTools.isperm(p) || throw(ArgumentError(lazy"not a valid permutation: $p")) - @assert FusionStyle(I) isa UniqueFusion braid_style = BraidingStyle(I) - if braid_style isa SymmetricBraiding # this assumes Fsymbols are 1! + if FusionStyle(I) isa UniqueFusion && braid_style isa SymmetricBraiding # Fsymbols are 1 coeff = one(sectorscalartype(I)) for i in 1:N for j in 1:(i - 1) @@ -322,35 +321,20 @@ end @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 + f, p, levels = key + return _braid_single(f::FusionTreePair{I}, p, levels) else src, (p1, p2), (l1, l2) = key src = src::FusionTreeBlock{I} - p = linearizepermutation(p1, p2, numout(src), numin(src)) - levels = (l1..., reverse(l2)...) - if length(src) == 1 # then so is every intermediate block: the braid is a product of scalars - (f, f₀), coeff = _bend_single(_bendleft_single, only(fusiontrees(src)), Val(numin(src))) - for s in permutation2swaps(p) - _check_levels(BraidingStyle(I), levels, s) - inv = levels[s] > levels[s + 1] - f, c = artin_braid(f, s; inv) - coeff *= c - l = levels[s] - levels = TupleTools.setindex(levels, levels[s + 1], s) - levels = TupleTools.setindex(levels, l, s + 1) - end - (f₁′, f₂′), c = _bend_single(_bendright_single, (f, f₀), Val(N₂)) - return FusionTreeBlock([(f₁′, f₂′)]) => fill(coeff * c, 1, 1) + f′, coeff = _braid_single(only(fusiontrees(src)), (p1, p2), (l1, l2)) + return FusionTreeBlock([f′]) => fill(coeff, 1, 1) end + p = linearizepermutation(p1, p2, numout(src), numin(src)) + levels = (l1..., reverse(l2)...) + # the intermediate blocks keep their trees in the order in which they are produced, # and only the result is brought into canonical order dst, U = repartition_noncanonical(src, Val(numind(src))) @@ -383,6 +367,16 @@ end end end +# also for a tree pair that is the only one with its uncoupled sectors (`only` throws otherwise) +function _braid_single((f₁, f₂)::FusionTreePair, (p1, p2)::Index2Tuple{N₁, N₂}, (l1, l2)::Index2Tuple) where {N₁, N₂} + 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 + CacheStyle(::typeof(fsbraid), k::Union{FSPBraidKey{I}, FSBBraidKey{I}}) where {I} = FusionStyle(I) isa UniqueFusion ? NoCache() : GlobalLRUCache() diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index ca49a3a47..8a3812d1d 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -53,12 +53,8 @@ function _bendright_treepair((f₁, f₂)::FusionTreePair) return (a, b, c), (f₁′, f₂′) end +# also for a final splitting vertex without multiplicity (`only` throws otherwise) function bendright((f₁, f₂)::FusionTreePair) - @assert FusionStyle((f₁, f₂)) === UniqueFusion() - return _bendright_single((f₁, f₂)) -end -# also for a final fusion vertex without multiplicity (`only` throws otherwise) -function _bendright_single((f₁, f₂)::FusionTreePair) N₁ = numout((f₁, f₂)) (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) @@ -142,22 +138,9 @@ See also [`bendright`](@ref). """ bendleft function bendleft((f₁, f₂)::FusionTreePair) - @assert FusionStyle((f₁, f₂)) === UniqueFusion() (f₂′, f₁′), coeff = bendright((f₂, f₁)) return (f₁′, f₂′) => conj(coeff) end -function _bendleft_single((f₁, f₂)::FusionTreePair) - (f₂′, f₁′), coeff = _bendright_single((f₂, f₁)) - return (f₁′, f₂′) => conj(coeff) -end - -# `n` bends of a tree pair that stays a single tree pair, and the product of their coefficients -_bend_single(bend, f, ::Val{0}) = f => one(sectorscalartype(sectortype(f))) -function _bend_single(bend, f, ::Val{n}) where {n} - f′, c′ = bend(f) - f″, c″ = _bend_single(bend, f′, Val(n - 1)) - return f″ => c′ * c″ -end # !! note that this is more or less a copy of bendright through # (f1, f2) => conj(coeff) for ((f2, f1), coeff) in bendleft(src) @@ -499,7 +482,7 @@ function _repartition_body(N, canonical = true) if N == 0 ex = quote T = fusionscalartype(sectortype(src)) - if FusionStyle(src) === UniqueFusion() + if src isa FusionTreePair return src => one(T) else U = copyto!(zeros(T, length(src), length(src)), LinearAlgebra.I) From d397cb27f1f3ed08c4d89cd36599567fc010c6bb Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 16:08:53 -0400 Subject: [PATCH 13/13] Keep the trees of every fusion tree block manipulation in production order The destination blocks of bends, folds and braids are only used together with their transformation matrix, so they no longer construct and index the canonical destination block, but keep their trees in the order in which they are produced. Since `planar_trace` now also works on complete blocks, the canonical mode of `TreeRows`, `canonicalize`, `repartition_noncanonical` and `treeindex_map` are removed. The tests compare the results of manipulations up to the order of their trees. Co-Authored-By: Claude Opus 5.5 --- docs/src/Changelog.md | 2 +- src/fusiontrees/braiding_manipulations.jl | 23 +++---- src/fusiontrees/duality_manipulations.jl | 78 ++++++----------------- src/fusiontrees/fusiontrees.jl | 51 ++++----------- test/setup.jl | 13 +++- test/symmetries/doubletree.jl | 20 +++--- 6 files changed, 61 insertions(+), 126 deletions(-) diff --git a/docs/src/Changelog.md b/docs/src/Changelog.md index 459a557ff..07fdc3bcf 100644 --- a/docs/src/Changelog.md +++ b/docs/src/Changelog.md @@ -44,7 +44,7 @@ When releasing a new version, move the "Unreleased" changes to a new version sec - Recoupling matrices are stored such that they can leverage BLAS, also for mixed complex tensor with real recoupling coefficients. - Constructing a `FusionTreeBlock` only takes the fusion product of the shorter side, and `fusionblocks` groups the cached fusion trees of the space instead of constructing a block for every combination of sectors (internal) -- Braiding fusion tree blocks no longer constructs and indexes the canonical intermediate blocks of every step: these keep their trees in the order in which they are produced, their buffers are recycled, and blocks with a single tree are braided as a product of scalars (internal) +- Fusion tree block manipulations no longer construct and index the canonical destination block, but keep the trees of the result in the order in which they are produced; the buffers of the intermediate blocks of a braid are recycled, and blocks with a single tree are braided as their tree pair (internal) ## [0.17.2](https://github.com/QuantumKitHub/TensorKit.jl/compare/v0.17.1...v0.17.2) - 2026-09-20 diff --git a/src/fusiontrees/braiding_manipulations.jl b/src/fusiontrees/braiding_manipulations.jl index 70f634118..ceaa1882b 100644 --- a/src/fusiontrees/braiding_manipulations.jl +++ b/src/fusiontrees/braiding_manipulations.jl @@ -93,10 +93,9 @@ function artin_braid(src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false) where return artin_braid!(Matrix{sectorscalartype(I)}(undef, n, n), src, i; inv) # zeroed by `artin_braid!` end # `U` is overwritten with the braid matrix of `src`, a complete block; the trees of the result are -# in canonical order or in the order in which they are produced, which recycles `trees` if given +# stored in `trees` if given function artin_braid!( - U::AbstractMatrix, src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, - canonical::Val = Val(true), trees = nothing + U::AbstractMatrix, src::FusionTreeBlock{I, N, 0}, i; inv::Bool = false, trees = nothing ) where {I, N} 1 <= i < N || throw(ArgumentError("Cannot swap outputs i=$i and i+1 out of only $N outputs")) size(U) == (length(src), length(src)) || throw(DimensionMismatch(lazy"braid matrix of size $(size(U)) for a block of $(length(src)) trees")) @@ -111,7 +110,7 @@ function artin_braid!( isdual = src.isdual[1] isdual′ = TupleTools.setindex(isdual, isdual[i], i + 1) isdual′ = TupleTools.setindex(isdual′, isdual[i + 1], i) - rows = TreeRows(I, (uncoupled′, ()), (isdual′, ()), length(src), canonical, trees) + rows = TreeRows(fusiontreetype(I, N, 0), length(src), trees) oneT = one(sectorscalartype(I)) fill!(U, zero(oneT)) @@ -335,9 +334,7 @@ end p = linearizepermutation(p1, p2, numout(src), numin(src)) levels = (l1..., reverse(l2)...) - # the intermediate blocks keep their trees in the order in which they are produced, - # and only the result is brought into canonical order - dst, U = repartition_noncanonical(src, Val(numind(src))) + dst, U = repartition(src, Val(numind(src))) # the swaps keep the size of the block, so that the tree vector and the matrices of the # previous step are recycled (except for the trees of `src`, which are part of the key) @@ -348,7 +345,7 @@ end for s in permutation2swaps(p) _check_levels(braid_style, levels, s) inv = levels[s] > levels[s + 1] - dst′, _ = artin_braid!(U_tmp, dst, s; inv, canonical = Val(false), trees) + dst′, _ = artin_braid!(U_tmp, dst, s; inv, trees) trees = dst.trees === src.trees ? nothing : dst.trees dst = dst′ U_spare, U = U, U_spare isa typeof(U_tmp) ? mul!(U_spare, U_tmp, U) : U_tmp * U @@ -357,13 +354,9 @@ end levels = TupleTools.setindex(levels, l, s + 1) end - if N₂ == 0 - return canonicalize(dst, U) - else - dst, U_tmp = repartition_noncanonical(dst::FusionTreeBlock{I}, Val(N₁)) - U = U_tmp * U - return canonicalize(dst, U) - end + N₂ == 0 && return dst => U + dst, U_tmp = repartition(dst::FusionTreeBlock{I}, Val(N₁)) + return dst => U_tmp * U end end diff --git a/src/fusiontrees/duality_manipulations.jl b/src/fusiontrees/duality_manipulations.jl index 8a3812d1d..58e542988 100644 --- a/src/fusiontrees/duality_manipulations.jl +++ b/src/fusiontrees/duality_manipulations.jl @@ -65,24 +65,14 @@ function bendright((f₁, f₂)::FusionTreePair) return (f₁′, f₂′) => coeff end -function bendright(src::FusionTreeBlock, canonical::Val = Val(true)) +function bendright(src::FusionTreeBlock) I = sectortype(src) N₁ = numout(src) N₂ = numin(src) @assert N₁ > 0 - uncoupled_dst = ( - TupleTools.front(src.uncoupled[1]), - (src.uncoupled[2]..., dual(src.uncoupled[1][N₁])), - ) - isdual_dst = ( - TupleTools.front(src.isdual[1]), - (src.isdual[2]..., !(src.isdual[1][N₁])), - ) - - rows = TreeRows(I, uncoupled_dst, isdual_dst, length(src), canonical) - # the canonical destination block has more trees than `src` if `src` holds only some of the - # trees with its uncoupled sectors (as in `planar_trace`) - U = zeros(fusionscalartype(I), max(length(rows.trees), length(src)), length(src)) + + rows = TreeRows(fusiontreetype(I, N₁ - 1, N₂ + 1), length(src)) + U = zeros(fusionscalartype(I), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₁′, f₂′) = _bendright_treepair((f₁, f₂)) @@ -144,24 +134,14 @@ end # !! note that this is more or less a copy of bendright through # (f1, f2) => conj(coeff) for ((f2, f1), coeff) in bendleft(src) -function bendleft(src::FusionTreeBlock, canonical::Val = Val(true)) +function bendleft(src::FusionTreeBlock) I = sectortype(src) N₁ = numout(src) N₂ = numin(src) @assert N₂ > 0 - uncoupled_dst = ( - (src.uncoupled[1]..., dual(src.uncoupled[2][N₂])), - TupleTools.front(src.uncoupled[2]), - ) - isdual_dst = ( - (src.isdual[1]..., !(src.isdual[2][N₂])), - TupleTools.front(src.isdual[2]), - ) - - rows = TreeRows(I, uncoupled_dst, isdual_dst, length(src), canonical) - # the canonical destination block has more trees than `src` if `src` holds only some of the - # trees with its uncoupled sectors (as in `planar_trace`) - U = zeros(fusionscalartype(I), max(length(rows.trees), length(src)), length(src)) + + rows = TreeRows(fusiontreetype(I, N₁ + 1, N₂ - 1), length(src)) + U = zeros(fusionscalartype(I), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) (a, b, c), (f₂′, f₁′) = _bendright_treepair((f₂, f₁)) @@ -236,17 +216,11 @@ function foldright((f₁, f₂)::FusionTreePair) end function foldright(src::FusionTreeBlock) - uncoupled_dst = ( - Base.tail(src.uncoupled[1]), - (dual(first(src.uncoupled[1])), src.uncoupled[2]...), - ) - isdual_dst = (Base.tail(src.isdual[1]), (!first(src.isdual[1]), src.isdual[2]...)) I = sectortype(src) N₁ = numout(src) @assert N₁ > 0 - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) + rows = TreeRows(fusiontreetype(I, N₁ - 1, numin(src) + 1), length(src)) f₁, f₂ = first(fusiontrees(src)) a::I = f₁.uncoupled[1] @@ -260,7 +234,7 @@ function foldright(src::FusionTreeBlock) c::I = f₁.coupled cache₃ = Dict((b, c) => Asymbol(a, b, c)) - U = zeros(eltype(coeff₁), length(dst), length(src)) + U = zeros(eltype(coeff₁), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) f₁′s, coeffs₁ = get!(cache₁, f₁) do multi_Fmove(f₁) @@ -280,12 +254,12 @@ function foldright(src::FusionTreeBlock) if isduala coeff *= κₐ end - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] += coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end @@ -319,21 +293,12 @@ function foldleft((f₁, f₂)::FusionTreePair) end function foldleft(src::FusionTreeBlock) - uncoupled_dst = ( - (dual(first(src.uncoupled[2])), src.uncoupled[1]...), - Base.tail(src.uncoupled[2]), - ) - isdual_dst = ( - (!first(src.isdual[2]), src.isdual[1]...), - Base.tail(src.isdual[2]), - ) I = sectortype(src) N₁ = numin(src) N₂ = numout(src) @assert N₁ > 0 - dst = FusionTreeBlock{I}(uncoupled_dst, isdual_dst; sizehint = length(src)) - indexmap = treeindex_map(dst) + rows = TreeRows(fusiontreetype(I, N₂ + 1, N₁ - 1), length(src)) f₁, f₂ = first(fusiontrees(src)) a::I = f₂.uncoupled[1] @@ -347,7 +312,7 @@ function foldleft(src::FusionTreeBlock) c::I = f₂.coupled cache₃ = Dict((b, c) => Asymbol(a, b, c)) - U = zeros(eltype(coeff₂), length(dst), length(src)) + U = zeros(eltype(coeff₂), length(src), length(src)) for (col, (f₁, f₂)) in enumerate(fusiontrees(src)) f₂′s, coeffs₂ = get!(cache₂, f₂) do multi_Fmove(f₂) @@ -367,12 +332,12 @@ function foldleft(src::FusionTreeBlock) if isduala coeff *= conj(κₐ) end - row = indexmap[treeindex_data((f₁′, f₂′))] + row = row!(rows, (f₁′, f₂′)) @inbounds U[row, col] += coeff end end end - return dst => U + return FusionTreeBlock(rows.trees) => U end # clockwise cyclic permutation while preserving (N₁, N₂): foldright & bendleft @@ -478,7 +443,7 @@ return dst, U # it, so that it is visible to the generator's world age. Otherwise a precompile workload that # first triggers `repartition` fails with `UndefVarError: _repartition_body` (at runtime the # first call happens at a later world age, which masks the ordering issue). -function _repartition_body(N, canonical = true) +function _repartition_body(N) if N == 0 ex = quote T = fusionscalartype(sectortype(src)) @@ -491,14 +456,13 @@ function _repartition_body(N, canonical = true) end else f = N < 0 ? bendleft : bendright - args = canonical ? () : (Val(false),) ex_rep = Expr(:block) for _ in 1:(abs(N) - 1) - push!(ex_rep.args, :((dst, Utmp) = $f(dst, $(args...)))) + push!(ex_rep.args, :((dst, Utmp) = $f(dst))) push!(ex_rep.args, :(U = Utmp * U)) end ex = quote - dst, U = $f(src, $(args...)) + dst, U = $f(src) $ex_rep return dst => U end @@ -508,10 +472,6 @@ end @generated function repartition(src::Union{FusionTreePair, FusionTreeBlock}, ::Val{N}) where {N} return _repartition_body(numout(src) - N) end -# the same for a block, with the trees of the result in the order in which they are produced -@generated function repartition_noncanonical(src::FusionTreeBlock, ::Val{N}) where {N} - return _repartition_body(numout(src) - N, false) -end """ transpose((f₁, f₂)::FusionTreePair{I}, p::Index2Tuple{N₁, N₂}) where {I, N₁, N₂} diff --git a/src/fusiontrees/fusiontrees.jl b/src/fusiontrees/fusiontrees.jl index be1c889a1..fef21b76b 100644 --- a/src/fusiontrees/fusiontrees.jl +++ b/src/fusiontrees/fusiontrees.jl @@ -117,8 +117,10 @@ const FusionTreePair{I, N₁, N₂} = Tuple{FusionTree{I, N₁}, FusionTree{I, N Collection of fusion-splitting tree pairs that share the same uncoupled sectors. Mostly internal structure to speed up non-`UniqueFusion` fusiontree manipulation computations. -Braiding a block requires it to be complete, i.e. to contain all the fusion-splitting tree pairs -with its uncoupled sectors and duality flags, as `FusionTreeBlock{I}(uncoupled, isdual)` does. +Manipulating a block requires it to be complete, i.e. to contain all the fusion-splitting tree +pairs with its uncoupled sectors and duality flags, as `FusionTreeBlock{I}(uncoupled, isdual)` +does. The resulting block holds its trees in the order in which they are produced, rather than in +the canonical order of `FusionTreeBlock{I}(uncoupled, isdual)`. """ struct FusionTreeBlock{I, N₁, N₂, F <: FusionTreePair{I, N₁, N₂}} trees::Vector{F} @@ -273,19 +275,11 @@ function treeindex_data((f₁, f₂)) return false end end -function treeindex_map(fs::FusionTreeBlock) - I = sectortype(fs) - return fusiontreedict(I)(treeindex_data(f) => ind for (ind, f) in enumerate(fusiontrees(fs))) -end -# The trees of the destination block of a block manipulation, and the row of each of them in the -# transformation matrix. With `Val(true)`, these are the trees of the destination block in their -# canonical order; with `Val(false)`, the trees are added in the order in which they are produced, -# which avoids constructing the destination block. The latter is used for the intermediate blocks -# of a composite manipulation, of which only the result is brought into canonical order; the tree -# vector of an intermediate block that is no longer needed can then be passed on as `buffer`. -# With `Val(false)`, the source block has to be complete: the destination block then has as many -# trees (`n`), and an incomplete source block raises an error if its destination block has more. +# The trees of the destination block of a block manipulation, in the order in which they are +# produced, and the row of each of them in the transformation matrix. The source block has to be +# complete, such that the destination block has as many trees (`n`); the tree vector of a block +# that is no longer needed can be passed on as `buffer`. struct TreeRows{F, D} trees::Vector{F} index::Union{D, Nothing} # `nothing`: few trees, the rows are found by a linear search @@ -293,18 +287,11 @@ struct TreeRows{F, D} # no default outer constructor, of which `D` would be unbound if `index === nothing` TreeRows{F, D}(trees::Vector{F}, index, n::Int) where {F, D} = new{F, D}(trees, index, n) end -function TreeRows( - ::Type{I}, uncoupled, isdual, n::Int, ::Val{canonical}, buffer = nothing - ) where {I, canonical} - F = fusiontreetype(I, length(uncoupled[1]), length(uncoupled[2])) - trees = if canonical - fusiontrees(FusionTreeBlock{I}(uncoupled, isdual; sizehint = n)) - else - buffer isa Vector{F} ? empty!(buffer) : sizehint!(Vector{F}(), n) - end - K = Base.promote_op(treeindex_data, eltype(trees)) - index = n > 32 ? Dict{K, Int}(treeindex_data(f) => i for (i, f) in enumerate(trees)) : nothing - return TreeRows{eltype(trees), Dict{K, Int}}(trees, index, n) +function TreeRows(::Type{F}, n::Int, buffer = nothing) where {F} + trees = buffer isa Vector{F} ? empty!(buffer) : sizehint!(Vector{F}(), n) + K = Base.promote_op(treeindex_data, F) + index = n > 32 ? Dict{K, Int}() : nothing + return TreeRows{F, Dict{K, Int}}(trees, index, n) end function row!(rows::TreeRows, f) key = treeindex_data(f) @@ -315,22 +302,12 @@ function row!(rows::TreeRows, f) end isnothing(i) || return i length(rows.trees) < rows.n || - throw(ArgumentError("braiding requires a complete `FusionTreeBlock`")) + throw(ArgumentError("manipulating a block requires it to be complete")) push!(rows.trees, f) isnothing(rows.index) || (rows.index[key] = length(rows.trees)) return length(rows.trees) end -# bring a block whose trees are not in canonical order, together with the transformation matrix -# whose rows correspond to these trees, into canonical order -function canonicalize(src::FusionTreeBlock{I}, U::AbstractMatrix) where {I} - length(src) == 1 && return src => U - rows = TreeRows(I, src.uncoupled, src.isdual, length(src), Val(true)) - U′ = zeros(eltype(U), length(rows.trees), size(U, 2)) - U′[map(f -> row!(rows, f), fusiontrees(src)), :] = U - return FusionTreeBlock(rows.trees) => U′ -end - # Facilitate getting correct fusion tree types Base.@assume_effects :foldable function fusiontreetype(::Type{I}, N::Int) where {I <: Sector} diff --git a/test/setup.jl b/test/setup.jl index 2030228a5..e928fdd11 100644 --- a/test/setup.jl +++ b/test/setup.jl @@ -9,12 +9,12 @@ export sectorlist, fast_sectorlist export default_spacelist, factorization_spacelist, ad_spacelist export VIBM, VIBMRepA4 export test_ad_rrule -export _isunitary, _isone +export _isunitary, _isone, _inorder using Random using Test: @test using TensorKit -using TensorKit: ℙ, PlanarTrivial +using TensorKit: ℙ, PlanarTrivial, FusionTreeBlock using TensorKitSectors using TensorOperations: IndexTuple, Index2Tuple using Base.Iterators: take, product @@ -149,6 +149,15 @@ _isunitary(x::Number; kwargs...) = isapprox(x * x', one(x); kwargs...) _isunitary(x; kwargs...) = isunitary(x; kwargs...) _isone(x; kwargs...) = isapprox(x, one(x); kwargs...) +# the rows of the transformation matrix `U` onto the trees of `dst`, in the order of the trees of +# `ref`, as block manipulations keep their trees in the order in which they are produced +_inorder(U::Number, dst, ref) = (@test dst == ref; U) +function _inorder(U::AbstractMatrix, dst::FusionTreeBlock, ref::FusionTreeBlock) + rows = indexin(fusiontrees(ref), fusiontrees(dst)) + @test length(dst) == length(ref) && !any(isnothing, rows) + return U[something.(rows, 1), :] +end + uniquefusionsectorlist = ( Z2Irrep, Z3Irrep, Z4Irrep, Z3Irrep ⊠ Z4Irrep, U1Irrep, diff --git a/test/symmetries/doubletree.jl b/test/symmetries/doubletree.jl index 4c5987451..e6fd1ca95 100644 --- a/test/symmetries/doubletree.jl +++ b/test/symmetries/doubletree.jl @@ -57,15 +57,13 @@ using TensorKitSectors # single bend dst, U = @constinferred TK.bendright(src) dst2, U2 = @constinferred TK.bendleft(dst) - @test src == dst2 - @test _isone(U2 * U) + @test _isone(_inorder(U2 * U, dst2, src)) # double bend dst1, U1 = @constinferred TK.bendleft(src) dst2, U2 = @constinferred TK.bendleft(dst1) dst3, U3 = @constinferred TK.bendright(dst2) dst4, U4 = @constinferred TK.bendright(dst3) - @test src == dst4 - @test _isone(U4 * U3 * U2 * U1) + @test _isone(_inorder(U4 * U3 * U2 * U1, dst4, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = (ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...) @@ -88,15 +86,13 @@ using TensorKitSectors # single bend dst, U = @constinferred TK.foldleft(src) dst2, U2 = @constinferred TK.foldright(dst) - @test src == dst2 - @test _isone(U2 * U) + @test _isone(_inorder(U2 * U, dst2, src)) # double bend dst1, U1 = @constinferred TK.foldright(src) dst2, U2 = @constinferred TK.foldright(dst1) dst3, U3 = @constinferred TK.foldleft(dst2) dst4, U4 = @constinferred TK.foldleft(dst3) - @test src == dst4 - @test _isone(U4 * U3 * U2 * U1) + @test _isone(_inorder(U4 * U3 * U2 * U1, dst4, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = TupleTools.circshift((ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...), -2) @@ -121,7 +117,7 @@ using TensorKitSectors # @test _isunitary(U) dst′, U′ = repartition(dst, N) - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) all_inds = (ntuple(identity, numout(src))..., reverse(ntuple(i -> i + numout(src), numin(src)))...) @@ -152,11 +148,11 @@ using TensorKitSectors dst, U = @constinferred transpose(src, (p1, p2)) dst′, U′ = @constinferred transpose(dst, (ip1, ip2)) - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) if BraidingStyle(I) isa Bosonic dst″, U″ = permute(src, (p1, p2)) - @test U″ ≈ U + @test _inorder(U″, dst″, dst) ≈ U end if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I) @@ -217,7 +213,7 @@ using TensorKitSectors else dst′, U′ = @constinferred TensorKit.braid(dst, (ip1, ip2), (il1, il2)) end - @test _isone(U * U′) + @test _isone(U * _inorder(U′, dst′, src)) # check fusiontensor compatibility if (BraidingStyle(I) isa Bosonic) && hasfusiontensor(I)