diff --git a/src/auxiliary/blockarrays.jl b/src/auxiliary/blockarrays.jl index 97ffd24..a155eda 100644 --- a/src/auxiliary/blockarrays.jl +++ b/src/auxiliary/blockarrays.jl @@ -16,6 +16,17 @@ function copy_dense!(Adense, A) return Adense end +function copy_blocks!(A::BlockMatrix, Adense) + for block_index in Iterators.product(blockaxes(A)...) + a = view(A, block_index...) + isempty(a) && continue + indices = getindex.(axes(A), block_index) + a .= view(Adense, indices...) + end + return A +end +copy_blocks!(A, Adense) = copyto!(A, Adense) + const BlockBlasMat{T <: MAK.BlasFloat} = BlockMatrix{T} function MAK.zero!(A::BlockBlasMat) diff --git a/src/auxiliary/sparsetensorarray.jl b/src/auxiliary/sparsetensorarray.jl index 7831407..8dfc423 100644 --- a/src/auxiliary/sparsetensorarray.jl +++ b/src/auxiliary/sparsetensorarray.jl @@ -44,9 +44,9 @@ TensorKit.codomain(A::SparseTensorArray) = codomain(space(A)) TensorKit.domain(A::SparseTensorArray) = domain(space(A)) TensorKit.numout(A::SparseTensorArray) = numout(eltype(A)) -TensorKit.numout(::Type{T}) where {T <: SparseTensorArray} = numout(eltype(A)) +TensorKit.numout(::Type{T}) where {T <: SparseTensorArray} = numout(eltype(T)) TensorKit.numin(A::SparseTensorArray) = numin(eltype(A)) -TensorKit.numin(::Type{T}) where {T <: SparseTensorArray} = numin(eltype(A)) +TensorKit.numin(::Type{T}) where {T <: SparseTensorArray} = numin(eltype(T)) # AbstractArray interface # ----------------------- @@ -93,6 +93,7 @@ function Base.delete!(A::SparseTensorArray, I::Vararg{Int, N}) where {N} end Base.delete!(A::SparseTensorArray, I::CartesianIndex) = delete!(A.data, I) Base.empty!(A::SparseTensorArray) = empty!(A.data) +Base.filter!(f, A::SparseTensorArray) = (filter!(f, A.data); A) function Base.haskey(A::SparseTensorArray, I::Vararg{Int, N}) where {N} return haskey(A.data, CartesianIndex(I)) end diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index ce887f5..317e4d4 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -28,7 +28,7 @@ for f! in ( Fblocks′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg) # deal with the case where the output is not in-place for (b′, b) in zip(Fblocks′, Fblocks) - b === b′ || copyto!(b, b′) + b === b′ || copy_blocks!(b, b′) end return nothing end @@ -46,7 +46,7 @@ for f! in ( TensorKit.foreachblock(t, N) do _, (tblock, Nblock) Nblock′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg) # deal with the case where the output is not the same as the input - Nblock === Nblock′ || copyto!(Nblock, Nblock′) + Nblock === Nblock′ || copy_blocks!(Nblock, Nblock′) return nothing end return N @@ -186,12 +186,12 @@ end function TensorKit.Factorizations.truncate_domain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds) TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block) I = get(inds, c, nothing) - dst_dense = copy_dense!(similar_dense(dst_block), dst_block) + dst_dense = similar_dense(dst_block) src_dense = copy_dense!(similar_dense(src_block), src_block) @assert !isnothing(I) @views dst_dense .= src_dense[:, I] # deal with the case where the output is not in-place - dst_dense === dst_block || copyto!(dst_block, dst_dense) + dst_dense === dst_block || copy_blocks!(dst_block, dst_dense) return nothing end return tdst @@ -199,12 +199,12 @@ end function TensorKit.Factorizations.truncate_codomain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds) TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block) I = get(inds, c, nothing) - dst_dense = copy_dense!(similar_dense(dst_block), dst_block) + dst_dense = similar_dense(dst_block) src_dense = copy_dense!(similar_dense(src_block), src_block) @assert !isnothing(I) @views dst_dense .= src_dense[I, :] # deal with the case where the output is not in-place - dst_dense === dst_block || copyto!(dst_block, dst_dense) + dst_dense === dst_block || copy_blocks!(dst_block, dst_dense) return nothing end return tdst diff --git a/src/linalg/linalg.jl b/src/linalg/linalg.jl index 0abd9a3..f0296e8 100644 --- a/src/linalg/linalg.jl +++ b/src/linalg/linalg.jl @@ -157,10 +157,34 @@ for (TC, TA, TB) in Iterators.product( end end -function LinearAlgebra.norm(tA::BlockTensorMap, p::Real = 2) +function LinearAlgebra.norm(tA::AbstractBlockTensorMap, p::Real = 2) vals = nonzero_values(tA) isempty(vals) && return norm(zero(scalartype(tA)), p) - return LinearAlgebra.norm(norm.(vals), p) + return LinearAlgebra.norm(map(Base.Fix2(norm, p), vals), p) +end + +function Base.:(==)(t1::AbstractBlockTensorMap, t2::AbstractBlockTensorMap) + space(t1) == space(t2) || return false + for I in nonzero_keys(t1) + v1 = t1[I] + (haskey(t2, I) ? v1 == t2[I] : iszero(norm(v1))) || return false + end + if issparse(t1) + for I in nonzero_keys(t2) + haskey(t1, I) || iszero(norm(t2[I])) || return false + end + end + return true +end + +function LinearAlgebra.tr(t::AbstractBlockTensorMap) + domain(t) == codomain(t) || + throw(SpaceMismatch("Trace of a tensor only exist when domain == codomain")) + s = zero(scalartype(t)) * zero(TK.dimscalartype(sectortype(t))) + for I in nonzero_keys(t) + TT.getindices(I.I, codomainind(t)) == TT.getindices(I.I, domainind(t)) && (s += tr(t[I])) + end + return s end for f in (:real, :imag) @@ -173,7 +197,7 @@ for f in (:real, :imag) end return t′ else - msg = "`$f` has not been implemented for `BlockTensorMap{$(S)}`." + msg = "`$f` has not been implemented for `BlockTensorMap{$(spacetype(t))}`." throw(ArgumentError(msg)) end end @@ -208,7 +232,7 @@ function LinearAlgebra.lmul!(D::DiagonalTensorMap, t::AbstractBlockTensorMap) domain(D) == codomain(t) || throw(SpaceMismatch()) TensorKit.foreachblock(t, D) do c, (tblock, Dblock) tblock′ = lmul!(Dblock, copy_dense!(similar_dense(tblock), tblock)) - tblock === tblock′ || copyto!(tblock, tblock′) + tblock === tblock′ || copy_blocks!(tblock, tblock′) return tblock end return t @@ -218,7 +242,7 @@ function LinearAlgebra.rmul!(t::AbstractBlockTensorMap, D::DiagonalTensorMap) codomain(D) == domain(t) || throw(SpaceMismatch()) TensorKit.foreachblock(t, D) do c, (tblock, Dblock) tblock′ = rmul!(copy_dense!(similar_dense(tblock), tblock), Dblock) - tblock === tblock′ || copyto!(tblock, tblock′) + tblock === tblock′ || copy_blocks!(tblock, tblock′) return tblock end return t diff --git a/src/tensors/abstractblocktensor/abstracttensormap.jl b/src/tensors/abstractblocktensor/abstracttensormap.jl index 3a68f2e..49622ce 100644 --- a/src/tensors/abstractblocktensor/abstracttensormap.jl +++ b/src/tensors/abstractblocktensor/abstracttensormap.jl @@ -10,34 +10,14 @@ end eachspace(t::AbstractBlockTensorMap) = SumSpaceIndices(space(t)) -@inline function TensorKit.subblock( - t::AbstractBlockTensorMap, (f₁, f₂)::Tuple{FusionTree, FusionTree} - ) - sectortype(t) === sectortype(f₁) === sectortype(f₂) || - throw(SectorMismatch("Not a valid sectortype for this tensor")) - numout(t) == length(f₁) && numin(t) == length(f₂) || - throw(DimensionMismatch("Invalid number of fusiontree legs for this tensor")) - - subblocks = map(eachspace(t), parent(t)) do V, x - sz = (dims(codomain(V), f₁.uncoupled)..., dims(domain(V), f₂.uncoupled)...) - if prod(sz) == 0 - data = storagetype(t)(undef, 0) - return sreshape(StridedView(data), sz) - else - return x[f₁, f₂] - end - end - - return mortar(subblocks) -end @inline function Base.setindex!( t::AbstractBlockTensorMap, v::AbstractBlockArray, f₁::FusionTree, f₂::FusionTree ) for I in eachindex(t) - b = v[Block(I.I)] - if !isempty(b) - getindex!(t, I)[f₁, f₂] = v[Block(I.I)] - end + b = view(v, Block(I.I)) + isempty(b) && continue + x = issparse(t) ? getindex!(t, I) : t[I] + x[f₁, f₂] = b end return t end @@ -53,57 +33,124 @@ end return setindex!(t, v′, f₁, f₂) end -function TensorKit.block(t::AbstractBlockTensorMap, c::Sector) - sectortype(t) == typeof(c) || throw(SectorMismatch()) +# Blocks and subblocks +# -------------------- +mapentries(f, t::AbstractBlockTensorMap) = + issparse(t) ? Dict(I => f(x) for (I, x) in nonzero_pairs(t)) : map(f, parent(t)) - rows = prod(TT.getindices(size(t), codomainind(t))) - cols = prod(TT.getindices(size(t), domainind(t))) +blockcache(x) = x +blockcache(x::AbstractTensorMap) = (b = TK.blocks(x); b isa TK.BlockIterator ? b : x) +getblock(x, c::Sector) = block(x, c) +getblock(b::TK.BlockIterator, c::Sector) = b[c] - if rows == 0 || cols == 0 - allblocks = Matrix{TK.blocktype(eltype(t))}(undef, rows, cols) +subblockcache(x) = sectortype(x) === Trivial ? x : TK.subblocks(x) +getsubblock(x, f) = subblock(x, f) +getsubblock(s::TK.SubblockIterator, f) = s[f] - rowaxes = Int[] - if rows != 0 - W′ = codomain(t) ← zero(spacetype(t)) - for V in eachspace(W′) - push!(rowaxes, blockdim(codomain(V), c)) - end - end +TensorKit.blocks(t::AbstractBlockTensorMap) = + TK.BlockIterator(t, (blocksectors(t), mapentries(blockcache, t))) +TensorKit.block(t::AbstractBlockTensorMap, c::Sector) = _block(t, parent(t), c) + +TensorKit.subblocks(t::AbstractBlockTensorMap) = + TK.SubblockIterator(t, (TK.fusiontrees(t), mapentries(subblockcache, t))) +TensorKit.subblock(t::AbstractBlockTensorMap, f::Tuple{FusionTree, FusionTree}) = _subblock(t, parent(t), f) + +Base.length(iter::TK.BlockIterator{<:AbstractBlockTensorMap}) = length(first(iter.structure)) +Base.isdone(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...) = + Base.isdone(first(iter.structure), state...) +function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...) + next = iterate(first(iter.structure), state...) + isnothing(next) && return next + c, newstate = next + return c => iter[c], newstate +end +Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = + _block(iter.t, last(iter.structure), c) + +Base.length(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}) = length(first(iter.structure)) +Base.isdone(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int = 1) = i > length(iter) +function Base.iterate(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int = 1) + i > length(iter) && return nothing + f = TK.gettokenvalue(first(iter.structure), i) + return f => iter[f], i + 1 +end +Base.getindex(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, i::Int) = + iter[TK.gettokenvalue(first(iter.structure), i)] +Base.getindex(iter::TK.SubblockIterator{<:AbstractBlockTensorMap}, f::Tuple{FusionTree, FusionTree}) = + _subblock(iter.t, last(iter.structure), f) - colaxes = Int[] - if cols != 0 - W′ = zero(spacetype(t)) ← domain(t) - for V in eachspace(W′) - push!(colaxes, blockdim(domain(V), c)) - end +function _block(t::AbstractBlockTensorMap, entries, c::Sector) + sectortype(t) == typeof(c) || throw(SectorMismatch()) + rowdims = _blockdims(codomain(t), c) + coldims = _blockdims(domain(t), c) + allblocks = Matrix{_blockeltype(typeof(t))}(undef, length(rowdims), length(coldims)) + for (k, I) in enumerate(eachindex(IndexCartesian(), t)) + allblocks[k] = if !issparse(t) || haskey(entries, I) + getblock(entries[I], c) + else + i, j = Tuple(CartesianIndices(allblocks)[k]) + zeroblock(eltype(t), rowdims[i], coldims[j]) end + end + return mortar(allblocks, rowdims, coldims) +end - return mortar(allblocks, rowaxes, colaxes) +function _subblock(t::AbstractBlockTensorMap, entries, (f₁, f₂)::Tuple{FusionTree, FusionTree}) + sectortype(t) === sectortype(f₁) === sectortype(f₂) || + throw(SectorMismatch("Not a valid sectortype for this tensor")) + numout(t) == length(f₁) && numin(t) == length(f₂) || + throw(DimensionMismatch("Invalid number of fusiontree legs for this tensor")) + W = eachspace(t) + subblocks = map(eachindex(IndexCartesian(), t)) do I + V = W[I] + sz = (dims(codomain(V), f₁.uncoupled)..., dims(domain(V), f₂.uncoupled)...) + if prod(sz) == 0 || (issparse(t) && !haskey(entries, I)) + return sreshape(StridedView(zerovector!(similar(storagetype(t), prod(sz)))), sz) + else + return getsubblock(entries[I], (f₁, f₂)) + end end + return mortar(subblocks) +end + +function _blockdims(P::ProductSumSpace{S, N}, c::Sector) where {S, N} + return vec( + map(CartesianIndices(map(length, P.spaces))) do I + return blockdim(ProductSpace{S, N}(map(getindex, P.spaces, Tuple(I))), c) + end + ) +end - allblocks = map(Base.Fix2(block, c), parent(t)) - return mortar(reshape(allblocks, rows, cols)) +zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: AbstractTensorMap} = + zerovector!(similar(storagetype(TT), d₁, d₂)) +function zeroblock(::Type{TT}, d₁::Int, d₂::Int) where {TT <: TensorMap} + data = zerovector!(similar(storagetype(TT), d₁ * d₂)) + return reshape(view(data, 1:(d₁ * d₂)), (d₁, d₂)) end +zeroblock(::Type{<:AdjointTensorMap{<:Any, <:Any, <:Any, <:Any, TT}}, d₁::Int, d₂::Int) where {TT} = + zeroblock(TT, d₂, d₁)' # TODO: this might get fixed once new tensormap is implemented TensorKit.blocksectors(t::AbstractBlockTensorMap) = blocksectors(space(t)) TensorKit.hasblock(t::AbstractBlockTensorMap, c::Sector) = c in blocksectors(t) -TensorKit.blocks(t::AbstractBlockTensorMap) = TK.BlockIterator(t, blocksectors(t)) -Base.@assume_effects :foldable function TensorKit.blocktype(::Type{TT}) where {TT <: AbstractBlockTensorMap} +Base.@assume_effects :foldable function _blockeltype(::Type{TT}) where {TT <: AbstractBlockTensorMap} T = scalartype(TT) B = TK.blocktype(eltype(TT)) - (B <: AbstractMatrix{T}) || (B = AbstractMatrix{T}) # safeguard against type-instability + return B <: AbstractMatrix{T} ? B : AbstractMatrix{T} # safeguard against type-instability +end +function TensorKit.blocktype(::Type{TT}) where {TT <: AbstractBlockTensorMap} BS = NTuple{2, BlockedOneTo{Int, Vector{Int}}} - return BlockMatrix{T, Matrix{B}, BS} + return BlockMatrix{scalartype(TT), Matrix{_blockeltype(TT)}, BS} end -function Base.iterate(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, state...) - next = iterate(iter.structure, state...) - isnothing(next) && return next - c, newstate = next - return c => block(iter.t, c), newstate +function TensorKit.foreachblock(f, t::AbstractBlockTensorMap, ts...; scheduler = nothing) + tensors = (t, ts...) + iters = map(blockcache, tensors) + foreach(union(blocksectors.(tensors)...)) do c + return f(c, map(Base.Fix2(getblock, c), iters)) + end + return nothing end -Base.getindex(iter::TK.BlockIterator{<:AbstractBlockTensorMap}, c::Sector) = block(iter.t, c) TK.storagetype(::Type{TT}) where {TT <: AbstractBlockTensorMap} = storagetype(eltype(TT)) diff --git a/src/tensors/abstractblocktensor/conversion.jl b/src/tensors/abstractblocktensor/conversion.jl index e7e742b..a57dc92 100644 --- a/src/tensors/abstractblocktensor/conversion.jl +++ b/src/tensors/abstractblocktensor/conversion.jl @@ -1,16 +1,38 @@ # Conversion # ---------- -function _copy_subblocks!(tdst, tsrc) - N₁, N₂ = numout(tsrc), numin(tsrc) - offsets = ntuple(i -> _sumspace_offsets(i <= N₁ ? codomain(tsrc)[i] : domain(tsrc)[i - N₁]), N₁ + N₂) - # a single `SubblockIterator` holds the structure: `tdst[f₁, f₂]` would look it up in TensorKit's locked cache every time - dstblocks = subblocks(tdst) - for (k, v) in nonzero_pairs(tsrc), ((f₁, f₂), src) in subblocks(v) - ranges = map(offsets, (f₁.uncoupled..., f₂.uncoupled...), Tuple(k)) do o, c, kᵢ - return (o[c][kᵢ] + 1):o[c][kᵢ + 1] +function _subblock_ranges(offsets, (f₁, f₂), I::CartesianIndex) + return map(offsets, (f₁.uncoupled..., f₂.uncoupled...), Tuple(I)) do o, c, kᵢ + return (o[c][kᵢ] + 1):o[c][kᵢ + 1] + end +end + +function _subblock_pairs(t::AbstractTensorMap) + sectortype(t) === Trivial || return subblocks(t) + f = TK.trivial_fusiontree(t) + return (f => subblock(t, f),) +end + +function _copy_subblocks!(tdst::TensorMap, tsrc::AbstractBlockTensorMap) + offsets = map(_sumspace_offsets, eachspace(tsrc).sumspaces) + dst = Base.Fix1(getsubblock, subblockcache(tdst)) + for (I, v) in nonzero_pairs(tsrc), (f, src) in _subblock_pairs(v) + copy!(view(dst(f), _subblock_ranges(offsets, f, I)...), src) + end + return tdst +end + +function _copy_subblocks!(tdst::AbstractBlockTensorMap, tsrc::AbstractTensorMap) + offsets = map(_sumspace_offsets, eachspace(tdst).sumspaces) + src = Base.Fix1(getsubblock, subblockcache(tsrc)) + for (I, V) in pairs(eachspace(tdst)) + v = similar(eltype(tdst), V) + vblocks = _subblock_pairs(v) + issparse(tdst) && all(isempty ∘ last, vblocks) && continue + for (f, dst) in vblocks + copy!(dst, view(src(f), _subblock_ranges(offsets, f, I)...)) end - copy!(dstblocks[(f₁, f₂)][ranges...], src) + @inbounds tdst[I] = v end return tdst end diff --git a/src/tensors/adjointblocktensor.jl b/src/tensors/adjointblocktensor.jl index 7e17dec..97059aa 100644 --- a/src/tensors/adjointblocktensor.jl +++ b/src/tensors/adjointblocktensor.jl @@ -47,3 +47,4 @@ end function nonzero_length(t::AdjointBlockTensorMap) return nonzero_length(t') end +Base.haskey(t::AdjointBlockTensorMap, I::CartesianIndex) = haskey(t', permute_adjointindices(t', I)) diff --git a/src/tensors/blocktensor.jl b/src/tensors/blocktensor.jl index 3d2b6d9..58146a9 100644 --- a/src/tensors/blocktensor.jl +++ b/src/tensors/blocktensor.jl @@ -112,16 +112,7 @@ end # AbstractTensorMap -> BlockTensorMap function BlockTensorMap(t::AbstractTensorMap, space::TensorMapSumSpace) TT = tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t)) - tdst = BlockTensorMap{TT}(undef, space) - @inbounds for (f₁, f₂) in fusiontrees(t) - dst = tdst[f₁, f₂] - src = t[f₁, f₂] - for block_index in Iterators.product(blockaxes(dst)...) - indices = getindex.(axes(dst), block_index) - dst[block_index...] .= @view src[indices...] - end - end - return tdst + return _copy_subblocks!(BlockTensorMap{TT}(undef_blocks, space), t) end # Convenience constructors diff --git a/src/tensors/indexmanipulations.jl b/src/tensors/indexmanipulations.jl index 1194120..9a0fce9 100644 --- a/src/tensors/indexmanipulations.jl +++ b/src/tensors/indexmanipulations.jl @@ -1,65 +1,4 @@ -@propagate_inbounds function TK.add_transform!( - tdst::BlockTensorMap, tsrc::BlockTensorMap, p::Index2Tuple, transformer, - α::Number, β::Number, backend, allocator - ) - @boundscheck TK.spacecheck_transform(permute, tdst, tsrc, p) - - dstdata = parent(tdst) - srcdata = permutedims(StridedView(parent(tsrc)), (p[1]..., p[2]...)) - - @inbounds for I in eachindex(dstdata, srcdata) - dstdata[I] = TK.add_transform!( - dstdata[I], srcdata[I], p, transformer, α, β, backend, allocator - ) - end - return tdst -end -@propagate_inbounds function TK.add_transform!( - tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, - p::Index2Tuple, transformer, α::Number, β::Number, backend, allocator - ) - @boundscheck TK.spacecheck_transform(permute, tdst, tsrc, p) - scale!(tdst, β) - p_lin = (p[1]..., p[2]...) - @inbounds for (I, v) in nonzero_pairs(tsrc) - I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.add_transform!( - tdst[I′], v, p_lin, transformer, α, One(), backend, allocator - ) - end - return tdst -end -function TK.add_transform!( - tdst::AbstractBlockTensorMap, tsrc::AdjointTensorMap{T, S, N₁, N₂, TT}, - p::Index2Tuple, transformer, α::Number, β::Number, backend, allocator - ) where {T, S, N₁, N₂, TT <: AbstractBlockTensorMap} - @boundscheck TK.spacecheck_transform(permute, tdst, tsrc, p) - scale!(tdst, β) - p_lin = (p[1]..., p[2]...) - @inbounds for (I, v) in nonzero_pairs(tsrc) - I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.add_transform!( - tdst[I′], v, p, transformer, α, One(), backend, allocator - ) - end - return tdst -end -function TK.add_transform!( - tdst::TensorMap, tsrc::BlockTensorMap, p::Index2Tuple, transformer, - α::Number, β::Number, backend, allocator - ) - return TK.add_transform!( - tdst, only(tsrc), p, transformer, α, β, backend, allocator - ) -end -function TK.add_transform!( - tdst::BlockTensorMap, tsrc::TensorMap, p::Index2Tuple, transformer, - α::Number, β::Number, backend, allocator - ) - return TK.add_transform!( - only(tdst), tsrc, p, transformer, α, β, backend, allocator - ) -end +const BlockTensorOrAdjoint = Union{AbstractBlockTensorMap, AdjointBlockTensorMap} # we need to capture the other functions earlier to enjoy the fast transformers... for f in (:permute, :transpose) @@ -80,28 +19,17 @@ for f in (:permute, :transpose) return tdst end function TK.$f!( - tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, + tdst::AbstractBlockTensorMap, tsrc::BlockTensorOrAdjoint, p::Index2Tuple, α::Number, β::Number, backend::AbstractBackend, allocator ) @boundscheck TK.spacecheck_transform(TK.$f, tdst, tsrc, p) - scale!(tdst, β) p_lin = (p[1]..., p[2]...) + _scale_untouched!(tdst, tsrc, β, p_lin) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.$f!(tdst[I′], v, p, α, One(), backend, allocator) - end - return tdst - end - function TK.$f!( - tdst::AbstractBlockTensorMap, tsrc::AdjointTensorMap{T, S, N₁, N₂, TT}, - p::Index2Tuple, α::Number, β::Number, backend::AbstractBackend, allocator - ) where {T, S, N₁, N₂, TT <: AbstractBlockTensorMap} - @boundscheck TK.spacecheck_transform(TK.$f, tdst, tsrc, p) - scale!(tdst, β) - p_lin = (p[1]..., p[2]...) - @inbounds for (I, v) in nonzero_pairs(tsrc) - I′ = CartesianIndex(TT.getindices(I.I, p)) - tdst[I′] = TK.$f!(tdst[I′], v, (p₁, p₂), α, One(), backend, allocator) + _addblock!(tdst, I′, β, allocator) do c, γ + TK.$f!(c, v, p, α, γ, backend, allocator) + end end return tdst end @@ -138,29 +66,18 @@ end return tdst end @propagate_inbounds function TK.braid!( - tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, + tdst::AbstractBlockTensorMap, tsrc::BlockTensorOrAdjoint, p::Index2Tuple, levels::IndexTuple, α::Number, β::Number, backend::AbstractBackend, allocator ) @boundscheck TK.spacecheck_transform(braid, tdst, tsrc, p, levels) - scale!(tdst, β) - p_lin = (p[1]..., p[2]...) - @inbounds for (I, v) in nonzero_pairs(tsrc) - I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.braid!(tdst[I′], v, p, levels, α, One(), backend, allocator) - end - return tdst -end -function TK.braid!( - tdst::AbstractBlockTensorMap, tsrc::AdjointTensorMap{T, S, N₁, N₂, TT}, - p::Index2Tuple, levels::IndexTuple, α::Number, β::Number, backend::AbstractBackend, allocator - ) where {T, S, N₁, N₂, TT <: AbstractBlockTensorMap} - @boundscheck TK.spacecheck_transform(braid, tdst, tsrc, p, levels) - scale!(tdst, β) p_lin = (p[1]..., p[2]...) + _scale_untouched!(tdst, tsrc, β, p_lin) @inbounds for (I, v) in nonzero_pairs(tsrc) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - tdst[I′] = TK.braid!(tdst[I′], v, p, levels, α, One(), backend, allocator) + _addblock!(tdst, I′, β, allocator) do c, γ + TK.braid!(c, v, p, levels, α, γ, backend, allocator) + end end return tdst end @@ -169,7 +86,7 @@ function TK.braid!( p::Index2Tuple, levels::IndexTuple, α::Number, β::Number, backend::AbstractBackend, allocator ) - return TK.add_braid!(tdst, only(tsrc), p, levels, α, β, backend, allocator) + return TK.braid!(tdst, only(tsrc), p, levels, α, β, backend, allocator) end function TK.braid!( tdst::BlockTensorMap, tsrc::TensorMap, diff --git a/src/tensors/sparseblocktensor.jl b/src/tensors/sparseblocktensor.jl index d919caf..e429159 100644 --- a/src/tensors/sparseblocktensor.jl +++ b/src/tensors/sparseblocktensor.jl @@ -85,11 +85,7 @@ end # AbstractTensorMap -> SparseBlockTensorMap function SparseBlockTensorMap(t::AbstractTensorMap, space::TensorMapSumSpace) TT = tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t)) - tdst = SparseBlockTensorMap{TT}(undef, space) - for (f₁, f₂) in fusiontrees(tdst) - tdst[f₁, f₂] = t[f₁, f₂] - end - return tdst + return _copy_subblocks!(SparseBlockTensorMap{TT}(undef_blocks, space), t) end # Utility constructors @@ -189,6 +185,43 @@ function Base.delete!(t::SparseBlockTensorMap{TT}, I::Vararg{Int, N}) where {TT, return delete!(t, CartesianIndex(I...)) end +# Blocks of missing entries are detached zeros, so the in-place fillers that write through +# `blocks` first store the (nonempty) entries they write to. +function storeentries!(t::SparseBlockTensorMap, keep = Returns(true)) + for (I, V) in pairs(eachspace(t)) + !haskey(t, I) && dim(V) > 0 && keep(I) && (t[I] = zerovector!(similar(eltype(t), V))) + end + return t +end +# the entries that overlap the diagonal of a block, i.e. the ones `one!` writes nonzeros to +function diagonalentries(t::SparseBlockTensorMap) + inds = reshape(collect(eachindex(IndexCartesian(), t)), :, prod(TT.getindices(size(t), domainind(t)))) + entries = Set{eltype(inds)}() + for c in blocksectors(t) + rows = cumsum(pushfirst!(_blockdims(codomain(t), c), 0)) + cols = cumsum(pushfirst!(_blockdims(domain(t), c), 0)) + for j in axes(inds, 2), i in axes(inds, 1) + max(rows[i], cols[j]) < min(rows[i + 1], cols[j + 1]) && push!(entries, inds[i, j]) + end + end + return entries +end + +Base.fill!(t::SparseBlockTensorMap, value::Number) = + iszero(value) ? zerovector!(t) : invoke(fill!, Tuple{AbstractTensorMap, Number}, storeentries!(t), value) +for f in (:(TK.one!), :(TK.isomorphism!), :(TK.isometry!)) + @eval function $f(t::SparseBlockTensorMap) + diagonal = diagonalentries(t) + invoke($f, Tuple{AbstractTensorMap}, storeentries!(t, in(diagonal))) + filter!(in(diagonal) ∘ first, t.data) + return t + end +end +for f in (:(Random.rand!), :(Random.randn!), :(Random.randexp!), :(TK.randisometry!)) + @eval $f(rng::Random.AbstractRNG, t::SparseBlockTensorMap) = + invoke($f, Tuple{Random.AbstractRNG, AbstractTensorMap}, rng, storeentries!(t)) +end + # Show # ---- function Base.showarg(io::IO, t::SparseBlockTensorMap, toplevel::Bool) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index a0342c2..ad9c357 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -21,11 +21,13 @@ function TO.tensoradd!( C::AbstractBlockTensorMap, A::AbstractBlockTensorMap, pA::Index2Tuple, conjA::Bool, α::Number, β::Number, backend, allocator ) - scale!(C, β) p_lin = (pA[1]..., pA[2]...) + _scale_untouched!(C, A, β, p_lin) @inbounds for (I, v) in nonzero_pairs(A) I′ = CartesianIndex(TT.getindices(I.I, p_lin)) - C[I′] = TO.tensoradd!(C[I′], v, pA, conjA, α, One(), backend, allocator) + _addblock!(C, I′, β, allocator) do c, γ + TO.tensoradd!(c, v, pA, conjA, α, γ, backend, allocator) + end end return C end @@ -145,13 +147,13 @@ function TK.trace_permute!( q₁ = $(q₁), q₂ = $(q₂)")) end - scale!(tdst, β) + isone(β) || (iszero(β) ? zerovector!(tdst) : scale!(tdst, β)) @inbounds for (Isrc, vsrc) in nonzero_pairs(tsrc) TT.getindices(Isrc.I, q₁) == TT.getindices(Isrc.I, q₂) || continue Idst = CartesianIndex(TT.getindices(Isrc.I, (p₁..., p₂...))) - tdst[Idst] = TensorKit.trace_permute!( - tdst[Idst], vsrc, (p₁, p₂), (q₁, q₂), α, One(), backend - ) + _addblock!(tdst, Idst) do c, γ + TensorKit.trace_permute!(c, vsrc, (p₁, p₂), (q₁, q₂), α, γ, backend) + end end return tdst end diff --git a/src/tensors/vectorinterface.jl b/src/tensors/vectorinterface.jl index f7fbe66..e7d0b6e 100644 --- a/src/tensors/vectorinterface.jl +++ b/src/tensors/vectorinterface.jl @@ -10,15 +10,8 @@ function VI.scale!(t::AbstractBlockTensorMap, α::Number) return t end -function VI.scale!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number) - space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - # entries of ty that are structurally zero in tx have to be zeroed out - issparse(tx) && zerovector!(ty) - for (I, v) in nonzero_pairs(tx) - ty[I] = scale!!(ty[I], v, α) - end - return ty -end +VI.scale!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number) = + add!(ty, tx, α, Zero()) # add # --- @@ -31,7 +24,7 @@ function VI.add(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Numb # This would give issues for example with DiagonalTensorMap + TensorMap T = VectorInterface.promote_add(ty, tx, α, β) tdst = if typeof(ty) === typeof(tx) - zerovector(ty, T) + similar(ty, T) else M = TK.promote_storagetype(TK.similarstoragetype(ty, T), TK.similarstoragetype(tx, T)) if issparse(ty) && issparse(tx) @@ -46,13 +39,42 @@ end function VI.add!(ty::AbstractBlockTensorMap, tx::AbstractBlockTensorMap, α::Number, β::Number) space(ty) == space(tx) || throw(SpaceMismatch("$(space(ty)) ≠ $(space(tx))")) - isone(β) || scale!(ty, β) + _scale_untouched!(ty, tx, β) for (I, v) in nonzero_pairs(tx) - ty[I] = add!!(ty[I], v, α, One()) + _addblock!((y, γ) -> add!!(y, v, α, γ), ty, I, β) end return ty end +# blockwise updates +# ----------------- +@inline function _addblock!( + f, t::AbstractBlockTensorMap, I::CartesianIndex, β::Number = One(), allocator = TO.DefaultAllocator() + ) + t[I] = if haskey(t, I) + f(t[I], isone(β) ? One() : β) + else + f(TO.tensoralloc(eltype(t), eachspace(t)[I], Val(false), allocator), Zero()) + end + return t +end + +function _scale_untouched!( + t::AbstractBlockTensorMap, tsrc::AbstractTensorMap, β::Number, p::IndexTuple = ntuple(identity, numind(t)) + ) + (issparse(tsrc) && !isone(β)) || return t + q = TT.invperm(p) + istouched(I) = haskey(tsrc, CartesianIndex(TT.getindices(I.I, q))) + if issparse(t) && iszero(β) + filter!(istouched ∘ first, parent(t)) + else + for (I, v) in nonzero_pairs(t) + istouched(I) || scale!(v, β) + end + end + return t +end + # inner # ----- function VI.inner(x::AbstractBlockTensorMap, y::AbstractBlockTensorMap) diff --git a/src/vectorspaces/sumspace.jl b/src/vectorspaces/sumspace.jl index 1375e69..a9b30f5 100644 --- a/src/vectorspaces/sumspace.jl +++ b/src/vectorspaces/sumspace.jl @@ -75,12 +75,12 @@ Base.axes(S::SumSpace) = Base.OneTo(dim(S)) Base.axes(S::SumSpace, n::Int) = axes(S.spaces, n) function Base.axes(S::SumSpace, c::Sector) offset = 0 - a = [] + a = Int[] for s in S.spaces - a = push!(a, axes(s, c) .+ offset) + append!(a, axes(s, c) .+ offset) offset += dim(s) end - return collect(flatten(a)) + return a end Base.hash(S::SumSpace, h::UInt) = hash(S.spaces, h) @@ -134,6 +134,22 @@ function TensorKit._sectors(S::SumSpace, ::Type{I}) where {I} return sort!(collect(s)) end +_sectorsubset(S₁::SumSpace, S₂::SumSpace) = all(V -> all(s -> hassector(S₂, s), sectors(V)), S₁.spaces) +function TensorKit.sectorequal(S₁::SumSpace, S₂::SumSpace) + return isdual(S₁) == isdual(S₂) && _sectorsubset(S₁, S₂) && _sectorsubset(S₂, S₁) +end +# `sectorequal` compares the set of sectors, so the hash may only depend on that set. The min and +# max sector hash are independent of component order and duplicates and need no allocation, at the +# cost of more collisions than hashing the full sorted set; collisions only fall back to `sectorequal`. +function TensorKit.sectorhash(S::SumSpace, h::UInt) + lo, hi = typemax(UInt), typemin(UInt) + for V in S.spaces, s in sectors(V) + x = hash(s) + lo, hi = min(lo, x), max(hi, x) + end + return hash(lo, hash(hi, hash(isdual(S), h))) +end + TensorKit.dim(S::SumSpace, sector::Sector) = sum(v -> dim(v, sector), S.spaces; init = 0) # ambiguity fix: TensorKit.dim(S::SumSpace, ::Trivial) = sum(v -> dim(v, Trivial()), S.spaces; init = 0) diff --git a/test/abstracttensor/blocktensor.jl b/test/abstracttensor/blocktensor.jl index 2da67da..13ebf7f 100644 --- a/test/abstracttensor/blocktensor.jl +++ b/test/abstracttensor/blocktensor.jl @@ -99,6 +99,21 @@ end @test t1 ≈ t1″ end +@testset "TensorMap conversion: sectors and empty blocks" begin + Vs = ( + Vect[U1Irrep](0 => 2, 1 => 1) ⊞ Vect[U1Irrep](-1 => 1, 2 => 1) ⊞ Vect[U1Irrep](1 => 2), + Vect[SU2Irrep](0 => 1, 1 // 2 => 2) ⊞ Vect[SU2Irrep](1 => 1) ⊞ Vect[SU2Irrep](1 // 2 => 1, 3 // 2 => 1), + ℂ^0 ⊞ ℂ^2 ⊞ ℂ^3, + ) + for V in Vs, W in (V ⊗ V' ← V ⊗ V, V ⊗ V ⊗ V' ← one(V)), T in scalartypes + t = randn(T, W) + t′ = @constinferred TensorMap(t) + @test norm(t) ≈ norm(t′) + @test BlockTensorMap(t′, W) ≈ t + @test SparseBlockTensorMap(t′, W) ≈ t + end +end + @testset "Adapt" begin W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5 t1 = rand(Float32, W) @@ -170,6 +185,34 @@ end end end +@testset "norm, ==, tr: test via conversion" begin + W = V1 ⊗ V2 ← V1 ⊗ V2 + for T in (Float64, ComplexF64) + t = rand(T, W) + t′ = convert(TensorMap, t) + for p in (1, 2, 3, Inf) + @test norm(t, p) ≈ norm(t′, p) + end + @test tr(t) ≈ tr(t′) + @test t == copy(t) + ts, td = SparseBlockTensorMap(t), BlockTensorMap(t) + @test t == ts == td + @test td == ts + @test td != 2 * ts + @test ts != 2 * td + @test norm(ts, 1) ≈ norm(td, 1) + @test t != 2 * t + t2 = copy(t) + t2[1] = zero(t2[1]) + t2[1] + @test t2 == t + @test (t2 + t) != t + @test t != rand(T, V1 ⊗ V2 ← V1) + @test iszero(norm(spzeros(T, W))) + @test spzeros(T, W) == zero(t) + @test iszero(tr(spzeros(T, W))) + end +end + @testset "Basic linear algebra: test via conversion" begin W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5 for T in (Float32, ComplexF64) diff --git a/test/abstracttensor/sparseblocktensor.jl b/test/abstracttensor/sparseblocktensor.jl index ae0d56c..ec8004f 100644 --- a/test/abstracttensor/sparseblocktensor.jl +++ b/test/abstracttensor/sparseblocktensor.jl @@ -91,6 +91,23 @@ end end end +@testset "TensorMap conversion: sectors and empty blocks" begin + Vs = ( + Vect[U1Irrep](0 => 2, 1 => 1) ⊞ Vect[U1Irrep](-1 => 1, 2 => 1) ⊞ Vect[U1Irrep](1 => 2), + Vect[SU2Irrep](0 => 1, 1 // 2 => 2) ⊞ Vect[SU2Irrep](1 => 1) ⊞ Vect[SU2Irrep](1 // 2 => 1, 3 // 2 => 1), + ℂ^0 ⊞ ℂ^2 ⊞ ℂ^3, + ) + for V in Vs, W in (V ⊗ V' ← V ⊗ V, V ⊗ V ⊗ V' ← one(V)), T in scalartypes + t = sprand(T, W, 0.5) + t′ = @constinferred TensorMap(t) + @test norm(t) ≈ norm(t′) + t″ = SparseBlockTensorMap(t′, W) + @test t″ ≈ t + @test all(I -> dim(eachspace(t)[I]) > 0, nonzero_keys(t″)) + @test TensorMap(BlockTensorMap(t)) ≈ t′ + end +end + @testset "Adapt" begin W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5 t1 = sprand(Float32, W, 0.5) @@ -145,6 +162,35 @@ end end end +@testset "norm, ==, tr: test via conversion" begin + W = V1 ⊗ V2 ← V1 ⊗ V2 + for T in (Float64, ComplexF64) + t = sprand(T, W, 0.5) + t[1] = randn(T, space(t[1])) + t′ = convert(TensorMap, t) + for p in (1, 2, 3, Inf) + @test norm(t, p) ≈ norm(t′, p) + end + @test tr(t) ≈ tr(t′) + @test t == copy(t) + ts, td = SparseBlockTensorMap(t), BlockTensorMap(t) + @test t == ts == td + @test td == ts + @test td != 2 * ts + @test ts != 2 * td + @test norm(ts, 1) ≈ norm(td, 1) + @test t != 2 * t + t2 = copy(t) + t2[1] = zero(t2[1]) + t2[1] + @test t2 == t + @test (t2 + t) != t + @test t != rand(T, V1 ⊗ V2 ← V1) + @test iszero(norm(spzeros(T, W))) + @test spzeros(T, W) == zero(t) + @test iszero(tr(spzeros(T, W))) + end +end + @testset "Basic linear algebra: test via conversion" begin W = V1 ⊗ V2 ⊗ V3 ← V4 ⊗ V5 for T in (Float32, ComplexF64) @@ -247,3 +293,53 @@ end @test HrA12array ≈ convert(TensorMap, HrA12) end + +@testset "sparse block access" begin + Vu = Vect[U1Irrep](0 => 2, 1 => 2, -1 => 2) ⊞ Vect[U1Irrep](0 => 3, 1 => 1) ⊞ Vect[U1Irrep](0 => 1, 2 => 1) + for W in (Vu ⊗ Vu ← Vu, Vu ⊗ Vu' ← Vu ⊗ Vu) + for t in (sprand(W, 0.3), spzeros(W)) + td = BlockTensorMap(t) + @test length(TensorKit.blocks(t)) == length(blocksectors(t)) + for (c, b) in TensorKit.blocks(t) + @test b isa TensorKit.blocktype(t) + @test @constinferred(block(t, c)) == b == block(td, c) + @test size.(b.blocks) == size.(block(td, c).blocks) + end + TensorKit.foreachblock(t, td) do c, (b, bd) + @test b == bd + end + for (f₁, f₂) in fusiontrees(t) + @test t[f₁, f₂] == td[f₁, f₂] + end + sbs = TensorKit.subblocks(t) + @test length(sbs) == length(fusiontrees(t)) + for (i, (f, b)) in enumerate(sbs) + @test b == sbs[i] == TensorKit.subblock(td, f) + end + @test t == td + end + end +end + +@testset "in-place fillers store only the entries they fill" begin + Vu = Vect[U1Irrep](0 => 2, 1 => 2, -1 => 2) ⊞ Vect[U1Irrep](0 => 3, 1 => 1) ⊞ Vect[U1Irrep](0 => 1, 2 => 1) + filled(t) = Set(I for (I, v) in nonzero_pairs(t) if !iszero(norm(v))) + nonempty(t) = Set(I for (I, V) in pairs(eachspace(t)) if dim(V) > 0) + for W in (Vu ⊗ Vu ← Vu ⊗ Vu, Vu ⊗ Vu ← Vu) + fillers = (t -> fill!(t, 2.0), TensorKit.isometry!) + codomain(W) == domain(W) && (fillers = (fillers..., TensorKit.one!, TensorKit.isomorphism!)) + for f! in fillers + t = f!(sprand(W, 0.3)) + td = f!(BlockTensorMap(spzeros(W))) + @test t == td + @test Set(nonzero_keys(t)) == filled(td) + end + @test nonzero_length(fill!(sprand(W, 0.3), 0)) == 0 + for f! in (rand!, randn!, randexp!, randisometry!) + t = f!(spzeros(W)) + @test Set(nonzero_keys(t)) == nonempty(t) + end + end + t = randisometry!(spzeros(Vu ⊗ Vu ← Vu)) + @test t' * t ≈ id(domain(t)) +end diff --git a/test/linalg/factorizations.jl b/test/linalg/factorizations.jl index dce9da3..4de3304 100644 --- a/test/linalg/factorizations.jl +++ b/test/linalg/factorizations.jl @@ -5,6 +5,8 @@ using BlockTensorKit using Random using Combinatorics using LinearAlgebra +using Adapt +using JLArrays Vtr = ( SumSpace(ℂ^3), @@ -526,3 +528,28 @@ end # end # end # end + +@testset "In-place write-back into block storage" begin + V1 = SumSpace(ℂ^2, ℂ^2) + V2 = SumSpace(ℂ^3, ℂ^1) + D = DiagonalTensorMap(rand(dim(V1)), V1) + D′ = DiagonalTensorMap(D.data, ℂ^4) + t = rand(V1 ← V1 ⊗ V2) + t0 = TensorMap(t) + @test TensorMap(lmul!(D, t)) ≈ D′ * t0 + t = rand(V1 ⊗ V2 ← V1) + t0 = TensorMap(t) + @test TensorMap(rmul!(t, D)) ≈ t0 * D′ + + t = rand(V1 ← V1 ⊗ V2) + b = block(t, first(blocksectors(t))) + d = BlockTensorKit.copy_dense!(BlockTensorKit.similar_dense(b), b) .* 2 + BlockTensorKit.copy_blocks!(b, d) + @test BlockTensorKit.copy_dense!(BlockTensorKit.similar_dense(b), b) == d + + tg = adapt(JLVector{Float64}, t) + c = first(blocksectors(t)) + BlockTensorKit.copy_blocks!(block(t, c), d .* 3) + BlockTensorKit.copy_blocks!(block(tg, c), JLArray(d .* 3)) + @test adapt(Vector{Float64}, tg) ≈ t +end diff --git a/test/linalg/indexmanipulations.jl b/test/linalg/indexmanipulations.jl index 8b13789..45e01eb 100644 --- a/test/linalg/indexmanipulations.jl +++ b/test/linalg/indexmanipulations.jl @@ -1 +1,51 @@ +using Test, TestExtras +using TensorKit +using BlockTensorKit +using BlockTensorKit: sprand, nonzero_keys, nonzero_length +using VectorInterface: Zero, One +Vtr = (ℂ^2 ⊞ ℂ^1, (ℂ^1 ⊞ ℂ^2)', ℂ^2 ⊞ ℂ^3, ℂ^2 ⊞ ℂ^2 ⊞ ℂ^1, (ℂ^1 ⊞ ℂ^2)') +F = Vect[FermionParity] +Vf = ( + F(0 => 1, 1 => 1) ⊞ F(0 => 2), F(1 => 1)' ⊞ F(0 => 1, 1 => 2)', F(0 => 2, 1 => 1), + F(0 => 1) ⊞ F(1 => 1) ⊞ F(0 => 1, 1 => 1), F(0 => 1, 1 => 1)' ⊞ F(1 => 2)', +) + +maketensor(sparse, T, W) = sparse ? sprand(T, W, 0.4) : randn(T, W) +dense(t) = convert(TensorMap, t) + +# `(p₁..., reverse(p₂)...)` is a cyclic rotation of the planar order of a 3 ← 2 tensor +perms(f) = f === transpose! ? ((2, 3), (1, 4, 5)) : ((3, 1), (5, 2, 4)) +transform!(f, C, A, p, α, β) = f === braid! ? f(C, A, p, (1, 2, 3, 4, 5), α, β) : f(C, A, p, α, β) + +@testset "$f ($(sectortype(V[1])), C sparse = $Csparse, A sparse = $Asparse, adjoint = $adj)" for V in (Vtr, Vf), + f in (permute!, braid!, transpose!), Csparse in (false, true), Asparse in (false, true), adj in (false, true) + T = ComplexF64 + W = V[1] ⊗ V[2] ⊗ V[3] ← V[4] ⊗ V[5] + A = maketensor(Asparse, T, adj ? W' : W) + A = adj ? A' : A + Ad = adj ? dense(A')' : dense(A) + p = perms(f) + Wdst = permute(space(A), p) + for β in (Zero(), zero(T), One(), randn(T)) + α = randn(T) + C = maketensor(Csparse, T, Wdst) + Cd = transform!(f, dense(C), Ad, p, α, β) + @test transform!(f, C, A, p, α, β) === C + @test dense(C) ≈ Cd + if iszero(β) && Csparse && Asparse + @test nonzero_length(C) == nonzero_length(A) + end + end + C = scale!(maketensor(Csparse, T, Wdst), NaN) + @test dense(transform!(f, C, A, p, One(), Zero())) ≈ transform!(f, zerovector(dense(C)), Ad, p, One(), Zero()) +end + +@testset "permute of adjoint ($(sectortype(V[1])), sparse = $sparse)" for V in (Vtr, Vf), sparse in (false, true) + W = V[1] ⊗ V[2] ⊗ V[3] ← V[4] ⊗ V[5] + t = maketensor(sparse, ComplexF64, W) + p = ((3, 1), (5, 2, 4)) + @test dense(permute(t', p)) ≈ permute(dense(t)', p) + C = similar(t, permute(space(t'), p)) + @test dense(permute!(C, t', p)) ≈ permute(dense(t)', p) +end diff --git a/test/vectorspaces/sumspace.jl b/test/vectorspaces/sumspace.jl index 0e9eb0c..64c1fc6 100644 --- a/test/vectorspaces/sumspace.jl +++ b/test/vectorspaces/sumspace.jl @@ -266,3 +266,32 @@ end @test issetequal(@constinferred(blocksectors(v ← v)), blocksectors(v)) end end + +@testset "sectorhash and sectorequal" begin + V = Vect[U1Irrep](0 => 2, 1 => 2, -1 => 2) ⊞ Vect[U1Irrep](0 => 3, 1 => 1) ⊞ Vect[U1Irrep](2 => 1) + V′ = Vect[U1Irrep](0 => 1, 1 => 1, -1 => 3) ⊞ Vect[U1Irrep](0 => 1, 1 => 2) ⊞ Vect[U1Irrep](2 => 4) + @test TensorKit.sectorequal(V, V′) + @test TensorKit.sectorhash(V, UInt(1)) == TensorKit.sectorhash(V′, UInt(1)) + @test !TensorKit.sectorequal(V, V') + for V″ in (V ⊞ V, V[[3, 1, 2]], ⊞(⊕(V)), ⊞(Vect[U1Irrep](0 => 1, 1 => 1, -1 => 1, 2 => 1))) + @test TensorKit.sectorequal(V, V″) + @test TensorKit.sectorhash(V, UInt(1)) == TensorKit.sectorhash(V″, UInt(1)) + end + @test !TensorKit.sectorequal(V, V[1:2]) + + fallback_equal(a, b) = invoke(TensorKit.sectorequal, Tuple{ElementarySpace, ElementarySpace}, a, b) + randcomp() = Vect[U1Irrep](c => rand(0:1) for c in -2:2) + for _ in 1:500 + dual = rand(Bool) + a, b = (SumSpace([randcomp() for _ in 1:rand(1:4)]) for _ in 1:2) + a, b = dual ? (a', b') : (a, b) + @test TensorKit.sectorequal(a, b) == fallback_equal(a, b) + TensorKit.sectorequal(a, b) && + @test TensorKit.sectorhash(a, UInt(1)) == TensorKit.sectorhash(b, UInt(1)) + end + @test axes(V, U1Irrep(1)) == [3, 4, 10] + for W in (V ⊗ V ← V, V ⊗ V' ← V ⊗ V, (V ⊞ V) ⊗ V' ← V) + @test blocksectors(W) == blocksectors(convert(TensorMapSpace, W)) + end + @test blocksectors(V′ ⊗ V′' ← V′) == blocksectors(V ⊗ V' ← V) +end