Relax MPO element type restrictions (#244)

* Relax MPO element type restrictions

* Apply suggestions from code review

Co-authored-by: Lukas Devos <[email protected]>

* Add TODO for merging finite and infinite MPO multiplication

* Replace `similar` with appropriate `allocate_GL`

---------

Co-authored-by: Lukas Devos <[email protected]>

Author: leburgel

src/operators/abstractmpo.jl

@@ -1,11 +1,11 @@
 # Matrix Product Operators
 # ========================
 """
-    abstract type AbstractMPO{O<:MPOTensor} <: AbstractVector{O} end
+    abstract type AbstractMPO{O} <: AbstractVector{O} end
 
 Abstract supertype for Matrix Product Operators (MPOs).
 """
-abstract type AbstractMPO{O<:MPOTensor} <: AbstractVector{O} end
+abstract type AbstractMPO{O} <: AbstractVector{O} end
 
 # useful union types
 const SparseMPO{O<:SparseBlockTensorMap} = AbstractMPO{O}
@@ -16,7 +16,7 @@ Base.size(mpo::AbstractMPO, args...) = size(parent(mpo), args...)
 Base.length(mpo::AbstractMPO) = length(parent(mpo))
 
 @inline Base.getindex(mpo::AbstractMPO, i::Int) = getindex(parent(mpo), i)
-@inline function Base.setindex!(mpo::AbstractMPO, value::MPOTensor, i::Int)
+@inline function Base.setindex!(mpo::AbstractMPO, value, i::Int)
     setindex!(parent(mpo), value, i)
     return mpo
 end
@@ -194,20 +194,24 @@ end
 # Kernels
 # -------
 # TODO: diagram
+
+function _fuse_mpo_mpo(O1::MPOTensor, O2::MPOTensor, Fₗ, Fᵣ)
+    return @plansor O′[-1 -2; -3 -4] := Fₗ[-1; 1 2] *
+                                        O2[1 3; -3 5] *
+                                        O1[2 -2; 3 4] *
+                                        conj(Fᵣ[-4; 5 4])
+end
+
 """
     fuse_mul_mpo(O1, O2)
 
 Compute the mpo tensor that arises from multiplying MPOs.
 """
-function fuse_mul_mpo(O1::MPOTensor, O2::MPOTensor)
+function fuse_mul_mpo(O1, O2)
     T = promote_type(scalartype(O1), scalartype(O2))
     F_left = fuser(T, left_virtualspace(O2), left_virtualspace(O1))
     F_right = fuser(T, right_virtualspace(O2), right_virtualspace(O1))
-    @plansor O[-1 -2; -3 -4] := F_left[-1; 1 2] *
-                                O2[1 5; -3 3] *
-                                O1[2 -2; 5 4] *
-                                conj(F_right[-4; 3 4])
-    return O
+    return _fuse_mpo_mpo(O1, O2, F_left, F_right)
 end
 function fuse_mul_mpo(O1::BraidingTensor, O2::BraidingTensor)
     T = promote_type(scalartype(O1), scalartype(O2))

src/operators/mpo.jl

@@ -1,24 +1,24 @@
 """
-    struct MPO{O<:MPOTensor,V<:AbstractVector{O}} <: AbstractMPO{O}
+    struct MPO{O,V<:AbstractVector{O}} <: AbstractMPO{O}
 
 Matrix Product Operator (MPO) acting on a tensor product space with a linear order.
 
 See also: [`FiniteMPO`](@ref), [`InfiniteMPO`](@ref)
 """
-struct MPO{TO<:MPOTensor,V<:AbstractVector{TO}} <: AbstractMPO{TO}
+struct MPO{TO,V<:AbstractVector{TO}} <: AbstractMPO{TO}
     O::V
 end
 
 """
-    FiniteMPO(Os::Vector{<:MPOTensor}) -> FiniteMPO
-    FiniteMPO(O::AbstractTensorMap{S,N,N}) where {S,N} -> FiniteMPO
+    FiniteMPO(Os::Vector{O}) -> FiniteMPO{O}
+    FiniteMPO(O::AbstractTensorMap{S,N,N}) where {S,N} -> FiniteMPO{O<:MPOTensor}
 
 Matrix Product Operator (MPO) acting on a finite tensor product space with a linear order.
 """
-const FiniteMPO{O<:MPOTensor} = MPO{O,Vector{O}}
+const FiniteMPO{O} = MPO{O,Vector{O}}
 Base.isfinite(::Type{<:FiniteMPO}) = true
 
-function FiniteMPO(Os::AbstractVector{O}) where {O<:MPOTensor}
+function FiniteMPO(Os::AbstractVector{O}) where {O}
     for i in eachindex(Os)[1:(end - 1)]
         right_virtualspace(Os[i]) == left_virtualspace(Os[i + 1]) ||
             throw(SpaceMismatch("unmatching virtual spaces at site $i"))
@@ -31,14 +31,14 @@ function FiniteMPO(O::AbstractTensorMap{T,S,N,N}) where {T,S,N}
 end
 
 """
-    InfiniteMPO(Os::PeriodicVector{<:MPOTensor}) -> InfiniteMPO
+    InfiniteMPO(Os::PeriodicVector{O}) -> InfiniteMPO{O}
 
 Matrix Product Operator (MPO) acting on an infinite tensor product space with a linear order.
 """
-const InfiniteMPO{O<:MPOTensor} = MPO{O,PeriodicVector{O}}
+const InfiniteMPO{O} = MPO{O,PeriodicVector{O}}
 Base.isfinite(::Type{<:InfiniteMPO}) = false
 
-function InfiniteMPO(Os::AbstractVector{O}) where {O<:MPOTensor}
+function InfiniteMPO(Os::AbstractVector{O}) where {O}
     for i in eachindex(Os)
         right_virtualspace(Os[i]) == left_virtualspace(Os[mod1(i + 1, end)]) ||
             throw(SpaceMismatch("umatching virtual spaces at site $i"))
@@ -55,7 +55,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(map(copy, mpo))
 
-function Base.similar(mpo::MPO, ::Type{O}, L::Int) where {O<:MPOTensor}
+function Base.similar(mpo::MPO, ::Type{O}, L::Int) where {O}
     return MPO(similar(parent(mpo), O, L))
 end
 
@@ -95,7 +95,7 @@ function _mps_to_mpo(A::GenericMPSTensor{S,3}) where {S}
     return O
 end
 
-function Base.convert(::Type{TensorMap}, mpo::FiniteMPO)
+function Base.convert(::Type{TensorMap}, mpo::FiniteMPO{<:MPOTensor})
     N = length(mpo)
     # add trivial tensors to remove left and right trivial leg.
     V_left = left_virtualspace(mpo, 1)
@@ -120,7 +120,7 @@ end
 # VectorInterface.scalartype(::Type{FiniteMPO{O}}) where {O} = scalartype(O)
 
 Base.:+(mpo::MPO) = MPO(map(+, parent(mpo)))
-function Base.:+(mpo1::FiniteMPO{TO}, mpo2::FiniteMPO{TO}) where {TO}
+function Base.:+(mpo1::FiniteMPO{TO}, mpo2::FiniteMPO{TO}) where {TO<:MPOTensor}
     (N = length(mpo1)) == length(mpo2) || throw(ArgumentError("dimension mismatch"))
     @assert left_virtualspace(mpo1, 1) == left_virtualspace(mpo2, 1) &&
             right_virtualspace(mpo1, N) == right_virtualspace(mpo2, N)
@@ -207,7 +207,8 @@ function VectorInterface.scale!(mpo::MPO, α::Number)
     return mpo
 end
 
-function Base.:*(mpo1::FiniteMPO{TO}, mpo2::FiniteMPO{TO}) where {TO}
+# TODO: merge implementation with that of InfiniteMPO
+function Base.:*(mpo1::FiniteMPO{TO}, mpo2::FiniteMPO{TO}) where {TO<:MPOTensor}
     (N = length(mpo1)) == length(mpo2) || throw(ArgumentError("dimension mismatch"))
     S = spacetype(TO)
     if (left_virtualspace(mpo1, 1) != oneunit(S) ||
@@ -273,29 +274,16 @@ function _fuse_mpo_mps(O::MPOTensor, A::MPSTensor, Fₗ, Fᵣ)
 end
 
 function Base.:*(mpo1::InfiniteMPO, mpo2::InfiniteMPO)
-    L = check_length(mpo1, mpo2)
-
-    T = promote_type(scalartype(mpo1), scalartype(mpo2))
-    make_fuser(i) = fuser(T, left_virtualspace(mpo2, i), left_virtualspace(mpo1, i))
-    fusers = PeriodicArray(map(make_fuser, 1:L))
-
-    Os = map(1:L) do i
-        return _fuse_mpo_mpo(mpo1[i], mpo2[i], fusers[i], fusers[i + 1])
-    end
+    check_length(mpo1, mpo2)
+    Os = map(fuse_mul_mpo, parent(mpo1), parent(mpo2))
     return InfiniteMPO(Os)
 end
 
-function _fuse_mpo_mpo(O1::MPOTensor, O2::MPOTensor, Fₗ, Fᵣ)
-    return @plansor O′[-1 -2; -3 -4] := Fₗ[-1; 1 4] *
-                                        O2[1 2; -3 3] *
-                                        O1[4 -2; 2 5] *
-                                        conj(Fᵣ[-4; 3 5])
-end
-
 # TODO: I think the fastest order is to start from both ends, and take the overlap at the
 # largest virtual space cut, but it might be better to just multithread both sides and meet
 # in the middle
-function TensorKit.dot(bra::FiniteMPS{T}, mpo::FiniteMPO, ket::FiniteMPS{T}) where {T}
+function TensorKit.dot(bra::FiniteMPS{T}, mpo::FiniteMPO{<:MPOTensor},
+                       ket::FiniteMPS{T}) where {T}
     (N = length(bra)) == length(mpo) == length(ket) ||
         throw(ArgumentError("dimension mismatch"))
     Nhalf = N ÷ 2
@@ -320,17 +308,15 @@ function TensorKit.dot(bra::FiniteMPS{T}, mpo::FiniteMPO, ket::FiniteMPS{T}) whe
 end
 function TensorKit.dot(bra::InfiniteMPS, mpo::InfiniteMPO, ket::InfiniteMPS;
                        ishermitian=false, krylovdim=30, kwargs...)
-    ρ₀ = similar(bra.AL[1],
-                 left_virtualspace(ket, 1) * left_virtualspace(mpo, 1) ←
-                 left_virtualspace(bra, 1))
+    ρ₀ = allocate_GL(bra, mpo, ket, 1)
     randomize!(ρ₀)
 
     val, = fixedpoint(TransferMatrix(ket.AL, parent(mpo), bra.AL), ρ₀, :LM; ishermitian,
                       krylovdim, kwargs...)
     return val
 end
 
-function TensorKit.dot(mpo₁::FiniteMPO, mpo₂::FiniteMPO)
+function TensorKit.dot(mpo₁::FiniteMPO{TO}, mpo₂::FiniteMPO{TO}) where {TO<:MPOTensor}
     length(mpo₁) == length(mpo₂) || throw(ArgumentError("dimension mismatch"))
     N = length(mpo₁)
     Nhalf = N ÷ 2

src/operators/multilinempo.jl

@@ -13,10 +13,10 @@ See also: [`Multiline`](@ref), [`AbstractMPO`](@ref)
 """
 const MultilineMPO = Multiline{<:AbstractMPO}
 
-function MultilineMPO(Os::AbstractMatrix{T}) where {T<:MPOTensor}
+function MultilineMPO(Os::AbstractMatrix)
     return MultilineMPO(map(FiniteMPO, eachrow(Os)))
 end
-function MultilineMPO(Os::PeriodicMatrix{T}) where {T<:MPOTensor}
+function MultilineMPO(Os::PeriodicMatrix)
     return MultilineMPO(map(InfiniteMPO, eachrow(Os)))
 end
 MultilineMPO(mpos::AbstractVector{<:AbstractMPO}) = Multiline(mpos)