Add double-layer-pepo reduced densitymatrix

.github/workflows/Tests.yml

@@ -33,6 +33,7 @@ jobs:
           - utility
           - bondenv
           - timeevol
+          - toolbox
         os:
           - ubuntu-latest
           - macOS-latest

src/algorithms/contractions/localoperator.jl

@@ -171,6 +171,62 @@ function _contract_state_expr(rowrange, colrange, cartesian_inds = nothing)
     end
 end
 
+function _contract_pepo_state_expr(rowrange, colrange, cartesian_inds = nothing)
+    rmin, rmax = extrema(rowrange)
+    cmin, cmax = extrema(colrange)
+    gridsize = (rmax - rmin + 1, cmax - cmin + 1)
+
+    return map((:bra, :ket)) do side
+        return map(Iterators.product(1:gridsize[1], 1:gridsize[2])) do (i, j)
+            inds_id = if isnothing(cartesian_inds)
+                nothing
+            else
+                findfirst(==(CartesianIndex(rmin + i - 1, cmin + j - 1)), cartesian_inds)
+            end
+            physical_label_in = if isnothing(inds_id)
+                physicallabel(:in, i, j)
+            else
+                physicallabel(:Oopen, inds_id)
+            end
+            physical_label_out = if isnothing(inds_id)
+                physicallabel(:out, i, j)
+            else
+                physicallabel(:O, side, inds_id)
+            end
+            return tensorexpr(
+                if side == :ket
+                    :(twistdual(ket[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)], (1, 2)))
+                else
+                    :(bra[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)])
+                end,
+                (physical_label_out, physical_label_in),
+                (
+                    if i == 1
+                        virtuallabel(NORTH, side, j)
+                    else
+                        virtuallabel(:vertical, side, i - 1, j)
+                    end,
+                    if j == gridsize[2]
+                        virtuallabel(EAST, side, i)
+                    else
+                        virtuallabel(:horizontal, side, i, j)
+                    end,
+                    if i == gridsize[1]
+                        virtuallabel(SOUTH, side, j)
+                    else
+                        virtuallabel(:vertical, side, i, j)
+                    end,
+                    if j == 1
+                        virtuallabel(WEST, side, i)
+                    else
+                        virtuallabel(:horizontal, side, i, j - 1)
+                    end,
+                ),
+            )
+        end
+    end
+end
+
 @generated function _contract_local_operator(
         inds::NTuple{N, Val},
         O::AbstractTensorMap{T, S, N, N},
@@ -251,25 +307,40 @@ end
     return macroexpand(@__MODULE__, returnex)
 end
 
-"""
+@doc """
 $(SIGNATURES)
 
 Construct the reduced density matrix `ρ` of the PEPS `peps` with open indices `inds` using the environment `env`.
+Alternatively, construct the reduced density matrix `ρ` of a full density matrix PEPO with open indices `inds` using the environment `env`.
 
 This works by generating the appropriate contraction on a rectangular patch with its corners
 specified by `inds`. The result is normalized such that `tr(ρ) = 1`. 
-"""
+""" reduced_densitymatrix
+
 function reduced_densitymatrix(
         inds::NTuple{N, CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
     ) where {N}
     static_inds = Val.(inds)
     return _contract_densitymatrix(static_inds, ket, bra, env)
 end
+function reduced_densitymatrix(
+        inds::NTuple{N, CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env::CTMRGEnv
+    ) where {N}
+    size(ket) == size(bra) || throw(DimensionMismatch("incompatible bra and ket dimensions"))
+    size(ket, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported"))
+    static_inds = Val.(inds)
+    return _contract_densitymatrix(static_inds, ket, bra, env)
+end
 function reduced_densitymatrix(
         inds::NTuple{N, Tuple{Int, Int}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
     ) where {N}
     return reduced_densitymatrix(CartesianIndex.(inds), ket, bra, env)
 end
+function reduced_densitymatrix(
+        inds::NTuple{N, Tuple{Int, Int}}, ket::InfinitePEPO, bra::InfinitePEPO, env::CTMRGEnv
+    ) where {N}
+    return reduced_densitymatrix(CartesianIndex.(inds), ket, bra, env)
+end
 function reduced_densitymatrix(inds, ket::InfinitePEPS, env::CTMRGEnv)
     return reduced_densitymatrix(inds, ket, ket, env)
 end
@@ -461,3 +532,34 @@ end
         return ρ / str(ρ)
     end
 end
+
+@generated function _contract_densitymatrix(
+        inds::NTuple{N, Val}, ket::InfinitePEPO, bra::InfinitePEPO, env::CTMRGEnv
+    ) where {N}
+    cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) # weird hack to extract information from Val
+    allunique(cartesian_inds) ||
+        throw(ArgumentError("Indices should not overlap: $cartesian_inds."))
+    rowrange = getindex.(cartesian_inds, 1)
+    colrange = getindex.(cartesian_inds, 2)
+
+    corner_NW, corner_NE, corner_SE, corner_SW = _contract_corner_expr(rowrange, colrange)
+    edges_N, edges_E, edges_S, edges_W = _contract_edge_expr(rowrange, colrange)
+    result = tensorexpr(
+        :ρ,
+        ntuple(i -> physicallabel(:O, :ket, i), N),
+        ntuple(i -> physicallabel(:O, :bra, i), N),
+    )
+    bra, ket = _contract_pepo_state_expr(rowrange, colrange, cartesian_inds)
+
+    multiplication_ex = Expr(
+        :call, :*,
+        corner_NW, corner_NE, corner_SE, corner_SW,
+        edges_N..., edges_E..., edges_S..., edges_W...,
+        ket..., map(x -> Expr(:call, :conj, x), bra)...,
+    )
+    multex = :(@autoopt @tensor contractcheck = true $result := $multiplication_ex)
+    return quote
+        $(macroexpand(@__MODULE__, multex))
+        return ρ / str(ρ)
+    end
+end

src/algorithms/toolbox.jl

@@ -1,17 +1,25 @@
 """
-    expectation_value(peps::InfinitePEPS, O::LocalOperator, env::CTMRGEnv)
+    expectation_value(state, O::LocalOperator, env::CTMRGEnv)
+    expectation_value(bra, O::LocalOperator, ket, env::CTMRGEnv)
 
-Compute the expectation value ⟨peps|O|peps⟩ / ⟨peps|peps⟩ of a [`LocalOperator`](@ref) `O`
-for a PEPS `peps` using a given CTMRG environment `env`.
+Compute the expectation value ⟨bra|O|ket⟩ / ⟨bra|ket⟩ of a [`LocalOperator`](@ref) `O`.
+This can be done either for a PEPS, or alternatively for a density matrix PEPO.
+In the latter case the first signature corresponds to a single layer PEPO contraction, while
+the second signature yields a bilayer contraction instead.
 """
-function MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env::CTMRGEnv)
-    checklattice(peps, O)
+function MPSKit.expectation_value(
+        bra::Union{InfinitePEPS, InfinitePEPO}, O::LocalOperator,
+        ket::Union{InfinitePEPS, InfinitePEPO}, env::CTMRGEnv
+    )
+    checklattice(bra, O, ket)
     term_vals = dtmap([O.terms...]) do (inds, operator)  # OhMyThreads can't iterate over O.terms directly
-        ρ = reduced_densitymatrix(inds, peps, peps, env)
+        ρ = reduced_densitymatrix(inds, ket, bra, env)
         return trmul(operator, ρ)
     end
     return sum(term_vals)
 end
+MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env::CTMRGEnv) = expectation_value(peps, O, peps, env)
+
 """
     expectation_value(pf::InfinitePartitionFunction, inds => O, env::CTMRGEnv)
 

src/networks/tensors.jl

@@ -123,3 +123,23 @@ function physicalspace(t::PEPOTensor)
     return codomain_physicalspace(t)
 end
 virtualspace(t::PEPOTensor, dir) = space(t, dir + 2)
+
+"""
+    $(SIGNATURES)
+
+Fuse the physical indices of a PEPO tensor, obtaining a PEPS tensor.
+"""
+function fuse_physicalspaces(O::PEPOTensor)
+    F = isomorphism(Int, fuse(codomain(O)), codomain(O))
+    return F * O, F
+end
+
+"""
+    $(SIGNATURES)
+
+Trace out the physical indices of a PEPO tensor, obtaining a partition function tensor.
+"""
+function trace_physicalspaces(O::PEPOTensor)
+    @plansor t[W S; N E] := O[p p; N E S W]
+    return t
+end

src/operators/infinitepepo.jl

@@ -153,6 +153,7 @@ domain_physicalspace(T::InfinitePEPO, r::Int, c::Int) = domain_physicalspace(T[r
 function codomain_physicalspace(T::InfinitePEPO, r::Int, c::Int)
     return codomain_physicalspace(T[r, c, end])
 end
+physicalspace(T::InfinitePEPO) = physicalspace.(Ref(T), 1:size(T, 1), 1:size(T, 2))
 function physicalspace(T::InfinitePEPO, r::Int, c::Int)
     codomain_physicalspace(T, r, c) == domain_physicalspace(T, r, c) || throw(
         SpaceMismatch(
@@ -173,6 +174,23 @@ function InfiniteSquareNetwork(top::InfinitePEPS, mid::InfinitePEPO, bot::Infini
     )
 end
 
+## Conversion to states
+
+function InfinitePartitionFunction(ρ::InfinitePEPO)
+    size(ρ, 3) == 1 || throw(DimensionMismatch("Only single-layer density matrices can be converted to partition functions"))
+    return InfinitePartitionFunction(
+        trace_physicalspaces.(reshape(unitcell(ρ), size(ρ, 1), size(ρ, 2)))
+    )
+end
+
+function InfinitePEPS(ρ::InfinitePEPO)
+    size(ρ, 3) == 1 || throw(DimensionMismatch("Only single-layer density matrices can be converted to partition functions"))
+    return InfinitePEPS(
+        first.(fuse_physicalspaces.(reshape(unitcell(ρ), size(ρ, 1), size(ρ, 2))))
+    )
+end
+
+
 ## Vector interface
 
 VI.scalartype(::Type{NT}) where {NT <: InfinitePEPO} = scalartype(eltype(NT))

src/operators/localoperator.jl

@@ -74,6 +74,10 @@ while the second version throws an error if the lattices do not match.
 function checklattice(args...)
     return checklattice(Bool, args...) || throw(ArgumentError("Lattice mismatch."))
 end
+checklattice(::Type{Bool}, arg) = true
+function checklattice(::Type{Bool}, arg1, arg2, args...)
+    return checklattice(Bool, arg1, arg2) && checklattice(Bool, arg2, args...)
+end
 function checklattice(::Type{Bool}, H1::LocalOperator, H2::LocalOperator)
     return H1.lattice == H2.lattice
 end
@@ -83,6 +87,12 @@ end
 function checklattice(::Type{Bool}, H::LocalOperator, peps::InfinitePEPS)
     return checklattice(Bool, peps, H)
 end
+function checklattice(::Type{Bool}, pepo::InfinitePEPO, O::LocalOperator)
+    return size(pepo, 3) == 1 && reshape(physicalspace(pepo), size(pepo, 1), size(pepo, 2)) == physicalspace(O)
+end
+function checklattice(::Type{Bool}, O::LocalOperator, pepo::InfinitePEPO)
+    return checklattice(Bool, pepo, O)
+end
 @non_differentiable checklattice(args...)
 
 function Base.repeat(O::LocalOperator, m::Int, n::Int)

src/states/infinitepeps.jl

@@ -137,7 +137,9 @@ end
 ## Spaces
 
 TensorKit.spacetype(::Type{T}) where {T <: InfinitePEPS} = spacetype(eltype(T))
+virtualspace(n::InfinitePEPS, dir) = virtualspace.(unitcell(n), dir)
 virtualspace(n::InfinitePEPS, r::Int, c::Int, dir) = virtualspace(n[r, c], dir)
+physicalspace(n::InfinitePEPS) = physicalspace.(unitcell(n))
 physicalspace(n::InfinitePEPS, r::Int, c::Int) = physicalspace(n[r, c])
 
 ## InfiniteSquareNetwork interface

test/runtests.jl

@@ -65,6 +65,11 @@ end
             include("timeevol/sitedep_truncation.jl")
         end
     end
+    if GROUP == "ALL" || GROUP == "TOOLBOX"
+        @time @safetestset "Density matrix from double-layer PEPO" begin
+            include("toolbox/densitymatrices.jl")
+        end
+    end
     if GROUP == "ALL" || GROUP == "UTILITY"
         @time @safetestset "LocalOperator" begin
             include("utility/localoperator.jl")

test/toolbox/densitymatrices.jl

@@ -0,0 +1,41 @@
+using TensorKit
+using PEPSKit
+using Test
+using TestExtras
+
+ds = Dict(Trivial => ℂ^2, U1Irrep => U1Space(i => d for (i, d) in zip(-1:1, (1, 1, 2))), FermionParity => Vect[FermionParity](0 => 2, 1 => 1))
+Ds = Dict(Trivial => ℂ^3, U1Irrep => U1Space(i => D for (i, D) in zip(-1:1, (1, 2, 2))), FermionParity => Vect[FermionParity](0 => 3, 1 => 2))
+χs = Dict(Trivial => ℂ^4, U1Irrep => U1Space(i => χ for (i, χ) in zip(-2:2, (1, 3, 2))), FermionParity => Vect[FermionParity](0 => 3, 1 => 2))
+
+@testset "Double-layer densitymatrix contractions ($I)" for I in keys(ds)
+    d = ds[I]
+    D = Ds[I]
+    χ = χs[I]
+    ρ = InfinitePEPO(d, D)
+
+    ρ_peps = @constinferred InfinitePEPS(ρ)
+    env = CTMRGEnv(ρ_peps, χ)
+
+    O = rand(d, d)
+    F = isomorphism(fuse(d ⊗ d'), d ⊗ d')
+    @tensor O_doubled[-1; -2] := F[-1; 1 2] * O[1; 3] * twist(F', 2)[3 2; -2]
+
+    # Single site
+    O_singlesite = LocalOperator(reshape(physicalspace(ρ), 1, 1), ((1, 1),) => O)
+    E1 = expectation_value(ρ, O_singlesite, ρ, env)
+    O_doubled_singlesite = LocalOperator(physicalspace(ρ_peps), ((1, 1),) => O_doubled)
+    E2 = expectation_value(ρ_peps, O_doubled_singlesite, ρ_peps, env)
+    @test E1 ≈ E2
+
+    # two sites
+    for inds in zip(
+            [(1, 1), (1, 1), (1, 1), (1, 2), (1, 1)],
+            [(2, 1), (1, 2), (2, 2), (2, 1), (3, 1)]
+        )
+        O_twosite = LocalOperator(reshape(physicalspace(ρ), 1, 1), inds => O ⊗ O)
+        E1 = expectation_value(ρ, O_twosite, ρ, env)
+        O_doubled_twosite = LocalOperator(physicalspace(ρ_peps), inds => O_doubled ⊗ O_doubled)
+        E2 = expectation_value(ρ_peps, O_doubled_twosite, ρ_peps, env)
+        @test E1 ≈ E2
+    end
+end