generalize observable evaluation

Based on PR QuantumKitHub#117 of PEPSKit, this generalizes the expectation value for InfinitePEPOs and InfinitePartitionFunctions

Author: sanderdemeyer

src/PEPSKit.jl

@@ -32,9 +32,9 @@ include("states/infinitepeps.jl")
 include("states/infiniteweightpeps.jl")
 include("states/infinitepartitionfunction.jl")
 
+include("operators/localoperator.jl")
 include("operators/infinitepepo.jl")
 include("operators/transfermatrix.jl")
-include("operators/localoperator.jl")
 include("operators/lattices/squarelattice.jl")
 include("operators/models.jl")
 

src/algorithms/contractions/localoperator.jl

@@ -126,6 +126,42 @@ function _contract_edge_expr(rowrange, colrange)
     return edges_N, edges_E, edges_S, edges_W
 end
 
+function _contract_pf_edge_expr(rowrange, colrange)
+    rmin, rmax = extrema(rowrange)
+    cmin, cmax = extrema(colrange)
+    gridsize = (rmax - rmin + 1, cmax - cmin + 1)
+
+    edges_N = map(1:gridsize[2]) do i
+        E_N = :(env.edges[NORTH, mod1($(rmin - 1), end), mod1($(cmin + i - 1), end)])
+        return tensorexpr(
+            E_N, (envlabel(NORTH, i - 1), virtuallabel(NORTH, i)), envlabel(NORTH, i)
+        )
+    end
+
+    edges_E = map(1:gridsize[1]) do i
+        E_E = :(env.edges[EAST, mod1($(rmin + i - 1), end), mod1($(cmax + 1), end)])
+        return tensorexpr(
+            E_E, (envlabel(EAST, i - 1), virtuallabel(EAST, i)), envlabel(EAST, i)
+        )
+    end
+
+    edges_S = map(1:gridsize[2]) do i
+        E_S = :(env.edges[SOUTH, mod1($(rmax + 1), end), mod1($(cmin + i - 1), end)])
+        return tensorexpr(
+            E_S, (envlabel(SOUTH, i), virtuallabel(SOUTH, i)), envlabel(SOUTH, i - 1)
+        )
+    end
+
+    edges_W = map(1:gridsize[1]) do i
+        E_W = :(env.edges[WEST, mod1($(rmin + i - 1), end), mod1($(cmin - 1), end)])
+        return tensorexpr(
+            E_W, (envlabel(WEST, i), virtuallabel(WEST, i)), envlabel(WEST, i - 1)
+        )
+    end
+
+    return edges_N, edges_E, edges_S, edges_W
+end
+
 function _contract_state_expr(rowrange, colrange, cartesian_inds=nothing)
     rmin, rmax = extrema(rowrange)
     cmin, cmax = extrema(colrange)
@@ -173,6 +209,96 @@ function _contract_state_expr(rowrange, colrange, cartesian_inds=nothing)
     end
 end
 
+function pf_tensor_expr(tensor_label, label_west, label_south, label_north, label_east, (len,place))
+    if len == 1
+        a = tensorexpr(
+            tensor_label,
+            ((label_west),(label_south),),
+            ((label_north),(label_east),),
+        )
+    else
+        if place == 1
+            a = tensorexpr(
+                tensor_label,
+                ((label_west), (label_south),),
+                ((label_north), (label_east), (virtuallabel(:op, len, place)),),
+            )
+        elseif place == len
+            a = tensorexpr(
+                tensor_label,
+                ((virtuallabel(:op, len, place - 1)), (label_west), (label_south),),
+                ((label_north), (label_east),),
+            )
+        else
+            a = tensorexpr(
+                tensor_label,
+                ((virtuallabel(:op, len, place - 1)), (label_west), (label_south),),
+                ((label_north), (label_east), (virtuallabel(:op, len, place)),),
+            )
+        end
+    end
+    return a
+end
+
+function _contract_tensor_expr(O, pos, rowrange, colrange, cartesian_inds=nothing)
+    rmin, rmax = extrema(rowrange)
+    cmin, cmax = extrema(colrange)
+    gridsize = (rmax - rmin + 1, cmax - cmin + 1)
+    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
+        tensor_label = if isnothing(inds_id) || O <: Nothing
+            :(pf[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)])
+        else
+            :(O[$inds_id])
+        end
+
+        label_west = if j == 1
+            virtuallabel(WEST, i)
+        else
+            virtuallabel(:horizontal, i, j - 1)
+        end
+        label_south = if i == gridsize[1]
+            virtuallabel(SOUTH, j)
+        else
+            virtuallabel(:vertical, i, j)
+        end
+        label_north = if i == 1
+            virtuallabel(NORTH, j)
+        else
+            virtuallabel(:vertical, i - 1, j)
+        end
+        label_east = if j == gridsize[2]
+            virtuallabel(EAST, i)
+        else
+            virtuallabel(:horizontal, i, j)
+        end
+        if isnothing(inds_id) || O <: Nothing
+            pf_tensor_expr(
+                tensor_label,
+                label_west,
+                label_south,
+                label_north,
+                label_east,
+                (1,1),
+            )
+
+        else
+            pf_tensor_expr(
+                tensor_label,
+                label_west,
+                label_south,
+                label_north,
+                label_east,
+                pos[inds_id],
+            )
+        end
+    end
+end
+
 @generated function _contract_local_operator(
     inds::NTuple{N,Val},
     O::AbstractTensorMap{T,S,N,N},
@@ -270,3 +396,70 @@ end
     end
     return macroexpand(@__MODULE__, returnex)
 end
+
+# Partition function contractions
+
+"""
+    contract_local_tensors(inds, [O], pf, env)
+
+Contract a local tensor `O` inserted into a partition function `pf` at position `inds`,
+using the environment `env`.
+"""
+function contract_local_tensors(
+    inds::NTuple{N,CartesianIndex{2}},
+    O::Union{Nothing,NTuple{N,AbstractTensorMap{T,S}}},
+    pos::Union{Nothing,NTuple{N,Tuple{Int,Int}}},
+    pf::InfinitePartitionFunction,
+    env::CTMRGEnv{C,<:CTMRG_PF_EdgeTensor},
+) where {N,T,S,C}
+    static_inds = Val.(inds)
+    static_pos = Val.(pos)
+    return _contract_local_tensors(static_inds, O, static_pos, pf, env)
+end
+function contract_local_tensors(
+    inds::NTuple{N,Tuple{Int,Int}},
+    O::Union{Nothing,NTuple{N,AbstractTensorMap{T,S}}},
+    pos::Union{Nothing,NTuple{N,Tuple{Int,Int}}},
+    pf::InfinitePartitionFunction,
+    env::CTMRGEnv{C,<:CTMRG_PF_EdgeTensor},
+) where {N,T,S,C}
+    return contract_local_tensors(CartesianIndex.(inds), O, pos, pf, env)
+end
+@generated function _contract_local_tensors(
+    inds::NTuple{N,Val},
+    O::Union{Nothing,NTuple{N,AbstractTensorMap{T,S}}},
+    pos::Union{Nothing,NTuple{N,Val}},
+    pf::InfinitePartitionFunction,
+    env::CTMRGEnv{C,<:CTMRG_PF_EdgeTensor},
+) where {N,T,S,C}
+    # weird hack to extract information from Val
+    cartesian_inds = collect(CartesianIndex{2}, map(x -> x.parameters[1], inds.parameters)) 
+    cartesian_pos = collect(Tuple{Int,Int}, map(x -> x.parameters[1], pos.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_pf_edge_expr(rowrange, colrange)
+    tensors = _contract_tensor_expr(O, Tuple(cartesian_pos), rowrange, colrange, cartesian_inds)
+
+    multiplication_ex = Expr(
+        :call,
+        :*,
+        corner_NW,
+        corner_NE,
+        corner_SE,
+        corner_SW,
+        edges_N...,
+        edges_E...,
+        edges_S...,
+        edges_W...,
+        tensors...,
+    )
+
+    returnex = quote
+        @autoopt @tensor opt = $multiplication_ex
+    end
+    return macroexpand(@__MODULE__, returnex)
+end

src/algorithms/toolbox.jl

@@ -25,16 +25,24 @@ convention.
 """
 function MPSKit.expectation_value(
     pf::InfinitePartitionFunction,
-    op::Pair{CartesianIndex{2},<:AbstractTensorMap{T,S,2,2}},
-    env::CTMRGEnv,
-) where {T,S}
-    return contract_local_tensor(op[1], op[2], env) /
-           contract_local_tensor(op[1], pf[op[1]], env)
-end
-function MPSKit.expectation_value(
-    pf::InfinitePartitionFunction, op::Pair{Tuple{Int,Int}}, env::CTMRGEnv
-)
-    return expectation_value(pf, CartesianIndex(op[1]) => op[2], env)
+    op::NTuple{N,Array{<:Pair{CartesianIndex{2},<:AbstractTensorMap{T,S}}}},
+    envs::CTMRGEnv,
+) where {N,T,S}
+    inds = CartesianIndex{2}[]
+    ops = AbstractTensorMap{T,S}[]
+    positions = Tuple{Int,Int}[]
+    for (n, term) = enumerate(op)
+        for (i,(ind, op)) in enumerate(term)
+            push!(inds, ind)
+            push!(ops, op)
+            push!(positions, (length(term), i))
+        end
+    end
+    # inds = first.(op)
+    # ops = last.(op)
+    pos = Tuple(positions)
+    return contract_local_tensors(Tuple(inds), Tuple(ops), pos, pf, envs) /
+           contract_local_tensors(Tuple(inds), nothing, pos, pf, envs)
 end
 
 """

src/operators/infinitepepo.jl

@@ -202,6 +202,55 @@ function Base.adjoint(O::InfinitePEPO)
     return InfinitePEPO(_dag.(unitcell(O)))
 end
 
+function trace_out(O::InfinitePEPO)
+    function _trace_out(O::PEPOTensor)
+        @tensor O_tr[-4 -3; -1 -2] := O[1 1; -1 -2 -3 -4]
+        return O_tr
+    end
+    return InfinitePartitionFunction(_trace_out.(unitcell(O))[:,:,1])
+end
+
+function MPSKit.expectation_value(O::InfinitePEPO, H::InfinitePEPO, χ, ctm_alg)
+    network_OH = InfiniteSquareNetwork(_stack_pepos((PEPSKit.unitcell(O), PEPSKit.unitcell(H))))
+    env0_OH = CTMRGEnv(network_OH, χ)
+    env_OH, = leading_boundary(env0_OH, network_OH, ctm_alg)
+
+    network_OO = InfiniteSquareNetwork(_stack_pepos((PEPSKit.unitcell(O), PEPSKit.unitcell(O))))
+    env0_OO = CTMRGEnv(network_OO, χ)
+    env_OO, = leading_boundary(env0_OO, network_OO, ctm_alg)
+    return network_value(network_OH, env_OH) / network_value(network_OO, env_OO)
+end
+
+function _expectation_value(O::InfinitePEPO, gate::Pair{NTuple{1, CartesianIndex{2}}, T}, χ, ctm_alg) where {T<:AbstractTensorMap}
+    (Nr, Nc) = size(O)
+    site = (mod1(gate[1][1][1], Nr), mod1(gate[1][1][2], Nc), 1)
+
+    @tensor t[-4 -3; -1 -2] := O[site...][1 2; -1 -2 -3 -4] * gate[2][2; 1]
+    network = trace_out(O)
+    env0 = CTMRGEnv(network, χ)
+    env, = leading_boundary(env0, network, ctm_alg)
+    return expectation_value(network, ([gate[1][1] => t],), env)
+end
+
+function _expectation_value(O::InfinitePEPO, gate::Pair{NTuple{2, CartesianIndex{2}}, T}, χ, ctm_alg) where {T<:AbstractTensorMap}
+    (Nr, Nc) = size(O)
+    left = (mod1(gate[1][1][1], Nr), mod1(gate[1][1][2], Nc), 1)
+    right = (mod1(gate[1][2][1], Nr), mod1(gate[1][2][2], Nc), 1)
+    @tensor t[-1 -2 -3 -4; -5 -6 -7 -8] := O[left...][1 2; -1 -2 -3 -4] * O[right...][3 4; -5 -6 -7 -8] * gate[2][2 4; 1 3]
+    U, S, V = tsvd(t)
+    L = permute(U * sqrt(S), ((4,3),(1,2,5)))
+    R = permute(sqrt(S) * V, ((1,5,4),(2,3)))
+
+    network = trace_out(O)
+    env0 = CTMRGEnv(network, χ)
+    env, = leading_boundary(env0, network, ctm_alg)
+    return expectation_value(network, ([gate[1][1] => L, gate[1][2] => R],), env)
+end
+
+function MPSKit.expectation_value(O::InfinitePEPO, H::LocalOperator, χ, ctm_alg)
+    return sum([_expectation_value(O, t, χ, ctm_alg) for t = H.terms])
+end
+
 ## Vector interface
 
 function VectorInterface.scalartype(::Type{NT}) where {NT<:InfinitePEPO}
@@ -270,3 +319,14 @@ function _stack_tuples(A::Matrix{NTuple{N,T}}) where {N,T}
     end
     return out
 end
+
+function _stack_pepos(A::NTuple{N,Array{T}}) where {N,T}
+    size_out = size(A[1])
+    out = Array{T}(undef, size_out[1], size_out[2], N)
+    for h = 1:N
+        for (r, c) in Iterators.product(axes(A[1])...)
+            out[r, c, h] = A[h][r, c]
+        end    
+    end
+    return InfinitePEPO(out)
+end

test/ctmrg/partition_function.jl

@@ -101,8 +101,8 @@ projector_algs = [:halfinfinite, :fullinfinite]
 
     # check observables
     λ = network_value(Z, env)
-    m = expectation_value(Z, (1, 1) => M, env)
-    e = expectation_value(Z, (1, 1) => E, env)
+    m = expectation_value(Z, ((1, 1) => M,), env)
+    e = expectation_value(Z, ((1, 1) => E,), env)
     f_exact, m_exact, e_exact = classical_ising_exact(; beta)
 
     # should be real-ish
@@ -115,3 +115,43 @@ projector_algs = [:halfinfinite, :fullinfinite]
     @test abs(m) ≈ abs(m_exact) rtol = 1e-4
     @test e ≈ e_exact rtol = 1e-1 # accuracy limited by bond dimension and maxiter
 end
+
+@testset "Classical Ising correlation functions" begin
+    ctm_alg = SimultaneousCTMRG(; maxiter=300)
+    β = [0.1, log(1 + sqrt(2)) / 2 + 0.05, 2.0]
+
+    # contract at high, critical and low temperature
+    O_high, M_high, = classical_ising(; beta=β[1])
+    Z_high = InfinitePartitionFunction(O_high)
+    env0_high = CTMRGEnv(Z_high, χenv)
+    env_high = leading_boundary(env0_high, Z_high, ctm_alg)
+
+    O_crit, M_crit, = classical_ising(; beta=β[2])
+    Z_crit = InfinitePartitionFunction(O_crit)
+    env0_crit = CTMRGEnv(Z_crit, χenv)
+    env_crit = leading_boundary(env0_crit, Z_crit, ctm_alg)
+
+    O_low, M_low, = classical_ising(; beta=β[3])
+    Z_low = InfinitePartitionFunction(O_low)
+    env0_low = CTMRGEnv(Z_low, χenv)
+    env_low = leading_boundary(env0_low, Z_low, ctm_alg)
+
+    # compute correlators
+    corr_zz_high = expectation_value(Z_high, ((1, 1) => M_high, (2, 1) => M_high), env_high)
+    corr_zz_crit = expectation_value(Z_crit, ((1, 1) => M_crit, (2, 1) => M_crit), env_crit)
+    corr_zz_low = expectation_value(Z_low, ((1, 1) => M_low, (2, 1) => M_low), env_low)
+    @test abs(corr_zz_high) < abs(corr_zz_crit) < abs(corr_zz_low)
+    @test abs(corr_zz_low) ≈ 1.0 rtol = 1e-6
+
+    corr_zzz_high = expectation_value(
+        Z_high, ((1, 1) => M_high, (2, 1) => M_high, (1, 2) => M_high), env_high
+    )
+    corr_zzz_crit = expectation_value(
+        Z_crit, ((1, 1) => M_crit, (2, 1) => M_crit, (1, 2) => M_crit), env_crit
+    )
+    corr_zzz_low = expectation_value(
+        Z_low, ((1, 1) => M_low, (2, 1) => M_low, (1, 2) => M_low), env_low
+    )
+    @test abs(corr_zzz_high) ≈ 0.0 atol = 1e-6
+    @test abs(corr_zzz_low) ≈ 1.0 rtol = 1e-6
+end