Environment initialization for CTMRG + ProductStateEnv

src/PEPSKit.jl

@@ -50,6 +50,7 @@ include("utility/rotations.jl")
 include("utility/hook_pullback.jl")
 include("utility/autoopt.jl")
 include("utility/retractions.jl")
+include("utility/initialization.jl")
 
 include("networks/tensors.jl")
 include("networks/local_sandwich.jl")
@@ -71,6 +72,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 +105,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,6 +137,8 @@ using .Defaults: set_scheduler!
 export set_scheduler!
 export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
+export initialize_ctmrg_environment,
+    RandomInitialization, ProductStateInitialization, ApplicationInitialization, IdentityInitialization
 export corner, edge, setcorner!, setedge!
 export FixedSpaceTruncation, SiteDependentTruncation
 export HalfInfiniteProjector, FullInfiniteProjector

src/algorithms/ctmrg/initialization.jl

@@ -0,0 +1,90 @@
+"""
+    initialize_ctmrg_environment([elt::Type{<:Number},] n::InfiniteSquareNetwork, alg::RandomInitialization, virtual_spaces...)
+
+Initialize a fully random `CTMRGEnv` using the given environment virtual spaces. See
+[`CTMRGEnv`](@ref) for details on the expected format of the virtual spaces.
+"""
+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
+
+"""
+    initialize_ctmrg_environment([elt::Type{<:Number},] n::InfiniteSquareNetwork, alg::RandomInitialization)
+
+Initialize a `CTMRGEnv` corresponding to a product state with trivial virtual spaces and
+corners. The product state edge tensors are initialized as `alg.f(elt, V::ProductSpace)`.
+"""
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        alg::ProductStateInitialization,
+    )
+    env = CTMRGEnv(ProductStateEnv(alg.f, elt, n))
+    return env
+end
+
+"""
+    initialize_ctmrg_environment([elt::Type{<:Number},] n::InfiniteSquareNetwork, alg::RandomInitialization, [env0])
+
+Initialize a `CTMRGEnv` by applying a single untruncated iteration of
+[`SimultaneousCTMRG`](@ref) to a given initial environment. By default, the starting
+environment is chosen as a random product state.
+"""
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        alg::ApplicationInitialization,
+        env0 = ProductStateEnv(alg.f, elt, n)
+    )
+    dummy_alg = SimultaneousCTMRG(trunc = (; alg = :notrunc))
+    env, = ctmrg_iteration(n, CTMRGEnv(env0), dummy_alg)
+    return env
+end
+
+_check_two_layer(::InfiniteSquareNetwork) = false
+_check_two_layer(::InfiniteSquareNetwork{<:PEPSSandwich}) = true
+
+"""
+    initialize_ctmrg_environment([elt::Type{<:Number},] n::InfiniteSquareNetwork, alg::RandomInitialization, [env0])
+
+Initialize a `CTMRGEnv` corresponding to a product state acting as an identity between the
+virtual spaces of a two-layer network, for example
+```
+         ╱      
+┌-----ket----- 
+|    ╱ |        
+|      |   
+|      | ╱      
+└-----bra----- 
+     ╱          
+```
+"""
+function initialize_ctmrg_environment(
+        elt::Type{<:Number},
+        n::InfiniteSquareNetwork,
+        ::IdentityInitialization,
+    )
+    _check_two_layer(n) ||
+        throw(ArgumentError("Identity initialization is only defined for two-layer networks."))
+    bp_env = BPEnv(isomorphism, elt, n)
+    env = CTMRGEnv(bp_env)
+    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/bp_environments.jl

@@ -51,7 +51,7 @@ end
 
 """
     BPEnv(
-        [f=randn, T=ComplexF64], Ds_north::A, Ds_east::A; posdef::Bool = true
+        [f=isomorphism, T=ComplexF64], Ds_north::A, Ds_east::A; posdef::Bool = true
     ) where {A <: AbstractMatrix{<:ProductSpace}}
 
 Construct a BP environment by specifying matrices of north and east virtual spaces of the
@@ -64,15 +64,9 @@ Each entry of the `Ds_north` and `Ds_east` matrices corresponds to an effective
 of the network, and can be represented as a `ProductSpace` (e.g.
 for the case of a network representing overlaps of PEPSs).
 """
-function BPEnv(
-        Ds_north::A, Ds_east::A; posdef::Bool = true
-    ) where {A <: AbstractMatrix{<:ProductSpace}}
-    return BPEnv(randn, ComplexF64, N, Ds_north, Ds_east; posdef)
-end
 function BPEnv(
         f, T, Ds_north::A, Ds_east::A; posdef::Bool = true
     ) where {A <: AbstractMatrix{<:ProductSpace}}
-    # no recursive broadcasting?
     Ds_south = _elementwise_dual.(circshift(Ds_north, (-1, 0)))
     Ds_west = _elementwise_dual.(circshift(Ds_east, (0, 1)))
     messages = map(Iterators.product(1:4, axes(Ds_north, 1), axes(Ds_north, 2))) do (dir, r, c)
@@ -93,32 +87,31 @@ function BPEnv(
     normalize!.(messages)
     return BPEnv(messages)
 end
+function BPEnv(
+        D_north::P, args...; kwargs...
+    ) where {P <: Union{Matrix{ProductSpace}, ProductSpace}}
+    return BPEnv(isomorphism, ComplexF64, D_north, args...; kwargs...)
+end
 
 """
     BPEnv(
-        [f=randn, T=ComplexF64], D_north::P, D_east::P;
+        [f=isomorphism, T=ComplexF64], D_north::P, D_east::P;
         unitcell::Tuple{Int, Int} = (1, 1), posdef::Bool = true
     ) where {P <: ProductSpace}
 
 Construct a BP 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 BPEnv(
-        D_north::P, D_east::P;
-        unitcell::Tuple{Int, Int} = (1, 1), posdef::Bool = true
-    ) where {P <: ProductSpace}
-    return BPEnv(randn, ComplexF64, D_north, D_east; unitcell, posdef)
-end
 function BPEnv(
         f, T, D_north::P, D_east::P;
         unitcell::Tuple{Int, Int} = (1, 1), posdef::Bool = true
     ) where {P <: ProductSpace}
-    return BPEnv(f, T, N, fill(D_north, unitcell), fill(D_east, unitcell); posdef)
+    return BPEnv(f, T, N, _fill_edge_physical_spaces(D_north, D_east; unitcell)...; posdef)
 end
 
 """
-    BPEnv([f=randn, T=ComplexF64], network::InfiniteSquareNetwork; posdef::Bool = true)
+    BPEnv([f=isomorphism, T=ComplexF64], network::InfiniteSquareNetwork; posdef::Bool = true)
 
 Construct a BP environment by specifying a corresponding [`InfiniteSquareNetwork`](@ref).
 """
@@ -127,17 +120,8 @@ function BPEnv(f, T, network::InfiniteSquareNetwork; posdef::Bool = true)
     Ds_east = _east_edge_physical_spaces(network)
     return BPEnv(f, T, Ds_north, Ds_east; posdef)
 end
-function BPEnv(network::InfiniteSquareNetwork; posdef::Bool = true)
-    return BPEnv(randn, scalartype(network), network; posdef)
-end
-
-function BPEnv(state::InfinitePartitionFunction, args...; kwargs...)
-    return BPEnv(InfiniteSquareNetwork(state), args...; kwargs...)
-end
-function BPEnv(state::Union{InfinitePEPS, InfinitePEPO}, args...; kwargs...)
-    bp_env = BPEnv(InfiniteSquareNetwork(state), args...; kwargs...)
-    TensorKit.id!.(bp_env.messages)
-    return bp_env
+function BPEnv(network::Union{InfiniteSquareNetwork, InfinitePartitionFunction, InfinitePEPS, InfinitePEPO}, args...; kwargs...)
+    return BPEnv(isomorphism, scalartype(network), network, args...; kwargs...)
 end
 function BPEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS, InfinitePEPO}, args...; kwargs...)
     return BPEnv(f, T, InfiniteSquareNetwork(state), args...; kwargs...)

src/environments/ctmrg_environments.jl

@@ -230,6 +230,9 @@ function CTMRGEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS}, a
     return CTMRGEnv(f, T, InfiniteSquareNetwork(state), args...)
 end
 
+# copy-like constructor
+CTMRGEnv(env::CTMRGEnv) = CTMRGEnv(env.corners, env.edges)
+
 @non_differentiable CTMRGEnv(state::Union{InfinitePartitionFunction, InfinitePEPS}, args...)
 
 # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation

src/environments/product_state_environments.jl

@@ -0,0 +1,117 @@
+"""
+$(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}
+    ProductStateEnv{T}(edges::Array{T, 3}) where {T} = new{T}(edges)
+    function ProductStateEnv(edges::Array{T, 3}) where {T}
+        foreach(Iterators.product(axes(edges)[2:3]...)) do (d, w)
+            codomain(edges[NORTH, d, w]) == _elementwise_dual(codomain(edges[SOUTH, _prev(d, end), w])) ||
+                throw(
+                SpaceMismatch("North virtual space at site $((d, w)) does not match: $(space(edges[NORTH, d, w])) vs $(space(edges[SOUTH, _prev(d, end), w])).")
+            )
+            codomain(edges[EAST, d, w]) == _elementwise_dual(codomain(edges[WEST, d, _next(w, end)])) ||
+                throw(SpaceMismatch("East virtual space at site $((d, w)) does not match: $(space(edges[EAST, d, w])) vs $(space(edges[WEST, d, _next(w, end)]))."))
+        end
+        foreach(Iterators.product(axes(edges)...)) do (dir, d, w)
+            dim(space(edges[dir, d, w])) > 0 || @warn "no fusion channels for edge ($dir, $d, $w)"
+        end
+        return new{T}(edges)
+    end
+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(
+        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
+function ProductStateEnv(Ds_north::A, args...; kwargs...) where {A <: AbstractMatrix{<:VectorSpace}}
+    return ProductStateEnv(randn, ComplexF64, Ds_north, args...; kwargs...)
+end
+
+"""
+    ProductStateEnv([f=randn, 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::Union{InfiniteSquareNetwork, InfinitePartitionFunction, InfinitePEPS})
+    return ProductStateEnv(randn, scalartype(network), network)
+end
+function ProductStateEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS}, args...)
+    return ProductStateEnv(f, T, InfiniteSquareNetwork(state), args...)
+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.eachindex(index_style, env::ProductStateEnv) = eachindex(index_style, env.edges)
+VectorInterface.scalartype(::Type{ProductStateEnv{T}}) where {T} = scalartype(T)
+TensorKit.storagetype(::Type{ProductStateEnv{T}}) where {T} = storagetype(T)
+TensorKit.spacetype(::Type{ProductStateEnv{T}}) where {T} = spacetype(T)
+
+# 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(eachindex(IndexCartesian(), prod_env)) do idx
+        return insertleftunit(insertleftunit(prod_env[idx]), 1)
+    end
+    corners = map(eachindex(IndexCartesian(), prod_env)) do _
+        return TensorKit.id(storagetype(prod_env), oneunit(spacetype(prod_env)))
+    end
+    return CTMRGEnv(corners, edges)
+end