[Fix] correctly handle eachsite for Multiline objects (#281)

* correctly handle `eachsite` for `Multiline` objects

* Bump v0.13.1

Author: lkdvos

Project.toml

@@ -1,7 +1,7 @@
 name = "MPSKit"
 uuid = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 authors = "Lukas Devos, Maarten Van Damme and contributors"
-version = "0.13.0"
+version = "0.13.1"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

src/algorithms/approximate/vomps.jl

@@ -58,11 +58,11 @@ end
 function localupdate_step!(::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any,<:Tuple},
                            ::SerialScheduler)
     alg_orth = QRpos()
-    eachsite = 1:length(state.mps)
+
     ACs = similar(state.mps.AC)
     dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
 
-    foreach(eachsite) do site
+    foreach(eachsite(state.mps)) do site
         AC = circshift([AC_projection(CartesianIndex(row, site), state.mps, state.operator,
                                       state.envs)
                         for row in 1:size(state.mps, 1)], 1)
@@ -78,12 +78,11 @@ end
 function localupdate_step!(::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any,<:Tuple},
                            scheduler)
     alg_orth = QRpos()
-    eachsite = 1:length(state.mps)
 
     ACs = similar(state.mps.AC)
     dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
 
-    tforeach(eachsite; scheduler) do site
+    tforeach(eachsite(state.mps); scheduler) do site
         local AC, C
         @sync begin
             Threads.@spawn begin

src/algorithms/groundstate/vumps.jl

@@ -104,13 +104,12 @@ function localupdate_step!(it::IterativeSolver{<:VUMPS}, state,
     alg_orth = QRpos()
 
     mps = state.mps
-    eachsite = 1:length(mps)
     src_Cs = mps isa Multiline ? eachcol(mps.C) : mps.C
     src_ACs = mps isa Multiline ? eachcol(mps.AC) : mps.AC
     ACs = similar(mps.AC)
     dst_ACs = mps isa Multiline ? eachcol(ACs) : ACs
 
-    tforeach(eachsite, src_ACs, src_Cs; scheduler) do site, AC₀, C₀
+    tforeach(eachsite(mps), src_ACs, src_Cs; scheduler) do site, AC₀, C₀
         dst_ACs[site] = _localupdate_vumps_step!(site, mps, state.operator, state.envs,
                                                  AC₀, C₀; parallel=false, alg_orth,
                                                  state.which, alg_eigsolve)

src/algorithms/statmech/vomps.jl

@@ -100,13 +100,12 @@ function localupdate_step!(::IterativeSolver{<:VOMPS}, state,
                            scheduler=Defaults.scheduler[])
     alg_orth = QRpos()
     mps = state.mps
-    eachsite = 1:length(mps)
     src_Cs = mps isa Multiline ? eachcol(mps.C) : mps.C
     src_ACs = mps isa Multiline ? eachcol(mps.AC) : mps.AC
     ACs = similar(mps.AC)
     dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
 
-    tforeach(eachsite, src_ACs, src_Cs; scheduler) do site, AC₀, C₀
+    tforeach(eachsite(mps), src_ACs, src_Cs; scheduler) do site, AC₀, C₀
         dst_ACs[site] = _localupdate_vomps_step!(site, mps, state.operator, state.envs,
                                                  AC₀, C₀; alg_orth, parallel=false)
         return nothing

src/states/abstractmps.jl

@@ -197,3 +197,10 @@ physicalspace(A::MPSTensor) = space(A, 2)
 physicalspace(A::GenericMPSTensor) = prod(x -> space(A, x), 2:(numind(A) - 1))
 physicalspace(O::MPOTensor) = space(O, 2)
 physicalspace(O::AbstractBlockTensorMap{<:Any,<:Any,2,2}) = only(space(O, 2))
+
+"""
+    eachsite(state::AbstractMPS)
+
+Return an iterator over the sites of the MPS `state`.
+"""
+eachsite(ψ::AbstractMPS) = eachindex(ψ)

src/utility/multiline.jl

@@ -30,6 +30,8 @@ function Base.axes(m::Multiline, i::Int)
 end
 Base.eachindex(m::Multiline) = CartesianIndices(size(m))
 
+eachsite(m::Multiline) = eachsite(first(parent(m)))
+
 Base.getindex(m::Multiline, i::Int) = getindex(parent(m), i)
 Base.setindex!(m::Multiline, v, i::Int) = (setindex!(parent(m), v, i); m)