Add block iterator for TensorMap

Author: lkdvos

src/TensorKit.jl

@@ -185,7 +185,7 @@ include("spaces/vectorspaces.jl")
 #-------------------------------------
 # general definitions
 include("tensors/abstracttensor.jl")
-# include("tensors/tensortreeiterator.jl")
+include("tensors/blockiterator.jl")
 include("tensors/tensor.jl")
 include("tensors/adjoint.jl")
 include("tensors/linalg.jl")

src/tensors/blockiterator.jl

@@ -0,0 +1,15 @@
+"""
+    struct BlockIterator{T<:AbstractTensorMap,S}
+
+Iterator over the blocks of type `T`, possibly holding some pre-computed data of type `S`
+"""
+struct BlockIterator{T<:AbstractTensorMap,S}
+    t::T
+    structure::S
+end
+
+Base.IteratorSize(::BlockIterator) = Base.HasLength()
+Base.IteratorEltype(::BlockIterator) = Base.HasEltype()
+Base.eltype(::Type{<:BlockIterator{T}}) where {T} = blocktype(T)
+Base.length(iter::BlockIterator) = length(iter.structure)
+Base.isdone(iter::BlockIterator, state...) = Base.isdone(iter.structure, state...)

src/tensors/tensor.jl

@@ -421,28 +421,35 @@ end
 
 # Getting and setting the data at the block level
 #-------------------------------------------------
-function block(t::TensorMap, s::Sector)
-    sectortype(t) == typeof(s) || throw(SectorMismatch())
-    structure = fusionblockstructure(t).blockstructure
-    (d₁, d₂), r = get(structure, s) do
+block(t::TensorMap, c::Sector) = blocks(t)[c]
+
+blocks(t::TensorMap) = BlockIterator(t, fusionblockstructure(t).blockstructure)
+
+function blocktype(::Type{TT}) where {TT<:TensorMap}
+    A = storagetype(TT)
+    T = eltype(A)
+    return Base.ReshapedArray{T,2,SubArray{T,1,A,Tuple{UnitRange{Int}},true},Tuple{}}
+end
+
+function Base.iterate(iter::BlockIterator{<:TensorMap}, state...)
+    next = iterate(iter.structure, state...)
+    isnothing(next) && return next
+    (c, (sz, r)), newstate = next
+    return c => reshape(view(iter.t.data, r), sz), newstate
+end
+
+function Base.getindex(iter::BlockIterator{<:TensorMap}, c::Sector)
+    sectortype(iter.t) === typeof(c) || throw(SectorMismatch())
+    (d₁, d₂), r = get(iter.structure, c) do
         # is s is not a key, at least one of the two dimensions will be zero:
         # it then does not matter where exactly we construct a view in `t.data`,
         # as it will have length zero anyway
-        d₁′ = blockdim(codomain(t), s)
-        d₂′ = blockdim(domain(t), s)
+        d₁′ = blockdim(codomain(iter.t), c)
+        d₂′ = blockdim(domain(iter.t), c)
         l = d₁′ * d₂′
         return (d₁′, d₂′), 1:l
     end
-    return reshape(view(t.data, r), (d₁, d₂))
-end
-
-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
+    return reshape(view(iter.t.data, r), (d₁, d₂))
 end
 
 # Indexing and getting and setting the data at the subblock level
@@ -559,18 +566,22 @@ function Base.show(io::IO, t::TensorMap)
     if sectortype(t) == Trivial
         Base.print_array(io, t[])
         println(io)
-    elseif FusionStyle(sectortype(t)) isa UniqueFusion
-        for (f₁, f₂) in fusiontrees(t)
-            println(io, "* Data for sector ", f₁.uncoupled, " ← ", f₂.uncoupled, ":")
-            Base.print_array(io, t[f₁, f₂])
-            println(io)
-        end
     else
-        for (f₁, f₂) in fusiontrees(t)
-            println(io, "* Data for fusiontree ", f₁, " ← ", f₂, ":")
-            Base.print_array(io, t[f₁, f₂])
-            println(io)
+        for (c, b) in blocks(t)
+            println(io, "* Data for sector $c:")
+            Base.print_array(io, b)
         end
+        #     for (f₁, f₂) in fusiontrees(t)
+        #         println(io, "* Data for sector ", f₁.uncoupled, " ← ", f₂.uncoupled, ":")
+        #         Base.print_array(io, t[f₁, f₂])
+        #         println(io)
+        #     end
+        # else
+        #     for (f₁, f₂) in fusiontrees(t)
+        #         println(io, "* Data for fusiontree ", f₁, " ← ", f₂, ":")
+        #         Base.print_array(io, t[f₁, f₂])
+        #         println(io)
+        #     end
     end
 end