Add unified contractible network interface (#131)

* Add contractible network interface

* Rename + fix typo

* Some variable renames, fix typo in `_codomain_space`, fix bug in PEPOSandwich map

* Fix typo and format

* Apply suggestions from code review

Co-authored-by: Lukas Devos <[email protected]>

* Fix merge conflict

* Update virtual space handling in sparse environments

* Format

* Update physicalspace handling for `InfinitePEPO`

* Rename `value -> network_value`

* Avoid type piracy

* Format

* Implement interface for trivial `PartitionFunctionTensor` local sandwich

* Remove physical network supertype, implement interface for each type individually

* Fix typo

* Improve handling of local sandwich adjoints, mark some unused things for potential deletion

* Format

* Fix typos

* Remove type restrictions where possible

* Format

* Export `InfiniteSquareNetwork`

* Remove `leading_boundary` method with trivial default environment

* Remove unused chainrules

* Apply suggestions from code review

Co-authored-by: Lukas Devos <[email protected]>

* Update src/networks/infinitesquarenetwork.jl

Co-authored-by: Lukas Devos <[email protected]>

* Update src/networks/infinitesquarenetwork.jl

Co-authored-by: Lukas Devos <[email protected]>

* Apply more suggestions

* Implement suggested physical space handling

* Use flat tuples for all local sandwiches

* Update transfer PEPO constructor to make flat tuples

* Forgotten sandwich

* Re-un-tuple partition function `InfiniteSquareNetwork`

---------

Co-authored-by: Lukas Devos <[email protected]>

src/PEPSKit.jl

@@ -23,6 +23,7 @@ include("utility/hook_pullback.jl")
 include("utility/autoopt.jl")
 
 include("networks/tensors.jl")
+include("networks/local_sandwich.jl")
 include("networks/infinitesquarenetwork.jl")
 
 include("states/infinitepeps.jl")
@@ -207,14 +208,15 @@ export SVDAdjoint, IterSVD
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
 export FixedSpaceTruncation, HalfInfiniteProjector, FullInfiniteProjector
 export LocalOperator
-export expectation_value, cost_function, product_peps, correlation_length
+export expectation_value, cost_function, product_peps, correlation_length, network_value
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver, EigSolver
 export fixedpoint
 
 export absorb_weight
 export su_iter, simpleupdate, SimpleUpdate
 
+export InfiniteSquareNetwork
 export InfinitePartitionFunction
 export InfinitePEPS, InfiniteTransferPEPS
 export SUWeight, InfiniteWeightPEPS

src/algorithms/contractions/vumps_contractions.jl

@@ -319,10 +319,10 @@ end
 # PEPS derivative
 
 # sandwich with the bottom dropped out...
-const ∂PEPOSandwich{N,T<:PEPSTensor,P<:PEPOTensor} = Tuple{T,Tuple{Vararg{P,N}}}
+const ∂PEPOSandwich{N,T<:PEPSTensor,P<:PEPOTensor} = Tuple{T,Vararg{P,N}}
 ket(p::∂PEPOSandwich) = p[1]
-pepo(p::∂PEPOSandwich) = p[2]
-pepo(p::∂PEPOSandwich, i::Int) = p[2][i]
+pepo(p::∂PEPOSandwich) = p[2:end]
+pepo(p::∂PEPOSandwich, i::Int) = p[1 + i]
 
 # specialize simple case
 function ∂peps(
@@ -358,7 +358,7 @@ end
     GR_e = _pepo_rightenv_expr(:GR, :E, H)
     ket_e = _pepo_pepstensor_expr(:(ket(O)), :top, 1)
     pepo_es = map(1:H) do h
-        return _pepo_pepotensor_expr(:(pepo(O)[$h]), h)
+        return _pepo_pepotensor_expr(:(pepo(O, $h)), h)
     end
 
     rhs = Expr(:call, :*, AC_e, Expr(:call, :conj, ĀC_e), GL_e, GR_e, ket_e, pepo_es...)

src/algorithms/ctmrg/ctmrg.jl

@@ -27,52 +27,52 @@ Different levels of output information are printed depending on `alg.verbosity`,
 suppresses all output, `1` only prints warnings, `2` gives information at the start and
 end, and `3` prints information every iteration.
 """
-function MPSKit.leading_boundary(state, alg::CTMRGAlgorithm)
-    return MPSKit.leading_boundary(CTMRGEnv(state, oneunit(spacetype(state))), state, alg)
-end
-function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
+function MPSKit.leading_boundary(
+    envinit, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm
+)
     CS = map(x -> tsvd(x)[2], envinit.corners)
     TS = map(x -> tsvd(x)[2], envinit.edges)
 
-    η = one(real(scalartype(state)))
+    η = one(real(scalartype(network)))
     env = deepcopy(envinit)
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
 
     return LoggingExtras.withlevel(; alg.verbosity) do
-        ctmrg_loginit!(log, η, state, envinit)
+        ctmrg_loginit!(log, η, network, envinit)
         local info
         for iter in 1:(alg.maxiter)
-            env, info = ctmrg_iteration(state, env, alg)  # Grow and renormalize in all 4 directions
+            env, info = ctmrg_iteration(network, env, alg)  # Grow and renormalize in all 4 directions
             η, CS, TS = calc_convergence(env, CS, TS)
 
             if η ≤ alg.tol && iter ≥ alg.miniter
-                ctmrg_logfinish!(log, iter, η, state, env)
+                ctmrg_logfinish!(log, iter, η, network, env)
                 break
             end
             if iter == alg.maxiter
-                ctmrg_logcancel!(log, iter, η, state, env)
+                ctmrg_logcancel!(log, iter, η, network, env)
             else
-                ctmrg_logiter!(log, iter, η, state, env)
+                ctmrg_logiter!(log, iter, η, network, env)
             end
         end
         return env, info
     end
 end
-
-# network-specific objective functions
-ctmrg_objective(state::InfinitePEPS, env::CTMRGEnv) = norm(state, env)
-ctmrg_objective(state::InfinitePartitionFunction, env::CTMRGEnv) = value(state, env)
+function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
+    return MPSKit.leading_boundary(envinit, InfiniteSquareNetwork(state), alg)
+end
 
 # custom CTMRG logging
-ctmrg_loginit!(log, η, state, env) = @infov 2 loginit!(log, η, ctmrg_objective(state, env))
-function ctmrg_logiter!(log, iter, η, state, env)
-    @infov 3 logiter!(log, iter, η, ctmrg_objective(state, env))
+function ctmrg_loginit!(log, η, network, env)
+    @infov 2 loginit!(log, η, network_value(network, env))
+end
+function ctmrg_logiter!(log, iter, η, network, env)
+    @infov 3 logiter!(log, iter, η, network_value(network, env))
 end
-function ctmrg_logfinish!(log, iter, η, state, env)
-    @infov 2 logfinish!(log, iter, η, ctmrg_objective(state, env))
+function ctmrg_logfinish!(log, iter, η, network, env)
+    @infov 2 logfinish!(log, iter, η, network_value(network, env))
 end
-function ctmrg_logcancel!(log, iter, η, state, env)
-    @warnv 1 logcancel!(log, iter, η, ctmrg_objective(state, env))
+function ctmrg_logcancel!(log, iter, η, network, env)
+    @warnv 1 logcancel!(log, iter, η, network_value(network, env))
 end
 
 @non_differentiable ctmrg_loginit!(args...)

src/algorithms/ctmrg/sequential.jl

@@ -30,88 +30,80 @@ function SequentialCTMRG(;
     )
 end
 
-function ctmrg_iteration(state, env::CTMRGEnv, alg::SequentialCTMRG)
-    truncation_error = zero(real(scalartype(state)))
-    condition_number = zero(real(scalartype(state)))
+function ctmrg_iteration(network, env::CTMRGEnv, alg::SequentialCTMRG)
+    truncation_error = zero(real(scalartype(network)))
+    condition_number = zero(real(scalartype(network)))
     for _ in 1:4 # rotate
-        for col in 1:size(state, 2) # left move column-wise
-            projectors, info = sequential_projectors(col, state, env, alg.projector_alg)
-            env = renormalize_sequentially(col, projectors, state, env)
+        for col in 1:size(network, 2) # left move column-wise
+            projectors, info = sequential_projectors(col, network, env, alg.projector_alg)
+            env = renormalize_sequentially(col, projectors, network, env)
             truncation_error = max(truncation_error, info.truncation_error)
             condition_number = max(condition_number, info.condition_number)
         end
-        state = rotate_north(state, EAST)
+        network = rotate_north(network, EAST)
         env = rotate_north(env, EAST)
     end
 
     return env, (; truncation_error, condition_number)
 end
 
 """
-    sequential_projectors(col::Int, state::InfinitePEPS, env::CTMRGEnv, alg::ProjectorAlgorithm)
-    sequential_projectors(coordinate::NTuple{3,Int}, state::InfinitePEPS, env::CTMRGEnv, alg::ProjectorAlgorithm)
+    sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm)
+    sequential_projectors(coordinate::NTuple{3,Int}, network::InfiniteSquareNetwork, env::CTMRGEnv, alg::ProjectorAlgorithm)
 
 Compute CTMRG projectors in the `:sequential` scheme either for an entire column `col` or
 for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequential` scheme).
 """
-function sequential_projectors(
-    col::Int, state::InfiniteSquareNetwork, env::CTMRGEnv, alg::ProjectorAlgorithm
-)
+function sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm)
     coordinates = eachcoordinate(env)[:, col]
     proj_and_info = dtmap(coordinates) do (r, c)
         trscheme = truncation_scheme(alg, env.edges[WEST, _prev(r, size(env, 2)), c])
         proj, info = sequential_projectors(
-            (WEST, r, c), state, env, @set(alg.trscheme = trscheme)
+            (WEST, r, c), network, env, @set(alg.trscheme = trscheme)
         )
         return proj, info
     end
     return _split_proj_and_info(proj_and_info)
 end
 function sequential_projectors(
-    coordinate::NTuple{3,Int},
-    state::InfiniteSquareNetwork,
-    env::CTMRGEnv,
-    alg::HalfInfiniteProjector,
+    coordinate::NTuple{3,Int}, network, env::CTMRGEnv, alg::HalfInfiniteProjector
 )
     _, r, c = coordinate
     r′ = _prev(r, size(env, 2))
-    Q1 = TensorMap(EnlargedCorner(state, env, (SOUTHWEST, r, c)), SOUTHWEST)
-    Q2 = TensorMap(EnlargedCorner(state, env, (NORTHWEST, r′, c)), NORTHWEST)
+    Q1 = TensorMap(EnlargedCorner(network, env, (SOUTHWEST, r, c)), SOUTHWEST)
+    Q2 = TensorMap(EnlargedCorner(network, env, (NORTHWEST, r′, c)), NORTHWEST)
     return compute_projector((Q1, Q2), coordinate, alg)
 end
 function sequential_projectors(
-    coordinate::NTuple{3,Int},
-    state::InfiniteSquareNetwork,
-    env::CTMRGEnv,
-    alg::FullInfiniteProjector,
+    coordinate::NTuple{3,Int}, network, env::CTMRGEnv, alg::FullInfiniteProjector
 )
     rowsize, colsize = size(env)[2:3]
     coordinate_nw = _next_coordinate(coordinate, rowsize, colsize)
     coordinate_ne = _next_coordinate(coordinate_nw, rowsize, colsize)
     coordinate_se = _next_coordinate(coordinate_ne, rowsize, colsize)
     ec = (
-        TensorMap(EnlargedCorner(state, env, coordinate_se), SOUTHEAST),
-        TensorMap(EnlargedCorner(state, env, coordinate), SOUTHWEST),
-        TensorMap(EnlargedCorner(state, env, coordinate_nw), NORTHWEST),
-        TensorMap(EnlargedCorner(state, env, coordinate_ne), NORTHEAST),
+        TensorMap(EnlargedCorner(network, env, coordinate_se), SOUTHEAST),
+        TensorMap(EnlargedCorner(network, env, coordinate), SOUTHWEST),
+        TensorMap(EnlargedCorner(network, env, coordinate_nw), NORTHWEST),
+        TensorMap(EnlargedCorner(network, env, coordinate_ne), NORTHEAST),
     )
     return compute_projector(ec, coordinate, alg)
 end
 
 """
-    renormalize_sequentially(col::Int, projectors, state, env)
+    renormalize_sequentially(col::Int, projectors, network, env)
 
 Renormalize one column of the CTMRG environment.
 """
-function renormalize_sequentially(col::Int, projectors, state, env)
+function renormalize_sequentially(col::Int, projectors, network, env)
     corners = Zygote.Buffer(env.corners)
     edges = Zygote.Buffer(env.edges)
 
-    for (dir, r, c) in eachcoordinate(state, 1:4)
+    for (dir, r, c) in eachcoordinate(network, 1:4)
         (c == col && dir in [SOUTHWEST, NORTHWEST]) && continue
         corners[dir, r, c] = env.corners[dir, r, c]
     end
-    for (dir, r, c) in eachcoordinate(state, 1:4)
+    for (dir, r, c) in eachcoordinate(network, 1:4)
         (c == col && dir == WEST) && continue
         edges[dir, r, c] = env.edges[dir, r, c]
     end
@@ -124,7 +116,7 @@ function renormalize_sequentially(col::Int, projectors, state, env)
         C_northwest = renormalize_top_corner((row, col), env, projectors)
         corners[NORTHWEST, row, col] = C_northwest / norm(C_northwest)
 
-        E_west = renormalize_west_edge((row, col), env, projectors, state)
+        E_west = renormalize_west_edge((row, col), env, projectors, network)
         edges[WEST, row, col] = E_west / norm(E_west)
     end
 

src/algorithms/ctmrg/simultaneous.jl

@@ -30,12 +30,12 @@ function SimultaneousCTMRG(;
     )
 end
 
-function ctmrg_iteration(state, env::CTMRGEnv, alg::SimultaneousCTMRG)
-    enlarged_corners = dtmap(eachcoordinate(state, 1:4)) do idx
-        return TensorMap(EnlargedCorner(state, env, idx), idx[1])
+function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG)
+    enlarged_corners = dtmap(eachcoordinate(network, 1:4)) do idx
+        return TensorMap(EnlargedCorner(network, env, idx), idx[1])
     end  # expand environment
     projectors, info = simultaneous_projectors(enlarged_corners, env, alg.projector_alg)  # compute projectors on all coordinates
-    env′ = renormalize_simultaneously(enlarged_corners, projectors, state, env)  # renormalize enlarged corners
+    env′ = renormalize_simultaneously(enlarged_corners, projectors, network, env)  # renormalize enlarged corners
     return env′, info
 end
 
@@ -96,26 +96,26 @@ function simultaneous_projectors(
 end
 
 """
-    renormalize_simultaneously(enlarged_corners, projectors, state, env)
+    renormalize_simultaneously(enlarged_corners, projectors, network, env)
 
 Renormalize all enlarged corners and edges simultaneously.
 """
-function renormalize_simultaneously(enlarged_corners, projectors, state, env)
+function renormalize_simultaneously(enlarged_corners, projectors, network, env)
     P_left, P_right = projectors
     coordinates = eachcoordinate(env, 1:4)
     corners_edges = dtmap(coordinates) do (dir, r, c)
         if dir == NORTH
             corner = renormalize_northwest_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_north_edge((r, c), env, P_left, P_right, state)
+            edge = renormalize_north_edge((r, c), env, P_left, P_right, network)
         elseif dir == EAST
             corner = renormalize_northeast_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_east_edge((r, c), env, P_left, P_right, state)
+            edge = renormalize_east_edge((r, c), env, P_left, P_right, network)
         elseif dir == SOUTH
             corner = renormalize_southeast_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_south_edge((r, c), env, P_left, P_right, state)
+            edge = renormalize_south_edge((r, c), env, P_left, P_right, network)
         elseif dir == WEST
             corner = renormalize_southwest_corner((r, c), enlarged_corners, P_left, P_right)
-            edge = renormalize_west_edge((r, c), env, P_left, P_right, state)
+            edge = renormalize_west_edge((r, c), env, P_left, P_right, network)
         end
         return corner / norm(corner), edge / norm(edge)
     end