Add block iterator for DiagonalTensorMap

Author: lkdvos

src/spaces/homspace.jl

@@ -349,3 +349,18 @@ function fusionblockstructure(W::HomSpace, ::GlobalLRUCache)
     end
     return structure
 end
+
+# Diagonal ranges
+#----------------
+# TODO: is this something we want to cache?
+function diagonalblockstructure(W::HomSpace)
+    structure = SectorDict{sectortype(W),UnitRange{Int}}() # range
+    offset = 0
+    dom = domain(W)[1]
+    for c in blocksectors(W)
+        d = dim(dom, c)
+        structure[c] = offset .+ (1:d)
+        offset + d
+    end
+    return structure
+end

src/tensors/diagonal.jl

@@ -110,7 +110,21 @@ function block(d::DiagonalTensorMap, s::Sector)
     return Diagonal(view(d.data, 1:0))
 end
 
-# TODO: is relying on generic AbstractTensorMap blocks sufficient?
+blocks(t::DiagonalTensorMap) = BlockIterator(t, diagonalblockstructure(space(t)))
+blocktype(::Type{TT}) where {TT<:DiagonalTensorMap} = Diagonal{eltype(TT),storagetype(TT)}
+
+function Base.iterate(iter::BlockIterator{<:DiagonalTensorMap}, state...)
+    next = iterate(iter.structure, state...)
+    isnothing(next) && return next
+    (c, r), newstate = next
+    return c => Diagonal(view(iter.t.data, r)), newstate
+end
+
+function Base.getindex(iter::BlockIterator{<:DiagonalTensorMap}, c::Sector)
+    sectortype(iter.t) === typeof(c) || throw(SectorMismatch())
+    r = get(iter.structure, c, 1:0)
+    return Diagonal(view(iter.t.data, r))
+end
 
 # Indexing and getting and setting the data at the subblock level
 #-----------------------------------------------------------------