Refactor projector steps for 2D and 3D HOTRG (#107)

* Refactor `_step_hotrg(3d)_x(y)`

* Fix test

Author: Yue-Zhengyuan

src/schemes/hotrg.jl

@@ -40,7 +40,7 @@ Hence both are written explicitly.
 =#
 
 function _get_hotrg_xproj(
-        A1::TensorMap{E, S, 2, 2}, A2::TensorMap{E, S, 2, 2},
+        A1::AbstractTensorMap{E, S, 2, 2}, A2::AbstractTensorMap{E, S, 2, 2},
         trunc::TensorKit.TruncationScheme
     ) where {E, S}
     #= join in y-direction, keep x-indices open (A1 below A2)
@@ -72,7 +72,7 @@ function _get_hotrg_xproj(
 end
 
 function _get_hotrg_yproj(
-        A1::TensorMap{E, S, 2, 2}, A2::TensorMap{E, S, 2, 2},
+        A1::AbstractTensorMap{E, S, 2, 2}, A2::AbstractTensorMap{E, S, 2, 2},
         trunc::TensorKit.TruncationScheme
     ) where {E, S}
     #= join in x-direction, keep y-indices open (A1 on the left of A2)
@@ -104,50 +104,50 @@ function _get_hotrg_yproj(
 end
 
 function _step_hotrg_y(
-        A1::TensorMap{E, S, 2, 2}, A2::TensorMap{E, S, 2, 2},
-        trunc::TensorKit.TruncationScheme
+        A1::AbstractTensorMap{E, S, 2, 2}, A2::AbstractTensorMap{E, S, 2, 2},
+        Ux::AbstractTensorMap{E, S, 2, 1}
     ) where {E, S}
     #= compression along the y-direction
                     -3
                     |
             ┌---1---A2---3--┐
             |       |       |
-        -1--U†      5       U-- -4
+        -1--Ux†      5      Ux-- -4
             |       |       |
             └---2---A1---4--┘
                     |
                     -2
     =#
-    U, = _get_hotrg_xproj(A1, A2, trunc)
     @tensor T[-1 -2; -3 -4] :=
-        conj(U[1 2; -1]) * U[3 4; -4] * A2[1 5; -3 3] * A1[2 -2; 5 4]
+        conj(Ux[1 2; -1]) * Ux[3 4; -4] * A2[1 5; -3 3] * A1[2 -2; 5 4]
     return T
 end
 
 function _step_hotrg_x(
-        A1::TensorMap{E, S, 2, 2}, A2::TensorMap{E, S, 2, 2},
-        trunc::TensorKit.TruncationScheme
+        A1::AbstractTensorMap{E, S, 2, 2}, A2::AbstractTensorMap{E, S, 2, 2},
+        Uy::AbstractTensorMap{E, S, 2, 1}
     ) where {E, S}
     #= compression along the x-direction
                 -3
                 |
-            ┌3--U--4┐
+            ┌3--Uy-4┐
             |       |
         -1--A1--5---A2-- -4
             |       |
-            └1--U†-2┘
+            └1-Uy†-2┘
                 |
                 -2
     =#
-    U, = _get_hotrg_yproj(A1, A2, trunc)
     @tensor T[-1 -2; -3 -4] :=
-        A1[-1 1; 3 5] * A2[5 2; 4 -4] * conj(U[1 2; -2]) * U[3 4; -3]
+        A1[-1 1; 3 5] * A2[5 2; 4 -4] * conj(Uy[1 2; -2]) * Uy[3 4; -3]
     return T
 end
 
 function step!(scheme::HOTRG, trunc::TensorKit.TruncationScheme)
-    scheme.T = _step_hotrg_y(scheme.T, scheme.T, trunc)
-    scheme.T = _step_hotrg_x(scheme.T, scheme.T, trunc)
+    Ux, = _get_hotrg_xproj(scheme.T, scheme.T, trunc)
+    scheme.T = _step_hotrg_y(scheme.T, scheme.T, Ux)
+    Uy, = _get_hotrg_yproj(scheme.T, scheme.T, trunc)
+    scheme.T = _step_hotrg_x(scheme.T, scheme.T, Uy)
     return scheme
 end
 

src/schemes/hotrg3d.jl

@@ -34,7 +34,7 @@ mutable struct HOTRG_3D <: TNRScheme
 end
 
 function _get_hotrg3d_xproj(
-        A1::TensorMap{E, S, 2, 4}, A2::TensorMap{E, S, 2, 4},
+        A1::AbstractTensorMap{E, S, 2, 4}, A2::AbstractTensorMap{E, S, 2, 4},
         trunc::TensorKit.TruncationScheme
     ) where {E, S}
     # join in z-direction, keep x-indices open (A1 below A2)
@@ -60,24 +60,30 @@ function _get_hotrg3d_xproj(
     return U, s, ε
 end
 
-function _step_hotrg3d(
-        A1::TensorMap{E, S, 2, 4}, A2::TensorMap{E, S, 2, 4},
+function _get_hotrg3d_yproj(
+        A1::AbstractTensorMap{E, S, 2, 4}, A2::AbstractTensorMap{E, S, 2, 4},
         trunc::TensorKit.TruncationScheme
     ) where {E, S}
-    Ux, = _get_hotrg3d_xproj(A1, A2, trunc)
-    # switch x/y directions
     perm = ((1, 2), (4, 3, 6, 5))
-    Uy, = _get_hotrg3d_xproj(permute(A1, perm), permute(A2, perm), trunc)
-    # apply the truncation
+    return _get_hotrg3d_xproj(permute(A1, perm), permute(A2, perm), trunc)
+end
+
+function _step_hotrg3d(
+        A1::AbstractTensorMap{E, S, 2, 4}, A2::AbstractTensorMap{E, S, 2, 4},
+        Ux::AbstractTensorMap{E, S, 2, 1}, Uy::AbstractTensorMap{E, S, 2, 1},
+    ) where {E, S}
     @tensoropt T[-1 -2; -3 -4 -5 -6] :=
         conj(Ux[x1 x2; -6]) * Ux[x1′ x2′; -4] *
         conj(Uy[y1 y2; -5]) * Uy[y1′ y2′; -3] *
         A1[-1 z; y1′ x1′ y1 x1] * A2[z -2; y2′ x2′ y2 x2]
     return T
 end
 
+# HOTRG step to compress along z direction
 function _step!(scheme::HOTRG_3D, trunc::TensorKit.TruncationScheme)
-    scheme.T = _step_hotrg3d(scheme.T, scheme.T, trunc)
+    Ux, = _get_hotrg3d_xproj(scheme.T, scheme.T, trunc)
+    Uy, = _get_hotrg3d_yproj(scheme.T, scheme.T, trunc)
+    scheme.T = _step_hotrg3d(scheme.T, scheme.T, Ux, Uy)
     return scheme
 end