Add support for fermionic systems (#47)

* ctmrg for fermions * NNanisotropic * Merge remote-tracking branch 'origin' * fix onsites * Add ChainRulesTestUtils test dependency * Format * Add FiniteDifferences test dependency * Move rotation functions * Make tests their own testset * Attempt GMRES for p-wave * Fix ugly "fix" for gaugefix * Fiddle with some settings * Refactor transfer matrix eigensolver * Fix typo * Add sum over onsite expectation values * Add nontrivial unit cell and change some settings to 'force' optimization to proceed * Some cleanup * Fix docstring AnisotropicNNOperator * Tweak some more settings to reduce runtime * Formatter [no ci] * Cleanup gradparts tests [no ci] --------- Co-authored-by: leburgel <[email protected]> Co-authored-by: lkdvos <[email protected]>
QuantumKitHub · Jun 26, 2024 · 4a1e5fb · 4a1e5fb
1 parent 996893a
commit 4a1e5fb
Show file tree

Hide file tree

Showing 17 changed files with 409 additions and 127 deletions.
diff --git a/Project.toml b/Project.toml
@@ -36,6 +36,8 @@ julia = "1.6"
 [extras]
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 
 [targets]
-test = ["Test", "SafeTestsets"]
+test = ["Test", "SafeTestsets", "ChainRulesTestUtils", "FiniteDifferences"]
diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl
@@ -6,18 +6,9 @@ using PEPSKit, KrylovKit
 # We use the parameters (J₁, J₂, J₃) = (-1, 1, -1) by default to capture
 # the ground state in a single-site unit cell. This can be seen from
 # sublattice rotating H from parameters (1, 1, 1) to (-1, 1, -1).
-function square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
-    physical_space = ComplexSpace(2)
-    T = ComplexF64
-    σ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 = (Jx * σx ⊗ σx) + (Jy * σy ⊗ σy) + (Jz * σz ⊗ σz)
-    return NLocalOperator{NearestNeighbor}(H / 4)
-end
+H = square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
 
 # Parameters
-H = square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
 χbond = 2
 χenv = 20
 ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, verbosity=1)

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
@@ -26,6 +26,7 @@ include("mpskit_glue/transferpepo_environments.jl")
 
 include("environments/ctmrgenv.jl")
 include("operators/localoperator.jl")
+include("operators/models.jl")
 
 include("algorithms/ctmrg.jl")
 include("algorithms/peps_opt.jl")
@@ -58,7 +59,7 @@ module Defaults
 end
 
 export CTMRG, CTMRGEnv
-export NLocalOperator, OnSite, NearestNeighbor
+export NLocalOperator, AnisotropicNNOperator, OnSite, NearestNeighbor
 export expectation_value, costfun
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolve
@@ -68,5 +69,6 @@ export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS
 export symmetrize, None, Depth, Full
 export showtypeofgrad
+export square_lattice_heisenberg, square_lattice_pwave
 
 end # module
diff --git a/src/algorithms/ctmrg.jl b/src/algorithms/ctmrg.jl
@@ -39,30 +39,30 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
 
     for i in 1:(alg.maxiter)
         env, ϵ = ctmrg_iter(state, env, alg)  # Grow and renormalize in all 4 directions
+        conv_condition, normold, CSold, TSold, ϵ = ignore_derivatives() do
+            # Compute convergence criteria and take max (TODO: How should we handle logging all of this?)
+            Δϵ = abs((ϵold - ϵ) / ϵold)
+            normnew = norm(state, env)
+            Δnorm = abs(normold - normnew) / abs(normold)
+            CSnew = map(c -> tsvd(c; alg=TensorKit.SVD())[2], env.corners)
+            ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new)
+                # only compute the difference on the smallest part of the spaces
+                smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new))
+                e_old = isometry(MPSKit._firstspace(c_old), smallest)
+                e_new = isometry(MPSKit._firstspace(c_new), smallest)
+                return norm(e_new' * c_new * e_new - e_old' * c_old * e_old)
+            end
+            TSnew = map(t -> tsvd(t; alg=TensorKit.SVD())[2], env.edges)
 
-        # Compute convergence criteria and take max (TODO: How should we handle logging all of this?)
-        Δϵ = abs((ϵold - ϵ) / ϵold)
-        normnew = norm(state, env)
-        Δnorm = abs(normold - normnew) / abs(normold)
-        CSnew = map(c -> tsvd(c; alg=TensorKit.SVD())[2], env.corners)
-        ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new)
-            # only compute the difference on the smallest part of the spaces
-            smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new))
-            e_old = isometry(MPSKit._firstspace(c_old), smallest)
-            e_new = isometry(MPSKit._firstspace(c_new), smallest)
-            return norm(e_new' * c_new * e_new - e_old' * c_old * e_old)
-        end
-        TSnew = map(t -> tsvd(t; alg=TensorKit.SVD())[2], env.edges)
+            ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new)
+                MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) ||
+                    return scalartype(t_old)(Inf)
+                # TODO: implement when spaces aren't the same
+                return norm(t_new - t_old)
+            end
 
-        ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new)
-            MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) ||
-                return scalartype(t_old)(Inf)
-            # TODO: implement when spaces aren't the same
-            return norm(t_new - t_old)
-        end
+            conv_condition = max(Δnorm, ΔCS, ΔTS) < alg.tol && i > alg.miniter
 
-        conv_condition = max(Δnorm, ΔCS, ΔTS) < alg.tol && i > alg.miniter
-        ignore_derivatives() do # Print verbose info
             if alg.verbosity > 1 || (alg.verbosity == 1 && (i == 1 || conv_condition))
                 @printf(
                     "CTMRG iter: %3d   norm: %.2e   Δnorm: %.2e   ΔCS: %.2e   ΔTS: %.2e   ϵ: %.2e   Δϵ: %.2e\n",
@@ -80,14 +80,9 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
                 @warn(
                     "CTMRG reached maximal number of iterations at (Δnorm=$Δnorm, ΔCS=$ΔCS, ΔTS=$ΔTS)"
                 )
+            return conv_condition, normnew, CSnew, TSnew, ϵ
         end
         conv_condition && break  # Converge if maximal Δ falls below tolerance
-
-        # Update convergence criteria
-        normold = normnew
-        CSold = CSnew
-        TSold = TSnew
-        ϵold = ϵ
     end
 
     # Do one final iteration that does not change the spaces
@@ -96,7 +91,7 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
     )
     env′, = ctmrg_iter(state, env, alg_fixed)
     envfix = gauge_fix(env, env′)
-    check_elementwise_convergence(env, envfix; atol=alg.tol^(3 / 4)) ||
+    check_elementwise_convergence(env, envfix; atol=alg.tol^(1 / 2)) ||
         @warn "CTMRG did not converge elementwise."
     return envfix
 end
@@ -145,12 +140,13 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
             scalartype(T),
             MPSKit._lastspace(Tsfinal[end])' ← MPSKit._lastspace(M[end])',
         )
-        ρprev = eigsolve(TransferMatrix(Tsprev, M), ρinit, 1, :LM)[2][1]
-        ρfinal = eigsolve(TransferMatrix(Tsfinal, M), ρinit, 1, :LM)[2][1]
+
+        ρ_prev = transfermatrix_fixedpoint(Tsprev, M, ρinit)
+        ρ_final = transfermatrix_fixedpoint(Tsfinal, M, ρinit)
 
         # Decompose and multiply
-        Up, _, Vp = tsvd(ρprev)
-        Uf, _, Vf = tsvd(ρfinal)
+        Up, _, Vp = tsvd!(ρ_prev)
+        Uf, _, Vf = tsvd!(ρ_final)
         Qprev = Up * Vp
         Qfinal = Uf * Vf
         σ = Qprev * Qfinal'
@@ -161,17 +157,25 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
     cornersfix, edgesfix = fix_relative_phases(envfinal, signs)
 
     # Fix global phase
-    cornersgfix = map(zip(envprev.corners, cornersfix)) do (Cprev, Cfix)
-        φ = dot(Cprev, Cfix)
-        φ' * Cfix
+    cornersgfix = map(envprev.corners, cornersfix) do Cprev, Cfix
+        return dot(Cfix, Cprev) * Cfix
     end
-    edgesgfix = map(zip(envprev.edges, edgesfix)) do (Tprev, Tfix)
-        φ = dot(Tprev, Tfix)
-        φ' * Tfix
+    edgesgfix = map(envprev.edges, edgesfix) do Tprev, Tfix
+        return dot(Tfix, Tprev) * Tfix
     end
-    envfix = CTMRGEnv(cornersgfix, edgesgfix)
+    return CTMRGEnv(cornersgfix, edgesgfix)
+end
 
-    return envfix
+# this is a bit of a hack to get the fixed point of the mixed transfer matrix
+# because MPSKit is not compatible with AD
+function transfermatrix_fixedpoint(tops, bottoms, ρinit)
+    _, vecs, info = eigsolve(ρinit, 1, :LM, Arnoldi()) do ρ
+        return foldr(zip(tops, bottoms); init=ρ) do (top, bottom), ρ
+            return @tensor ρ′[-1; -2] := top[-1 4 3; 1] * conj(bottom[-2 4 3; 2]) * ρ[1; 2]
+        end
+    end
+    info.converged > 0 || @warn "eigsolve did not converge"
+    return first(vecs)
 end
 
 # Explicit fixing of relative phases (doing this compactly in a loop is annoying)
@@ -329,7 +333,8 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}
             else
                 alg.trscheme
             end
-            U, S, V, ϵ_local = tsvd!(Q_sw * Q_nw; trunc=trscheme, alg=TensorKit.SVD())  # TODO: Add field in CTMRG to choose SVD function
+            @tensor QQ[-1 -2 -3; -4 -5 -6] := Q_sw[-1 -2 -3; 1 2 3] * Q_nw[1 2 3; -4 -5 -6]
+            U, S, V, ϵ_local = tsvd!(QQ; trunc=trscheme, alg=TensorKit.SVD())
             ϵ = max(ϵ, ϵ_local / norm(S))
             # TODO: check if we can just normalize enlarged corners s.t. trunc behaves a bit better
 
@@ -409,8 +414,8 @@ function build_projectors(
     U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q_sw, Q_nw
 ) where {E<:ElementarySpace}
     isqS = sdiag_inv_sqrt(S)
-    @tensor Pl[-1 -2 -3; -4] := Q_nw[-1 -2 -3; 1 2 3] * conj(V[4; 1 2 3]) * isqS[4; -4]
-    @tensor Pr[-1; -2 -3 -4] := isqS[-1; 1] * conj(U[2 3 4; 1]) * Q_sw[2 3 4; -2 -3 -4]
+    Pl = Q_nw * V' * isqS
+    Pr = isqS * U' * Q_sw
     return Pl, Pr
 end
 
@@ -448,12 +453,10 @@ function LinearAlgebra.norm(peps::InfinitePEPS, env::CTMRGEnv)
             peps[r, c][17; 6 10 14 2] *
             conj(peps[r, c][17; 7 11 15 3])
 
-        total *= tr(
-            env.corners[NORTHWEST, r, c] *
-            env.corners[NORTHEAST, r, mod1(c - 1, end)] *
-            env.corners[SOUTHEAST, mod1(r - 1, end), mod1(c - 1, end)] *
-            env.corners[SOUTHWEST, mod1(r - 1, end), c],
-        )
+        total *= @tensor env.corners[NORTHWEST, r, c][1; 2] *
+            env.corners[NORTHEAST, r, mod1(c - 1, end)][2; 3] *
+            env.corners[SOUTHEAST, mod1(r - 1, end), mod1(c - 1, end)][3; 4] *
+            env.corners[SOUTHWEST, mod1(r - 1, end), c][4; 1]
 
         total /= @tensor env.edges[WEST, r, c][1 10 11; 4] *
             env.corners[NORTHWEST, r, c][4; 5] *

diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl
@@ -112,6 +112,7 @@ function _rrule(
     alg::CTMRG,
 )
     envs = leading_boundary(envinit, state, alg)
+    #TODO: fixed space for unit cells
 
     function leading_boundary_pullback(Δenvs′)
         Δenvs = unthunk(Δenvs′)

diff --git a/src/operators/infinitepepo.jl b/src/operators/infinitepepo.jl
@@ -121,8 +121,6 @@ Base.getindex(T::InfinitePEPO, args...) = Base.getindex(T.A, args...)
 Base.axes(T::InfinitePEPO, args...) = axes(T.A, args...)
 TensorKit.space(T::InfinitePEPO, i, j) = space(T[i, j, end], 1)
 
-Base.rotl90(T::InfinitePEPO) = InfinitePEPO(rotl90(rotl90.(T.A)));
-
 function initializePEPS(
     T::InfinitePEPO{<:PEPOTensor{S}}, vspace::S
 ) where {S<:ElementarySpace}

diff --git a/src/operators/localoperator.jl b/src/operators/localoperator.jl
@@ -1,3 +1,5 @@
+# TODO: change this implementation to a type-stable one
+
 abstract type AbstractInteraction end
 
 """
@@ -24,6 +26,29 @@ struct NLocalOperator{I<:AbstractInteraction}
     op::AbstractTensorMap
 end
 
+"""
+    struct AnisotropicNNOperator{I<:AbstractInteraction}
+    
+Operator which includes an on-site term and two nearest-neighbor terms, vertical and horizontal.
+"""
+struct AnisotropicNNOperator
+    h0::NLocalOperator{OnSite}
+    hx::NLocalOperator{NearestNeighbor}
+    hy::NLocalOperator{NearestNeighbor}
+end
+function AnisotropicNNOperator(
+    h0::AbstractTensorMap{S,1,1},
+    hx::AbstractTensorMap{S,2,2},
+    hy::AbstractTensorMap{S,2,2}=hx,
+) where {S}
+    return AnisotropicNNOperator(
+        NLocalOperator{OnSite}(h0),
+        NLocalOperator{NearestNeighbor}(hx),
+        NLocalOperator{NearestNeighbor}(hy),
+    )
+end
+# TODO: include the on-site term in the two-site terms, to reduce number of contractions.
+
 @doc """
     operator_env(peps::InfinitePEPS, env::CTMRGEnv, ::AbstractInteraction)
 
@@ -178,3 +203,17 @@ function costfun(peps::InfinitePEPS, env, op::NLocalOperator{NearestNeighbor})
     ov = sum(expectation_value(rotl90(peps), rotl90(env), op))
     return real(oh + ov)
 end
+
+"""
+    costfun(peps::InfinitePEPS, env, op::AnisotropicNNOperator)
+    
+Compute the expectation value of an on-site and an anisotropic nearest-neighbor operator.
+This is used to evaluate and differentiate the energy in ground-state PEPS optimizations.
+"""
+function costfun(peps::InfinitePEPS, env, op::AnisotropicNNOperator)
+    oos = sum(expectation_value(peps, env, op.h0))
+    oh = sum(expectation_value(peps, env, op.hx))
+    ov = sum(expectation_value(rotr90(peps), rotate_north(env, WEST), op.hy))
+    #ov = sum(expectation_value(rotl90(peps), rotl90(env), op.hy))
+    return real(oos + oh + ov)
+end
diff --git a/src/operators/models.jl b/src/operators/models.jl
@@ -0,0 +1,45 @@
+## Model Hamiltonians
+# -------------------
+
+"""
+    square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
+
+Square lattice Heisenberg model.
+By default, this implements a single site unit cell via a sublattice rotation.
+"""
+function square_lattice_heisenberg(
+    ::Type{T}=ComplexF64; Jx=-1, Jy=1, Jz=-1
+) where {T<:Number}
+    physical_space = ComplexSpace(2)
+    σ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 = (Jx * σx ⊗ σx) + (Jy * σy ⊗ σy) + (Jz * σz ⊗ σz)
+    return NLocalOperator{NearestNeighbor}(H / 4)
+end
+
+"""
+    square_lattice_pwave(; t=1, μ=2, Δ=1)
+
+Square lattice p-wave superconductor model.
+"""
+function square_lattice_pwave(
+    ::Type{T}=ComplexF64; t::Number=1, μ::Number=2, Δ::Number=1
+) where {T<:Number}
+    physical_space = Vect[FermionParity](0 => 1, 1 => 1)
+    # on-site
+    h0 = TensorMap(zeros, T, physical_space ← physical_space)
+    block(h0, FermionParity(1)) .= -μ
+
+    # two-site (x-direction)
+    hx = TensorMap(zeros, T, physical_space^2 ← physical_space^2)
+    block(hx, FermionParity(0)) .= [0 -Δ; -Δ 0]
+    block(hx, FermionParity(1)) .= [0 -t; -t 0]
+
+    # two-site (y-direction)
+    hy = TensorMap(zeros, T, physical_space^2 ← physical_space^2)
+    block(hy, FermionParity(0)) .= [0 Δ*im; -Δ*im 0]
+    block(hy, FermionParity(1)) .= [0 -t; -t 0]
+
+    return AnisotropicNNOperator(h0, hx, hy)
+end
diff --git a/src/states/abstractpeps.jl b/src/states/abstractpeps.jl
@@ -60,6 +60,3 @@ abstract type AbstractPEPS end
 Abstract supertype for a 2D projected entangled-pair operator.
 """
 abstract type AbstractPEPO end
-
-Base.rotl90(t::PEPSTensor) = permute(t, ((1,), (3, 4, 5, 2)))
-Base.rotl90(t::PEPOTensor) = permute(t, ((1, 2), (4, 5, 6, 3)))
diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
@@ -5,7 +5,7 @@ Represents an infinite projected entangled-pair state on a 2D square lattice.
 """
 struct InfinitePEPS{T<:PEPSTensor} <: AbstractPEPS
     A::Matrix{T}
-
+    InfinitePEPS{T}(A::Matrix{T}) where {T<:PEPSTensor} = new{T}(A)
     function InfinitePEPS(A::Array{T,2}) where {T<:PEPSTensor}
         for (d, w) in Tuple.(CartesianIndices(A))
             space(A[d, w], 2) == space(A[_prev(d, end), w], 4)' || throw(
@@ -106,9 +106,6 @@ Base.setindex!(T::InfinitePEPS, args...) = (Base.setindex!(T.A, args...); T)
 Base.axes(T::InfinitePEPS, args...) = axes(T.A, args...)
 TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1)
 
-Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A)))
-Base.rotr90(t::InfinitePEPS) = InfinitePEPS(rotr90(rotr90.(t.A)))
-
 ## Math
 Base.:+(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = InfinitePEPS(ψ₁.A + ψ₂.A)
 Base.:-(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = InfinitePEPS(ψ₁.A - ψ₂.A)
@@ -160,3 +157,11 @@ function ChainRulesCore.rrule(::typeof(rotl90), peps::InfinitePEPS)
     end
     return peps′, rotl90_pullback
 end
+
+function ChainRulesCore.rrule(::typeof(rotr90), peps::InfinitePEPS)
+    peps′ = rotr90(peps)
+    function rotr90_pullback(Δpeps)
+        return NoTangent(), rotl90(Δpeps)
+    end
+    return peps′, rotr90_pullback
+end
diff --git a/src/utility/symmetrization.jl b/src/utility/symmetrization.jl
@@ -60,10 +60,22 @@ function herm_height_inv(x::Union{PEPSTensor,PEPOTensor})
 end
 
 # rotation invariance
-
+Base.rotl90(t::PEPSTensor) = permute(t, ((1,), (3, 4, 5, 2)))
 Base.rotr90(t::PEPSTensor) = permute(t, ((1,), (5, 2, 3, 4)))
 Base.rot180(t::PEPSTensor) = permute(t, ((1,), (4, 5, 2, 3)))
 
+Base.rotl90(t::PEPOTensor) = permute(t, ((1, 2), (4, 5, 6, 3)))
+Base.rotr90(t::PEPOTensor) = permute(t, ((1, 2), (6, 3, 4, 5)))
+Base.rot180(t::PEPOTensor) = permute(t, ((1, 2), (5, 6, 3, 4)))
+
+Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A)))
+Base.rotr90(t::InfinitePEPS) = InfinitePEPS(rotr90(rotr90.(t.A)))
+Base.rot180(t::InfinitePEPS) = InfinitePEPS(rot180(rot180.(t.A)))
+
+Base.rotl90(T::InfinitePEPO) = InfinitePEPO(stack(rotl90, eachslice(T.A; dims=3)))
+Base.rotr90(T::InfinitePEPO) = InfinitePEPO(stack(rotr90, eachslice(T.A; dims=3)))
+Base.rot180(T::InfinitePEPO) = InfinitePEPO(stack(rot180, eachslice(T.A; dims=3)))
+
 function rot_inv(x)
     return 0.25 * (
         x +