Extend CTMRGEnv constructor for nontrivial unit cells

README.md

@@ -59,7 +59,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, Heisenberg_hamiltonian, opt_alg, ctm)
 
 @show result.E # -0.6625...

docs/src/index.md

@@ -44,7 +44,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, Heisenberg_hamiltonian, opt_alg, ctm)
 
 @show result.E # -0.6625...

examples/boundary_mps.jl

@@ -33,7 +33,7 @@ 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, fixedspace=true), CTMRGEnv(peps; Venv=ComplexSpace(20))
+    peps, CTMRG(; verbosity=1, fixedspace=true), CTMRGEnv(peps, ComplexSpace(20))
 )
 N´ = abs(norm(peps, ctm))
 
@@ -56,7 +56,7 @@ mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS())
 N2 = abs(prod(expectation_value(mps2, T2)))
 
 ctm2 = leading_boundary(
-    peps2, CTMRG(; verbosity=1, fixedspace=true), CTMRGEnv(peps2; Venv=ComplexSpace(20))
+    peps2, CTMRG(; verbosity=1, fixedspace=true), CTMRGEnv(peps2, ComplexSpace(20))
 )
 N2´ = abs(norm(peps2, ctm2))
 

examples/heisenberg.jl

@@ -27,6 +27,6 @@ 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), ψ₀, ctmalg)
+env₀ = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctmalg)
 result = fixedpoint(ψ₀, H, alg, env₀)
 @show result.E

src/algorithms/ctmrg.jl

@@ -28,7 +28,7 @@ Contract `state` using CTMRG and return the CTM environment.
 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,44 +8,272 @@ struct CTMRGEnv{C,T}
     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<:ElementarySpace}
+    return TensorMap(f, T, left_vspace ← right_vspace)
+end
+function _corner_tensor(
+    f, ::Type{T}, left_vspace::Int, right_vspace::Int=left_vspace
+) where {T}
+    return _corner_tensor(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<:ElementarySpace}
+    return TensorMap(f, T, left_vspace ⊗ top_pspace ⊗ dual(bot_pspace) ← right_vspace)
+end
+function _edge_tensor(
+    f,
+    ::Type{T},
+    left_vspace::Int,
+    top_pspace::Int,
+    bot_pspace::Int,
+    right_vspace::Int=left_vspace,
+) where {T}
+    return _edge_tensor(
+        f,
+        T,
+        _to_space(left_vspace),
+        _to_space(top_pspace),
+        _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)...)
+    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(peps)...)
-        for dir in 1:4
-            corners[dir, r, c] = TensorMap(randn, scalartype(P), Venv, Venv)
-        end
-        edges[NORTH, _prev(r, end), c] = TensorMap(
-            randn, scalartype(P), Venv * space(peps[r, c], 2)' * space(peps[r, c], 2), Venv
+    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[EAST, r, _next(c, end)] = TensorMap(
-            randn, scalartype(P), Venv * space(peps[r, c], 3)' * space(peps[r, c], 3), Venv
+        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[SOUTH, _next(r, end), c] = TensorMap(
-            randn, scalartype(P), Venv * space(peps[r, c], 4)' * space(peps[r, c], 4), Venv
+        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]
         )
-        edges[WEST, r, _prev(c, end)] = TensorMap(
-            randn, scalartype(P), Venv * space(peps[r, c], 5)' * space(peps[r, c], 5), Venv
+        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
 
     corners[:, :, :] ./= norm.(corners[:, :, :])
     edges[:, :, :] ./= norm.(edges[:, :, :])
 
     return CTMRGEnv(corners, edges)
 end
-@non_differentiable CTMRGEnv(peps::InfinitePEPS)
+
+"""
+    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)

test/boundarymps/vumps.jl

@@ -18,7 +18,7 @@ const vumps_alg = VUMPS(; alg_eigsolve=MPSKit.Defaults.alg_eigsolve(; ishermitia
     N = abs(sum(expectation_value(mps, T)))
 
     ctm = leading_boundary(
-        CTMRGEnv(psi; Venv=ComplexSpace(20)), psi, CTMRG(; verbosity=1, fixedspace=true)
+        CTMRGEnv(psi, ComplexSpace(20)), psi, CTMRG(; verbosity=1, fixedspace=true)
     )
     N´ = abs(norm(psi, ctm))
 
@@ -34,7 +34,7 @@ end
     N = abs(prod(expectation_value(mps, T)))
 
     ctm = leading_boundary(
-        CTMRGEnv(psi; Venv=ComplexSpace(20)), psi, CTMRG(; verbosity=1, fixedspace=true)
+        CTMRGEnv(psi, ComplexSpace(20)), psi, CTMRG(; verbosity=1, fixedspace=true)
     )
     N´ = abs(norm(psi, ctm))
 

test/ctmrg/gaugefix.jl

@@ -42,7 +42,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(;
@@ -67,7 +67,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(;