Add docstrings, clean up J1-J2 test

src/PEPSKit.jl

@@ -85,7 +85,7 @@ export fixedpoint
 export InfinitePEPS, InfiniteTransferPEPS
 export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS
-export symmetrize!, ReflectDepth, ReflectWidth, RotateReflect, symmetrize_callback
+export ReflectDepth, ReflectWidth, RotateReflect, symmetrize!, symmetrize_callback
 export showtypeofgrad
 export square_lattice_tf_ising, square_lattice_heisenberg, square_lattice_j1j2
 export square_lattice_pwave

src/utility/symmetrization.jl

@@ -1,11 +1,32 @@
 abstract type SymmetrizationStyle end
 
+"""
+    struct ReflectDepth <: SymmetrizationStyle
+
+Reflection symmmetrization along the horizontal axis, such that north and south are mirrored.
+"""
 struct ReflectDepth <: SymmetrizationStyle end
+
+"""
+    struct ReflectWidth <: SymmetrizationStyle
+
+Reflection symmmetrization along the vertical axis, such that east and west are mirrored.
+"""
 struct ReflectWidth <: SymmetrizationStyle end
+
+# TODO
 struct Rotate <: SymmetrizationStyle end
+
+"""
+    struct RotateReflect <: SymmetrizationStyle
+
+Full reflection and rotation symmmetrization, such that reflection along the horizontal and
+vertical axis as well as π/2 rotations leave the object invariant.
+"""
 struct RotateReflect <: SymmetrizationStyle end
 
 # some rather shady definitions for 'hermitian conjugate' at the level of a single tensor
+
 function herm_depth(x::PEPSTensor)
     return permute(x', ((5,), (3, 2, 1, 4)))
 end
@@ -67,47 +88,54 @@ function rot_inv(x)
     )
 end
 
-## PEPS unit cell symmetrization
+# PEPS unit cell symmetrization
 
-const PEPSLike = Union{InfinitePEPS,AbstractArray{<:PEPSTensor,2}}
+"""
+    symmetrize!(peps::InfinitePEPS, ::SymmetrizationStyle)
 
-symmetrize!(p::PEPSLike, ::Nothing) = p
+Symmetrize a PEPS using the given `SymmetrizationStyle` in-place.
+"""
+symmetrize!(peps::InfinitePEPS, ::Nothing) = peps
 
-function symmetrize!(p::PEPSLike, ::ReflectDepth)
-    depth, width = size(p)
+function symmetrize!(peps::InfinitePEPS, ::ReflectDepth)
+    depth, width = size(peps)
     if mod(depth, 2) == 1
         for w in 1:width
-            p[ceil(Int, depth / 2), w] = herm_depth_inv(p[ceil(Int, depth / 2), w])
+            peps[ceil(Int, depth / 2), w] = herm_depth_inv(peps[ceil(Int, depth / 2), w])
         end
     end
     for d in 1:floor(Int, depth / 2)
         for w in 1:width
-            p[depth - d + 1, w] = _fit_spaces(herm_depth(p[d, w]), p[depth - d + 1, w])
+            peps[depth - d + 1, w] = _fit_spaces(
+                herm_depth(peps[d, w]), peps[depth - d + 1, w]
+            )
         end
     end
-    return p
+    return peps
 end
 
-function symmetrize!(p::PEPSLike, ::ReflectWidth)
-    depth, width = size(p)
+function symmetrize!(peps::InfinitePEPS, ::ReflectWidth)
+    depth, width = size(peps)
     if mod(width, 2) == 1
         for d in 1:depth
-            p[d, ceil(Int, width / 2)] = herm_width_inv(p[d, ceil(Int, width / 2)])
+            peps[d, ceil(Int, width / 2)] = herm_width_inv(peps[d, ceil(Int, width / 2)])
         end
     end
     for w in 1:floor(Int, width / 2)
         for d in 1:depth
-            p[d, width - w + 1] = _fit_spaces(herm_width(p[d, w]), p[d, width - w + 1])
+            peps[d, width - w + 1] = _fit_spaces(
+                herm_width(peps[d, w]), peps[d, width - w + 1]
+            )
         end
     end
-    return p
+    return peps
 end
 
-function symmetrize!(p::PEPSLike, ::Rotate)
+function symmetrize!(peps::InfinitePEPS, ::Rotate)
     return error("TODO")
 end
 
-function symmetrize!(p::PEPSLike, ::RotateReflect)
+function symmetrize!(peps::InfinitePEPS, symm::RotateReflect)
     # TODO: clean up this mess...
 
     # some auxiliary transformations
@@ -119,56 +147,71 @@ function symmetrize!(p::PEPSLike, ::RotateReflect)
         return x + _fit_spaces(permute(x', ((5,), (3, 2, 1, 4))), x)
     end
 
-    depth, width = size(p)
-    depth == width || error("This only works for square unit cells.")
+    depth, width = size(peps)
+    depth == width || ArgumentError("$symm can only be applied to square unit cells")
 
     odd = mod(depth, 2)
     if odd == 1
-        p[ceil(Int, depth / 2), ceil(Int, width / 2)] = symmetrize_center(
-            p[ceil(Int, depth / 2), ceil(Int, width / 2)]
+        peps[ceil(Int, depth / 2), ceil(Int, width / 2)] = symmetrize_center(
+            peps[ceil(Int, depth / 2), ceil(Int, width / 2)]
         )
     end
     for d in 1:ceil(Int, depth / 2)
         for w in 1:floor(Int, width / 2)
             if d == w
-                p[d, w] = symmetrize_corner(p[d, w])
-                p[d, width - w + 1] = _fit_spaces(rotr90(p[d, w]), p[d, width - w + 1])
-                p[depth - d + 1, w] = _fit_spaces(herm_depth(p[d, w]), p[depth - d + 1, w])
-                p[depth - d + 1, width - w + 1] = _fit_spaces(
-                    herm_depth(rotr90(p[d, w])), p[depth - d + 1, width - w + 1]
+                peps[d, w] = symmetrize_corner(peps[d, w])
+                peps[d, width - w + 1] = _fit_spaces(
+                    rotr90(peps[d, w]), peps[d, width - w + 1]
+                )
+                peps[depth - d + 1, w] = _fit_spaces(
+                    herm_depth(peps[d, w]), peps[depth - d + 1, w]
+                )
+                peps[depth - d + 1, width - w + 1] = _fit_spaces(
+                    herm_depth(rotr90(peps[d, w])), peps[depth - d + 1, width - w + 1]
                 )
 
             elseif odd == 1 && d == ceil(Int, depth / 2)
-                p[d, w] = symmetrize_mid_depth(p[d, w])
-                p[w, d] = _fit_spaces(rotr90(p[d, w]), p[w, d])
-                p[d, width - w + 1] = _fit_spaces(rot180(p[d, w]), p[d, width - w + 1])
-                p[width - w + 1, d] = _fit_spaces(
-                    herm_depth(rotr90(p[d, w])), p[width - w + 1, d]
+                peps[d, w] = symmetrize_mid_depth(peps[d, w])
+                peps[w, d] = _fit_spaces(rotr90(peps[d, w]), peps[w, d])
+                peps[d, width - w + 1] = _fit_spaces(
+                    rot180(peps[d, w]), peps[d, width - w + 1]
+                )
+                peps[width - w + 1, d] = _fit_spaces(
+                    herm_depth(rotr90(peps[d, w])), peps[width - w + 1, d]
                 )
 
             else
-                p[depth - d + 1, w] = _fit_spaces(herm_depth(p[d, w]), p[depth - d + 1, w])
-                p[w, depth - d + 1] = _fit_spaces(rotr90(p[d, w]), p[w, depth - d + 1])
-                p[width - w + 1, depth - d + 1] = _fit_spaces(
-                    herm_depth(rotr90(p[d, w])), [width - w + 1, depth - d + 1]
+                peps[depth - d + 1, w] = _fit_spaces(
+                    herm_depth(peps[d, w]), peps[depth - d + 1, w]
                 )
-                p[w, d] = _fit_spaces(rotr90(herm_depth(p[d, w])), p[w, d])
-                p[width - w + 1, d] = _fit_spaces(
-                    herm_depth(rotr90(herm_depth(p[d, w]))), p[width - w + 1, d]
+                peps[w, depth - d + 1] = _fit_spaces(
+                    rotr90(peps[d, w]), peps[w, depth - d + 1]
                 )
-                p[d, width - w + 1] = _fit_spaces(
-                    rotr90(rotr90(herm_depth(p[d, w]))), p[d, width - w + 1]
+                peps[width - w + 1, depth - d + 1] = _fit_spaces(
+                    herm_depth(rotr90(peps[d, w])), [width - w + 1, depth - d + 1]
                 )
-                p[depth - d + 1, width - w + 1] = _fit_spaces(
-                    herm_depth(rotr90(rotr90(herm_depth(p[d, w])))),
-                    p[depth - d + 1, width - w + 1],
+                peps[w, d] = _fit_spaces(rotr90(herm_depth(peps[d, w])), peps[w, d])
+                peps[width - w + 1, d] = _fit_spaces(
+                    herm_depth(rotr90(herm_depth(peps[d, w]))), peps[width - w + 1, d]
+                )
+                peps[d, width - w + 1] = _fit_spaces(
+                    rotr90(rotr90(herm_depth(peps[d, w]))), peps[d, width - w + 1]
+                )
+                peps[depth - d + 1, width - w + 1] = _fit_spaces(
+                    herm_depth(rotr90(rotr90(herm_depth(peps[d, w])))),
+                    peps[depth - d + 1, width - w + 1],
                 )
             end
         end
     end
-    return p
+    return peps
 end
 
+"""
+    symmetrize_callback(peps, envs, E, grad, symm::SymmetrizationStyle)
+
+Callback function symmetrizing both the `peps` and `grad` tensors.
+"""
 function symmetrize_callback(peps, envs, E, grad, symm::SymmetrizationStyle)
     peps_symm = symmetrize!(deepcopy(peps), symm)
     grad_symm = symmetrize!(deepcopy(grad), symm)

test/j1j2_model.jl

@@ -27,17 +27,14 @@ opt_alg = PEPSOptimize(;
 Random.seed!(91283219347)
 H = square_lattice_j1j2(; J2=0.25)
 psi_init = InfinitePEPS(2, χbond)
-psi_init = PEPSKit.symmetrize!(psi_init, PEPSKit.Full())
-env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
+psi_init = PEPSKit.symmetrize!(psi_init, PEPSKit.RotateReflect())
+env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg);
 
 # find fixedpoint
-cb = (args...) -> PEPSKit.symmetrize_callback(args..., RotateReflect())
-result = fixedpoint(psi_init, H, opt_alg, env_init)#; callback=cb)
+callback = (args...) -> PEPSKit.symmetrize_callback(args..., RotateReflect())
+result = fixedpoint(psi_init, H, opt_alg, env_init; callback)
 ξ_h, ξ_v, = correlation_length(result.peps, result.env)
 
-p = symmetrize!(deepcopy(result.peps), RotateReflect())
-ξ_h, ξ_v, = correlation_length(p, leading_boundary(result.env, p, ctm_alg))
-
 # compare against Juraj Hasik's data:
 # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.25/state_1s_A1_j20.25_D2_chi_opt48.dat
 ξ_ref = -1 / log(0.2723596743547324)