Rewrite local expectation values as a single contraction, switch `opt…

…=true`
QuantumKitHub · Jun 19, 2024 · f5ee89c · f5ee89c
1 parent f242f9a
commit f5ee89c
Showing 1 changed file with 59 additions and 9 deletions.
diff --git a/src/operators/localoperator.jl b/src/operators/localoperator.jl
@@ -35,7 +35,7 @@ operator_env
 
 function operator_env(peps::InfinitePEPS, env::CTMRGEnv, ::OnSite)
     return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
-        @tensor ρ[-1; -2] :=
+        @tensor opt = true ρ[-1; -2] :=
             env.corners[NORTHWEST, r, c][1; 2] *
             env.edges[NORTH, r, c][2 3 4; 5] *
             env.corners[NORTHEAST, r, c][5; 6] *
@@ -52,7 +52,7 @@ end
 function operator_env(peps::InfinitePEPS, env::CTMRGEnv, ::NearestNeighbor)
     return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
         cnext = _next(c, size(peps, 2))
-        @tensor ρ[-11 -20; -12 -18] :=
+        @tensor opt = true ρ[-12 -18; -11 -20] :=
             env.corners[NORTHWEST, r, c][1; 3] *
             env.edges[NORTH, r, c][3 5 8; 13] *
             env.edges[NORTH, r, cnext][13 16 22; 23] *
@@ -78,20 +78,70 @@ of `peps` and `env`.
 """ MPSKit.expectation_value(::InfinitePEPS, ::Any, ::NLocalOperator)
 
 # 1-site operator expectation values on unit cell
-function MPSKit.expectation_value(peps::InfinitePEPS, env, O::NLocalOperator{OnSite})
-    return map(operator_env(peps, env, OnSite())) do ρ
-        o = @tensor ρ[-1; -2] * O.op[-1; -2]
-        n = @tensor ρ[-1; -1]
+function MPSKit.expectation_value(peps::InfinitePEPS, env::CTMRGEnv, O::NLocalOperator{OnSite})
+    return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
+        o = @tensor opt = true  env.corners[NORTHWEST, r, c][1; 2] *
+                                env.edges[NORTH, r, c][2 3 4; 5] *
+                                env.corners[NORTHEAST, r, c][5; 6] *
+                                env.edges[EAST, r, c][6 7 8; 9] *
+                                env.corners[SOUTHEAST, r, c][9; 10] *
+                                env.edges[SOUTH, r, c][10 11 12; 13] *
+                                env.corners[SOUTHWEST, r, c][13; 14] *
+                                env.edges[WEST, r, c][14 15 16; 1] *
+                                peps[r, c][17; 3 7 11 15] *
+                                conj(peps[r, c][18; 4 8 12 16]) * 
+                                O.op[18; 17]
+
+        n = @tensor opt = true  env.corners[NORTHWEST, r, c][1; 2] *
+                                env.edges[NORTH, r, c][2 3 4; 5] *
+                                env.corners[NORTHEAST, r, c][5; 6] *
+                                env.edges[EAST, r, c][6 7 8; 9] *
+                                env.corners[SOUTHEAST, r, c][9; 10] *
+                                env.edges[SOUTH, r, c][10 11 12; 13] *
+                                env.corners[SOUTHWEST, r, c][13; 14] *
+                                env.edges[WEST, r, c][14 15 16; 1] *
+                                peps[r, c][17; 3 7 11 15] *
+                                conj(peps[r, c][17; 4 8 12 16])
+
         o / n
     end
 end
 
 function MPSKit.expectation_value(
     peps::InfinitePEPS, env, O::NLocalOperator{NearestNeighbor}
 )
-    return map(operator_env(peps, env, NearestNeighbor())) do ρ
-        o = @tensor ρ[1 2; 3 4] * O.op[1 2; 3 4]
-        n = @tensor ρ[1 2; 1 2]
+    return map(Iterators.product(axes(env.corners, 2), axes(env.corners, 3))) do (r, c)
+        cnext = _next(c, size(peps, 2))
+        o = @tensor opt = true  env.corners[NORTHWEST, r, c][1; 3] *
+                                env.edges[NORTH, r, c][3 5 8; 13] *
+                                env.edges[NORTH, r, cnext][13 16 22; 23] *
+                                env.corners[NORTHEAST, r, cnext][23; 24] *
+                                env.edges[EAST, r, cnext][24 25 26; 27] *
+                                env.corners[SOUTHEAST, r, cnext][27; 28] *
+                                env.edges[SOUTH, r, cnext][28 17 21; 14] *
+                                env.edges[SOUTH, r, c][14 6 10; 4] *
+                                env.corners[SOUTHWEST, r, c][4; 2] *
+                                env.edges[WEST, r, c][2 7 9; 1] *
+                                peps[r, c][29; 5 15 6 7] *
+                                conj(peps[r, c][31; 8 19 10 9]) *
+                                peps[r, cnext][30; 16 25 17 15] *
+                                conj(peps[r, cnext][32; 22 26 21 19]) * 
+                                O.op[31 32; 29 30]
+
+        n = @tensor opt = true  env.corners[NORTHWEST, r, c][1; 3] *
+                                env.edges[NORTH, r, c][3 5 8; 13] *
+                                env.edges[NORTH, r, cnext][13 16 22; 23] *
+                                env.corners[NORTHEAST, r, cnext][23; 24] *
+                                env.edges[EAST, r, cnext][24 25 26; 27] *
+                                env.corners[SOUTHEAST, r, cnext][27; 28] *
+                                env.edges[SOUTH, r, cnext][28 17 21; 14] *
+                                env.edges[SOUTH, r, c][14 6 10; 4] *
+                                env.corners[SOUTHWEST, r, c][4; 2] *
+                                env.edges[WEST, r, c][2 7 9; 1] *
+                                peps[r, c][29; 5 15 6 7] *
+                                conj(peps[r, c][29; 8 19 10 9]) *
+                                peps[r, cnext][30; 16 25 17 15] *
+                                conj(peps[r, cnext][30; 22 26 21 19])
         o / n
     end
 end