make blocks iterator adn use where possible

Author: Jutho

src/TensorKit.jl

@@ -146,23 +146,23 @@ const FusionTreeDict{K,V} = Dict{K,V}
 abstract type TensorException <: Exception end
 
 # Exception type for all errors related to sector mismatch
-struct SectorMismatch{S<:Union{Nothing,String}} <: TensorException
+struct SectorMismatch{S<:Union{Nothing,AbstractString}} <: TensorException
     message::S
 end
 SectorMismatch() = SectorMismatch{Nothing}(nothing)
 Base.show(io::IO, ::SectorMismatch{Nothing}) = print(io, "SectorMismatch()")
 Base.show(io::IO, e::SectorMismatch) = print(io, "SectorMismatch(\"", e.message, "\")")
 
 # Exception type for all errors related to vector space mismatch
-struct SpaceMismatch{S<:Union{Nothing,String}} <: TensorException
+struct SpaceMismatch{S<:Union{Nothing,AbstractString}} <: TensorException
     message::S
 end
 SpaceMismatch() = SpaceMismatch{Nothing}(nothing)
 Base.show(io::IO, ::SpaceMismatch{Nothing}) = print(io, "SpaceMismatch()")
 Base.show(io::IO, e::SpaceMismatch) = print(io, "SpaceMismatch(\"", e.message, "\")")
 
 # Exception type for all errors related to invalid tensor index specification.
-struct IndexError{S<:Union{Nothing,String}} <: TensorException
+struct IndexError{S<:Union{Nothing,AbstractString}} <: TensorException
     message::S
 end
 IndexError() = IndexError{Nothing}(nothing)

src/tensors/abstracttensor.jl

@@ -286,19 +286,24 @@ end
 # tensor data: block access
 #---------------------------
 @doc """
-    blocks(t::AbstractTensorMap) -> SectorDict{<:Sector,<:DenseMatrix}
+    blocks(t::AbstractTensorMap)
 
 Return an iterator over all blocks of a tensor, i.e. all coupled sectors and their
-corresponding blocks.
+corresponding matrix blocks.
 
 See also [`block`](@ref), [`blocksectors`](@ref), [`blockdim`](@ref) and [`hasblock`](@ref).
 """
-blocks(t::AbstractTensorMap) = SectorDict(c => block(t, c) for c in blocksectors(t)) # TODO: make iterator
+function blocks(t::AbstractTensorMap)
+    iter = Base.Iterators.map(blocksectors(t)) do c
+        return c => block(t, c)
+    end
+    return iter
+end
 
 @doc """
-    block(t::AbstractTensorMap, c::Sector) -> DenseMatrix
+    block(t::AbstractTensorMap, c::Sector)
 
-Return the block of a tensor corresponding to a coupled sector `c`.
+Return the matrix block of a tensor corresponding to a coupled sector `c`.
 
 See also [`blocks`](@ref), [`blocksectors`](@ref), [`blockdim`](@ref) and [`hasblock`](@ref).
 """ block

src/tensors/adjoint.jl

@@ -23,9 +23,15 @@ storagetype(::Type{AdjointTensorMap{T,S,N₁,N₂,TT}}) where {T,S,N₁,N₂,TT}
 
 # Blocks and subblocks
 #----------------------
-blocksectors(t::AdjointTensorMap) = blocksectors(parent(t))
 block(t::AdjointTensorMap, s::Sector) = block(parent(t), s)'
 
+function blocks(t::AdjointTensorMap)
+    iter = Base.Iterators.map(blocks(parent(t))) do (c, b)
+        return c => b'
+    end
+    return iter
+end
+
 function Base.getindex(t::AdjointTensorMap{T,S,N₁,N₂},
                        f₁::FusionTree{I,N₁}, f₂::FusionTree{I,N₂}) where {T,S,N₁,N₂,I}
     tp = parent(t)

src/tensors/linalg.jl

@@ -174,7 +174,7 @@ function LinearAlgebra.diagm(codom::VectorSpace, dom::VectorSpace, v::SectorDict
                                                          blockdim(dom, c), b)
                                 for (c, b) in v), codom ← dom)
 end
-LinearAlgebra.isdiag(t::AbstractTensorMap) = all(LinearAlgebra.isdiag, values(blocks(t)))
+LinearAlgebra.isdiag(t::AbstractTensorMap) = all(LinearAlgebra.isdiag ∘ last, blocks(t))
 
 # In-place methods
 #------------------
@@ -184,8 +184,8 @@ LinearAlgebra.isdiag(t::AbstractTensorMap) = all(LinearAlgebra.isdiag, values(bl
 # Copy, adjoint and fill:
 function Base.copy!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap)
     space(tdst) == space(tsrc) || throw(SpaceMismatch("$(space(tdst)) ≠ $(space(tsrc))"))
-    for c in blocksectors(tdst)
-        copy!(StridedView(block(tdst, c)), StridedView(block(tsrc, c)))
+    for ((c, bdst), (_, bsrc)) in zip(blocks(tdst), blocks(tsrc))
+        copy!(StridedView(bdst), StridedView(bsrc))
     end
     return tdst
 end
@@ -284,20 +284,44 @@ function LinearAlgebra.mul!(tC::AbstractTensorMap,
                             tA::AbstractTensorMap,
                             tB::AbstractTensorMap, α=true, β=false)
     compose(space(tA), space(tB)) == space(tC) ||
-        throw(SpaceMismatch("$(space(tC)) ≠ $(space(tA)) * $(space(tB))"))
-
-    for c in blocksectors(tC)
-        if hasblock(tA, c) # then also tB should have such a block
-            A = block(tA, c)
-            B = block(tB, c)
-            C = block(tC, c)
-            mul!(C, A, B, α, β)
-        elseif β != one(β)
-            rmul!(block(tC, c), β)
+        throw(SpaceMismatch(lazy"$(space(tC)) ≠ $(space(tA)) * $(space(tB))"))
+
+    iterC = blocks(tC)
+    iterA = blocks(tA)
+    iterB = blocks(tB)
+    nextA = iterate(iterA)
+    nextB = iterate(iterB)
+    nextC = iterate(iterC)
+    while !isnothing(nextC)
+        (cC, C), stateC = nextC
+        if !isnothing(nextA) && !isnothing(nextB)
+            (cA, A), stateA = nextA
+            (cB, B), stateB = nextB
+            if cA == cC && cB == cC
+                mul!(C, A, B, α, β)
+                nextA = iterate(iterA, stateA)
+                nextB = iterate(iterB, stateB)
+                nextC = iterate(iterC, stateC)
+            elseif cA < cC
+                nextA = iterate(iterA, stateA)
+            elseif cB < cC
+                nextB = iterate(iterB, stateB)
+            else
+                if β != one(β)
+                    rmul!(C, β)
+                end
+                nextC = iterate(iterC, stateC)
+            end
+        else
+            if β != one(β)
+                rmul!(C, β)
+            end
+            nextC = iterate(iterC, stateC)
         end
     end
     return tC
 end
+
 # TODO: consider spawning threads for different blocks, support backends
 
 # TensorMap inverse

src/tensors/tensor.jl

@@ -426,7 +426,14 @@ function block(t::TensorMap, s::Sector)
     return reshape(view(t.data, r), (d₁, d₂))
 end
 
-blocks(t::TensorMap) = SectorDict(c => block(t, c) for c in blocksectors(t)) # TODO: make iterator
+function blocks(t::TensorMap)
+    structure = fusionblockstructure(t).blockstructure
+    iter = Base.Iterators.map(structure) do (c, ((d₁, d₂), r))
+        b = reshape(view(t.data, r), (d₁, d₂))
+        return c => b
+    end
+    return iter
+end
 
 # Indexing and getting and setting the data at the subblock level
 #-----------------------------------------------------------------

src/tensors/tensortreeiterator.jl

@@ -17,15 +17,15 @@ function force_planar(tsrc::TensorMap{<:Any,ComplexSpace})
     tdst = TensorMap{scalartype(tsrc)}(undef,
                                        force_planar(codomain(tsrc)) ←
                                        force_planar(domain(tsrc)))
-    copyto!(blocks(tdst)[PlanarTrivial()], blocks(tsrc)[Trivial()])
+    copyto!(block(tdst, PlanarTrivial()), block(tsrc, Trivial()))
     return tdst
 end
 function force_planar(tsrc::TensorMap{<:Any,<:GradedSpace})
     tdst = TensorMap{scalartype(tsrc)}(undef,
                                        force_planar(codomain(tsrc)) ←
                                        force_planar(domain(tsrc)))
     for (c, b) in blocks(tsrc)
-        copyto!(blocks(tdst)[c ⊠ PlanarTrivial()], b)
+        copyto!(block(tdst, c ⊠ PlanarTrivial()), b)
     end
     return tdst
 end