MPO utility and fixes (#256)

* improve `tensorexpr` flexibility

* Refactor `convert(TensorMap, ::FiniteMPO(Hamiltonian))`

* Bump minimal tensorkit/blocktensorkit versions because `removeunit`

* implement `FiniteMPO(Hamiltonian) * TensorMap`

* expand operator testset

Author: lkdvos

Project.toml

@@ -24,7 +24,7 @@ VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 [compat]
 Accessors = "0.1"
 Aqua = "0.8.9"
-BlockTensorKit = "0.1.1"
+BlockTensorKit = "0.1.4"
 Compat = "3.47, 4.10"
 DocStringExtensions = "0.9.3"
 HalfIntegers = "1.6.0"
@@ -38,7 +38,7 @@ Plots = "1.40"
 Printf = "1"
 Random = "1"
 RecipesBase = "1.1"
-TensorKit = "0.13, 0.14"
+TensorKit = "0.14"
 TensorKitManifolds = "0.7"
 TensorOperations = "5"
 Test = "1"

src/operators/abstractmpo.jl

@@ -269,3 +269,36 @@ function add_physical_charge(O::AbstractBlockTensorMap{<:Any,<:Any,2,2}, charge:
     end
     return Odst
 end
+
+# Contractions
+# ------------
+# This function usually does not require to be specified for many N, so @generated function is fine?
+@generated function _instantiate_finitempo(L::AbstractTensorMap{<:Any,S,1,2},
+                                           O::NTuple{N,MPOTensor{S}},
+                                           R::AbstractTensorMap{<:Any,S,2,1}) where {N,S}
+    sites = N + 2
+    t_out = tensorexpr(:T, -(1:sites), -(1:sites) .- sites)
+    t_left = tensorexpr(:L, -1, (-1 - sites, 1))
+    t_mid = ntuple(N) do n
+        return tensorexpr(:(O[$n]), (n, -n - 1), (-n - sites - 1, n + 1))
+    end
+    t_right = tensorexpr(:R, (sites - 1, -sites), -2sites)
+    ex = :(@plansor $t_out ≔ *($t_left, $t_right, $(t_mid...)))
+    return macroexpand(@__MODULE__, ex)
+end
+
+@generated function _apply_finitempo(x::AbstractTensorMap{<:Any,S,M,A},
+                                     L::AbstractTensorMap{<:Any,S,1,2},
+                                     O::NTuple{N,MPOTensor{S}},
+                                     R::AbstractTensorMap{<:Any,S,2,1}) where {N,M,S,A}
+    M == N + 2 || throw(ArgumentError("Incompatible number of spaces"))
+    t_out = tensorexpr(:y, -(1:M), -(1:A) .- M)
+    t_in = tensorexpr(:x, 1:2:(2M - 1), -(1:A) .- M)
+    t_left = tensorexpr(:L, -1, (1, 2))
+    t_mid = ntuple(N) do n
+        return tensorexpr(:(O[$n]), (2n, -n - 1), (2n + 1, 2n + 2))
+    end
+    t_right = tensorexpr(:R, (2N + 2, -M), 2N + 3)
+    ex = :(@plansor $t_out ≔ *($t_in, $t_left, $t_right, $(t_mid...)))
+    return macroexpand(@__MODULE__, ex)
+end

src/operators/mpo.jl

@@ -96,23 +96,10 @@ function _mps_to_mpo(A::GenericMPSTensor{S,3}) where {S}
 end
 
 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)
-    @assert V_left == oneunit(V_left)
-    U_left = ones(scalartype(mpo), V_left)'
-
-    V_right = right_virtualspace(mpo, length(mpo))
-    @assert V_right == oneunit(V_right)
-    U_right = ones(scalartype(mpo), V_right)
-
-    tensors = vcat(U_left, parent(mpo), U_right)
-    indices = [[i, -i, -(2N - i + 1), i + 1] for i in 1:length(mpo)]
-    pushfirst!(indices, [1])
-    push!(indices, [N + 1])
-    O = ncon(tensors, indices)
-
-    return repartition(O, N, N)
+    L = removeunit(mpo[1], 1)
+    R = removeunit(mpo[end], 4)
+    M = Tuple(mpo[2:(end - 1)])
+    return convert(TensorMap, _instantiate_finitempo(L, M, R))
 end
 
 # Linear Algebra
@@ -279,6 +266,16 @@ function Base.:*(mpo1::InfiniteMPO, mpo2::InfiniteMPO)
     return InfiniteMPO(Os)
 end
 
+function Base.:*(mpo::FiniteMPO{<:MPOTensor}, x::AbstractTensorMap)
+    @assert length(mpo) > 1
+    @assert numout(x) == length(mpo)
+    L = removeunit(mpo[1], 1)
+    M = Tuple(mpo[2:(end - 1)])
+    R = removeunit(mpo[end], 4)
+    y = _apply_finitempo(x, L, M, R)
+    return convert(TensorMap, y)
+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

src/operators/mpohamiltonian.jl

@@ -359,23 +359,10 @@ function isemptylevel(H::InfiniteMPOHamiltonian, i::Int)
 end
 
 function Base.convert(::Type{TensorMap}, H::FiniteMPOHamiltonian)
-    N = length(H)
-    # add trivial tensors to remove left and right trivial leg.
-    V_left = left_virtualspace(H, 1)
-    @assert V_left == oneunit(V_left)
-    U_left = ones(scalartype(H), V_left)'
-
-    V_right = right_virtualspace(H, length(H))
-    @assert V_right == oneunit(V_right)
-    U_right = ones(scalartype(H), V_right)
-
-    tensors = vcat(U_left, parent(H), U_right)
-    indices = [[i, -i, -(i + N), i + 1] for i in 1:length(H)]
-    pushfirst!(indices, [1])
-    push!(indices, [N + 1])
-    O = convert(TensorMap, ncon(tensors, indices))
-
-    return transpose(O, (ntuple(identity, N), ntuple(i -> i + N, N)))
+    L = removeunit(H[1], 1)
+    R = removeunit(H[end], 4)
+    M = Tuple(H[2:(end - 1)])
+    return convert(TensorMap, _instantiate_finitempo(L, M, R))
 end
 
 function add_physical_charge(H::MPOHamiltonian, charges::AbstractVector{<:Sector})
@@ -556,6 +543,15 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS)
     return FiniteMPS(A′)
 end
 
+function Base.:*(H::FiniteMPOHamiltonian{<:MPOTensor}, x::AbstractTensorMap)
+    @assert length(H) > 1
+    @assert numout(x) == length(H)
+    L = removeunit(H[1], 1)
+    M = Tuple(H[2:(end - 1)])
+    R = removeunit(H[end], 4)
+    return convert(TensorMap, _apply_finitempo(x, L, M, R))
+end
+
 function TensorKit.dot(H₁::FiniteMPOHamiltonian, H₂::FiniteMPOHamiltonian)
     check_length(H₁, H₂)
 

src/utility/utility.jl

@@ -120,15 +120,18 @@ function randomize!(a::AbstractBlockTensorMap)
     return a
 end
 
+_totuple(t) = t isa Tuple ? t : Tuple(t)
+
 """
-    tensorexpr(name::Symbol, ind_out, [ind_in])
+    tensorexpr(name, ind_out, [ind_in])
 
 Generates expressions for use within [`@tensor`](@extref TensorOperations.@tensor) environments
 of the form `name[ind_out...; ind_in]`.
 """
-tensorexpr(name::Symbol, inds) = Expr(:ref, name, inds...)
-function tensorexpr(name::Symbol, indout, indin)
-    return Expr(:typed_vcat, name, Expr(:row, indout...), Expr(:row, indin...))
+tensorexpr(name, inds) = Expr(:ref, name, _totuple(inds)...)
+function tensorexpr(name, indout, indin)
+    return Expr(:typed_vcat, name, Expr(:row, _totuple(indout)...),
+                Expr(:row, _totuple(indin)...))
 end
 
 function check_length(a, b...)

test/operators.jl

@@ -45,6 +45,7 @@ vspaces = (ℙ^10, Rep[U₁]((0 => 20)), Rep[SU₂](1 // 2 => 10, 3 // 2 => 5, 5
         mps₁ = FiniteMPS(ψ₁)
 
         @test convert(TensorMap, mpo₁ * mps₁) ≈ O₁ * ψ₁
+        @test mpo₁ * ψ₁ ≈ O₁ * ψ₁
 
         @test dot(mps₁, mpo₁, mps₁) ≈ dot(ψ₁, O₁, ψ₁)
         @test dot(mps₁, mpo₁, mps₁) ≈ dot(mps₁, mpo₁ * mps₁)
@@ -60,66 +61,70 @@ end
 
 @testset "Finite MPOHamiltonian" begin
     L = 3
-    lattice = fill(ℂ^2, L)
-    O₁ = rand(ComplexF64, ℂ^2, ℂ^2)
-    E = id(storagetype(O₁), domain(O₁))
-    O₂ = rand(ComplexF64, ℂ^2 * ℂ^2, ℂ^2 * ℂ^2)
-
-    H1 = FiniteMPOHamiltonian(lattice, i => O₁ for i in 1:L)
-    H2 = FiniteMPOHamiltonian(lattice, (i, i + 1) => O₂ for i in 1:(L - 1))
-    H3 = FiniteMPOHamiltonian(lattice, 1 => O₁, (2, 3) => O₂, (1, 3) => O₂)
-
-    # check if constructor works by converting back to tensormap
-    H1_tm = convert(TensorMap, H1)
-    operators = vcat(fill(E, L - 1), O₁)
-    @test H1_tm ≈ mapreduce(+, 1:L) do i
-        return reduce(⊗, circshift(operators, i))
-    end
-    operators = vcat(fill(E, L - 2), O₂)
-    @test convert(TensorMap, H2) ≈ mapreduce(+, 1:(L - 1)) do i
-        return reduce(⊗, circshift(operators, i))
+    T = ComplexF64
+    for V in (ℂ^2, U1Space(0 => 1, 1 => 1))
+        lattice = fill(V, L)
+        O₁ = rand(T, V, V)
+        E = id(storagetype(O₁), domain(O₁))
+        O₂ = rand(T, V^2 ← V^2)
+
+        H1 = FiniteMPOHamiltonian(lattice, i => O₁ for i in 1:L)
+        H2 = FiniteMPOHamiltonian(lattice, (i, i + 1) => O₂ for i in 1:(L - 1))
+        H3 = FiniteMPOHamiltonian(lattice, 1 => O₁, (2, 3) => O₂, (1, 3) => O₂)
+
+        # check if constructor works by converting back to tensormap
+        H1_tm = convert(TensorMap, H1)
+        operators = vcat(fill(E, L - 1), O₁)
+        @test H1_tm ≈ mapreduce(+, 1:L) do i
+            return reduce(⊗, circshift(operators, i))
+        end
+        operators = vcat(fill(E, L - 2), O₂)
+        @test convert(TensorMap, H2) ≈ mapreduce(+, 1:(L - 1)) do i
+            return reduce(⊗, circshift(operators, i))
+        end
+        @test convert(TensorMap, H3) ≈
+              O₁ ⊗ E ⊗ E + E ⊗ O₂ + permute(O₂ ⊗ E, ((1, 3, 2), (4, 6, 5)))
+
+        # check if adding terms on the same site works
+        single_terms = Iterators.flatten(Iterators.repeated((i => O₁ / 2 for i in 1:L), 2))
+        H4 = FiniteMPOHamiltonian(lattice, single_terms)
+        @test H4 ≈ H1 atol = 1e-6
+        double_terms = Iterators.flatten(Iterators.repeated(((i, i + 1) => O₂ / 2
+                                                             for i in 1:(L - 1)), 2))
+        H5 = FiniteMPOHamiltonian(lattice, double_terms)
+        @test H5 ≈ H2 atol = 1e-6
+
+        # test linear algebra
+        @test H1 ≈
+              FiniteMPOHamiltonian(lattice, 1 => O₁) +
+              FiniteMPOHamiltonian(lattice, 2 => O₁) +
+              FiniteMPOHamiltonian(lattice, 3 => O₁)
+        @test 0.8 * H1 + 0.2 * H1 ≈ H1 atol = 1e-6
+        @test convert(TensorMap, H1 + H2) ≈ convert(TensorMap, H1) + convert(TensorMap, H2) atol = 1e-6
+
+        # test dot and application
+        state = rand(T, prod(lattice))
+        mps = FiniteMPS(state)
+
+        @test convert(TensorMap, H1 * mps) ≈ H1_tm * state
+        @test H1 * state ≈ H1_tm * state
+        @test dot(mps, H2, mps) ≈ dot(mps, H2 * mps)
+
+        # test constructor from dictionary with mixed linear and Cartesian lattice indices as keys
+        grid = square = fill(V, 3, 3)
+
+        local_operators = Dict((I,) => O₁ for I in eachindex(grid))
+        I_vertical = CartesianIndex(1, 0)
+        vertical_operators = Dict((I, I + I_vertical) => O₂
+                                  for I in eachindex(IndexCartesian(), square)
+                                  if I[1] < size(square, 1))
+        operators = merge(local_operators, vertical_operators)
+        H4 = FiniteMPOHamiltonian(grid, operators)
+
+        @test H4 ≈
+              FiniteMPOHamiltonian(grid, local_operators) +
+              FiniteMPOHamiltonian(grid, vertical_operators)
     end
-    @test convert(TensorMap, H3) ≈
-          O₁ ⊗ E ⊗ E + E ⊗ O₂ + permute(O₂ ⊗ E, ((1, 3, 2), (4, 6, 5)))
-
-    # check if adding terms on the same site works
-    single_terms = Iterators.flatten(Iterators.repeated((i => O₁ / 2 for i in 1:L), 2))
-    H4 = FiniteMPOHamiltonian(lattice, single_terms)
-    @test H4 ≈ H1 atol = 1e-6
-    double_terms = Iterators.flatten(Iterators.repeated(((i, i + 1) => O₂ / 2
-                                                         for i in 1:(L - 1)), 2))
-    H5 = FiniteMPOHamiltonian(lattice, double_terms)
-    @test H5 ≈ H2 atol = 1e-6
-
-    # test linear algebra
-    @test H1 ≈
-          FiniteMPOHamiltonian(lattice, 1 => O₁) +
-          FiniteMPOHamiltonian(lattice, 2 => O₁) +
-          FiniteMPOHamiltonian(lattice, 3 => O₁)
-    @test 0.8 * H1 + 0.2 * H1 ≈ H1 atol = 1e-6
-    @test convert(TensorMap, H1 + H2) ≈ convert(TensorMap, H1) + convert(TensorMap, H2) atol = 1e-6
-
-    # test dot and application
-    state = rand(ComplexF64, prod(lattice))
-    mps = FiniteMPS(state)
-
-    @test convert(TensorMap, H1 * mps) ≈ H1_tm * state
-    @test dot(mps, H2, mps) ≈ dot(mps, H2 * mps)
-
-    # test constructor from dictionary with mixed linear and Cartesian lattice indices as keys
-    grid = square = fill(ℂ^2, 3, 3)
-
-    local_operators = Dict((I,) => O₁ for I in eachindex(grid))
-    I_vertical = CartesianIndex(1, 0)
-    vertical_operators = Dict((I, I + I_vertical) => O₂
-                              for I in eachindex(IndexCartesian(), square)
-                              if I[1] < size(square, 1))
-    operators = merge(local_operators, vertical_operators)
-    H4 = FiniteMPOHamiltonian(grid, operators)
-
-    @test H4 ≈
-          FiniteMPOHamiltonian(grid, local_operators) +
-          FiniteMPOHamiltonian(grid, vertical_operators)
 end
 
 @testset "InfiniteMPOHamiltonian $(sectortype(pspace))" for (pspace, Dspace) in