Extend CTMRGEnv constructor for nontrivial unit cells

README.md

@@ -52,7 +52,7 @@ opt_alg = PEPSOptimize(;
 
 # ground state search
 state = InfinitePEPS(2, D)
-ctm = leading_boundary(CTMRGEnv(state; Venv=ComplexSpace(chi)), state, ctm_alg)
+ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
 result = fixedpoint(state, H, opt_alg, ctm)
 
 @show result.E # -0.6625...

docs/src/index.md

@@ -37,8 +37,13 @@ opt_alg = PEPSOptimize(;
 
 # ground state search
 state = InfinitePEPS(2, D)
+<<<<<<< HEAD
+ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
+result = fixedpoint(state, Heisenberg_hamiltonian, opt_alg, ctm)
+=======
 ctm = leading_boundary(CTMRGEnv(state; Venv=ComplexSpace(chi)), state, ctm_alg)
 result = fixedpoint(state, H, opt_alg, ctm)
+>>>>>>> master
 
 @show result.E # -0.6625...
 ```

examples/boundary_mps.jl

@@ -32,7 +32,7 @@ mps, envs, ϵ = leading_boundary(mps, T, VUMPS())
 N = abs(prod(expectation_value(mps, T)))
 
 # This can be compared to the result obtained using the CTMRG algorithm
-ctm = leading_boundary(peps, CTMRG(; verbosity=1), CTMRGEnv(peps; Venv=ComplexSpace(20)))
+ctm = leading_boundary(peps, CTMRG(; verbosity=1), CTMRGEnv(peps, ComplexSpace(20)))
 N´ = abs(norm(peps, ctm))
 
 @show abs(N - N´) / N
@@ -53,7 +53,7 @@ mps2 = PEPSKit.initializeMPS(T2, fill(ComplexSpace(20), 2, 2))
 mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS())
 N2 = abs(prod(expectation_value(mps2, T2)))
 
-ctm2 = leading_boundary(peps2, CTMRG(; verbosity=1), CTMRGEnv(peps2; Venv=ComplexSpace(20)))
+ctm2 = leading_boundary(peps2, CTMRG(; verbosity=1), CTMRGEnv(peps2, ComplexSpace(20)))
 N2´ = abs(norm(peps2, ctm2))
 
 @show abs(N2 - N2´) / N2

examples/heisenberg.jl

@@ -27,6 +27,6 @@ opt_alg = PEPSOptimize(;
 # E/N = −0.6694421, which is a QMC estimate from https://arxiv.org/abs/1101.3281.
 # Of course there is a noticable bias for small χbond and χenv.
 ψ₀ = InfinitePEPS(2, χbond)
-env₀ = leading_boundary(CTMRGEnv(ψ₀; Venv=ℂ^χenv), ψ₀, ctm_alg)
+env₀ = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctmalg)
 result = fixedpoint(ψ₀, H, opt_alg, env₀)
 @show result.E

src/algorithms/ctmrg.jl

@@ -64,7 +64,7 @@
 Per default, a random initial environment is used.
 """
 function MPSKit.leading_boundary(state, alg::CTMRG)
-    return MPSKit.leading_boundary(CTMRGEnv(state), state, alg)
+    return MPSKit.leading_boundary(CTMRGEnv(state, oneunit(spacetype(state))), state, alg)
 end
 function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
     normold = 1.0

src/algorithms/peps_opt.jl

@@ -60,7 +60,7 @@ in `alg`. The initial environment `env₀` serves as an initial guess for the fi
 By default, a random initial environment is used.
 """
 function fixedpoint(
-    ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(ψ₀; Venv=field(T)^20)
+    ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(T)^20)
 ) where {T}
     (peps, env), E, ∂E, info = optimize(
         (ψ₀, env₀), alg.optimizer; retract=my_retract, inner=my_inner

src/environments/ctmrgenv.jl

@@ -8,36 +8,256 @@
     edges::Array{T,3}
 end
 
+_spacetype(::Int) = ComplexSpace
+_spacetype(::S) where {S<:ElementarySpace} = S
+
+_to_space(χ::Int) = ℂ^χ
+_to_space(χ::ElementarySpace) = χ
+
+function _corner_tensor(
+    f, ::Type{T}, left_vspace::S, right_vspace::S=left_vspace
+) where {T,S<:Union{Int,ElementarySpace}}
+    return TensorMap(f, T, _to_space(left_vspace) ← _to_space(right_vspace))
+end
+
+function _edge_tensor(
+    f,
+    ::Type{T},
+    left_vspace::S,
+    top_pspace::S,
+    bot_pspace::S=top_pspace,
+    right_vspace::S=left_vspace,
+) where {T,S<:Union{Int,ElementarySpace}}
+    return TensorMap(
+        f,
+        T,
+        _to_space(left_vspace) ⊗ _to_space(top_pspace) ⊗ dual(_to_space(bot_pspace)) ←
+        _to_space(right_vspace),
+    )
+end
+
 """
-    CTMRGEnv(peps::InfinitePEPS{P}; Venv=oneunit(spacetype(P)))
+    CTMRGEnv(
+        [f=randn, ComplexF64], Ds_east::A, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+    ) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
 
-Create a random CTMRG environment from a PEPS tensor. The environment bond dimension
-defaults to one and can be specified using the `Venv` space.
+TODO: docstring.
 """
-function CTMRGEnv(peps::InfinitePEPS{P}; Venv=oneunit(spacetype(P))) where {P}
-    C_type = tensormaptype(spacetype(P), 1, 1, storagetype(P))
-    T_type = tensormaptype(spacetype(P), 3, 1, storagetype(P))
+function CTMRGEnv(
+    Ds_north::A, Ds_east::A, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+    return CTMRGEnv(
+        randn, ComplexF64, Ds_north, Ds_east, chis_north, chis_east, chis_south, chis_west
+    )
+end
+function CTMRGEnv(
+    f, T, Ds_north::A, Ds_east::A, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+    Ds_south = adjoint.(circshift(Ds_north, (-1, 0)))
+    Ds_west = adjoint.(circshift(Ds_east, (0, 1)))
+
+    # do the whole thing
+    st = _spacetype(first(Ds_north))
+    C_type = tensormaptype(st, 1, 1, T)
+    T_type = tensormaptype(st, 3, 1, T)
 
     # First index is direction
-    corners = Array{C_type}(undef, 4, size(peps)...)
-    edges = Array{T_type}(undef, 4, size(peps)...)
-
-    for dir in 1:4, i in 1:size(peps, 1), j in 1:size(peps, 2)
-        @diffset corners[dir, i, j] = TensorMap(randn, scalartype(P), Venv, Venv)
-        @diffset edges[dir, i, j] = TensorMap(
-            randn,
-            scalartype(P),
-            Venv * space(peps[i, j], dir + 1)' * space(peps[i, j], dir + 1),
-            Venv,
+    corners = Array{C_type}(undef, 4, size(Ds_north)...)
+    edges = Array{T_type}(undef, 4, size(Ds_north)...)
+
+    for (r, c) in Iterators.product(axes(Ds_north)...)
+        edges[NORTH, r, c] = _edge_tensor(
+            f,
+            T,
+            chis_north[r, _prev(c, end)],
+            Ds_north[_next(r, end), c],
+            Ds_north[_next(r, end), c],
+            chis_north[r, c],
+        )
+        edges[EAST, r, c] = _edge_tensor(
+            f,
+            T,
+            chis_east[r, c],
+            Ds_east[r, _prev(c, end)],
+            Ds_east[r, _prev(c, end)],
+            chis_east[_next(r, end), c],
+        )
+        edges[SOUTH, r, c] = _edge_tensor(
+            f,
+            T,
+            chis_south[r, c],
+            Ds_south[_prev(r, end), c],
+            Ds_south[_prev(r, end), c],
+            chis_south[r, _prev(c, end)],
         )
+        edges[WEST, r, c] = _edge_tensor(
+            f,
+            T,
+            chis_west[_next(r, end), c],
+            Ds_west[r, _next(c, end)],
+            Ds_west[r, _next(c, end)],
+            chis_west[r, c],
+        )
+
+        corners[NORTHWEST, r, c] = _corner_tensor(
+            f, T, chis_west[_next(r, end), c], chis_north[r, c]
+        )
+        corners[NORTHEAST, r, c] = _corner_tensor(
+            f, T, chis_north[r, _prev(c, end)], chis_east[_next(r, end), c]
+        )
+        corners[SOUTHEAST, r, c] = _corner_tensor(
+            f, T, chis_east[r, c], chis_south[r, _prev(c, end)]
+        )
+        corners[SOUTHWEST, r, c] = _corner_tensor(f, T, chis_south[r, c], chis_west[r, c])
     end
 
-    @diffset corners[:, :, :] ./= norm.(corners[:, :, :])
-    @diffset edges[:, :, :] ./= norm.(edges[:, :, :])
+    corners[:, :, :] ./= norm.(corners[:, :, :])
+    edges[:, :, :] ./= norm.(edges[:, :, :])
 
     return CTMRGEnv(corners, edges)
 end
 
+"""
+    CTMRGEnv(
+        [f=randn, ComplexF64], D_north::S, D_south::S, chi_north::S, chi_east::S, chi_south::S, chi_west::S; unitcell::Tuple{Int,Int}=(1, 1),
+    ) where {S<:Union{Int,ElementarySpace}}
+
+TODO: docstring.
+"""
+function CTMRGEnv(
+    D_north::S,
+    D_south::S,
+    chi_north::S,
+    chi_east::S,
+    chi_south::S,
+    chi_west::S;
+    unitcell::Tuple{Int,Int}=(1, 1),
+) where {S<:Union{Int,ElementarySpace}}
+    return CTMRGEnv(
+        randn,
+        ComplexF64,
+        fill(D_north, unitcell),
+        fill(D_south, unitcell),
+        fill(chi_north, unitcell),
+        fill(chi_east, unitcell),
+        fill(chi_south, unitcell),
+        fill(chi_west, unitcell),
+    )
+end
+function CTMRGEnv(
+    f,
+    T,
+    D_north::S,
+    D_south::S,
+    chi_north::S,
+    chi_east::S,
+    chi_south::S,
+    chi_west::S;
+    unitcell::Tuple{Int,Int}=(1, 1),
+) where {S<:Union{Int,ElementarySpace}}
+    return CTMRGEnv(
+        f,
+        T,
+        fill(D_north, unitcell),
+        fill(D_south, unitcell),
+        fill(chi_north, unitcell),
+        fill(chi_east, unitcell),
+        fill(chi_south, unitcell),
+        fill(chi_west, unitcell),
+    )
+end
+
+"""
+    function CTMRGEnv(
+        [f=randn, T=ComplexF64], peps::InfinitePEPS, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+    ) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+
+TODO: docstring.
+"""
+function CTMRGEnv(
+    peps::InfinitePEPS, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+    Ds_north = map(peps.A) do t
+        return adjoint(space(t, 2))
+    end
+    Ds_east = map(peps.A) do t
+        return adjoint(space(t, 3))
+    end
+    return CTMRGEnv(
+        randn,
+        ComplexF64,
+        Ds_north,
+        Ds_east,
+        _to_space.(chis_north),
+        _to_space.(chis_east),
+        _to_space.(chis_south),
+        _to_space.(chis_west),
+    )
+end
+function CTMRGEnv(
+    f, T, peps::InfinitePEPS, chis_north::A, chis_east::A, chis_south::A, chis_west::A
+) where {A<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+    Ds_north = map(peps.A) do t
+        return adjoint(space(t, 2))
+    end
+    Ds_east = map(peps.A) do t
+        return adjoint(space(t, 3))
+    end
+    return CTMRGEnv(
+        f,
+        T,
+        Ds_north,
+        Ds_east,
+        _to_space.(chis_north),
+        _to_space.(chis_east),
+        _to_space.(chis_south),
+        _to_space.(chis_west),
+    )
+end
+
+"""
+function CTMRGEnv(
+    peps::InfinitePEPS, chi_north::S, chi_east::S=chi_north, chi_south::S=chi_north, chi_west::S=chi_north,
+) where {S<:Union{Int,ElementarySpace}}
+
+TODO: docstring.
+"""
+function CTMRGEnv(
+    peps::InfinitePEPS,
+    chi_north::S,
+    chi_east::S=chi_north,
+    chi_south::S=chi_north,
+    chi_west::S=chi_north,
+) where {S<:Union{Int,ElementarySpace}}
+    return CTMRGEnv(
+        peps,
+        fill(chi_north, size(peps)),
+        fill(chi_east, size(peps)),
+        fill(chi_south, size(peps)),
+        fill(chi_west, size(peps)),
+    )
+end
+function CTMRGEnv(
+    f,
+    T,
+    peps::InfinitePEPS,
+    chi_north::S,
+    chi_east::S=chi_north,
+    chi_south::S=chi_north,
+    chi_west::S=chi_north,
+) where {S<:Union{Int,ElementarySpace}}
+    return CTMRGEnv(
+        f,
+        T,
+        peps,
+        fill(chi_north, size(peps)),
+        fill(chi_east, size(peps)),
+        fill(chi_south, size(peps)),
+        fill(chi_west, size(peps)),
+    )
+end
+@non_differentiable CTMRGEnv(peps::InfinitePEPS, args...)
+
 # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation
 function ChainRulesCore.rrule(::Type{CTMRGEnv}, corners, edges)
     ctmrgenv_pullback(ē) = NoTangent(), ē.corners, ē.edges

test/boundarymps/vumps.jl

@@ -17,7 +17,7 @@ const vumps_alg = VUMPS(; alg_eigsolve=MPSKit.Defaults.alg_eigsolve(; ishermitia
     mps, envs, ϵ = leading_boundary(mps, T, vumps_alg)
     N = abs(sum(expectation_value(mps, T)))
 
-    ctm = leading_boundary(CTMRGEnv(psi; Venv=ComplexSpace(20)), psi, CTMRG(; verbosity=1))
+    ctm = leading_boundary(CTMRGEnv(psi, ComplexSpace(20)), psi, CTMRG(; verbosity=1))
     N´ = abs(norm(psi, ctm))
 
     @test N ≈ N´ atol = 1e-3
@@ -31,7 +31,7 @@ end
     mps, envs, ϵ = leading_boundary(mps, T, vumps_alg)
     N = abs(prod(expectation_value(mps, T)))
 
-    ctm = leading_boundary(CTMRGEnv(psi; Venv=ComplexSpace(20)), psi, CTMRG(; verbosity=1))
+    ctm = leading_boundary(CTMRGEnv(psi, ComplexSpace(20)), psi, CTMRG(; verbosity=1))
     N´ = abs(norm(psi, ctm))
 
     @test N ≈ N´ rtol = 1e-2

test/ctmrg/gaugefix.jl

@@ -43,7 +43,7 @@ end
     psi = _make_symmetric(psi)
 
     Random.seed!(987654321)  # Seed RNG to make random environment consistent
-    ctm = CTMRGEnv(psi; Venv=ctm_space)
+    ctm = CTMRGEnv(psi, ctm_space)
 
     verbosity = 1
     alg = CTMRG(;
@@ -68,7 +68,7 @@ end
     psi = _make_symmetric(psi)
 
     Random.seed!(123456789)  # Seed RNG to make random environment consistent
-    ctm = CTMRGEnv(psi; Venv=ctm_space)
+    ctm = CTMRGEnv(psi, ctm_space)
 
     verbosity = 1
     alg = CTMRG(;

test/ctmrg/gradients.jl

@@ -36,7 +36,7 @@ steps = -0.01:0.005:0.01
     @testset "$alg_rrule" for alg_rrule in gradmodes
         dir = InfinitePEPS(Pspace, Vspace, Vspace)
         psi = InfinitePEPS(Pspace, Vspace, Vspace)
-        env = leading_boundary(CTMRGEnv(psi; Venv=Espace), psi, boundary_alg)
+        env = leading_boundary(CTMRGEnv(psi, Espace), psi, boundary_alg)
         alphas, fs, dfs1, dfs2 = OptimKit.optimtest(
             (psi, env),
             dir;

test/ctmrg/gradparts.jl

@@ -52,7 +52,7 @@ end
 title = "Reverse rules for composite parts of the CTMRG fixed point with spacetype"
 @testset title * "$(Vspaces[i])" for i in eachindex(Pspaces)
     psi = InfinitePEPS(Pspaces[i], Vspaces[i], Vspaces[i])
-    env = CTMRGEnv(psi; Venv=Espaces[i])
+    env = CTMRGEnv(psi, Espaces[i])
 
     @testset "$f" for f in functions
         atol = f == leading_boundary ? sqrt(tol) : tol