[Utility] Be more consistent about what space functions work for what arguments (#310)

lkdvos · web-flow · commit 0aa451fb1d76 · 2025-08-03T12:48:46.000-04:00
* Consistently implement `physicalspace(x)`

* Consistently implement `left`- and `right_virtualspace(x)`

* Consistently implement `eachsite`

* Add tests

* Add spaces for multilinemps

* some small fixes

* format

* use functions correctly

* improve test coverage

* add missing constructor

* ensure infinite spaces are periodic
diff --git a/src/algorithms/expval.jl b/src/algorithms/expval.jl
@@ -42,7 +42,7 @@ function expectation_value(ψ::AbstractMPS, (inds, O)::Pair)
         dual(_lastspace(last(local_mpo)))
     for (site, o) in zip(sites, local_mpo)
         if o isa MPOTensor
-            physicalspace(ψ)[site] == physicalspace(o) ||
+            physicalspace(ψ, site) == physicalspace(o) ||
                 throw(SpaceMismatch("physical space does not match at site $site"))
         end
     end
diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl
@@ -207,7 +207,7 @@ function variance(
             state.AC[i], envs.GLs[i], H[i][:, :, :, end], envs.GRs[i][end]
         )
     end
-    lattice = physicalspace.(Ref(state), 1:length(state))
+    lattice = physicalspace(state)
     H_renormalized = InfiniteMPOHamiltonian(
         lattice, i => e * id(storagetype(eltype(H)), lattice[i]) for (i, e) in enumerate(e_local)
     )
@@ -234,7 +234,7 @@ function variance(state::InfiniteQP, H::InfiniteMPOHamiltonian, envs = environme
         GR = rightenv(envs, i, gs)
         return contract_mpo_expval(gs.AC[i], GL, H[i][:, :, :, end], GR[end])
     end
-    lattice = physicalspace.(Ref(gs), 1:length(state))
+    lattice = physicalspace(gs)
     H_regularized = H - InfiniteMPOHamiltonian(
         lattice, i => e * id(storagetype(eltype(H)), lattice[i]) for (i, e) in enumerate(e_local)
     )
diff --git a/src/operators/abstractmpo.jl b/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))
diff --git a/src/operators/mpo.jl b/src/operators/mpo.jl
@@ -54,6 +54,7 @@ DenseMPO(mpo::MPO) = mpo isa DenseMPO ? copy(mpo) : MPO(map(TensorMap, parent(mp
 # -------
 Base.parent(mpo::MPO) = mpo.O
 Base.copy(mpo::MPO) = MPO(copy.(parent(mpo)))
+eachsite(mpo::InfiniteMPO) = PeriodicArray(eachindex(mpo))
 
 function Base.similar(mpo::MPO{<:MPOTensor}, ::Type{O}, L::Int) where {O <: MPOTensor}
     return MPO(similar(parent(mpo), O, L))
diff --git a/src/states/abstractmps.jl b/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)
diff --git a/src/states/finitemps.jl b/src/states/finitemps.jl
@@ -394,7 +394,6 @@ function right_virtualspace(ψ::FiniteMPS, n::Integer)
         _firstspace(ψ.C[n])
 end
 
-physicalspace(ψ::FiniteMPS) = physicalspace.(Ref(ψ), 1:length(ψ))
 function physicalspace(ψ::FiniteMPS{<:GenericMPSTensor{<:Any, N}}, n::Integer) where {N}
     N == 1 && return ProductSpace{spacetype(ψ)}()
     return physicalspace(coalesce(ψ.ALs[n], ψ.ARs[n], ψ.ACs[n]))
diff --git a/src/states/infinitemps.jl b/src/states/infinitemps.jl
@@ -265,17 +265,15 @@ end
 
 Base.eachindex(ψ::InfiniteMPS) = eachindex(ψ.AL)
 Base.eachindex(l::IndexStyle, ψ::InfiniteMPS) = eachindex(l, ψ.AL)
+eachsite(ψ::InfiniteMPS) = PeriodicArray(eachindex(ψ))
 
 Base.checkbounds(::Type{Bool}, ψ::InfiniteMPS, i::Integer) = true
 
 site_type(::Type{<:InfiniteMPS{A}}) where {A} = A
 bond_type(::Type{<:InfiniteMPS{<:Any, B}}) where {B} = B
 
-left_virtualspace(ψ::InfiniteMPS) = left_virtualspace.(ψ.AL)
 left_virtualspace(ψ::InfiniteMPS, n::Integer) = left_virtualspace(ψ.AL[n])
-right_virtualspace(ψ::InfiniteMPS) = right_virtualspace.(ψ.AL)
 right_virtualspace(ψ::InfiniteMPS, n::Integer) = right_virtualspace(ψ.AL[n])
-physicalspace(ψ::InfiniteMPS) = physicalspace.(ψ.AL)
 physicalspace(ψ::InfiniteMPS, n::Integer) = physicalspace(ψ.AL[n])
 
 # TensorKit.space(ψ::InfiniteMPS{<:MPSTensor}, n::Integer) = space(ψ.AC[n], 2)
diff --git a/src/states/multilinemps.jl b/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
diff --git a/src/states/quasiparticle_state.jl b/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}
diff --git a/src/states/windowmps.jl b/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)
diff --git a/src/utility/windowarray.jl b/src/utility/windowarray.jl
@@ -15,9 +15,9 @@ struct WindowArray{T} <: AbstractVector{T}
     right::PeriodicVector{T}
 end
 function WindowArray(
-        left::PeriodicVector{T}, middle::AbstractVector{T}, right::PeriodicVector{T}
+        left::AbstractVector{T}, middle::AbstractVector{T}, right::AbstractVector{T}
     ) where {T}
-    return WindowArray{T}(left, convert(Vector{T}, middle), right)
+    return WindowArray{T}(left, middle, right)
 end
 
 # these definitions are a bit iffy, but will do for now
diff --git a/test/operators.jl b/test/operators.jl
@@ -31,6 +31,15 @@ module TestOperators
             mpo₁ = FiniteMPO(O₁) # type-unstable for now!
             mpo₂ = FiniteMPO(O₂)
             mpo₃ = FiniteMPO(O₃)
+
+            @test @constinferred physicalspace(mpo₁) == fill(V, L)
+            Vleft = @constinferred left_virtualspace(mpo₁)
+            Vright = @constinferred right_virtualspace(mpo₂)
+            for i in 1:L
+                @test Vleft[i] == left_virtualspace(mpo₁, i)
+                @test Vright[i] == right_virtualspace(mpo₁, i)
+            end
+
             @test convert(TensorMap, mpo₁) ≈ O₁
             @test convert(TensorMap, -mpo₂) ≈ -O₂
             @test convert(TensorMap, @constinferred complex(mpo₃)) ≈ complex(O₃)
diff --git a/test/states.jl b/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,15 @@ module TestStates
             ϕ₁ = LeftGaugedQP(rand, ψ)
             ϕ₂ = LeftGaugedQP(rand, ψ)
 
+            @test @constinferred physicalspace(ϕ₁) == physicalspace(ψ)
+            @test @constinferred left_virtualspace(ϕ₁) == left_virtualspace(ψ)
+            @test @constinferred right_virtualspace(ϕ₁) == right_virtualspace(ψ)
+            for i in 1:period
+                @test physicalspace(ψ, i) == physicalspace(ϕ₁, i)
+                @test left_virtualspace(ψ, i) == left_virtualspace(ϕ₁, i)
+                @test right_virtualspace(ψ, i) == right_virtualspace(ϕ₁, i)
+            end
+
             @test norm(axpy!(1, ϕ₁, copy(ϕ₂))) ≤ norm(ϕ₁) + norm(ϕ₂)
             @test norm(ϕ₁) * 3 ≈ norm(ϕ₁ * 3)
 

Original file line number	Diff line number	Diff line change
`@@ -207,7 +207,7 @@ function variance(`
`207`	`207`	`state.AC[i], envs.GLs[i], H[i][:, :, :, end], envs.GRs[i][end]`
`208`	`208`	`)`
`209`	`209`	`end`
`210`		`- lattice = physicalspace.(Ref(state), 1:length(state))`
	`210`	`+ lattice = physicalspace(state)`
`211`	`211`	`H_renormalized = InfiniteMPOHamiltonian(`
`212`	`212`	`lattice, i => e * id(storagetype(eltype(H)), lattice[i]) for (i, e) in enumerate(e_local)`
`213`	`213`	`)`
`@@ -234,7 +234,7 @@ function variance(state::InfiniteQP, H::InfiniteMPOHamiltonian, envs = environme`
`234`	`234`	`GR = rightenv(envs, i, gs)`
`235`	`235`	`return contract_mpo_expval(gs.AC[i], GL, H[i][:, :, :, end], GR[end])`
`236`	`236`	`end`
`237`		`- lattice = physicalspace.(Ref(gs), 1:length(state))`
	`237`	`+ lattice = physicalspace(gs)`
`238`	`238`	`H_regularized = H - InfiniteMPOHamiltonian(`
`239`	`239`	`lattice, i => e * id(storagetype(eltype(H)), lattice[i]) for (i, e) in enumerate(e_local)`
`240`	`240`	`)`