[Utility] Be more consistent about what space functions work for what arguments

src/operators/abstractmpo.jl

@@ -14,6 +14,7 @@ Base.isfinite(O::AbstractMPO) = isfinite(typeof(O))
 # By default, define things in terms of parent
 Base.size(mpo::AbstractMPO, args...) = size(parent(mpo), args...)
 Base.length(mpo::AbstractMPO) = length(parent(mpo))
+eachsite(mpo::AbstractMPO) = eachindex(mpo)
 
 @inline Base.getindex(mpo::AbstractMPO, i::Int) = getindex(parent(mpo), i)
 @inline function Base.setindex!(mpo::AbstractMPO, value, i::Int)
@@ -24,11 +25,11 @@ end
 # Properties
 # ----------
 left_virtualspace(mpo::AbstractMPO, site::Int) = left_virtualspace(mpo[site])
-left_virtualspace(mpo::AbstractMPO) = map(left_virtualspace, parent(mpo))
+left_virtualspace(mpo::AbstractMPO) = map(Base.Fix1(left_virtualspace, mpo), eachsite(mpo))
 right_virtualspace(mpo::AbstractMPO, site::Int) = right_virtualspace(mpo[site])
-right_virtualspace(mpo::AbstractMPO) = map(right_virtualspace, parent(mpo))
+right_virtualspace(mpo::AbstractMPO) = map(Base.Fix1(right_virtualspace, mpo), eachsite(mpo))
 physicalspace(mpo::AbstractMPO, site::Int) = physicalspace(mpo[site])
-physicalspace(mpo::AbstractMPO) = map(physicalspace, mpo)
+physicalspace(mpo::AbstractMPO) = map(Base.Fix1(physicalspace, mpo), eachsite(mpo))
 
 for ftype in (:spacetype, :sectortype, :storagetype)
     @eval TensorKit.$ftype(mpo::AbstractMPO) = $ftype(typeof(mpo))

src/states/abstractmps.jl

@@ -187,6 +187,7 @@ Return the virtual space of the bond to the left of sites `pos`.
 function left_virtualspace end
 left_virtualspace(A::GenericMPSTensor) = space(A, 1)
 left_virtualspace(O::MPOTensor) = space(O, 1)
+left_virtualspace(ψ::AbstractMPS) = map(Base.Fix1(left_virtualspace, ψ), eachsite(ψ))
 
 """
     right_virtualspace(ψ::AbstractMPS, [pos=1:length(ψ)])
@@ -200,6 +201,7 @@ Return the virtual space of the bond to the right of site(s) `pos`.
 function right_virtualspace end
 right_virtualspace(A::GenericMPSTensor) = space(A, numind(A))'
 right_virtualspace(O::MPOTensor) = space(O, 4)'
+right_virtualspace(ψ::AbstractMPS) = map(Base.Fix1(right_virtualspace, ψ), eachsite(ψ))
 
 """
     physicalspace(ψ::AbstractMPS, [pos=1:length(ψ)])
@@ -211,6 +213,7 @@ 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))
+physicalspace(ψ::AbstractMPS) = map(Base.Fix1(physicalspace, ψ), eachsite(ψ))
 
 """
     eachsite(state::AbstractMPS)

src/states/multilinemps.jl

@@ -101,5 +101,8 @@ Base.convert(::Type{InfiniteMPS}, st::MultilineMPS) = only(st)
 Base.eltype(t::MultilineMPS) = eltype(t[1])
 Base.copy!(ψ::MultilineMPS, ϕ::MultilineMPS) = (copy!.(parent(ψ), parent(ϕ)); ψ)
 
-left_virtualspace(t::MultilineMPS, i::Int, j::Int) = left_virtualspace(t[i], j)
-right_virtualspace(t::MultilineMPS, i::Int, j::Int) = right_virtualspace(t[i], j)
+for f_space in (:physicalspace, :left_virtualspace, :right_virtualspace)
+    @eval $f_space(t::MultilineMPS, i::Int, j::Int) = $f_space(t[i], j)
+    @eval $f_space(t::MultilineMPS, I::CartesianIndex{2}) = $f_space(t, Tuple(I)...)
+    @eval $f_space(t::MultilineMPS) = map(Base.Fix1($f_space, t), eachindex(t))
+end

src/states/quasiparticle_state.jl

@@ -215,10 +215,16 @@ const MultilineQP{Q <: QP} = Multiline{Q}
 TensorKit.spacetype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = spacetype(S)
 TensorKit.sectortype(::Union{QP{S}, Type{<:QP{S}}}) where {S} = sectortype(S)
 
+physicalspace(state::QP, i::Int) = physicalspace(state.left_gs, i)
+physicalspace(state::QP) = physicalspace(state.left_gs)
 left_virtualspace(state::QP, i::Int) = left_virtualspace(state.left_gs, i)
+left_virtualspace(state::QP) = map(Base.Fix1(left_virtualspace, state), eachsite(state))
 right_virtualspace(state::QP, i::Int) = right_virtualspace(state.right_gs, i)
+right_virtualspace(state::QP) = map(Base.Fix1(right_virtualspace, state), eachsite(state))
 auxiliaryspace(state::QP) = space(state.Xs[1], 2)
 
+eachsite(state::QP) = eachsite(state.left_gs)
+
 Base.copy(a::QP) = copy!(similar(a), a)
 Base.copyto!(a::QP, b::QP) = copy!(a, b)
 function Base.copy!(a::T, b::T) where {T <: QP}

src/states/windowmps.jl

@@ -150,11 +150,7 @@ for f in (:physicalspace, :left_virtualspace, :right_virtualspace)
     @eval $f(ψ::WindowMPS, n::Integer) = n < 1 ? $f(ψ.left_gs, n) :
         n > length(ψ) ? $f(ψ.right_gs, n - length(ψ)) :
         $f(ψ.window, n)
-end
-function physicalspace(ψ::WindowMPS)
-    return WindowArray(
-        physicalspace(ψ.left_gs), physicalspace(ψ.window), physicalspace(ψ.right_gs)
-    )
+    @eval $f(ψ::WindowMPS) = WindowArray($f(ψ.left_gs), $f(ψ.window), $f(ψ.right_gs))
 end
 r_RR(ψ::WindowMPS) = r_RR(ψ.right_gs, length(ψ))
 l_LL(ψ::WindowMPS) = l_LL(ψ.left_gs, 1)

test/states.jl

@@ -21,7 +21,12 @@ module TestStates
                 ComplexF32,
             ),
         ]
-        ψ = FiniteMPS(rand, elt, rand(3:20), d, D)
+        L = rand(3:20)
+        ψ = FiniteMPS(rand, elt, L, d, D)
+
+        @test @constinferred physicalspace(ψ) == fill(d, L)
+        @test all(x -> x ≾ D, @constinferred left_virtualspace(ψ))
+        @test all(x -> x ≾ D, @constinferred right_virtualspace(ψ))
 
         ovl = dot(ψ, ψ)
 
@@ -94,6 +99,10 @@ module TestStates
 
         ψ = InfiniteMPS([rand(elt, D * d, D), rand(elt, D * d, D)]; tol)
 
+        @test physicalspace(ψ) == fill(d, 2)
+        @test all(x -> x ≾ D, left_virtualspace(ψ))
+        @test all(x -> x ≾ D, right_virtualspace(ψ))
+
         for i in 1:length(ψ)
             @plansor difference[-1 -2; -3] := ψ.AL[i][-1 -2; 1] * ψ.C[i][1; -3] -
                 ψ.C[i - 1][-1; 1] * ψ.AR[i][1 -2; -3]
@@ -121,6 +130,10 @@ module TestStates
             ]; tol
         )
 
+        @test physicalspace(ψ) == fill(d, 2, 2)
+        @test all(x -> x ≾ D, left_virtualspace(ψ))
+        @test all(x -> x ≾ D, right_virtualspace(ψ))
+
         for i in 1:size(ψ, 1), j in 1:size(ψ, 2)
             @plansor difference[-1 -2; -3] := ψ.AL[i, j][-1 -2; 1] * ψ.C[i, j][1; -3] -
                 ψ.C[i, j - 1][-1; 1] * ψ.AR[i, j][1 -2; -3]
@@ -155,6 +168,16 @@ module TestStates
         # constructor 2 - used to take a "slice" from an infinite mps
         window_2 = WindowMPS(gs, 10)
 
+        P = @constinferred physicalspace(window_2)
+        Vleft = @constinferred left_virtualspace(window_2)
+        Vright = @constinferred right_virtualspace(window_2)
+
+        for i in -3:13
+            @test physicalspace(window_2, i) == P[i]
+            @test left_virtualspace(window_2, i) == Vleft[i]
+            @test right_virtualspace(window_2, i) == Vright[i]
+        end
+
         # we should logically have that window_1 approximates window_2
         ovl = dot(window_1, window_2)
         @test ovl ≈ 1 atol = 1.0e-8
@@ -206,6 +229,10 @@ module TestStates
             ϕ₁ = LeftGaugedQP(rand, ψ)
             ϕ₂ = LeftGaugedQP(rand, ψ)
 
+            @test @constinferred physicalspace(ϕ₁) == physicalspace(ψ)
+            @test @constinferred left_virtualspace(ϕ₁) == left_virtualspace(ψ)
+            @test @constinferred right_virtualspace(ϕ₁) == right_virtualspace(ψ)
+
             @test norm(axpy!(1, ϕ₁, copy(ϕ₂))) ≤ norm(ϕ₁) + norm(ϕ₂)
             @test norm(ϕ₁) * 3 ≈ norm(ϕ₁ * 3)
 
@@ -237,6 +264,10 @@ module TestStates
             ϕ₁ = LeftGaugedQP(rand, ψ)
             ϕ₂ = LeftGaugedQP(rand, ψ)
 
+            @test @constinferred physicalspace(ϕ₁) == physicalspace(ψ)
+            @test @constinferred left_virtualspace(ϕ₁) == left_virtualspace(ψ)
+            @test @constinferred right_virtualspace(ϕ₁) == right_virtualspace(ψ)
+
             @test norm(axpy!(1, ϕ₁, copy(ϕ₂))) ≤ norm(ϕ₁) + norm(ϕ₂)
             @test norm(ϕ₁) * 3 ≈ norm(ϕ₁ * 3)