Remove support for integer space specifiers in state and environment constructors (#266)

* Remove for integer space specifiers, simplify(?) constructors

* Typo

* Typo

* Switch optional to kwarg

* Forgot about tests

* Do nothing if already properly filled

* Typo

* Forgot the kwargs

* Missed a few more

* Fix(?) ambiguity

* If it's not used, probably don't need it

* Update src/environments/ctmrg_environments.jl

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: leburgel

README.md

@@ -56,7 +56,7 @@ ctmrg_tol = 1e-10
 grad_tol = 1e-4
 
 # initialize a PEPS and CTMRG environment
-peps₀ = InfinitePEPS(2, D)
+peps₀ = InfinitePEPS(ComplexSpace(2), ComplexSpace(D))
 env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χ)), peps₀; tol=ctmrg_tol)
 
 # ground state search

docs/src/index.md

@@ -41,7 +41,7 @@ ctmrg_tol = 1e-10
 grad_tol = 1e-4
 
 # initialize a PEPS and CTMRG environment
-peps₀ = InfinitePEPS(2, D)
+peps₀ = InfinitePEPS(ComplexSpace(2), ComplexSpace(D))
 env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χ)), peps₀; tol=ctmrg_tol)
 
 # ground state search

src/environments/ctmrg_environments.jl

@@ -31,37 +31,22 @@ struct CTMRGEnv{C, T}
     edges::Array{T, 3}
 end
 
-const ProductSpaceLike{N} = Union{
-    NTuple{N, Int}, NTuple{N, <:ElementarySpace}, ProductSpace{<:ElementarySpace, N},
-}
-const SpaceLike = Union{ElementarySpaceLike, ProductSpaceLike}
-
-_elementwise_dual(s::NTuple{N, <:ElementarySpace}) where {N} = dual.(s)
-
-_spacetype(::Int) = ComplexSpace
-_spacetype(::S) where {S <: ElementarySpace} = S
-_spacetype(s::ProductSpaceLike) = _spacetype(first(s))
-
-_to_space(χ::Int) = ℂ^χ
-_to_space(χ::ElementarySpace) = χ
-_to_space(χ::ProductSpaceLike) = prod(_to_space, χ)
-
 function _corner_tensor(
         f, ::Type{T}, left_vspace::S, right_vspace::S = left_vspace
-    ) where {T, S <: ElementarySpaceLike}
-    return f(T, _to_space(left_vspace) ← _to_space(right_vspace))
+    ) where {T, S <: ElementarySpace}
+    return f(T, left_vspace ← right_vspace)
 end
 
 function _edge_tensor(
         f, ::Type{T}, left_vspace::S, pspaces::P, right_vspace::S = left_vspace
-    ) where {T, S <: ElementarySpaceLike, P <: ProductSpaceLike}
-    return f(T, _to_space(left_vspace) ⊗ _to_space(pspaces), _to_space(right_vspace))
+    ) where {T, S <: ElementarySpace, P <: ProductSpace}
+    return f(T, left_vspace ⊗ pspaces, right_vspace)
 end
 
 """
     CTMRGEnv(
         [f=randn, T=ComplexF64], Ds_north::A, Ds_east::A, chis_north::B, [chis_east::B], [chis_south::B], [chis_west::B]
-    ) where {A<:AbstractMatrix{<:SpaceLike}, B<:AbstractMatrix{<:ElementarySpaceLike}}
+    ) where {A<:AbstractMatrix{<:VectorSpace}, B<:AbstractMatrix{<:ElementarySpace}}
 
 Construct a CTMRG environment by specifying matrices of north and east virtual spaces of the
 corresponding partition function and the north, east, south and west virtual spaces of the
@@ -76,32 +61,27 @@ of the corresponding edge tensor for each direction. Specifically, for a given s
 `chis_west[r, c]` corresponds to the north space of the west edge tensor.
 
 Each entry of the `Ds_north` and `Ds_east` matrices corresponds to an effective local space
-of the partition function, represented as a tuple of elementary spaces encoding a product
-space. This can either contain a single elementary space for the case of a partition
-function defined in terms of local rank-4 tensors, or a tuple of elementary spaces
-representing a product space for the case of a partition function representing overlaps of
-PEPSs and PEPOs.
+of the partition function, and can be represented as an `ElementarySpace` (e.g. for the case
+of a partition function defined in terms of local rank-4 tensors) or a `ProductSpace` (e.g.
+for the case of a network representing overlaps of PEPSs and PEPOs).
 """
 function CTMRGEnv(
-        Ds_north::A, Ds_east::A,
-        chis_north::B, chis_east::B = chis_north, chis_south::B = chis_north, chis_west::B = chis_north,
-    ) where {A <: AbstractMatrix{<:ProductSpaceLike}, B <: AbstractMatrix{<:ElementarySpaceLike}}
-    return CTMRGEnv(
-        randn, ComplexF64, N, 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::B, chis_east::B = chis_north, chis_south::B = chis_north, chis_west::B = chis_north,
-    ) where {A <: AbstractMatrix{<:ProductSpaceLike}, B <: AbstractMatrix{<:ElementarySpaceLike}}
+        f, T, Ds_north::A, Ds_east::A, chis_north::B, chis_east::B = chis_north,
+        chis_south::B = chis_north, chis_west::B = chis_north,
+    ) where {
+        A <: AbstractMatrix{<:ProductSpace}, B <: AbstractMatrix{<:ElementarySpace},
+    }
+    # check all of the sizes
+    size(Ds_north) == size(Ds_east) == size(chis_north) == size(chis_east) ==
+        size(chis_south) == size(chis_west) || throw(ArgumentError("Input spaces should have equal sizes."))
+
     # no recursive broadcasting?
     Ds_south = _elementwise_dual.(circshift(Ds_north, (-1, 0)))
     Ds_west = _elementwise_dual.(circshift(Ds_east, (0, 1)))
 
     # do the whole thing
     N = length(first(Ds_north))
-    st = _spacetype(first(Ds_north))
+    st = spacetype(first(Ds_north))
     C_type = tensormaptype(st, 1, 1, T)
     T_type = tensormaptype(st, N + 1, 1, T)
 
@@ -140,12 +120,37 @@ function CTMRGEnv(
     edges[:, :, :] ./= norm.(edges[:, :, :])
     return CTMRGEnv(corners, edges)
 end
+function CTMRGEnv(D_north::P, args...; kwargs...) where {P <: Union{Matrix{VectorSpace}, VectorSpace}}
+    return CTMRGEnv(randn, ComplexF64, D_north, args...; kwargs...)
+end
+
+# expand physical edge spaces to unit cell size
+function _fill_edge_physical_spaces(
+        D_north::S, D_east::S = D_north; unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {S <: VectorSpace}
+    return fill(ProductSpace(D_north), unitcell), fill(ProductSpace(D_east), unitcell)
+end
+
+# expand virtual environment spaces to unit cell size
+function _fill_environment_virtual_spaces(
+        chis_north::S, chis_east::S = chis_north, chis_south::S = chis_north, chis_west::S = chis_north;
+        unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {S <: ElementarySpace}
+    return fill(chis_north, unitcell), fill(chis_east, unitcell), fill(chis_south, unitcell), fill(chis_west, unitcell)
+end
+function _fill_environment_virtual_spaces(
+        chis_north::M, chis_east::M = chis_north, chis_south::M = chis_north, chis_west::M = chis_north;
+        unitcell::Tuple{Int, Int} = (1, 1)
+    ) where {M <: AbstractMatrix{<:ElementarySpace}}
+    @assert size(chis_north) == size(chis_east) == size(chis_south) == size(chis_west) == unitcell "Incompatible size"
+    return chis_north, chis_east, chis_south, chis_west
+end
 
 """
     CTMRGEnv(
         [f=randn, T=ComplexF64], D_north::P, D_east::P, chi_north::S, [chi_east::S], [chi_south::S], [chi_west::S];
         unitcell::Tuple{Int,Int}=(1, 1),
-    ) where {P<:ProductSpaceLike,S<:ElementarySpaceLike}
+    ) where {P<:VectorSpace,S<:ElementarySpace}
 
 Construct a CTMRG environment by specifying the north and east virtual spaces of the
 corresponding [`InfiniteSquareNetwork`](@ref) and the north, east, south and west virtual
@@ -156,122 +161,55 @@ same.
 The environment virtual spaces for each site correspond to virtual space of the
 corresponding edge tensor for each direction.
 """
-function CTMRGEnv(
-        D_north::P, D_east::P,
-        chi_north::S, chi_east::S = chi_north, chi_south::S = chi_north, chi_west::S = chi_north;
-        unitcell::Tuple{Int, Int} = (1, 1),
-    ) where {P <: ProductSpaceLike, S <: ElementarySpaceLike}
-    return CTMRGEnv(
-        randn, ComplexF64,
-        fill(D_north, unitcell), fill(D_east, unitcell),
-        fill(chi_north, unitcell), fill(chi_east, unitcell),
-        fill(chi_south, unitcell), fill(chi_west, unitcell),
-    )
-end
-function CTMRGEnv(
-        f, T,
-        D_north::P, D_east::P,
-        chi_north::S, chi_east::S = chi_north,
-        chi_south::S = chi_north, chi_west::S = chi_north;
-        unitcell::Tuple{Int, Int} = (1, 1),
-    ) where {P <: ProductSpaceLike, S <: ElementarySpaceLike}
-    return CTMRGEnv(
-        f, T, N,
-        fill(D_north, unitcell), fill(D_east, unitcell),
-        fill(chi_north, unitcell), fill(chi_east, unitcell),
-        fill(chi_south, unitcell), fill(chi_west, unitcell),
-    )
-end
-
-"""
-    CTMRGEnv(
-        [f=randn, T=ComplexF64], network::InfiniteSquareNetwork, chis_north::A, [chis_east::A], [chis_south::A], [chis_west::A]
-    ) where {A<:AbstractMatrix{<:ElementarySpaceLike}}}
-
-Construct a CTMRG environment by specifying a corresponding [`InfiniteSquareNetwork`](@ref), and the
-north, east, south and west virtual spaces of the environment as matrices. Each respective
-matrix entry corresponds to a site in the unit cell. By default, the virtual spaces for all
-directions are taken to be the same.
-
-The environment virtual spaces for each site correspond to the north or east virtual space
-of the corresponding edge tensor for each direction. Specifically, for a given site
-`(r, c)`, `chis_north[r, c]` corresponds to the east space of the north edge tensor,
-`chis_east[r, c]` corresponds to the north space of the east edge tensor,
-`chis_south[r, c]` corresponds to the east space of the south edge tensor, and
-`chis_west[r, c]` corresponds to the north space of the west edge tensor.
-"""
-function CTMRGEnv(
-        network::InfiniteSquareNetwork,
-        chis_north::A, chis_east::A = chis_north, chis_south::A = chis_north, chis_west::A = chis_north,
-    ) where {A <: AbstractMatrix{<:ElementarySpaceLike}}
-    Ds_north = _north_env_spaces(network)
-    Ds_east = _east_env_spaces(network)
-    return CTMRGEnv(
-        randn, scalartype(network),
-        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,
-        network::InfiniteSquareNetwork,
-        chis_north::A, chis_east::A = chis_north, chis_south::A = chis_north, chis_west::A = chis_north,
-    ) where {A <: AbstractMatrix{<:ElementarySpaceLike}}
-    Ds_north = _north_env_spaces(network)
-    Ds_east = _east_env_spaces(network)
+        D_north::S, D_east::S, virtual_spaces...; unitcell::Tuple{Int, Int} = (1, 1),
+    ) where {S <: VectorSpace}
     return CTMRGEnv(
         f, T,
-        Ds_north, Ds_east,
-        _to_space.(chis_north), _to_space.(chis_east), _to_space.(chis_south), _to_space.(chis_west),
+        _fill_edge_physical_spaces(D_north, D_east; unitcell)...,
+        _fill_environment_virtual_spaces(virtual_spaces...; unitcell)...,
     )
 end
 
-function _north_env_spaces(network::InfiniteSquareNetwork)
+# get edge physical spaces from network
+function _north_edge_physical_spaces(network::InfiniteSquareNetwork)
     return map(ProductSpace ∘ _elementwise_dual ∘ north_virtualspace, unitcell(network))
 end
-function _east_env_spaces(network::InfiniteSquareNetwork)
+function _east_edge_physical_spaces(network::InfiniteSquareNetwork)
     return map(ProductSpace ∘ _elementwise_dual ∘ east_virtualspace, unitcell(network))
 end
 
 """
     CTMRGEnv(
-        [f=randn, T=ComplexF64,] network::InfiniteSquareNetwork, chi_north::S, [chi_east::S], [chi_south::S], [chi_west::S],
-    ) where {S<:ElementarySpaceLike}
+        [f=randn, T=ComplexF64], network::InfiniteSquareNetwork, chis_north::A, [chis_east::A], [chis_south::A], [chis_west::A]
+    ) where {A<:Union{AbstractMatrix{<:ElementarySpace}, ElementarySpace}}
 
 Construct a CTMRG environment by specifying a corresponding
 [`InfiniteSquareNetwork`](@ref), and the north, east, south and west virtual spaces of the
-environment. By default, the virtual spaces for all directions are taken to be the same.
+environment. The virtual spaces can either be specified as matrices of `ElementarySpace`s,
+or as individual `ElementarySpace`s which are then filled to match the size of the unit
+cell. Each respective matrix entry corresponds to a site in the unit cell. By default, the
+virtual spaces for all directions are taken to be the same.
 
-The environment virtual spaces for each site correspond to virtual space of the
-corresponding edge tensor for each direction.
+The environment virtual spaces for each site correspond to the north or east virtual space
+of the corresponding edge tensor for each direction. Specifically, for a given site
+`(r, c)`, `chis_north[r, c]` corresponds to the east space of the north edge tensor,
+`chis_east[r, c]` corresponds to the north space of the east edge tensor,
+`chis_south[r, c]` corresponds to the east space of the south edge tensor, and
+`chis_west[r, c]` corresponds to the north space of the west edge tensor.
 """
-function CTMRGEnv(
-        network::InfiniteSquareNetwork,
-        chi_north::S, chi_east::S = chi_north, chi_south::S = chi_north, chi_west::S = chi_north,
-    ) where {S <: ElementarySpaceLike}
-    return CTMRGEnv(
-        network,
-        fill(chi_north, size(network)), fill(chi_east, size(network)),
-        fill(chi_south, size(network)), fill(chi_west, size(network)),
-    )
+function CTMRGEnv(f, T, network::InfiniteSquareNetwork, virtual_spaces...)
+    Ds_north = _north_edge_physical_spaces(network)
+    Ds_east = _east_edge_physical_spaces(network)
+    virtual_spaces = _fill_environment_virtual_spaces(virtual_spaces...; unitcell = size(network))
+    return CTMRGEnv(f, T, Ds_north, Ds_east, virtual_spaces...)
 end
-function CTMRGEnv(
-        f, T,
-        network::InfiniteSquareNetwork,
-        chi_north::S, chi_east::S = chi_north, chi_south::S = chi_north, chi_west::S = chi_north,
-    ) where {S <: ElementarySpaceLike}
-    return CTMRGEnv(
-        f, T,
-        network,
-        fill(chi_north, size(network)), fill(chi_east, size(network)),
-        fill(chi_south, size(network)), fill(chi_west, size(network)),
-    )
+function CTMRGEnv(network::Union{InfiniteSquareNetwork, InfinitePartitionFunction, InfinitePEPS}, virtual_spaces...)
+    return CTMRGEnv(randn, scalartype(network), network, virtual_spaces...)
 end
 
 # allow constructing environments for implicitly defined contractible networks
-function CTMRGEnv(state::Union{InfinitePartitionFunction, InfinitePEPS}, args...)
-    return CTMRGEnv(InfiniteSquareNetwork(state), args...)
-end
 function CTMRGEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS}, args...)
     return CTMRGEnv(f, T, InfiniteSquareNetwork(state), args...)
 end

src/states/infinitepartitionfunction.jl

@@ -49,14 +49,9 @@ Create an `InfinitePartitionFunction` by specifying the physical, north virtual
 of the PEPS tensor at each site in the unit cell as a matrix. Each individual space can be
 specified as either an `Int` or an `ElementarySpace`.
 """
-function InfinitePartitionFunction(
-        Nspaces::A, Espaces::A
-    ) where {A <: AbstractMatrix{<:ElementarySpaceLike}}
-    return InfinitePartitionFunction(randn, ComplexF64, Nspaces, Espaces)
-end
 function InfinitePartitionFunction(
         f, T, Nspaces::M, Espaces::M = Nspaces
-    ) where {M <: AbstractMatrix{<:ElementarySpaceLike}}
+    ) where {M <: AbstractMatrix{<:ElementarySpace}}
     size(Nspaces) == size(Espaces) ||
         throw(ArgumentError("Input spaces should have equal sizes."))
 
@@ -69,6 +64,9 @@ function InfinitePartitionFunction(
 
     return InfinitePartitionFunction(A)
 end
+function InfinitePartitionFunction(Nspaces::A, args...) where {A <: Union{AbstractMatrix{<:ElementarySpace}, ElementarySpace}}
+    return InfinitePartitionFunction(randn, ComplexF64, Nspaces, args...)
+end
 
 """
     InfinitePartitionFunction(A; unitcell=(1, 1))
@@ -99,22 +97,15 @@ end
 """
     InfinitePartitionFunction(
         [f=randn, T=ComplexF64,] Pspace::S, Nspace::S, [Espace::S]; unitcell=(1,1)
-    ) where {S<:ElementarySpaceLike}
+    ) where {S<:ElementarySpace}
 
 Create an InfinitePartitionFunction by specifying its physical, north and east spaces and unit cell.
 Spaces can be specified either via `Int` or via `ElementarySpace`.
 """
-function InfinitePartitionFunction(
-        Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1)
-    ) where {S <: ElementarySpaceLike}
-    return InfinitePartitionFunction(
-        randn, ComplexF64, fill(Nspace, unitcell), fill(Espace, unitcell)
-    )
-end
 function InfinitePartitionFunction(
         f, T, Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1)
-    ) where {S <: ElementarySpaceLike}
-    return InfinitePartitionFunction(f, T, fill(Nspace, unitcell), fill(Espace, unitcell))
+    ) where {S <: ElementarySpace}
+    return InfinitePartitionFunction(f, T, _fill_state_virtual_spaces(Nspace, Espace; unitcell)...)
 end
 
 ## Unit cell interface