[WIP] Environment initialization for CTMRG

src/PEPSKit.jl

@@ -71,6 +71,7 @@ include("environments/ctmrg_environments.jl")
 include("environments/vumps_environments.jl")
 include("environments/suweight.jl")
 include("environments/bp_environments.jl")
+include("environments/product_state_environments.jl")
 
 include("algorithms/contractions/ctmrg/types.jl")
 include("algorithms/contractions/ctmrg/expr.jl")
@@ -103,6 +104,7 @@ include("algorithms/ctmrg/simultaneous.jl")
 include("algorithms/ctmrg/sequential.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 include("algorithms/ctmrg/c4v.jl")
+include("algorithms/ctmrg/initialization.jl")
 
 include("algorithms/truncation/truncationschemes.jl")
 include("algorithms/truncation/fullenv_truncation.jl")
@@ -134,13 +136,15 @@ using .Defaults: set_scheduler!
 export set_scheduler!
 export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
+export initialize_ctmrg_environment,
+    RandomInitialization, ProductStateInitialization, ApplicationInitialization
 export corner, edge, setcorner!, setedge!
 export FixedSpaceTruncation, SiteDependentTruncation
 export HalfInfiniteProjector, FullInfiniteProjector
 export C4vCTMRG, C4vEighProjector, C4vQRProjector
 export initialize_random_c4v_env, initialize_singlet_c4v_env
 export LocalOperator, physicalspace
-export product_peps
+export product_peps, bipartite_id
 export reduced_densitymatrix, expectation_value, network_value, cost_function
 export correlator, correlation_length
 export leading_boundary

src/algorithms/ctmrg/initialization.jl

@@ -0,0 +1,56 @@
+abstract type InitializationStyle end
+struct ProductStateInitialization{F} <: InitializationStyle
+    f::F
+    ProductStateInitialization(f::F = ones) where {F} = new{F}(f)
+end
+struct RandomInitialization{F} <: InitializationStyle
+    f::F
+    RandomInitialization(f::F = randn) where {F} = new{F}(f)
+end
+struct ApplicationInitialization{F} <: InitializationStyle
+    f::F
+    ApplicationInitialization(f::F = ones) where {F} = new{F}(f)
+end
+
+# initialize randomly, using same virtual space specification as the CTMRGEnv constructor
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        alg::RandomInitialization,
+        virtual_spaces... = oneunit(spacetype(n)),
+    )
+    return CTMRGEnv(alg.f, elt, n, virtual_spaces...)
+end
+
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        alg::ProductStateInitialization,
+    )
+    env = CTMRGEnv(ProductStateEnv(alg.f, elt, n))
+    return env
+end
+
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        alg::ApplicationInitialization,
+        env0::ProductStateEnv = ProductStateEnv(alg.f, elt, n)
+    )
+    dummy_alg = SimultaneousCTMRG(trunc = (; alg = :notrunc))
+    env, = ctmrg_iteration(n, CTMRGEnv(env0), dummy_alg)
+    return env
+end
+
+function initialize_ctmrg_environment(
+        A::Union{InfiniteSquareNetwork, InfinitePEPS, InfinitePartitionFunction}, args...;
+        kwargs...
+    )
+    return initialize_ctmrg_environment(scalartype(A), A, args...; kwargs...)
+end
+function initialize_ctmrg_environment(
+        elt::Type{<:Number}, A::Union{InfinitePEPS, InfinitePartitionFunction}, args...;
+        kwargs...
+    )
+    return initialize_ctmrg_environment(elt, InfiniteSquareNetwork(A), args...; kwargs...)
+end

src/environments/product_state_environments.jl

@@ -0,0 +1,141 @@
+"""
+$(TYPEDEF)
+
+Tensor product environment for an infinite square network, containing a 4 x rows x cols
+array of tensors, defined for each nearest neighbor bond in the network. 
+
+The product state tensors `p` connect to the network tensors 
+`P` at site `[r,c]` in the unit cell as:
+```
+                    p[1,r-1,c]
+                    |
+    p[4,r,c-1]------P[r,c]------p[2,r,c+1]
+                    |
+                    p[3,r+1,c]
+```
+## Fields
+
+$(TYPEDFIELDS)
+"""
+struct ProductStateEnv{T}
+    "4 x rows x cols array of edge tensors making up a product state environment, where the
+    first dimension specifies the spatial direction"
+    edges::Array{T, 3}
+end
+
+"""
+    ProductStateEnv(
+        [f=randn, T=ComplexF64], Ds_north::A, Ds_east::A
+    ) where {A <: AbstractMatrix{<:ProductSpace}}
+
+Construct a product state environment by specifying matrices of north and east virtual spaces of the
+corresponding [`InfiniteSquareNetwork`](@ref). Each matrix entry corresponds to a site in the unit cell.
+
+Each entry of the `Ds_north` and `Ds_east` matrices corresponds to an effective local space
+of the network, and can be represented as a `ProductSpace` (e.g.
+for the case of a network representing overlaps of PEPSs).
+"""
+function ProductStateEnv(
+        Ds_north::A, Ds_east::A
+    ) where {A <: AbstractMatrix{<:ProductSpace}}
+    return ProductStateEnv(randn, ComplexF64, N, Ds_north, Ds_east)
+end
+function ProductStateEnv(
+        f, T, Ds_north::A, Ds_east::A
+    ) where {A <: AbstractMatrix{<:ProductSpace}}
+    Ds_south = _elementwise_dual.(circshift(Ds_north, (-1, 0)))
+    Ds_west = _elementwise_dual.(circshift(Ds_east, (0, 1)))
+    edges = map(Iterators.product(1:4, axes(Ds_north, 1), axes(Ds_north, 2))) do (dir, r, c)
+        msg = if dir == NORTH
+            f(T, Ds_north[_next(r, end), c])
+        elseif dir == EAST
+            f(T, Ds_east[r, _prev(c, end)])
+        elseif dir == SOUTH
+            f(T, Ds_south[_prev(r, end), c])
+        else # WEST
+            f(T, Ds_west[r, _next(c, end)])
+        end
+        return msg
+    end
+    normalize!.(edges)
+    return ProductStateEnv(edges)
+end
+
+"""
+    ProductStateEnv(
+        [f=randn, T=ComplexF64], D_north::P, D_east::P;
+        unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {P <: ProductSpace}
+
+Construct a product state environment by specifying the north and east virtual spaces of the
+corresponding [`InfiniteSquareNetwork`](@ref). The network unit cell can be specified
+by the `unitcell` keyword argument.
+"""
+function ProductStateEnv(
+        D_north::P, D_east::P;
+        unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {P <: ProductSpace}
+    return ProductStateEnv(randn, ComplexF64, D_north, D_east; unitcell)
+end
+function ProductStateEnv(
+        f, T, D_north::P, D_east::P;
+        unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {P <: ProductSpace}
+    return ProductStateEnv(f, T, N, fill(D_north, unitcell), fill(D_east, unitcell))
+end
+
+"""
+    ProductStateEnv([f=ones, T=ComplexF64], network::InfiniteSquareNetwork)
+
+Construct a product state environment by specifying a corresponding [`InfiniteSquareNetwork`](@ref).
+"""
+function ProductStateEnv(f, T, network::InfiniteSquareNetwork)
+    Ds_north = _north_edge_physical_spaces(network)
+    Ds_east = _east_edge_physical_spaces(network)
+    return ProductStateEnv(f, T, Ds_north, Ds_east)
+end
+function ProductStateEnv(network::InfiniteSquareNetwork)
+    return ProductStateEnv(ones, scalartype(network), network) # TODO: do we want to use a different default function?
+end
+
+function ProductStateEnv(state::Union{InfinitePartitionFunction, InfinitePEPS, InfinitePEPO}, args...; kwargs...)
+    return ProductStateEnv(InfiniteSquareNetwork(state), args...; kwargs...)
+end
+function ProductStateEnv(state::Union{InfinitePEPS, InfinitePEPO}, args...; kwargs...)
+    return ProductStateEnv(InfiniteSquareNetwork(state), args...; kwargs...)
+end
+function ProductStateEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS, InfinitePEPO}, args...; kwargs...)
+    return ProductStateEnv(f, T, InfiniteSquareNetwork(state), args...; kwargs...)
+end
+
+Base.eltype(::Type{ProductStateEnv{T}}) where {T} = T
+Base.size(env::ProductStateEnv, args...) = size(env.edges, args...)
+Base.getindex(env::ProductStateEnv, args...) = Base.getindex(env.edges, args...)
+Base.axes(env::ProductStateEnv, args...) = Base.axes(env.edges, args...)
+Base.eachindex(env::ProductStateEnv) = eachindex(IndexCartesian(), env.edges)
+VectorInterface.scalartype(::Type{ProductStateEnv{T}}) where {T} = scalartype(T)
+TensorKit.spacetype(::Type{ProductStateEnv{T}}) where {T} = spacetype(T)
+
+function eachcoordinate(x::ProductStateEnv)
+    return collect(Iterators.product(axes(x, 2), axes(x, 3)))
+end
+function eachcoordinate(x::ProductStateEnv, dirs)
+    return collect(Iterators.product(dirs, axes(x, 2), axes(x, 3)))
+end
+
+# conversion to CTMRGEnv
+"""
+    CTMRGEnv(prod_env::ProductStateEnv)
+
+Construct a CTMRG environment with a trivial virtual space of bond dimension χ = 1 
+from the product state environment `prod_env`.
+"""
+function CTMRGEnv(prod_env::ProductStateEnv)
+    edges = map(CartesianIndices(prod_env.edges)) do idx
+        return insertleftunit(insertleftunit(prod_env.edges[idx]), 1)
+    end
+    corners = map(CartesianIndices(edges)) do _
+        return TensorKit.id(scalartype(prod_env), oneunit(spacetype(prod_env)))
+    end
+    return CTMRGEnv(corners, edges)
+end

src/utility/util.jl

@@ -208,3 +208,17 @@ function _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N}
     new_axes = (ntuple(i -> axes[biperm[1][i]], N - 1)..., ax)
     return new_axes, biperm
 end
+
+"""
+    bipartite_id([T::Type=Float64], V::ProductSpace{S, 2}) where {S}
+
+Constructs a tensor corresonding to the (permutation of) the identity operator between
+two componens of a bipartite product space.
+"""
+bipartite_id(V::ProductSpace) = bipartite_id(Float64, V)
+function bipartite_id(T::Type, V::ProductSpace)
+    throw(ArgumentError("bipartite_id is only defined for ProductSpace with 2 components, but got $(length(V)) components."))
+end
+function bipartite_id(T::Type, V::ProductSpace{S, 2}) where {S}
+    return permute(isomorphism(T, V[1] ← V[2]'), ((1, 2), ()))
+end

test/ctmrg/initialization.jl

@@ -0,0 +1,86 @@
+using Test
+using TensorKit
+using PEPSKit
+using Random
+
+using MPSKitModels: classical_ising
+
+sd = 12345
+
+# toggle symmetry, but same issue for both
+symmetries = [Z2Irrep, Trivial]
+make_space(::Type{Z2Irrep}, d::Int) = Z2Space(0 => d / 2, 1 => d / 2)
+make_space(::Type{Trivial}, d::Int) = ComplexSpace(d)
+
+d = 2
+D = 4
+χ = 20
+tol = 1.0e-4
+maxiter = 1000
+verbosity = 2
+trunc = truncrank(χ)
+boundary_alg = (; alg = :SimultaneousCTMRG, tol, verbosity, trunc, maxiter)
+
+@testset "CTMRG environment initialization for critical ising with $S symmetry (#255)" for S in symmetries
+    # initialize
+    T = classical_ising(S)
+    O = T[1]
+    n = InfinitePartitionFunction([O O; O O])
+    Venv = make_space(S, χ)
+    P = space(O, 2)
+
+    # random, doesn't converge
+    Random.seed!(sd)
+    env0_rand = initialize_ctmrg_environment(n, RandomInitialization(), Venv)
+    env_rand, info = leading_boundary(env0_rand, n; boundary_alg...)
+    @test_broken info.convergence_error ≤ tol
+
+    # embedded random product state, converges
+    Random.seed!(sd)
+    env0_prod = initialize_ctmrg_environment(n, ProductStateInitialization())
+    env_prod, info = leading_boundary(env0_prod, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+
+    # grown product state, converges
+    Random.seed!(sd)
+    env0_appl = initialize_ctmrg_environment(n, ApplicationInitialization())
+    env_appl, info = leading_boundary(env0_appl, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+
+    # PEPS-specific identity initialization; should throw when used on partition functions
+    Random.seed!(sd)
+    @test_throws ArgumentError env0_prod_id = initialize_ctmrg_environment(n, ProductStateInitialization(bipartite_id))
+end
+
+@testset "CTMRG environment initialization for PEPS with $S symmetry (#255)" for S in symmetries
+    # initialize
+    P = make_space(S, d)
+    Vpeps = make_space(S, D)
+    Venv = make_space(S, χ)
+    peps = InfinitePEPS(P, Vpeps; unitcell = (2, 2))
+    n = InfiniteSquareNetwork(peps)
+
+    # random, converges
+    Random.seed!(sd)
+    env0_rand = initialize_ctmrg_environment(n, RandomInitialization(), Venv)
+    env_rand, info = leading_boundary(env0_rand, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+
+    # embedded random product state, converges
+    Random.seed!(sd)
+    env0_prod = initialize_ctmrg_environment(n, ProductStateInitialization())
+    env_prod, info = leading_boundary(env0_prod, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+
+    # embedded product state as identity from ket to bra, converges
+    Random.seed!(sd)
+    env0_prod = initialize_ctmrg_environment(n, ProductStateInitialization(bipartite_id))
+    env_prod, info = leading_boundary(env0_prod, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+
+    # grown product state, converges
+    Random.seed!(sd)
+    env0_appl = initialize_ctmrg_environment(n, ApplicationInitialization())
+    env_appl, info = leading_boundary(env0_appl, n; boundary_alg...)
+    @test info.convergence_error ≤ tol
+end

test/timeevol/j1j2_finiteT.jl

@@ -10,8 +10,8 @@ using PEPSKit
 bm = [-0.1235, -0.213]
 
 function converge_env(state, χ::Int)
+    env0 = initialize_ctmrg_environment(state, ProductStateInitialization())
     trunc1 = truncrank(χ) & truncerror(; atol = 1.0e-12)
-    env0 = CTMRGEnv(ones, Float64, state, Vect[SU2Irrep](0 => 1))
     env, = leading_boundary(env0, state; alg = :SequentialCTMRG, trunc = trunc1, tol = 1.0e-10)
     return env
 end

test/timeevol/tf_ising_finiteT.jl

@@ -11,7 +11,7 @@ bm_2β = [0.5297, 0.8265]
 
 function converge_env(state, χ::Int)
     trunc1 = truncrank(4) & truncerror(; atol = 1.0e-12)
-    env0 = CTMRGEnv(rand, Float64, state, ℂ^2)
+    env0 = initialize_ctmrg_environment(state, ProductStateInitialization())
     env, = leading_boundary(env0, state; alg = :SequentialCTMRG, trunc = trunc1, tol = 1.0e-10)
     trunc2 = truncrank(χ) & truncerror(; atol = 1.0e-12)
     env, = leading_boundary(env, state; alg = :SequentialCTMRG, trunc = trunc2, tol = 1.0e-10)