Reduced density matrices (#230)

* Switch expectation_value to reduced densitymatrix

* Minor improvement

* Use AD compatible normalization

* FFF: Fix the f'ing fermions

* Fix bug in contraction expressions

* Special case 1x1 density matrix

* Special case 2x1 density matrix

* Safeguard fermionic non-canonical arrows

* Special case 1x2 density matrix

* Add utility syntax bra=ket

* Export `reduced_densitymatrix`

Author: lkdvos

src/PEPSKit.jl

@@ -90,8 +90,9 @@ export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
 export FixedSpaceTruncation, SiteDependentTruncation
 export HalfInfiniteProjector, FullInfiniteProjector
 export LocalOperator, physicalspace
-export expectation_value, cost_function, product_peps, correlation_length, network_value
-export correlator
+export product_peps
+export reduced_densitymatrix, expectation_value, network_value, cost_function
+export correlator, correlation_length
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver, EigSolver
 export fixedpoint

src/algorithms/contractions/localoperator.jl

@@ -144,7 +144,7 @@ function _contract_state_expr(rowrange, colrange, cartesian_inds=nothing)
                 physicallabel(:O, side, inds_id)
             end
             return tensorexpr(
-                :(bra[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)]),
+                :($(side)[mod1($(rmin + i - 1), end), mod1($(cmin + j - 1), end)]),
                 (physical_label,),
                 (
                     if i == 1
@@ -212,7 +212,7 @@ end
     )
 
     returnex = quote
-        @autoopt @tensor opt = $multiplication_ex
+        @autoopt @tensor $multiplication_ex
     end
     return macroexpand(@__MODULE__, returnex)
 end
@@ -266,7 +266,226 @@ end
     )
 
     returnex = quote
-        @autoopt @tensor opt = $multiplication_ex
+        @autoopt @tensor $multiplication_ex
     end
     return macroexpand(@__MODULE__, returnex)
 end
+
+"""
+$(SIGNATURES)
+
+Construct the reduced density matrix `ρ` of the PEPS `peps` 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`. 
+"""
+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,Tuple{Int,Int}}, ket::InfinitePEPS, bra::InfinitePEPS, 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
+
+# Special case 1x1 density matrix:
+# Keep contraction order but try to optimize intermediate permutations:
+# EE_SWA is largest object so keep largest legs to the front there
+function reduced_densitymatrix(
+    inds::Tuple{CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
+)
+    row, col = Tuple(inds[1])
+
+    # Unpack variables and absorb corners
+    A = ket[mod1(row, end), mod1(col, end)]
+    Ā = bra[mod1(row, end), mod1(col, end)]
+
+    E_north =
+        env.edges[NORTH, mod1(row - 1, end), mod1(col, end)] *
+        twistdual(env.corners[NORTHEAST, mod1(row - 1, end), mod1(col + 1, end)], 1)
+    E_east =
+        env.edges[EAST, mod1(row, end), mod1(col + 1, end)] *
+        twistdual(env.corners[SOUTHEAST, mod1(row + 1, end), mod1(col + 1, end)], 1)
+    E_south =
+        env.edges[SOUTH, mod1(row + 1, end), mod1(col, end)] *
+        twistdual(env.corners[SOUTHWEST, mod1(row + 1, end), mod1(col - 1, end)], 1)
+    E_west =
+        env.edges[WEST, mod1(row, end), mod1(col - 1, end)] *
+        twistdual(env.corners[NORTHWEST, mod1(row - 1, end), mod1(col - 1, end)], 1)
+
+    @tensor EE_SW[χSE χNW DSb DWb; DSt DWt] :=
+        E_south[χSE DSt DSb; χSW] * E_west[χSW DWt DWb; χNW]
+
+    @tensor EE_SWA[χSE χNW DNt DEt; dt DSb DWb] :=
+        EE_SW[χSE χNW DSb DWb; DSt DWt] * A[dt; DNt DEt DSt DWt]
+
+    @tensor EE_NE[DNb DEb; χSE χNW DNt DEt] :=
+        E_north[χNW DNt DNb; χNE] * E_east[χNE DEt DEb; χSE]
+
+    @tensor EEAEE[dt; DNb DEb DSb DWb] :=
+        EE_NE[DNb DEb; χSE χNW DNt DEt] * EE_SWA[χSE χNW DNt DEt; dt DSb DWb]
+
+    @tensor ρ[dt; db] := EEAEE[dt; DNb DEb DSb DWb] * conj(Ā[db; DNb DEb DSb DWb])
+
+    return ρ / str(ρ)
+end
+
+function reduced_densitymatrix(
+    inds::NTuple{2,CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
+)
+    if inds[2] - inds[1] == CartesianIndex(1, 0)
+        return reduced_densitymatrix2x1(inds[1], ket, bra, env)
+    elseif inds[2] - inds[1] == CartesianIndex(0, 1)
+        return reduced_densitymatrix1x2(inds[1], ket, bra, env)
+    else
+        static_inds = Val.(inds)
+        return _contract_densitymatrix(static_inds, ket, bra, env)
+    end
+end
+
+# Special case 2x1 density matrix:
+# Keep contraction order but try to optimize intermediate permutations:
+function reduced_densitymatrix2x1(
+    ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
+)
+    row, col = Tuple(ind)
+
+    # Unpack variables and absorb corners
+    A_north = ket[mod1(row, end), mod1(col, end)]
+    Ā_north = bra[mod1(row, end), mod1(col, end)]
+    A_south = ket[mod1(row + 1, end), mod1(col, end)]
+    Ā_south = bra[mod1(row + 1, end), mod1(col, end)]
+
+    E_north =
+        env.edges[NORTH, mod1(row - 1, end), mod1(col, end)] *
+        twistdual(env.corners[NORTHEAST, mod1(row - 1, end), mod1(col + 1, end)], 1)
+    E_northeast = env.edges[EAST, mod1(row, end), mod1(col + 1, end)]
+    E_southeast =
+        env.edges[EAST, mod1(row + 1, end), mod1(col + 1, end)] *
+        twistdual(env.corners[SOUTHEAST, mod1(row + 2, end), mod1(col + 1, end)], 1)
+    E_south =
+        env.edges[SOUTH, mod1(row + 2, end), mod1(col, end)] *
+        twistdual(env.corners[SOUTHWEST, mod1(row + 2, end), mod1(col - 1, end)], 1)
+    E_southwest = env.edges[WEST, mod1(row + 1, end), mod1(col - 1, end)]
+    E_northwest =
+        env.edges[WEST, mod1(row, end), mod1(col - 1, end)] *
+        twistdual(env.corners[NORTHWEST, mod1(row - 1, end), mod1(col - 1, end)], 1)
+
+    @tensor EE_NW[χW χNE DNWt DNt; DNWb DNb] :=
+        E_northwest[χW DNWt DNWb; χNW] * E_north[χNW DNt DNb; χNE]
+    @tensor EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] :=
+        EE_NW[χW χNE DNWt DNt; DNWb DNb] * conj(Ā_north[dNb; DNb DNEb DMb DNWb])
+    @tensor EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] :=
+        EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] * A_north[dNt; DNt DNEt DMt DNWt]
+    @tensor EEEAA_N[dNt dNb; χW DMt DMb χE] :=
+        EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] * E_northeast[χNE DNEt DNEb; χE]
+
+    @tensor EE_SE[χE χSW DSEt DSt; DSEb DSb] :=
+        E_southeast[χE DSEt DSEb; χSE] * E_south[χSE DSt DSb; χSW]
+    @tensor EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] :=
+        EE_SE[χE χSW DSEt DSt; DSEb DSb] * conj(Ā_south[dSb; DMb DSEb DSb DSWb])
+    @tensor EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] :=
+        EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] * A_south[dSt; DMt DSEt DSt DSWt]
+    @tensor EEEAA_S[χW DMt DMb χE; dSt dSb] :=
+        EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] * E_southwest[χSW DSWt DSWb; χW]
+
+    @tensor ρ[dNt dSt; dNb dSb] :=
+        EEEAA_N[dNt dNb; χW DMt DMb χE] * EEEAA_S[χW DMt DMb χE; dSt dSb]
+
+    return ρ / str(ρ)
+end
+
+function reduced_densitymatrix1x2(
+    ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv
+)
+    row, col = Tuple(ind)
+
+    # Unpack variables and absorb corners
+    A_west = ket[mod1(row, end), mod1(col, end)]
+    Ā_west = bra[mod1(row, end), mod1(col, end)]
+    A_east = ket[mod1(row, end), mod1(col + 1, end)]
+    Ā_east = bra[mod1(row, end), mod1(col + 1, end)]
+
+    E_northwest = env.edges[NORTH, mod1(row - 1, end), mod1(col, end)]
+    E_northeast =
+        env.edges[NORTH, mod1(row - 1, end), mod1(col + 1, end)] *
+        twistdual(env.corners[NORTHEAST, mod1(row - 1, end), mod1(col + 2, end)], 1)
+    E_east =
+        env.edges[EAST, mod1(row, end), mod1(col + 2, end)] *
+        twistdual(env.corners[SOUTHEAST, mod1(row + 1, end), mod1(col + 2, end)], 1)
+    E_southeast = env.edges[SOUTH, mod1(row + 1, end), mod1(col + 1, end)]
+    E_southwest =
+        env.edges[SOUTH, mod1(row + 1, end), mod1(col, end)] *
+        twistdual(env.corners[SOUTHWEST, mod1(row + 1, end), mod1(col - 1, end)], 1)
+    E_west =
+        env.edges[WEST, mod1(row, end), mod1(col - 1, end)] *
+        twistdual(env.corners[NORTHWEST, mod1(row - 1, end), mod1(col - 1, end)], 1)
+
+    @tensor EE_SW[χS χNW DSWt DWt; DSWb DWb] :=
+        E_southwest[χS DSWt DSWb; χSW] * E_west[χSW DWt DWb; χNW]
+    @tensor EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] :=
+        EE_SW[χS χNW DSWt DWt; DSWb DWb] * conj(Ā_west[dWb; DNWb DMb DSWb DWb])
+    @tensor EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] :=
+        EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] * A_west[dWt; DNWt DMt DSWt DWt]
+    @tensor EEEAA_W[dWt dWb; χS DMt DMb χN] :=
+        EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] * E_northwest[χNW DNWt DNWb; χN]
+
+    @tensor EE_NE[χN χSE DNEt DEt; DNEb DEb] :=
+        E_northeast[χN DNEt DNEb; χNE] * E_east[χNE DEt DEb; χSE]
+    @tensor EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] :=
+        EE_NE[χN χSE DNEt DEt; DNEb DEb] * conj(Ā_east[dEb; DNEb DEb DSEb DMb])
+    @tensor EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] :=
+        EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] * A_east[dEt; DNEt DEt DSEt DMt]
+    @tensor EEEAA_E[χS DMt DMb χN; dEt dEb] :=
+        EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] * E_southeast[χSE DSEt DSEb; χS]
+
+    @tensor ρ[dWt dEt; dWb dEb] :=
+        EEEAA_W[dWt dWb; χS DMt DMb χN] * EEEAA_E[χS DMt DMb χN; dEt dEb]
+
+    return ρ / str(ρ)
+end
+
+@generated function _contract_densitymatrix(
+    inds::NTuple{N,Val}, ket::InfinitePEPS, bra::InfinitePEPS, 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_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 $result := $multiplication_ex)
+    return quote
+        $(macroexpand(@__MODULE__, multex))
+        return ρ / str(ρ)
+    end
+end

src/algorithms/toolbox.jl

@@ -7,8 +7,8 @@ for a PEPS `peps` using a given CTMRG environment `env`.
 function MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env::CTMRGEnv)
     checklattice(peps, O)
     term_vals = dtmap([O.terms...]) do (inds, operator)  # OhMyThreads can't iterate over O.terms directly
-        contract_local_operator(inds, operator, peps, peps, env) /
-        contract_local_norm(inds, peps, peps, env)
+        ρ = reduced_densitymatrix(inds, peps, peps, env)
+        return trmul(operator, ρ)
     end
     return sum(term_vals)
 end

src/utility/util.jl

@@ -131,6 +131,31 @@ function twistdual!(t::AbstractTensorMap, is)
 end
 twistdual(t::AbstractTensorMap, is) = twistdual!(copy(t), is)
 
+"""
+    str(t)
+
+Fermionic supertrace by using `@tensor`.
+"""
+str(t::AbstractTensorMap) = _str(BraidingStyle(sectortype(t)), t)
+_str(::Bosonic, t::AbstractTensorMap) = tr(t)
+@generated function _str(::Fermionic, t::AbstractTensorMap{<:Any,<:Any,N,N}) where {N}
+    tex = tensorexpr(:t, ntuple(identity, N), ntuple(identity, N))
+    return macroexpand(@__MODULE__, :(@tensor $tex))
+end
+
+"""
+    trmul(H, ρ)
+
+Compute `tr(H * ρ)` without forming `H * ρ`.
+"""
+@generated function trmul(
+    H::AbstractTensorMap{<:Any,S,N,N}, ρ::AbstractTensorMap{<:Any,S,N,N}
+) where {S,N}
+    Hex = tensorexpr(:H, ntuple(identity, N), ntuple(i -> i + N, N))
+    ρex = tensorexpr(:ρ, ntuple(i -> i + N, N), ntuple(identity, N))
+    return macroexpand(@__MODULE__, :(@tensor $Hex * $ρex))
+end
+
 # Check whether diagonals contain degenerate values up to absolute or relative tolerance
 function is_degenerate_spectrum(
     S; atol::Real=0, rtol::Real=atol > 0 ? 0 : sqrt(eps(scalartype(S)))