Clean up symmetrization and add optimization callback

src/PEPSKit.jl

@@ -85,8 +85,9 @@ export fixedpoint
 export InfinitePEPS, InfiniteTransferPEPS
 export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS
-export symmetrize, None, Depth, Full
+export ReflectDepth, ReflectWidth, RotateReflect, symmetrize!, symmetrize_finalize!
 export showtypeofgrad
-export square_lattice_tf_ising, square_lattice_heisenberg, square_lattice_pwave
+export square_lattice_tf_ising, square_lattice_heisenberg, square_lattice_j1j2
+export square_lattice_pwave
 
 end # module

src/algorithms/peps_opt.jl

@@ -121,12 +121,17 @@ function PEPSOptimize(;
 end
 
 """
-    fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv]) where {T}
+    fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv];
+               finalize!=OptimKit._finalize!) where {T}
 
 Optimize `ψ₀` with respect to the Hamiltonian `H` according to the parameters supplied
 in `alg`. The initial environment `env₀` serves as an initial guess for the first CTMRG run.
 By default, a random initial environment is used.
 
+The `finalize!` kwarg can be used to insert a function call after each optimization step
+by utilizing the `finalize!` kwarg of `OptimKit.optimize`.
+The function maps `(peps, envs), f, g = finalize!((peps, envs), f, g, numiter)`.
+
 The function returns a `NamedTuple` which contains the following entries:
 - `peps`: final `InfinitePEPS`
 - `env`: `CTMRGEnv` corresponding to the final PEPS
@@ -137,12 +142,16 @@ The function returns a `NamedTuple` which contains the following entries:
 - `numfg`: total number of calls to the energy function
 """
 function fixedpoint(
-    ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(T)^20)
+    ψ₀::InfinitePEPS{T},
+    H,
+    alg::PEPSOptimize,
+    env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(T)^20);
+    (finalize!)=OptimKit._finalize!,
 ) where {T}
     (peps, env), E, ∂E, numfg, convhistory = optimize(
-        (ψ₀, env₀), alg.optimizer; retract=my_retract, inner=my_inner
+        (ψ₀, env₀), alg.optimizer; retract=my_retract, inner=my_inner, finalize!
     ) do (peps, envs)
-        E, g = withgradient(peps) do ψ
+        E, gs = withgradient(peps) do ψ
             envs´ = hook_pullback(
                 leading_boundary,
                 envs,
@@ -155,8 +164,8 @@ function fixedpoint(
             end
             return costfun(ψ, envs´, H)
         end
-        # withgradient returns tuple of gradients `g`
-        return E, only(g)
+        g = only(gs)  # `withgradient` returns tuple of gradients `gs`
+        return E, g
     end
     return (;
         peps,

src/operators/models.jl

@@ -38,6 +38,38 @@ function square_lattice_heisenberg(
     return repeat(nearest_neighbour_hamiltonian(lattice, H / 4), unitcell...)
 end
 
+"""
+    square_lattice_j1j2(::Type{T}=ComplexF64; J1=1, J2=1, unitcell=(1, 1), sublattice=true)
+
+
+Square lattice J₁-J₂ model. The `sublattice` kwarg enables a single site unit cell via a
+sublattice rotation.
+"""
+function square_lattice_j1j2(
+    ::Type{T}=ComplexF64; J1=1, J2=1, unitcell::Tuple{Int,Int}=(1, 1), sublattice=true
+) where {T<:Number}
+    physical_space = ComplexSpace(2)
+    lattice = fill(physical_space, 1, 1)
+    σx = TensorMap(T[0 1; 1 0], physical_space, physical_space)
+    σy = TensorMap(T[0 im; -im 0], physical_space, physical_space)
+    σz = TensorMap(T[1 0; 0 -1], physical_space, physical_space)
+    h_AA = σx ⊗ σx + σy ⊗ σy + σz ⊗ σz
+    h_AB = sublattice ? -σx ⊗ σx + σy ⊗ σy - σz ⊗ σz : h_AA  # Apply sublattice rotation
+
+    terms = []
+    for I in eachindex(IndexCartesian(), lattice)
+        nearest_x = I + CartesianIndex(1, 0)
+        nearest_y = I + CartesianIndex(0, 1)
+        next_xy = I + CartesianIndex(1, 1)
+        push!(terms, (I, nearest_x) => J1 / 4 * h_AB)
+        push!(terms, (I, nearest_y) => J1 / 4 * h_AB)
+        push!(terms, (I, next_xy) => J2 / 4 * h_AA)
+        push!(terms, (nearest_x, nearest_y) => J2 / 4 * h_AA)
+    end
+
+    return repeat(LocalOperator(lattice, terms...), unitcell...)
+end
+
 """
     square_lattice_pwave(::Type{T}=ComplexF64; t=1, μ=2, Δ=1, unitcell=(1, 1))
 

src/states/infinitepeps.jl

@@ -133,6 +133,18 @@
 LinearAlgebra.dot(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = dot(ψ₁.A, ψ₂.A)
 LinearAlgebra.norm(ψ::InfinitePEPS) = norm(ψ.A)
 
+## (Approximate) equality
+function Base.:(==)(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS)
+    return all(zip(ψ₁.A, ψ₂.A)) do (p₁, p₂)
+        return p₁ == p₂
+    end
+end
+function Base.isapprox(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS; kwargs...)
+    return all(zip(ψ₁.A, ψ₂.A)) do (p₁, p₂)
+        return isapprox(p₁, p₂; kwargs...)
+    end
+end
+
 # Used in _scale during OptimKit.optimize
 function LinearAlgebra.rmul!(ψ::InfinitePEPS, α::Number)
     rmul!(ψ.A, α)

src/utility/symmetrization.jl

@@ -1,14 +1,32 @@
-# some basic symmetrization routines for PEPS
-
 abstract type SymmetrizationStyle end
 
-struct None <: SymmetrizationStyle end
-struct Depth <: SymmetrizationStyle end
-struct Width <: SymmetrizationStyle end
-struct Rot <: SymmetrizationStyle end
-struct Full <: 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
@@ -32,7 +50,7 @@
 # make two TensorMap's have the same spaces, by force if necessary
 # this is definitely not what you would want to do, but it circumvents having to think
 # about what hermiticity means at the level of transfer operators, which is something
-function _make_it_fit(
+function _fit_spaces(
     y::AbstractTensorMap{S,N₁,N₂}, x::AbstractTensorMap{S,N₁,N₂}
 ) where {S<:IndexSpace,N₁,N₂}
     for i in 1:(N₁ + N₂)
@@ -46,128 +64,158 @@
     end
     return y
 end
+_fit_spaces(y::InfinitePEPS, x::InfinitePEPS) = InfinitePEPS(map(_fit_spaces, y.A, x.A))
 
 function herm_depth_inv(x::Union{PEPSTensor,PEPOTensor})
-    return 0.5 * (x + _make_it_fit(herm_depth(x), x))
+    return 0.5 * (x + _fit_spaces(herm_depth(x), x))
 end
 
 function herm_width_inv(x::Union{PEPSTensor,PEPOTensor})
-    return 0.5 * (x + _make_it_fit(herm_width(x), x))
+    return 0.5 * (x + _fit_spaces(herm_width(x), x))
 end
 
-function herm_height_inv(x::Union{PEPSTensor,PEPOTensor})
-    return 0.5 * (x + _make_it_fit(herm_height(x), x))
+function herm_height_inv(x::PEPOTensor)
+    return 0.5 * (x + _fit_spaces(herm_height(x), x))
 end
 
 # rotation invariance
 
 function rot_inv(x)
     return 0.25 * (
         x +
-        _make_it_fit(rotl90(x), x) +
-        _make_it_fit(rot180(x), x) +
-        _make_it_fit(rotr90(x), x)
+        _fit_spaces(rotl90(x), x) +
+        _fit_spaces(rot180(x), x) +
+        _fit_spaces(rotr90(x), x)
     )
 end
 
-## PEPS unit cell symmetrization
+# PEPS unit cell symmetrization
 
-PEPSLike = Union{InfinitePEPS,AbstractArray{<:PEPSTensor,2}}
+"""
+    symmetrize!(peps::InfinitePEPS, ::SymmetrizationStyle)
 
-symmetrize(p::PEPSLike, ::None) = p
+Symmetrize a PEPS using the given `SymmetrizationStyle` in-place.
+"""
+symmetrize!(peps::InfinitePEPS, ::Nothing) = peps
 
-function symmetrize(p::PEPSLike, ::Depth)
-    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] = _make_it_fit(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, ::Width)
-    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), h])
+            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] = _make_it_fit(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, ::Rot)
+function symmetrize!(peps::InfinitePEPS, ::Rotate)
     return error("TODO")
 end
 
-function symmetrize(p::PEPSLike, ::Full)
+function symmetrize!(peps::InfinitePEPS, symm::RotateReflect)
     # TODO: clean up this mess...
 
     # some auxiliary transformations
     function symmetrize_corner(x::PEPSTensor)
-        return 0.5 * (x + _make_it_fit(permute(x', ((5,), (4, 3, 2, 1))), x))
+        return 0.5 * (x + _fit_spaces(permute(x', ((5,), (4, 3, 2, 1))), x))
     end
     symmetrize_center(x::PEPSTensor) = herm_depth_inv(rot_inv(x))
     function symmetrize_mid_depth(x::PEPSTensor)
-        return x + _make_it_fit(permute(x', ((5,), (3, 2, 1, 4))), x)
+        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] = _make_it_fit(rotr90(p[d, w]), p[d, width - w + 1])
-                p[depth - d + 1, w] = _make_it_fit(herm_depth(p[d, w]), p[depth - d + 1, w])
-                p[depth - d + 1, width - w + 1] = _make_it_fit(
-                    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] = _make_it_fit(rotr90(p[d, w]), p[w, d])
-                p[d, width - w + 1] = _make_it_fit(rot180(p[d, w]), p[d, width - w + 1])
-                p[width - w + 1, d] = _make_it_fit(
-                    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] = _make_it_fit(herm_depth(p[d, w]), p[depth - d + 1, w])
-                p[w, depth - d + 1] = _make_it_fit(rotr90(p[d, w]), p[w, depth - d + 1])
-                p[width - w + 1, depth - d + 1] = _make_it_fit(
-                    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] = _make_it_fit(rotr90(herm_depth(p[d, w])), p[w, d])
-                p[width - w + 1, d] = _make_it_fit(
-                    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] = _make_it_fit(
-                    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] = _make_it_fit(
-                    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_finalize!(symm::SymmetrizationStyle)
+
+Return `finalize!` function for symmetrizing the `peps` and `grad` tensors in-place,
+which maps `(peps_symm, envs), E, grad_symm = symmetrize_finalize!((peps, envs), E, grad, numiter)`.
+"""
+function symmetrize_finalize!(symm::SymmetrizationStyle)
+    function symmetrize_finalize!((peps, envs), E, grad, _)
+        peps_symm = symmetrize!(peps, symm)
+        grad_symm = symmetrize!(grad, symm)
+        return (peps_symm, envs), E, grad_symm
+    end
 end