From 8ec07423c99f102ab2c5b6f7723072260ae38350 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 30 Oct 2026 16:04:28 +0100 Subject: [PATCH 1/2] Allow CTMRG methods to use an allocator --- Project.toml | 3 + src/PEPSKit.jl | 2 +- src/algorithms/contractions/ctmrg/gaugefix.jl | 36 ++--- .../contractions/ctmrg/network_value.jl | 44 +++--- .../contractions/ctmrg/projector.jl | 70 +++++----- .../contractions/ctmrg/renormalize_corner.jl | 125 +++++++++--------- .../contractions/ctmrg/renormalize_edge.jl | 84 ++++++------ .../local_patch/densitymatrix/bp.jl | 14 +- .../local_patch/densitymatrix/ctmrg.jl | 52 ++++---- .../local_patch/densitymatrix/generic.jl | 4 +- .../contractions/local_patch/network_expr.jl | 4 +- .../local_patch/patch_contractions.jl | 8 +- src/algorithms/ctmrg/c4v.jl | 20 +-- src/algorithms/ctmrg/ctmrg.jl | 34 ++--- src/algorithms/ctmrg/gaugefix.jl | 74 ++++++----- src/algorithms/ctmrg/initialization.jl | 6 +- .../ctmrg/projectors/fullinfinite.jl | 4 +- .../ctmrg/projectors/halfinfinite.jl | 6 +- src/algorithms/ctmrg/sequential.jl | 38 +++--- src/algorithms/ctmrg/simultaneous.jl | 42 +++--- src/algorithms/ctmrg/sparse_environments.jl | 36 ++--- .../expectation_value/expectation_value.jl | 62 +++++---- .../expectation_value/network_value.jl | 28 ++-- .../expectation_value/patch_contractions.jl | 35 ++--- .../reduced_densitymatrix.jl | 44 +++--- .../optimization/implicit_differentiation.jl | 32 +++-- .../optimization/peps_optimization.jl | 5 +- .../optimization/preconditioning.jl | 14 +- src/algorithms/time_evolution/simpleupdate.jl | 3 +- test/testsuite/boundarymps/vumps.jl | 4 +- test/testsuite/ctmrg/c4v_contractions.jl | 22 +-- test/testsuite/ctmrg/c4v_sign_branches.jl | 13 +- test/testsuite/ctmrg/contractions.jl | 19 +-- test/testsuite/ctmrg/fixed_iterscheme.jl | 51 +++---- test/testsuite/ctmrg/gaugefix.jl | 12 +- test/testsuite/ctmrg/jacobian_real_linear.jl | 12 +- test/testsuite/ctmrg/partition_function.jl | 4 +- test/testsuite/ctmrg/pepo.jl | 12 +- test/testsuite/ctmrg/suweight.jl | 5 +- test/testsuite/ctmrg/unitcell.jl | 9 +- test/testsuite/toolbox/densitymatrices.jl | 50 +++---- 41 files changed, 611 insertions(+), 531 deletions(-) diff --git a/Project.toml b/Project.toml index a718d73fd..219270598 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" +[sources] +TensorKit = {url = "https://github.com/QuantumKitHub/TensorKit.jl", rev = "main"} + [extensions] PEPSKitAdaptExt = "Adapt" PEPSKitGPUArraysExt = "GPUArrays" diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index b715dc2ec..9ce61b28d 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -52,7 +52,7 @@ import MPSKitModels: transverse_field_ising, heisenberg_XYZ, heisenberg_XXZ, hubbard_model, bose_hubbard_model, tj_model using FiniteDifferences -using OhMyThreads: tmap, tmap!, SerialScheduler +using OhMyThreads: tmap, tmapreduce, tmap!, SerialScheduler using DocStringExtensions include("Defaults.jl") # Include first to allow for docstring interpolation with Defaults values diff --git a/src/algorithms/contractions/ctmrg/gaugefix.jl b/src/algorithms/contractions/ctmrg/gaugefix.jl index 41f9222eb..01a8bddcc 100644 --- a/src/algorithms/contractions/ctmrg/gaugefix.jl +++ b/src/algorithms/contractions/ctmrg/gaugefix.jl @@ -17,9 +17,9 @@ Multiply corner tensor with incoming and outgoing gauge signs. ``` """ function fix_gauge_corner( - corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor + corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor, allocator ) - return @autoopt @tensor corner_fix[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner_fix[χ_out; χ_in] := σ_out[χ_out; χ1] * corner[χ1; χ2] * conj(σ_in[χ_in; χ2]) end @@ -28,11 +28,12 @@ $(SIGNATURES) Apply `fix_gauge_corner` to the northwest corner with appropriate row and column indices. """ -function fix_gauge_northwest_corner((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_northwest_corner((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_corner( corner(env, NORTHWEST, row, col), signs[WEST, row, col], signs[NORTH, row, _next(col, end)], + allocator, ) end @@ -41,11 +42,12 @@ $(SIGNATURES) Apply `fix_gauge_corner` to the northeast corner with appropriate row and column indices. """ -function fix_gauge_northeast_corner((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_northeast_corner((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_corner( corner(env, NORTHEAST, row, col), signs[NORTH, row, col], signs[EAST, _next(row, end), col], + allocator, ) end @@ -54,11 +56,12 @@ $(SIGNATURES) Apply `fix_gauge_corner` to the southeast corner with appropriate row and column indices. """ -function fix_gauge_southeast_corner((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_southeast_corner((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_corner( corner(env, SOUTHEAST, row, col), signs[EAST, row, col], signs[SOUTH, row, _prev(col, end)], + allocator, ) end @@ -67,11 +70,12 @@ $(SIGNATURES) Apply `fix_gauge_corner` to the southwest corner with appropriate row and column indices. """ -function fix_gauge_southwest_corner((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_southwest_corner((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_corner( corner(env, SOUTHWEST, row, col), signs[SOUTH, row, col], signs[WEST, _prev(row, end), col], + allocator, ) end @@ -88,7 +92,7 @@ Multiply edge tensor with incoming and outgoing gauge signs. ``` """ @generated function fix_gauge_edge( - edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor + edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor, allocator ) where {T, S, N} edge_fix_e = tensorexpr( :edge_fix, @@ -102,7 +106,7 @@ Multiply edge tensor with incoming and outgoing gauge signs. σ_in_e = tensorexpr(:σ_in, (envlabel(:in),), (envlabel(2),)) return macroexpand( @__MODULE__, - :(return @autoopt @tensor $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)), + :(return @autoopt @tensor allocator = allocator $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)), ) end @@ -111,9 +115,9 @@ $(SIGNATURES) Apply `fix_gauge_edge` to the north edge with appropriate row and column indices. """ -function fix_gauge_north_edge((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_north_edge((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_edge( - edge(env, NORTH, row, col), signs[NORTH, row, col], signs[NORTH, row, _next(col, end)], + edge(env, NORTH, row, col), signs[NORTH, row, col], signs[NORTH, row, _next(col, end)], allocator ) end @@ -122,9 +126,9 @@ $(SIGNATURES) Apply `fix_gauge_edge` to the east edge with appropriate row and column indices. """ -function fix_gauge_east_edge((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_east_edge((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_edge( - edge(env, EAST, row, col), signs[EAST, row, col], signs[EAST, _next(row, end), col] + edge(env, EAST, row, col), signs[EAST, row, col], signs[EAST, _next(row, end), col], allocator ) end @@ -133,9 +137,9 @@ $(SIGNATURES) Apply `fix_gauge_edge` to the south edge with appropriate row and column indices. """ -function fix_gauge_south_edge((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_south_edge((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_edge( - edge(env, SOUTH, row, col), signs[SOUTH, row, col], signs[SOUTH, row, _prev(col, end)], + edge(env, SOUTH, row, col), signs[SOUTH, row, col], signs[SOUTH, row, _prev(col, end)], allocator ) end @@ -144,9 +148,9 @@ $(SIGNATURES) Apply `fix_gauge_edge` to the west edge with appropriate row and column indices. """ -function fix_gauge_west_edge((row, col), env::AbstractCTMRGEnv, signs) +function fix_gauge_west_edge((row, col), env::AbstractCTMRGEnv, signs, allocator) return fix_gauge_edge( - edge(env, WEST, row, col), signs[WEST, row, col], signs[WEST, _prev(row, end), col] + edge(env, WEST, row, col), signs[WEST, row, col], signs[WEST, _prev(row, end), col], allocator ) end diff --git a/src/algorithms/contractions/ctmrg/network_value.jl b/src/algorithms/contractions/ctmrg/network_value.jl index eb757f4d3..90f02a1d5 100644 --- a/src/algorithms/contractions/ctmrg/network_value.jl +++ b/src/algorithms/contractions/ctmrg/network_value.jl @@ -10,9 +10,9 @@ function _contract_site( C_northwest, C_northeast, C_southeast, C_southwest, E_north::CTMRG_PEPS_EdgeTensor, E_east::CTMRG_PEPS_EdgeTensor, E_south::CTMRG_PEPS_EdgeTensor, E_west::CTMRG_PEPS_EdgeTensor, - O::PEPSSandwich, + O::PEPSSandwich, allocator ) - return @autoopt @tensor E_west[χ_WSW D_W_above D_W_below; χ_WNW] * + return @autoopt @tensor allocator = allocator E_west[χ_WSW D_W_above D_W_below; χ_WNW] * C_northwest[χ_WNW; χ_NNW] * E_north[χ_NNW D_N_above D_N_below; χ_NNE] * C_northeast[χ_NNE; χ_ENE] * @@ -27,9 +27,9 @@ function _contract_site( C_northwest, C_northeast, C_southeast, C_southwest, E_north::CTMRG_PF_EdgeTensor, E_east::CTMRG_PF_EdgeTensor, E_south::CTMRG_PF_EdgeTensor, E_west::CTMRG_PF_EdgeTensor, - O::PFTensor, + O::PFTensor, allocator ) - return @autoopt @tensor E_west[χ_WSW D_W; χ_WNW] * + return @autoopt @tensor allocator = allocator E_west[χ_WSW D_W; χ_WNW] * C_northwest[χ_WNW; χ_NNW] * E_north[χ_NNW D_N; χ_NNE] * C_northeast[χ_NNE; χ_ENE] * @@ -43,7 +43,7 @@ end @generated function _contract_site( C_northwest, C_northeast, C_southeast, C_southwest, E_north::TE, E_east::TE, E_south::TE, E_west::TE, - O::PEPOSandwich{H}, + O::PEPOSandwich{H}, allocator ) where {TE <: CTMRGEdgeTensor, H} @assert numout(TE) == H + 3 @@ -67,15 +67,15 @@ end pepo_es..., ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs)) end """ - _contract_site(ind::Tuple{Int,Int}, network::InfiniteSquareNetwork, env::AbstractCTMRGEnv) + _contract_site(ind::Tuple{Int,Int}, network::InfiniteSquareNetwork, env::AbstractCTMRGEnv, allocator) Contract around a single site `ind` of a square network using a given CTMRG environment. """ -function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::CTMRGEnv) +function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::CTMRGEnv, allocator) r, c = ind return _contract_site( corner(env, NORTHWEST, r - 1, c - 1), @@ -84,11 +84,11 @@ function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, en corner(env, SOUTHWEST, r + 1, c - 1), edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1), edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1), - network[r, c], + network[r, c], allocator ) end -function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::C4vCTMRGEnv) +function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, env::C4vCTMRGEnv, allocator) r, c = ind @assert r == c == 1 lazy"only trivial unit cells are supported for C4v-CTMRG, ind ($ind) must be (1, 1)" return _contract_site( @@ -98,7 +98,7 @@ function _contract_site(ind::Tuple{Int, Int}, network::InfiniteSquareNetwork, en corner(env, SOUTHWEST, r + 1, c - 1), edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1), edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1), - network[r, c], + network[r, c], allocator, ) end @@ -106,24 +106,26 @@ end function _contract_corners( C_northwest::CTMRGCornerTensor, C_northeast::CTMRGCornerTensor, C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor, + allocator, ) - return @tensor C_northwest[1; 2] * C_northeast[2; 3] * + return @tensor allocator = allocator C_northwest[1; 2] * C_northeast[2; 3] * C_southeast[3; 4] * C_southwest[4; 1] end """ - _contract_corners(ind::Tuple{Int,Int}, env::CTMRGEnv) + _contract_corners(ind::Tuple{Int,Int}, env::CTMRGEnv, allocator) Contract all corners around the south-east at position `ind` of the CTMRG environment `env`. """ -function _contract_corners(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv) +function _contract_corners(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator) r, c = ind return _contract_corners( corner(env, NORTHWEST, r - 1, c - 1), corner(env, NORTHEAST, r - 1, c), corner(env, SOUTHEAST, r, c), corner(env, SOUTHWEST, r, c - 1), + allocator, ) end @@ -132,6 +134,7 @@ end C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor, E_east::CTMRGEdgeTensor{T, S, N}, E_west::CTMRGEdgeTensor{T, S, N}, + allocator, ) where {T, S, N} C_northwest_e = tensorexpr(:C_northwest, (envlabel(:NW),), (envlabel(:N),)) C_northeast_e = tensorexpr(:C_northeast, (envlabel(:N),), (envlabel(:NE),)) @@ -150,7 +153,7 @@ end E_west_e, C_northwest_e, C_northeast_e, E_east_e, C_southeast_e, C_southwest_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs)) end """ @@ -159,7 +162,7 @@ end Contract the vertical edges and corners around the east edge at position `ind` of the CTMRG environment `env`. """ -function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv) +function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator) r, c = ind return _contract_vertical_edges( corner(env, NORTHWEST, r - 1, c - 1), @@ -168,6 +171,7 @@ function _contract_vertical_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv) corner(env, SOUTHWEST, r + 1, c - 1), edge(env, EAST, r, c), edge(env, WEST, r, c - 1), + allocator ) end @@ -175,6 +179,7 @@ end C_northwest::CTMRGCornerTensor, C_northeast::CTMRGCornerTensor, C_southeast::CTMRGCornerTensor, C_southwest::CTMRGCornerTensor, E_north::CTMRGEdgeTensor{T, S, N}, E_south::CTMRGEdgeTensor{T, S, N}, + allocator, ) where {T, S, N} C_northwest_e = tensorexpr(:C_northwest, (envlabel(:W),), (envlabel(:NW),)) C_northeast_e = tensorexpr(:C_northeast, (envlabel(:NE),), (envlabel(:E),)) @@ -193,16 +198,16 @@ end C_northwest_e, E_north_e, C_northeast_e, C_southeast_e, E_south_e, C_southwest_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $rhs)) end """ - _contract_horizontal_edges(ind::Tuple{Int,Int}, env::CTMRGEnv) + _contract_horizontal_edges(ind::Tuple{Int,Int}, env::CTMRGEnv, allocator) Contract the horizontal edges and corners around the south edge at position `ind` of the CTMRG environment `env`. """ -function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv) +function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv, allocator) r, c = ind return _contract_horizontal_edges( corner(env, NORTHWEST, r - 1, c - 1), @@ -211,5 +216,6 @@ function _contract_horizontal_edges(ind::Tuple{Int, Int}, env::AbstractCTMRGEnv) corner(env, SOUTHWEST, r, c - 1), edge(env, NORTH, r - 1, c), edge(env, SOUTH, r, c), + allocator, ) end diff --git a/src/algorithms/contractions/ctmrg/projector.jl b/src/algorithms/contractions/ctmrg/projector.jl index f8791d77f..831441d4a 100644 --- a/src/algorithms/contractions/ctmrg/projector.jl +++ b/src/algorithms/contractions/ctmrg/projector.jl @@ -14,23 +14,23 @@ Contract the CTMRG left projector with the higher-dimensional subspace facing to out ``` """ -function left_projector(E_1, C, E_2, V, isqS, A) +function left_projector(E_1, C, E_2, V, isqS, A, allocator) Vdt_isqS = twistdual(V', 1:(numind(V) - 1)) * isqS - return _left_projector(E_1, C, E_2, Vdt_isqS, A) + return _left_projector(E_1, C, E_2, Vdt_isqS, A, allocator) end -function _left_projector(E_1, C, E_2, Vd, A::PEPSSandwich) - return @autoopt @tensor P_left[χ_out D_outabove D_outbelow; χ_in] := +function _left_projector(E_1, C, E_2, Vd, A::PEPSSandwich, allocator) + return @autoopt @tensor allocator = allocator P_left[χ_out D_outabove D_outbelow; χ_in] := E_1[χ_out D1 D2; χ1] * C[χ1; χ2] * E_2[χ2 D3 D4; χ3] * ket(A)[d; D3 D5 D_outabove D1] * conj(bra(A)[d; D4 D6 D_outbelow D2]) * Vd[χ3 D5 D6; χ_in] end -function _left_projector(E_1, C, E_2, Vd, A::PFTensor) - return @autoopt @tensor P_left[χ_out D_out; χ_in] := +function _left_projector(E_1, C, E_2, Vd, A::PFTensor, allocator) + return @autoopt @tensor allocator = allocator P_left[χ_out D_out; χ_in] := E_1[χ_out D1; χ1] * C[χ1; χ2] * E_2[χ2 D2; χ3] * A[D1 D_out; D2 D3] * Vd[χ3 D3; χ_in] end @generated function _left_projector( - E_1, C, E_2, Vd, A::PEPOSandwich{H} + E_1, C, E_2, Vd, A::PEPOSandwich{H}, allocator ) where {H} E_west_e = _pepo_edge_expr(:E_1, :out, :WNW, :W, H) @@ -51,7 +51,7 @@ end ) return macroexpand( - @__MODULE__, :(return @autoopt @tensor $P_left_e := $rhs) + @__MODULE__, :(return @autoopt @tensor allocator = allocator $P_left_e := $rhs) ) end @@ -72,14 +72,14 @@ Contract the CTMRG left projector with the higher-dimensional subspace facing to in ``` """ -function left_projector(E_1, C_1, E_2, E_3, C_2, E_4, V, isqS, A_1, A_2) +function left_projector(E_1, C_1, E_2, E_3, C_2, E_4, V, isqS, A_1, A_2, allocator) Vdt_isqS = twistdual(V', 1:(numind(V) - 1)) * isqS - return _left_projector(E_1, C_1, E_2, E_3, C_2, E_4, Vdt_isqS, A_1, A_2) + return _left_projector(E_1, C_1, E_2, E_3, C_2, E_4, Vdt_isqS, A_1, A_2, allocator) end function _left_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PEPSSandwich, A_2::PEPSSandwich + E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PEPSSandwich, A_2::PEPSSandwich, allocator ) - return @autoopt @tensor P_left[χ_out D_outa D_outb; χ_in] := + return @autoopt @tensor allocator = allocator P_left[χ_out D_outa D_outb; χ_in] := E_1[χ_out D_W1a D_W1b; χ_WNW] * C_1[χ_WNW; χ_NNW] * E_2[χ_NNW D_N1a D_N1b; χ_N] * ket(A_1)[d1; D_N1a D_Ca D_outa D_W1a] * conj(bra(A_1)[d1; D_N1b D_Cb D_outb D_W1b]) * E_3[χ_N D_N2a D_N2b; χ_NNE] * C_2[χ_NNE; χ_ENE] * E_4[χ_ENE D_E2a D_E2b; χ_SE] * @@ -87,9 +87,9 @@ function _left_projector( Vd[χ_SE D_S2a D_S2b; χ_in] end function _left_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PFTensor, A_2::PFTensor + E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PFTensor, A_2::PFTensor, allocator ) - return @autoopt @tensor P_left[χ_out D_out; χ_in] := + return @autoopt @tensor allocator = allocator P_left[χ_out D_out; χ_in] := E_1[χ_out D_W1; χ_WNW] * C_1[χ_WNW; χ_NNW] * E_2[χ_NNW D_N1; χ_N] * A_1[D_W1 D_out; D_N1 D_C] * E_3[χ_N D_N2; χ_NNE] * C_2[χ_NNE; χ_ENE] * E_4[χ_ENE D_E2; χ_SE] * @@ -97,7 +97,7 @@ function _left_projector( Vd[χ_SE D_S2; χ_in] end @generated function _left_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PEPOSandwich{H}, A_2::PEPOSandwich{H} + E_1, C_1, E_2, E_3, C_2, E_4, Vd, A_1::PEPOSandwich{H}, A_2::PEPOSandwich{H}, allocator ) where {H} Vd_e = _pepo_domain_projector_expr(:Vd, :SE, :S, :in, H, 2) @@ -105,7 +105,7 @@ end P_left_e = _pepo_domain_projector_expr(:P_left, :SW, :S, :in, H, 1) return macroexpand( - @__MODULE__, :(return @autoopt @tensor $P_left_e := $proj_expr * $Vd_e) + @__MODULE__, :(return @autoopt @tensor allocator = allocator $P_left_e := $proj_expr * $Vd_e) ) end @@ -122,24 +122,24 @@ Contract the CTMRG right projector with the higher-dimensional subspace facing t in ``` """ -function right_projector(E_1, C, E_2, U, isqS, A) +function right_projector(E_1, C, E_2, U, isqS, A, allocator) isqS_Udt = isqS * twistnondual(U', 2:numind(U)) - return _right_projector(E_1, C, E_2, isqS_Udt, A) + return _right_projector(E_1, C, E_2, isqS_Udt, A, allocator) end -function _right_projector(E_1, C, E_2, Ud, A::PEPSSandwich) - return @autoopt @tensor P_right[χ_out; χ_in D_inabove D_inbelow] := - Ud[χ_out; χ2 D1 D2] * - ket(A)[d; D3 D5 D_inabove D1] * conj(bra(A)[d; D4 D6 D_inbelow D2]) * +function _right_projector(E_1, C, E_2, Ud, A::PEPSSandwich, allocator) + return @autoopt @tensor allocator = allocator P_right[χ_out; χ_in D_inabove D_inbelow] := + Ud[χ_out; χ2 D1 D2] * ket(A)[d; D3 D5 D_inabove D1] * + conj(bra(A)[d; D4 D6 D_inbelow D2]) * E_1[χ2 D3 D4; χ3] * C[χ3; χ4] * E_2[χ4 D5 D6; χ_in] end -function _right_projector(E_1, C, E_2, Ud, A::PFTensor) - return @autoopt @tensor P_right[χ_out; χ_in D_in] := +function _right_projector(E_1, C, E_2, Ud, A::PFTensor, allocator) + return @autoopt @tensor allocator = allocator P_right[χ_out; χ_in D_in] := Ud[χ_out; χ2 D1] * A[D1 D_in; D2 D3] * E_1[χ2 D2; χ3] * C[χ3; χ4] * E_2[χ4 D3; χ_in] end @generated function _right_projector( - E_1, C, E_2, Ud, A::PEPOSandwich{H} + E_1, C, E_2, Ud, A::PEPOSandwich{H}, allocator ) where {H} E_north_e = _pepo_edge_expr(:E_1, :NW, :NNE, :N, H) @@ -160,7 +160,7 @@ end ) return macroexpand( - @__MODULE__, :(return @autoopt @tensor $P_right_e := $rhs) + @__MODULE__, :(return @autoopt @tensor allocator = allocator $P_right_e := $rhs) ) end @@ -181,14 +181,14 @@ Contract the CTMRG left projector with the higher-dimensional subspace facing to out ``` """ -function right_projector(E_1, C_1, E_2, E_3, C_2, E_4, U, isqS, A_1, A_2) +function right_projector(E_1, C_1, E_2, E_3, C_2, E_4, U, isqS, A_1, A_2, allocator) isqS_Udt = isqS * twistnondual(U', 2:numind(U)) - return _right_projector(E_1, C_1, E_2, E_3, C_2, E_4, isqS_Udt, A_1, A_2) + return _right_projector(E_1, C_1, E_2, E_3, C_2, E_4, isqS_Udt, A_1, A_2, allocator) end function _right_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PEPSSandwich, A_2::PEPSSandwich + E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PEPSSandwich, A_2::PEPSSandwich, allocator ) - return @autoopt @tensor P_right[χ_out; χ_in D_ina D_inb] := + return @autoopt @tensor allocator = allocator P_right[χ_out; χ_in D_ina D_inb] := E_1[χ_SW D_W1a D_W1b; χ_WNW] * C_1[χ_WNW; χ_NNW] * E_2[χ_NNW D_N1a D_N1b; χ_N] * ket(A_1)[d1; D_N1a D_Ca D_S1a D_W1a] * conj(bra(A_1)[d1; D_N1b D_Cb D_S1b D_W1b]) * E_3[χ_N D_N2a D_N2b; χ_NNE] * C_2[χ_NNE; χ_ENE] * E_4[χ_ENE D_E2a D_E2b; χ_in] * @@ -196,9 +196,9 @@ function _right_projector( Ud[χ_out; χ_SW D_S1a D_S1b] end function _right_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PFTensor, A_2::PFTensor + E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PFTensor, A_2::PFTensor, allocator ) - return @autoopt @tensor P_right[χ_out; χ_in D_in] := + return @autoopt @tensor allocator = allocator P_right[χ_out; χ_in D_in] := E_1[χ_SW D_W1; χ_WNW] * C_1[χ_WNW; χ_NNW] * E_2[χ_NNW D_N1; χ_N] * A_1[D_W1 D_S1; D_N1 D_C] * E_3[χ_N D_N2; χ_NNE] * C_2[χ_NNE; χ_ENE] * E_4[χ_ENE D_E2; χ_in] * @@ -206,7 +206,7 @@ function _right_projector( Ud[χ_out; χ_SW D_S1] end @generated function _right_projector( - E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PEPOSandwich{H}, A_2::PEPOSandwich{H} + E_1, C_1, E_2, E_3, C_2, E_4, Ud, A_1::PEPOSandwich{H}, A_2::PEPOSandwich{H}, allocator ) where {H} Ud_e = _pepo_codomain_projector_expr(:Ud, :out, :SW, :S, H, 1) @@ -214,7 +214,7 @@ end P_right_e = _pepo_codomain_projector_expr(:P_right, :out, :SE, :S, H, 2) return macroexpand( - @__MODULE__, :(return @autoopt @tensor $P_right_e := $Ud_e * $proj_expr) + @__MODULE__, :(return @autoopt @tensor allocator = allocator $P_right_e := $Ud_e * $proj_expr) ) end @@ -238,7 +238,7 @@ Right projector: |~~| == |~~~| == ``` """ -function contract_projectors(U, S, V, Q, Q_next) +function contract_projectors(U, S, V, Q, Q_next, allocator) isqS = sdiag_pow(S, -0.5) P_left = Q_next * V' * isqS # use * to respect fermionic case P_right = isqS * U' * Q diff --git a/src/algorithms/contractions/ctmrg/renormalize_corner.jl b/src/algorithms/contractions/ctmrg/renormalize_corner.jl index b551c6283..50970bc9f 100644 --- a/src/algorithms/contractions/ctmrg/renormalize_corner.jl +++ b/src/algorithms/contractions/ctmrg/renormalize_corner.jl @@ -17,6 +17,7 @@ Apply projectors to each side of a quadrant. quadrant::AbstractTensorMap{<:Any, S, N, N}, P_left::AbstractTensorMap{<:Any, S, N, 1}, P_right::AbstractTensorMap{<:Any, S, 1, N}, + allocator, ) where {S, N} corner_e = tensorexpr(:corner, (envlabel(:out),), (envlabel(:in),)) P_right_e = tensorexpr( @@ -36,7 +37,7 @@ Apply projectors to each side of a quadrant. ) return macroexpand( @__MODULE__, - :(return @autoopt @tensor $corner_e := $P_right_e * $quadrant_e * $P_left_e), + :(return @autoopt @tensor allocator = allocator $corner_e := $P_right_e * $quadrant_e * $P_left_e), ) end @@ -44,9 +45,9 @@ end # ---------------- """ - renormalize_northwest_corner((row, col), enlarged_env, P_left, P_right) - renormalize_northwest_corner(quadrant, P_left, P_right) - renormalize_northwest_corner(E_west, C_northwest, E_north, P_left, P_right, A) + renormalize_northwest_corner((row, col), enlarged_env, P_left, P_right, allocator) + renormalize_northwest_corner(quadrant, P_left, P_right, allocator) + renormalize_northwest_corner(E_west, C_northwest, E_north, P_left, P_right, A, allocator) Apply `renormalize_corner` to the enlarged northwest corner. @@ -72,31 +73,32 @@ Alternatively, provide the constituent tensors and perform the complete contract out ``` """ -function renormalize_northwest_corner((row, col), enlarged_env, P_left, P_right) +function renormalize_northwest_corner((row, col), enlarged_env, P_left, P_right, allocator) return renormalize_northwest_corner( enlarged_env[NORTHWEST, row, col], P_left[NORTH, row, col], P_right[WEST, _next(row, end), col], + allocator, ) end function renormalize_northwest_corner( - quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right + quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right, allocator ) where {T, S, N} - return renormalize_corner(quadrant, P_left, P_right) + return renormalize_corner(quadrant, P_left, P_right, allocator) end function renormalize_northwest_corner( - E_west, C_northwest, E_north, P_left, P_right, A::PEPSSandwich + E_west, C_northwest, E_north, P_left, P_right, A::PEPSSandwich, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1 D2] * E_west[χ1 D3 D4; χ2] * C_northwest[χ2; χ3] * E_north[χ3 D5 D6; χ4] * ket(A)[d; D5 D7 D1 D3] * conj(bra(A)[d; D6 D8 D2 D4]) * P_left[χ4 D7 D8; χ_in] end function renormalize_northwest_corner( - E_west, C_northwest, E_north, P_left, P_right, A::PFTensor + E_west, C_northwest, E_north, P_left, P_right, A::PFTensor, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1] * E_west[χ1 D3; χ2] * C_northwest[χ2; χ3] * E_north[χ3 D5; χ4] * A[D3 D1; D5 D7] * @@ -104,7 +106,7 @@ function renormalize_northwest_corner( end @generated function renormalize_northwest_corner( - E_west, C_northwest, E_north, P_left, P_right, A::PEPOSandwich{H} + E_west, C_northwest, E_north, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {H} C_out_e = _corner_expr(:corner, :out, :in) @@ -124,16 +126,16 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $C_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $C_out_e := $rhs)) end # Northeast corner # ---------------- """ - renormalize_northeast_corner((row, col), enlarged_env, P_left, P_right) - renormalize_northeast_corner(quadrant, P_left, P_right) - renormalize_northeast_corner(E_north, C_northeast, E_east, P_left, P_right, A) + renormalize_northeast_corner((row, col), enlarged_env, P_left, P_right, allocator) + renormalize_northeast_corner(quadrant, P_left, P_right, allocator) + renormalize_northeast_corner(E_north, C_northeast, E_east, P_left, P_right, A, allocator) Apply `renormalize_corner` to the enlarged northeast corner. @@ -159,33 +161,34 @@ Alternatively, provide the constituent tensors and perform the complete contract in ``` """ -function renormalize_northeast_corner((row, col), enlarged_env, P_left, P_right) +function renormalize_northeast_corner((row, col), enlarged_env, P_left, P_right, allocator) return renormalize_northeast_corner( enlarged_env[NORTHEAST, row, col], P_left[EAST, row, col], P_right[NORTH, row, _prev(col, end)], + allocator, ) end function renormalize_northeast_corner( - quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right + quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right, allocator ) where {T, S, N} - return renormalize_corner(quadrant, P_left, P_right) + return renormalize_corner(quadrant, P_left, P_right, allocator) end function renormalize_northeast_corner( - E_north, C_northeast, E_east, P_left, P_right, A::PEPSSandwich + E_north, C_northeast, E_east, P_left, P_right, A::PEPSSandwich, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1 D2] * E_north[χ1 D3 D4; χ2] * C_northeast[χ2; χ3] * E_east[χ3 D5 D6; χ4] * ket(A)[d; D3 D5 D7 D1] * conj(bra(A)[d; D4 D6 D8 D2]) * P_left[χ4 D7 D8; χ_in] end function renormalize_northeast_corner( - E_north, C_northeast, E_east, P_left, P_right, A::PFTensor + E_north, C_northeast, E_east, P_left, P_right, A::PFTensor, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1] * E_north[χ1 D3; χ2] * C_northeast[χ2; χ3] * E_east[χ3 D5; χ4] * A[D1 D7; D3 D5] * @@ -193,7 +196,7 @@ function renormalize_northeast_corner( end @generated function renormalize_northeast_corner( - E_north, C_northeast, E_east, P_left, P_right, A::PEPOSandwich{H} + E_north, C_northeast, E_east, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {H} C_out_e = _corner_expr(:corner, :out, :in) @@ -213,16 +216,16 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $C_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $C_out_e := $rhs)) end # Southeast corner # ---------------- """ - renormalize_southeast_corner((row, col), enlarged_env, P_left, P_right) - renormalize_southeast_corner(quadrant, P_left, P_right) - renormalize_southeast_corner(E_east, C_southeast, E_south, P_left, P_right, A) + renormalize_southeast_corner((row, col), enlarged_env, P_left, P_right, allocator) + renormalize_southeast_corner(quadrant, P_left, P_right, allocator) + renormalize_southeast_corner(E_east, C_southeast, E_south, P_left, P_right, A, allocator) Apply `renormalize_corner` to the enlarged southeast corner. @@ -248,31 +251,32 @@ Alternatively, provide the constituent tensors and perform the complete contract |~~~~~~| -- E_south -- C_southeast ``` """ -function renormalize_southeast_corner((row, col), enlarged_env, P_left, P_right) +function renormalize_southeast_corner((row, col), enlarged_env, P_left, P_right, allocator) return renormalize_southeast_corner( enlarged_env[SOUTHEAST, row, col], P_left[SOUTH, row, col], P_right[EAST, _prev(row, end), col], + allocator, ) end function renormalize_southeast_corner( - quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right + quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right, allocator ) where {T, S, N} - return renormalize_corner(quadrant, P_left, P_right) + return renormalize_corner(quadrant, P_left, P_right, allocator) end function renormalize_southeast_corner( - E_east, C_southeast, E_south, P_left, P_right, A::PEPSSandwich + E_east, C_southeast, E_south, P_left, P_right, A::PEPSSandwich, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1 D2] * E_east[χ1 D3 D4; χ2] * C_southeast[χ2; χ3] * E_south[χ3 D5 D6; χ4] * ket(A)[d; D1 D3 D5 D7] * conj(bra(A)[d; D2 D4 D6 D8]) * P_left[χ4 D7 D8; χ_in] end function renormalize_southeast_corner( - E_east, C_southeast, E_south, P_left, P_right, A::PFTensor + E_east, C_southeast, E_south, P_left, P_right, A::PFTensor, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1] * E_east[χ1 D3; χ2] * C_southeast[χ2; χ3] * E_south[χ3 D5; χ4] * A[D7 D5; D1 D3] * @@ -280,7 +284,7 @@ function renormalize_southeast_corner( end @generated function renormalize_southeast_corner( - E_east, C_southeast, E_south, P_left, P_right, A::PEPOSandwich{H} + E_east, C_southeast, E_south, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {H} C_out_e = _corner_expr(:corner, :out, :in) @@ -300,16 +304,16 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $C_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $C_out_e := $rhs)) end # Southwest corner # ---------------- """ - renormalize_southwest_corner((row, col), enlarged_env, P_left, P_right) - renormalize_southwest_corner(quadrant, P_left, P_right) - renormalize_southwest_corner(E_south, C_southwest, E_west, P_left, P_right, A) + renormalize_southwest_corner((row, col), enlarged_env, P_left, P_right, allocator) + renormalize_southwest_corner(quadrant, P_left, P_right, allocator) + renormalize_southwest_corner(E_south, C_southwest, E_west, P_left, P_right, A, allocator) Apply `renormalize_corner` to the enlarged southwest corner. @@ -335,31 +339,32 @@ Alternatively, provide the constituent tensors and perform the complete contract C_southwest -- E_south -- |~~~~~~~| ``` """ -function renormalize_southwest_corner((row, col), enlarged_env, P_left, P_right) +function renormalize_southwest_corner((row, col), enlarged_env, P_left, P_right, allocator) return renormalize_southwest_corner( enlarged_env[SOUTHWEST, row, col], P_left[WEST, row, col], P_right[SOUTH, row, _next(col, end)], + allocator, ) end function renormalize_southwest_corner( - quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right + quadrant::AbstractTensorMap{T, S, N, N}, P_left, P_right, allocator ) where {T, S, N} - return renormalize_corner(quadrant, P_left, P_right) + return renormalize_corner(quadrant, P_left, P_right, allocator) end function renormalize_southwest_corner( - E_south, C_southwest, E_west, P_left, P_right, A::PEPSSandwich + E_south, C_southwest, E_west, P_left, P_right, A::PEPSSandwich, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1 D2] * E_south[χ1 D3 D4; χ2] * C_southwest[χ2; χ3] * E_west[χ3 D5 D6; χ4] * ket(A)[d; D7 D1 D3 D5] * conj(bra(A)[d; D8 D2 D4 D6]) * P_left[χ4 D7 D8; χ_in] end function renormalize_southwest_corner( - E_south, C_southwest, E_west, P_left, P_right, A::PFTensor + E_south, C_southwest, E_west, P_left, P_right, A::PFTensor, allocator ) - return @autoopt @tensor corner[χ_out; χ_in] := + return @autoopt @tensor allocator = allocator corner[χ_out; χ_in] := P_right[χ_out; χ1 D1] * E_south[χ1 D3; χ2] * C_southwest[χ2; χ3] * E_west[χ3 D5; χ4] * A[D5 D3; D7 D1] * @@ -367,7 +372,7 @@ function renormalize_southwest_corner( end @generated function renormalize_southwest_corner( - E_south, C_southwest, E_west, P_left, P_right, A::PEPOSandwich{H} + E_south, C_southwest, E_west, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {H} C_out_e = _corner_expr(:corner, :out, :in) @@ -387,15 +392,15 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $C_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $C_out_e := $rhs)) end # For SequentialCTMRG left move only # ---------------------------------- """ - renormalize_southwest_corner((r, c), env, projectors) - renormalize_southwest_corner(C_southwest, E_south, P_left) + renormalize_southwest_corner((r, c), env, projectors, allocator) + renormalize_southwest_corner(C_southwest, E_south, P_left, allocator) Apply left projector to southwest corner and south edge. ``` @@ -406,16 +411,17 @@ Apply left projector to southwest corner and south edge. C -- E --out ``` """ -function renormalize_southwest_corner((row, col), env::AbstractCTMRGEnv, projectors) +function renormalize_southwest_corner((row, col), env::AbstractCTMRGEnv, projectors, allocator) C_southwest = corner(env, SOUTHWEST, row, col - 1) E_south = edge(env, SOUTH, row, col) P_left = projectors[1][row] - return renormalize_southwest_corner(C_southwest, E_south, P_left) + return renormalize_southwest_corner(C_southwest, E_south, P_left, allocator) end @generated function renormalize_southwest_corner( C_southwest::CTMRGCornerTensor{<:Any, S}, E_south::CTMRGEdgeTensor{<:Any, S, N}, P_left::AbstractTensorMap{<:Any, S, N, 1}, + allocator, ) where {S, N} C_out_e = tensorexpr(:corner, (envlabel(:out),), (envlabel(:in),)) C_southwest_e = tensorexpr(:C_southwest, (envlabel(:SSW),), (envlabel(:WSW),)) @@ -431,13 +437,13 @@ end ) return macroexpand( @__MODULE__, - :(return @autoopt @tensor $C_out_e := $E_south_e * $C_southwest_e * $P_left_e), + :(return @autoopt @tensor allocator = allocator $C_out_e := $E_south_e * $C_southwest_e * $P_left_e), ) end """ - renormalize_northwest_corner((row, col), env, projectors) - renormalize_northwest_corner(C_northwest, E_north, P_right) + renormalize_northwest_corner((row, col), env, projectors, allocator) + renormalize_northwest_corner(C_northwest, E_north, P_right, allocator) Apply right projector to northwest corner and north edge. ``` @@ -448,16 +454,17 @@ Apply right projector to northwest corner and north edge. out ``` """ -function renormalize_northwest_corner((row, col), env::AbstractCTMRGEnv, projectors) +function renormalize_northwest_corner((row, col), env::AbstractCTMRGEnv, projectors, allocator) C_northwest = corner(env, NORTHWEST, row, col - 1) E_north = edge(env, NORTH, row, col) P_right = projectors[2][_next(row, end)] - return renormalize_northwest_corner(C_northwest, E_north, P_right) + return renormalize_northwest_corner(C_northwest, E_north, P_right, allocator) end @generated function renormalize_northwest_corner( C_northwest::CTMRGCornerTensor{<:Any, S}, E_north::CTMRGEdgeTensor{<:Any, S, N}, P_right::AbstractTensorMap{<:Any, S, 1, N}, + allocator, ) where {S, N} C_out_e = tensorexpr(:corner, (envlabel(:out),), (envlabel(:in),)) C_northwest_e = tensorexpr(:C_northwest, (envlabel(:WNW),), (envlabel(:NNW),)) @@ -469,6 +476,6 @@ end ) return macroexpand( @__MODULE__, - :(return @autoopt @tensor $C_out_e := $E_north_e * $C_northwest_e * $P_right_e), + :(return @autoopt @tensor allocator = allocator $C_out_e := $E_north_e * $C_northwest_e * $P_right_e), ) end diff --git a/src/algorithms/contractions/ctmrg/renormalize_edge.jl b/src/algorithms/contractions/ctmrg/renormalize_edge.jl index 5efd31cac..8f0ef7c92 100644 --- a/src/algorithms/contractions/ctmrg/renormalize_edge.jl +++ b/src/algorithms/contractions/ctmrg/renormalize_edge.jl @@ -2,8 +2,8 @@ # ---------- """ - renormalize_north_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}) - renormalize_north_edge(E_north, P_left, P_right, A::P) + renormalize_north_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}, allocator) + renormalize_north_edge(E_north, P_left, P_right, A::P, allocator) Absorb a local effective tensor `A` into the north edge using the given projectors and environment tensors. @@ -16,22 +16,22 @@ environment tensors. ``` """ function renormalize_north_edge( - (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork + (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork, allocator ) return renormalize_north_edge( edge(env, NORTH, row - 1, col), P_left[NORTH, row, col], P_right[NORTH, row, _prev(col, end)], - network[row, col], # so here it's fine + network[row, col], allocator ) end -function renormalize_north_edge(E_north, P_left, P_right, A) +function renormalize_north_edge(E_north, P_left, P_right, A, allocator) A_west = _rotl90_localsandwich(A) - return renormalize_west_edge(E_north, P_left, P_right, A_west) + return renormalize_west_edge(E_north, P_left, P_right, A_west, allocator) end # specialize PartitionFunction to avoid permute(A) -function renormalize_north_edge(E_north::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor) - return @tensor begin +function renormalize_north_edge(E_north::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor, allocator) + return @tensor allocator = allocator begin temp = permute(E_north, ((2, 1), (3,))) # impose D_N as 1st leg PE[D_N D_E; χNW χ_E] := temp[D_N χNW; χNE] * P_left[χNE D_E; χ_E] PEA[D_W χNW; D_S χ_E] := A[D_W D_S; D_N D_E] * PE[D_N D_E; χNW χ_E] @@ -41,7 +41,7 @@ function renormalize_north_edge(E_north::CTMRG_PF_EdgeTensor, P_left, P_right, A end @generated function renormalize_north_edge( - E_north::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H} + E_north::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {T, S, N, H} @assert N == H + 3 @@ -61,15 +61,15 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $E_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $E_out_e := $rhs)) end # East edge # --------- """ - renormalize_east_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}) - renormalize_east_edge(E_east, P_left, P_right, A::P) + renormalize_east_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}, allocator) + renormalize_east_edge(E_east, P_left, P_right, A::P, allocator) Absorb a local effective tensor into the east edge using the given projectors and environment tensors. @@ -87,22 +87,22 @@ environment tensors. ``` """ function renormalize_east_edge( - (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork + (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork, allocator ) return renormalize_east_edge( edge(env, EAST, row, col + 1), P_left[EAST, row, col], P_right[EAST, _prev(row, end), col], - network[row, col], + network[row, col], allocator ) end -function renormalize_east_edge(E_east, P_left, P_right, A) +function renormalize_east_edge(E_east, P_left, P_right, A, allocator) A_west = _rot180_localsandwich(A) - return renormalize_west_edge(E_east, P_left, P_right, A_west) + return renormalize_west_edge(E_east, P_left, P_right, A_west, allocator) end # specialize PartitionFunction to avoid permute(A) -function renormalize_east_edge(E_east::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor) - return @tensor begin +function renormalize_east_edge(E_east::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor, allocator) + return @tensor allocator = allocator begin temp = permute(P_right, ((3, 1), (2,))) # impose D_N as 1st leg PE[D_N D_E; χN χSE] := temp[D_N χN; χNE] * E_east[χNE D_E; χSE] PEA[D_W χN; χSE D_S] := A[D_W D_S; D_N D_E] * PE[D_N D_E; χN χSE] @@ -111,7 +111,7 @@ function renormalize_east_edge(E_east::CTMRG_PF_EdgeTensor, P_left, P_right, A:: end @generated function renormalize_east_edge( - E_east::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H} + E_east::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {T, S, N, H} @assert N == H + 3 @@ -131,15 +131,15 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $E_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $E_out_e := $rhs)) end # South edge # ---------- """ - renormalize_south_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}) - renormalize_south_edge(E_south, P_left, P_right, A::P) + renormalize_south_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}, allocator) + renormalize_south_edge(E_south, P_left, P_right, A::P, allocator) Absorb a local effective tensor into the south edge using the given projectors and environment tensors. @@ -152,22 +152,22 @@ environment tensors. ``` """ function renormalize_south_edge( - (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork + (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork, allocator ) return renormalize_south_edge( edge(env, SOUTH, row + 1, col), P_left[SOUTH, row, col], P_right[SOUTH, row, _next(col, end)], - network[row, col], + network[row, col], allocator ) end -function renormalize_south_edge(E_south, P_left, P_right, A) +function renormalize_south_edge(E_south, P_left, P_right, A, allocator) A_west = _rotr90_localsandwich(A) - return renormalize_west_edge(E_south, P_left, P_right, A_west) + return renormalize_west_edge(E_south, P_left, P_right, A_west, allocator) end # specialize PartitionFunction to avoid permute(A) -function renormalize_south_edge(E_south::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor) - return @tensor begin +function renormalize_south_edge(E_south::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor, allocator) + return @tensor allocator = allocator begin P_leftp = permute(P_left, ((3, 2), (1,))) # impose χ_W as 1st leg PE[χ_W χSE; D_W D_S] := P_leftp[χ_W D_W; χSW] * E_south[χSE D_S; χSW] PEA[χ_W D_N; χSE D_E] := PE[χ_W χSE; D_W D_S] * A[D_W D_S; D_N D_E] @@ -176,7 +176,7 @@ function renormalize_south_edge(E_south::CTMRG_PF_EdgeTensor, P_left, P_right, A end @generated function renormalize_south_edge( - E_south::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H} + E_south::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {T, S, N, H} @assert N == H + 3 @@ -196,15 +196,15 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $E_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $E_out_e := $rhs)) end # West edge # --------- """ - renormalize_west_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}) - renormalize_west_edge(E_west, P_left, P_right, A::P) + renormalize_west_edge((row, col), env, P_left, P_right, network::InfiniteSquareNetwork{P}, allocator) + renormalize_west_edge(E_west, P_left, P_right, A::P, allocator) Absorb a local effective tensor into the west edge using the given projectors and environment tensors. @@ -222,30 +222,30 @@ environment tensors. ``` """ function renormalize_west_edge( # For simultaneous CTMRG scheme - (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork, + (row, col), env::AbstractCTMRGEnv, P_left, P_right, network::InfiniteSquareNetwork, allocator ) return renormalize_west_edge( edge(env, WEST, row, col - 1), P_left[WEST, row, col], P_right[WEST, _next(row, end), col], - network[row, col], + network[row, col], allocator ) end function renormalize_west_edge( # For sequential CTMRG scheme - (row, col), env::AbstractCTMRGEnv, projectors, network::InfiniteSquareNetwork, + (row, col), env::AbstractCTMRGEnv, projectors, network::InfiniteSquareNetwork, allocator ) return renormalize_west_edge( edge(env, WEST, row, col - 1), projectors[1][row], projectors[2][_next(row, end)], - network[row, col], + network[row, col], allocator ) end function renormalize_west_edge( - E_west::CTMRG_PEPS_EdgeTensor, P_left, P_right, A::PEPSSandwich + E_west::CTMRG_PEPS_EdgeTensor, P_left, P_right, A::PEPSSandwich, allocator ) # starting with P_bottom to save one permute in the end - return @tensor begin + return @tensor allocator = allocator begin # already putting χE in front here to make next permute cheaper PE[χS χNW DSb DWb; DSt DWt] := P_right[χS; χSW DSt DSb] * E_west[χSW DWt DWb; χNW] PEket[χS χNW DNt DEt; DSb DWb d] := @@ -255,8 +255,8 @@ function renormalize_west_edge( edge[χS DEt DEb; χN] := corner[χS DEt DEb; χNW DNt DNb] * P_left[χNW DNt DNb; χN] end end -function renormalize_west_edge(E_west::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor) - return @tensor begin +function renormalize_west_edge(E_west::CTMRG_PF_EdgeTensor, P_left, P_right, A::PFTensor, allocator) + return @tensor allocator = allocator begin PE[χ_S χNW; D_W D_S] := P_right[χ_S; χSW D_S] * E_west[χSW D_W; χNW] PEA[χ_S D_E; χNW D_N] := PE[χ_S χNW; D_W D_S] * A[D_W D_S; D_N D_E] edge[χ_S D_E; χ_N] := PEA[χ_S D_E; χNW D_N] * P_left[χNW D_N; χ_N] @@ -264,7 +264,7 @@ function renormalize_west_edge(E_west::CTMRG_PF_EdgeTensor, P_left, P_right, A:: end @generated function renormalize_west_edge( - E_west::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H} + E_west::CTMRGEdgeTensor{T, S, N}, P_left, P_right, A::PEPOSandwich{H}, allocator ) where {T, S, N, H} @assert N == H + 3 @@ -284,5 +284,5 @@ end P_left_e, ) - return macroexpand(@__MODULE__, :(return @autoopt @tensor $E_out_e := $rhs)) + return macroexpand(@__MODULE__, :(return @autoopt @tensor allocator = allocator $E_out_e := $rhs)) end diff --git a/src/algorithms/contractions/local_patch/densitymatrix/bp.jl b/src/algorithms/contractions/local_patch/densitymatrix/bp.jl index 16c087b52..0d91d0369 100644 --- a/src/algorithms/contractions/local_patch/densitymatrix/bp.jl +++ b/src/algorithms/contractions/local_patch/densitymatrix/bp.jl @@ -4,7 +4,7 @@ # NOTE: currently restricted to 1x1, 2x1, and 1x2 patches, since evaluating larger patches # without including loop corrections tends to be a crude and not very useful appoximation. -function _contract_densitymatrix(inds::NTuple{N, Val}, state, env::BPEnv) where {N} +function _contract_densitymatrix(inds::NTuple{N, Val}, state, env::BPEnv, allocator) where {N} sites = _patch_inds(inds) return throw( ArgumentError( @@ -15,7 +15,7 @@ function _contract_densitymatrix(inds::NTuple{N, Val}, state, env::BPEnv) where end function reduced_densitymatrix1x1( - ind::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv + ind::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv, allocator ) row, col = Tuple(ind) M_north = env[NORTH, row - 1, col] @@ -23,7 +23,7 @@ function reduced_densitymatrix1x1( M_south = env[SOUTH, row + 1, col] M_west = env[WEST, row, col - 1] - @autoopt @tensor ρ[dt; db] := + @autoopt @tensor allocator = allocator ρ[dt; db] := ket[row, col][dt; DNt DEt DSt DWt] * conj(bra[row, col][db; DNb DEb DSb DWb]) * M_north[DNt; DNb] * @@ -35,7 +35,7 @@ function reduced_densitymatrix1x1( end function reduced_densitymatrix2x1( - coord::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv + coord::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv, allocator ) row, col = Tuple(coord) M_north = env[NORTH, row - 1, col] @@ -45,7 +45,7 @@ function reduced_densitymatrix2x1( M_southwest = env[WEST, row + 1, col - 1] M_northwest = env[WEST, row, col - 1] - @autoopt @tensor ρ[dNt dSt; dNb dSb] := + @autoopt @tensor allocator = allocator ρ[dNt dSt; dNb dSb] := ket[row, col][dNt; DNt DNEt DMt DNWt] * ket[row + 1, col][dSt; DMt DSEt DSt DSWt] * conj(bra[row, col][dNb; DNb DNEb DMb DNWb]) * @@ -61,7 +61,7 @@ function reduced_densitymatrix2x1( end function reduced_densitymatrix1x2( - coord::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv + coord::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv, allocator ) row, col = Tuple(coord) M_west = env[WEST, row, col - 1] @@ -75,7 +75,7 @@ function reduced_densitymatrix1x2( A_east = ket[row, col + 1] Ā_east = bra[row, col + 1] - @autoopt @tensor ρ[dWt dEt; dWb dEb] := + @autoopt @tensor allocator = allocator ρ[dWt dEt; dWb dEb] := A_west[dWt; DNWt DMt DSWt DWt] * A_east[dEt; DNEt DEt DSEt DMt] * conj(Ā_west[dWb; DNWb DMb DSWb DWb]) * diff --git a/src/algorithms/contractions/local_patch/densitymatrix/ctmrg.jl b/src/algorithms/contractions/local_patch/densitymatrix/ctmrg.jl index 882b95097..12cc7033d 100644 --- a/src/algorithms/contractions/local_patch/densitymatrix/ctmrg.jl +++ b/src/algorithms/contractions/local_patch/densitymatrix/ctmrg.jl @@ -4,7 +4,7 @@ # Keep contraction order but try to optimize intermediate permutations: # EE_SWA is largest object so keep largest legs to the front there function reduced_densitymatrix1x1( - inds::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv + inds::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv, allocator ) row, col = Tuple(inds) @@ -25,19 +25,19 @@ function reduced_densitymatrix1x1( edge(env, WEST, row, col - 1), corner(env, NORTHWEST, row - 1, col - 1) ) - @tensor EE_SW[χSE χNW DSb DWb; DSt DWt] := + @tensor allocator = allocator EE_SW[χSE χNW DSb DWb; DSt DWt] := E_south[χSE DSt DSb; χSW] * E_west[χSW DWt DWb; χNW] - @tensor EE_SWA[χSE χNW DNt DEt; dt DSb DWb] := + @tensor allocator = allocator EE_SWA[χSE χNW DNt DEt; dt DSb DWb] := EE_SW[χSE χNW DSb DWb; DSt DWt] * A[dt; DNt DEt DSt DWt] - @tensor EE_NE[DNb DEb; χSE χNW DNt DEt] := + @tensor allocator = allocator EE_NE[DNb DEb; χSE χNW DNt DEt] := E_north[χNW DNt DNb; χNE] * E_east[χNE DEt DEb; χSE] - @tensor EEAEE[dt; DNb DEb DSb DWb] := + @tensor allocator = allocator EEAEE[dt; DNb DEb DSb DWb] := EE_NE[DNb DEb; χSE χNW DNt DEt] * EE_SWA[χSE χNW DNt DEt; dt DSb DWb] - @tensor ρ[dt; db] := EEAEE[dt; DNb DEb DSb DWb] * conj(Ā[db; DNb DEb DSb DWb]) + @tensor allocator = allocator ρ[dt; db] := EEAEE[dt; DNb DEb DSb DWb] * conj(Ā[db; DNb DEb DSb DWb]) return ρ / str(ρ) end @@ -45,7 +45,7 @@ end # Special case 2x1 density matrix: # Keep contraction order but try to optimize intermediate permutations: function reduced_densitymatrix2x1( - ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv + ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv, allocator ) row, col = Tuple(ind) @@ -70,32 +70,32 @@ function reduced_densitymatrix2x1( edge(env, WEST, row, col - 1), corner(env, NORTHWEST, row - 1, col - 1) ) - @tensor EE_NW[χW χNE DNWt DNt; DNWb DNb] := + @tensor allocator = allocator EE_NW[χW χNE DNWt DNt; DNWb DNb] := E_northwest[χW DNWt DNWb; χNW] * E_north[χNW DNt DNb; χNE] - @tensor EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] := + @tensor allocator = allocator EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] := EE_NW[χW χNE DNWt DNt; DNWb DNb] * conj(Ā_north[dNb; DNb DNEb DMb DNWb]) - @tensor EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] := + @tensor allocator = allocator EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] := EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] * A_north[dNt; DNt DNEt DMt DNWt] - @tensor EEEAA_N[dNt dNb; χW DMt DMb χE] := + @tensor allocator = allocator EEEAA_N[dNt dNb; χW DMt DMb χE] := EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] * E_northeast[χNE DNEt DNEb; χE] - @tensor EE_SE[χE χSW DSEt DSt; DSEb DSb] := + @tensor allocator = allocator EE_SE[χE χSW DSEt DSt; DSEb DSb] := E_southeast[χE DSEt DSEb; χSE] * E_south[χSE DSt DSb; χSW] - @tensor EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] := + @tensor allocator = allocator EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] := EE_SE[χE χSW DSEt DSt; DSEb DSb] * conj(Ā_south[dSb; DMb DSEb DSb DSWb]) - @tensor EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] := + @tensor allocator = allocator EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] := EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] * A_south[dSt; DMt DSEt DSt DSWt] - @tensor EEEAA_S[χW DMt DMb χE; dSt dSb] := + @tensor allocator = allocator EEEAA_S[χW DMt DMb χE; dSt dSb] := EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] * E_southwest[χSW DSWt DSWb; χW] - @tensor ρ[dNt dSt; dNb dSb] := + @tensor allocator = allocator ρ[dNt dSt; dNb dSb] := EEEAA_N[dNt dNb; χW DMt DMb χE] * EEEAA_S[χW DMt DMb χE; dSt dSb] return ρ / str(ρ) end function reduced_densitymatrix1x2( - ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv + ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, env::AbstractCTMRGEnv, allocator ) row, col = Tuple(ind) @@ -120,25 +120,25 @@ function reduced_densitymatrix1x2( edge(env, WEST, row, col - 1), corner(env, NORTHWEST, row - 1, col - 1) ) - @tensor EE_SW[χS χNW DSWt DWt; DSWb DWb] := + @tensor allocator = allocator EE_SW[χS χNW DSWt DWt; DSWb DWb] := E_southwest[χS DSWt DSWb; χSW] * E_west[χSW DWt DWb; χNW] - @tensor EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] := + @tensor allocator = allocator EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] := EE_SW[χS χNW DSWt DWt; DSWb DWb] * conj(Ā_west[dWb; DNWb DMb DSWb DWb]) - @tensor EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] := + @tensor allocator = allocator EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] := EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] * A_west[dWt; DNWt DMt DSWt DWt] - @tensor EEEAA_W[dWt dWb; χS DMt DMb χN] := + @tensor allocator = allocator EEEAA_W[dWt dWb; χS DMt DMb χN] := EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] * E_northwest[χNW DNWt DNWb; χN] - @tensor EE_NE[χN χSE DNEt DEt; DNEb DEb] := + @tensor allocator = allocator EE_NE[χN χSE DNEt DEt; DNEb DEb] := E_northeast[χN DNEt DNEb; χNE] * E_east[χNE DEt DEb; χSE] - @tensor EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] := + @tensor allocator = allocator EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] := EE_NE[χN χSE DNEt DEt; DNEb DEb] * conj(Ā_east[dEb; DNEb DEb DSEb DMb]) - @tensor EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] := + @tensor allocator = allocator EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] := EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] * A_east[dEt; DNEt DEt DSEt DMt] - @tensor EEEAA_E[χS DMt DMb χN; dEt dEb] := + @tensor allocator = allocator EEEAA_E[χS DMt DMb χN; dEt dEb] := EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] * E_southeast[χSE DSEt DSEb; χS] - @tensor ρ[dWt dEt; dWb dEb] := + @tensor allocator = allocator ρ[dWt dEt; dWb dEb] := EEEAA_W[dWt dWb; χS DMt DMb χN] * EEEAA_E[χS DMt DMb χN; dEt dEb] return ρ / str(ρ) diff --git a/src/algorithms/contractions/local_patch/densitymatrix/generic.jl b/src/algorithms/contractions/local_patch/densitymatrix/generic.jl index 28f2c65ef..d625d5d14 100644 --- a/src/algorithms/contractions/local_patch/densitymatrix/generic.jl +++ b/src/algorithms/contractions/local_patch/densitymatrix/generic.jl @@ -13,7 +13,7 @@ legs become the indices of `ρ`: `physicallabel(:O, 1, k)` in its codomain and than `tr`, since the supertrace carries the fermionic signs. """ @generated function _contract_densitymatrix( - inds::NTuple{N, Val}, state, env + inds::NTuple{N, Val}, state, env, allocator ) where {N} sites = _patch_inds(inds) rowrange, colrange = _patch_ranges(sites) @@ -29,7 +29,7 @@ than `tr`, since the supertrace carries the fermionic signs. ntuple(i -> physicallabel(:O, 2, i), N), ) - multex = _tensor_expr(multiplication_ex, result) + multex = _tensor_expr(multiplication_ex, allocator, result) return quote $(macroexpand(@__MODULE__, multex)) return ρ / str(ρ) diff --git a/src/algorithms/contractions/local_patch/network_expr.jl b/src/algorithms/contractions/local_patch/network_expr.jl index 6e73b7b45..d4d8d319f 100644 --- a/src/algorithms/contractions/local_patch/network_expr.jl +++ b/src/algorithms/contractions/local_patch/network_expr.jl @@ -10,8 +10,8 @@ $(SIGNATURES) Wrap an assembled product in `@autoopt @tensor`, with `lhs` on the left if given and as a scalar otherwise. """ -function _tensor_expr(prod, lhs = nothing) - return isnothing(lhs) ? :(@autoopt @tensor $prod) : :(@autoopt @tensor $lhs := $prod) +function _tensor_expr(prod, allocator, lhs = nothing) + return isnothing(lhs) ? :(@autoopt @tensor allocator = allocator $prod) : :(@autoopt @tensor allocator = allocator $lhs := $prod) end diff --git a/src/algorithms/contractions/local_patch/patch_contractions.jl b/src/algorithms/contractions/local_patch/patch_contractions.jl index caa830d73..685580075 100644 --- a/src/algorithms/contractions/local_patch/patch_contractions.jl +++ b/src/algorithms/contractions/local_patch/patch_contractions.jl @@ -15,7 +15,7 @@ types of `env`, `state` and `operator` respectively, so this single method cover combination those three have methods for. """ @generated function _contract_local_operator( - inds::NTuple{N, Val}, operator, state, env + inds::NTuple{N, Val}, operator, state, env, allocator ) where {N} sites = _patch_inds(inds) rowrange, colrange = _patch_ranges(sites) @@ -27,7 +27,7 @@ combination those three have methods for. operator_contraction_expr(operator, N)..., ) - returnex = _tensor_expr(multiplication_ex) + returnex = _tensor_expr(multiplication_ex, allocator) return macroexpand(@__MODULE__, returnex) end @@ -44,7 +44,7 @@ leg is left open and no operator factor is needed. Note this is the norm of the the given environment, not the physical norm of the state. """ @generated function _contract_local_norm( - inds::NTuple{N, Val}, state, env + inds::NTuple{N, Val}, state, env, allocator ) where {N} sites = _patch_inds(inds) rowrange, colrange = _patch_ranges(sites) @@ -55,6 +55,6 @@ the given environment, not the physical norm of the state. bulk_contraction_expr(state, rowrange, colrange, nothing)..., # legs paired, not open ) - returnex = _tensor_expr(multiplication_ex) + returnex = _tensor_expr(multiplication_ex, allocator) return macroexpand(@__MODULE__, returnex) end diff --git a/src/algorithms/ctmrg/c4v.jl b/src/algorithms/ctmrg/c4v.jl index cc10c6216..63c8eaeb0 100644 --- a/src/algorithms/ctmrg/c4v.jl +++ b/src/algorithms/ctmrg/c4v.jl @@ -93,12 +93,12 @@ the corners are identical, and the edges have identical singular values. convergence_tensors(env::C4vCTMRGEnv, ::C4vCTMRG) = (view(env.corners, 1:1, :, :), view(env.edges, 1:1, :, :)) -function ctmrg_iteration(network, env::C4vCTMRGEnv, alg::C4vCTMRG) +function ctmrg_iteration(network, env::C4vCTMRGEnv, alg::C4vCTMRG, allocator) projector_alg = alg.projector_alg enlarged_corner = c4v_enlarge(network, env, projector_alg) projector, info = c4v_projector!(enlarged_corner, projector_alg) - edge′ = c4v_renormalize_edge(network, env, projector) - corner′ = c4v_renormalize_corner(edge′, projector, info, projector_alg) + edge′ = c4v_renormalize_edge(network, env, projector, allocator) + corner′ = c4v_renormalize_corner(edge′, projector, info, projector_alg, allocator) return C4vCTMRGEnv(corner′, edge′), info end @@ -116,9 +116,9 @@ Renormalize the single C₄ᵥ edge tensor. | ``` """ -function c4v_renormalize_edge(network, env::C4vCTMRGEnv, projector) +function c4v_renormalize_edge(network, env::C4vCTMRGEnv, projector, allocator) # TODO: possible missing twists for fermions - new_edge = renormalize_north_edge(env.edges[1], projector, projector', network[1, 1]) + new_edge = renormalize_north_edge(env.edges[1], projector, projector', network[1, 1], allocator) # additional Hermitian projection step for numerical stability return normalize(_project_hermitian(new_edge)) end @@ -153,28 +153,28 @@ to rewrite the renormalized corner as which reuses the renormalized edge `E′` (`new_edge`). (Credit: https://github.com/qiyang-ustc/QRCTM/blob/dd160116c3d7b02076691ceaf0a9833511ae532d/heisenberg.py#L80) """ -function c4v_renormalize_corner(new_edge::CTMRGEdgeTensor, projector, info, ::C4vEighProjector) +function c4v_renormalize_corner(new_edge::CTMRGEdgeTensor, projector, info, ::C4vEighProjector, allocator) return normalize(info.D) end -function c4v_renormalize_corner(new_edge::CTMRGEdgeTensor, projector, info, ::C4vQRProjector) +function c4v_renormalize_corner(new_edge::CTMRGEdgeTensor, projector, info, ::C4vQRProjector, allocator) # TODO: possible missing twists for fermions # contract updated edge and R edge′ = physical_flip(new_edge) ER = edge′ * twistdual(info.R, 1) # contract (edge′, R) with projector - new_corner = contract_edges(ER, projector) + new_corner = contract_edges(ER, projector, allocator) new_corner = project_hermitian(new_corner) return normalize(new_corner) end """Contract two C4v CTMRG edge tensors into a corner tensor.""" @generated function contract_edges( - EL::CTMRGEdgeTensor{T, S, N}, ER::CTMRGEdgeTensor{T, S, N} + EL::CTMRGEdgeTensor{T, S, N}, ER::CTMRGEdgeTensor{T, S, N}, allocator ) where {T, S, N} C´_e = tensorexpr(:C´, -1, -2) EL_e = tensorexpr(:EL, (-1, (2:N)...), 1) ER_e = tensorexpr(:ER, 1:N, -2) - return macroexpand(@__MODULE__, :(return @tensor $C´_e := $EL_e * $ER_e)) + return macroexpand(@__MODULE__, :(return @tensor allocator = allocator $C´_e := $EL_e * $ER_e)) end function C4vCTMRGEnv( diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 65f23a979..8814c4eab 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -43,11 +43,11 @@ function CTMRGAlgorithm(; end """ - ctmrg_iteration(network, env, alg::CTMRGAlgorithm) -> env′, info + ctmrg_iteration(network, env, alg::CTMRGAlgorithm, allocator) -> env′, info Perform a single CTMRG iteration in which all directions are being grown and renormalized. """ -function ctmrg_iteration(network, env, alg::CTMRGAlgorithm) end +function ctmrg_iteration(network, env, alg::CTMRGAlgorithm, allocator) end """ leading_boundary(env₀, network; kwargs...) -> env, info @@ -106,7 +106,9 @@ function leading_boundary(env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwo return leading_boundary(env₀, network, alg) end function leading_boundary( - env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm + env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm; + scheduler = Defaults.scheduler[], + allocator = default_allocator(network, scheduler) ) check_input(leading_boundary, network, env₀, alg) log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) @@ -116,23 +118,23 @@ function leading_boundary( return corner_spectra(first(convergence_tensors(env₀, alg)), alg), nothing end η = one(real(scalartype(network))) - ctmrg_loginit!(log, η, network, env₀) + ctmrg_loginit!(log, η, network, env₀, allocator) local info_iter converged = false for iter in 1:(alg.maxiter) env_prev = env - env, info_iter = ctmrg_iteration(network, env, alg) + env, info_iter = ctmrg_iteration(network, env, alg, allocator) η, CS, TS = calc_convergence_lazy(env, env_prev, CS, TS, alg; exact = iter == alg.maxiter) if η ≤ alg.tol && iter ≥ alg.miniter - ctmrg_logfinish!(log, iter, η, network, env) + ctmrg_logfinish!(log, iter, η, network, env, allocator) converged = true break end if iter == alg.maxiter - ctmrg_logcancel!(log, iter, η, network, env) + ctmrg_logcancel!(log, iter, η, network, env, allocator) else - ctmrg_logiter!(log, iter, η, network, env, isnothing(TS)) + ctmrg_logiter!(log, iter, η, network, env, allocator, isnothing(TS)) end end info = (; @@ -156,21 +158,21 @@ function check_input(::typeof(leading_boundary), network, env, alg::CTMRGAlgorit @non_differentiable check_input(args...) # custom CTMRG logging -function ctmrg_loginit!(log, η, network, env) - return @log_initialization loginit!(log, η, network_value(network, env)) +function ctmrg_loginit!(log, η, network, env, allocator) + return @log_initialization loginit!(log, η, network_value(network, env, allocator)) end -function ctmrg_logiter!(log, iter, η, network, env, lower_bound::Bool = false) +function ctmrg_logiter!(log, iter, η, network, env, allocator, lower_bound::Bool = false) return @log_iteration _iterlog_message( - logiter!(log, iter, η, network_value(network, env)), lower_bound + logiter!(log, iter, η, network_value(network, env, allocator)), lower_bound ) end _iterlog_message(log, lower_bound::Bool) = lower_bound ? replace(sprint(show, log), "err = " => "err ≥ ") : log -function ctmrg_logfinish!(log, iter, η, network, env) - return @log_convergence logfinish!(log, iter, η, network_value(network, env)) +function ctmrg_logfinish!(log, iter, η, network, env, allocator) + return @log_convergence logfinish!(log, iter, η, network_value(network, env, allocator)) end -function ctmrg_logcancel!(log, iter, η, network, env) - return @log_nonconvergence logcancel!(log, iter, η, network_value(network, env)) +function ctmrg_logcancel!(log, iter, η, network, env, allocator) + return @log_nonconvergence logcancel!(log, iter, η, network_value(network, env, allocator)) end @non_differentiable ctmrg_loginit!(args...) diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index f6bf3463f..150f6e586 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -19,14 +19,14 @@ Given that the CTMRG run is converged, the returned environment will be element-wise converged to `envprev`. """ function gauge_fix( - envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, alg::ScramblingEnvGauge + envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, alg::ScramblingEnvGauge, allocator ) where {C, T} - signs, corner_phases, edge_phases = compute_gauge_fix_gauge(envfinal, envprev, alg) - return fix_phases(envfinal, signs, corner_phases, edge_phases) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(envfinal, envprev, alg, allocator) + return fix_phases(envfinal, signs, corner_phases, edge_phases, allocator) end function compute_gauge_fix_gauge( - envfinal::CTMRGEnv{C, T}, envprev::CTMRGEnv{C, T}, alg::ScramblingEnvGauge + envfinal::CTMRGEnv{C, T}, envprev::CTMRGEnv{C, T}, alg::ScramblingEnvGauge, allocator ) where {C, T} # Check if spaces in envprev and envfinal are the same same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) @@ -36,18 +36,18 @@ function compute_gauge_fix_gauge( all(same_spaces) || throw(ArgumentError("Spaces of envfinal and envprev are not the same")) # find relative phases - signs = compute_relative_phases(envfinal, envprev, alg) + signs = compute_relative_phases(envfinal, envprev, alg, allocator) # find additional global phases corner_phases, edge_phases = compute_global_phases( - fix_relative_phases(envfinal, signs), envprev + fix_relative_phases(envfinal, signs, allocator), envprev ) return signs, corner_phases, edge_phases end function compute_gauge_fix_gauge( - envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, alg::ScramblingEnvGauge + envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, alg::ScramblingEnvGauge, allocator ) where {C, T} # Check if spaces in envprev and envfinal are the same same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) @@ -57,24 +57,24 @@ function compute_gauge_fix_gauge( all(same_spaces) || throw(ArgumentError("Spaces of envfinal and envprev are not the same")) # find relative phases - signs = compute_relative_phases(envfinal, envprev, alg) + signs = compute_relative_phases(envfinal, envprev, alg, allocator) # find additional global phases corner_phases, edge_phases = compute_global_phases( - fix_relative_phases(envfinal, signs), envprev + fix_relative_phases(envfinal, signs, allocator), envprev ) return signs, corner_phases, edge_phases end -function fix_phases(env::AbstractCTMRGEnv, signs, corner_phases, edge_phases) - env_fixed_relative = fix_relative_phases(env, signs) - env_fixed_global = fix_global_phases(env_fixed_relative, corner_phases, edge_phases) +function fix_phases(env::AbstractCTMRGEnv, signs, corner_phases, edge_phases, allocator) + env_fixed_relative = fix_relative_phases(env, signs, allocator) + env_fixed_global = fix_global_phases(env_fixed_relative, corner_phases, edge_phases, allocator) return env_fixed_global end function compute_relative_phases( - envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge + envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge, allocator ) where {C, T} signs = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) @@ -113,7 +113,7 @@ function compute_relative_phases( return signs end -function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, ::ScramblingEnvGauge) where {C, T} +function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, ::ScramblingEnvGauge, allocator) where {C, T} # "general" algorithm from https://arxiv.org/abs/2311.11894 Tprev = envprev.edges[1, 1, 1] Tfinal = envfinal.edges[1, 1, 1] @@ -146,13 +146,14 @@ end @generated function mps_transfer_right( v::AbstractTensorMap{<:Any, S, 1, N₁}, A::GenericMPSTensor{S, N₂}, Abar::GenericMPSTensor{S, N₂}, + allocator ) where {S, N₁, N₂} t_out = tensorexpr(:v, -1, -(2:(N₁ + 1))) t_top = tensorexpr(:A, (-1, reverse(3:(N₂ + 1))...), 1) t_bot = tensorexpr(:Abar, (-(N₁ + 1), reverse(3:(N₂ + 1))...), 2) t_in = tensorexpr(:v, 1, (-(2:N₁)..., 2)) return macroexpand( - @__MODULE__, :(return @tensor $t_out := $t_top * conj($t_bot) * $t_in) + @__MODULE__, :(return @tensor allocator = allocator $t_out := $t_top * conj($t_bot) * $t_in) ) end @@ -165,10 +166,11 @@ end function right_transfermatrix_fixedpoint( tops, bottoms, alg = Arnoldi(), ρ0 = initialize_right_fixedpoint(tops, bottoms), + allocator = default_allocator(ρ0, Defaults.scheduler[]) ) _, vecs, info = eigsolve(ρ0, 1, :LM, alg) do ρ return foldr(zip(tops, bottoms); init = ρ) do (top, bottom), ρ - return mps_transfer_right(ρ, top, bottom) + return mps_transfer_right(ρ, top, bottom, allocator) end end ignore_derivatives() do @@ -178,56 +180,56 @@ function right_transfermatrix_fixedpoint( end # Explicit fixing of relative phases (doing this compactly in a loop is annoying) -function fix_relative_phases(envfinal::CTMRGEnv, signs) +function fix_relative_phases(envfinal::CTMRGEnv, signs, allocator) corners_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) Cf = if dir == NORTHWEST - fix_gauge_northwest_corner((r, c), envfinal, signs) + fix_gauge_northwest_corner((r, c), envfinal, signs, allocator) elseif dir == NORTHEAST - fix_gauge_northeast_corner((r, c), envfinal, signs) + fix_gauge_northeast_corner((r, c), envfinal, signs, allocator) elseif dir == SOUTHEAST - fix_gauge_southeast_corner((r, c), envfinal, signs) + fix_gauge_southeast_corner((r, c), envfinal, signs, allocator) else # dir == SOUTHWEST - fix_gauge_southwest_corner((r, c), envfinal, signs) + fix_gauge_southwest_corner((r, c), envfinal, signs, allocator) end return Cf end edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) Ef = if dir == NORTHWEST - fix_gauge_north_edge((r, c), envfinal, signs) + fix_gauge_north_edge((r, c), envfinal, signs, allocator) elseif dir == NORTHEAST - fix_gauge_east_edge((r, c), envfinal, signs) + fix_gauge_east_edge((r, c), envfinal, signs, allocator) elseif dir == SOUTHEAST - fix_gauge_south_edge((r, c), envfinal, signs) + fix_gauge_south_edge((r, c), envfinal, signs, allocator) else # dir == SOUTHWEST - fix_gauge_west_edge((r, c), envfinal, signs) + fix_gauge_west_edge((r, c), envfinal, signs, allocator) end return Ef end return CTMRGEnv(corners_fixed, edges_fixed) end -function fix_relative_phases(envfinal::C4vCTMRGEnv, signs) +function fix_relative_phases(envfinal::C4vCTMRGEnv, signs, allocator) # all corners are the same corner_fixed = map(eachcoordinate(envfinal, 1:1)) do (dir, r, c) - fix_gauge_northwest_corner((r, c), envfinal, signs) + fix_gauge_northwest_corner((r, c), envfinal, signs, allocator) end edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) Ef = if dir == NORTHWEST - fix_gauge_north_edge((r, c), envfinal, signs) + fix_gauge_north_edge((r, c), envfinal, signs, allocator) elseif dir == NORTHEAST - fix_gauge_east_edge((r, c), envfinal, signs) + fix_gauge_east_edge((r, c), envfinal, signs, allocator) elseif dir == SOUTHEAST - fix_gauge_south_edge((r, c), envfinal, signs) + fix_gauge_south_edge((r, c), envfinal, signs, allocator) else # dir == SOUTHWEST - fix_gauge_west_edge((r, c), envfinal, signs) + fix_gauge_west_edge((r, c), envfinal, signs, allocator) end return Ef end return C4vCTMRGEnv(corner_fixed, edges_fixed) end function fix_relative_phases( - U::Array{Ut, 3}, V::Array{Vt, 3}, signs + U::Array{Ut, 3}, V::Array{Vt, 3}, signs, allocator ) where {Ut <: AbstractTensorMap, Vt <: AbstractTensorMap} U_fixed = map(eachindex(IndexCartesian(), U)) do I dir, r, c = Tuple(I) @@ -267,16 +269,16 @@ Fix global multiplicative phase of the environment tensors. To that end, the dot between all corners and all edges are computed to obtain the global phase which is then divided out. """ -function fix_global_phases(envfix::AbstractCTMRGEnv, envprev::AbstractCTMRGEnv) +function fix_global_phases(envfix::AbstractCTMRGEnv, envprev::AbstractCTMRGEnv, allocator) corner_phases, edge_phases = compute_global_phases(envfix, envprev) - return fix_global_phases(envfix, corner_phases, edge_phases) + return fix_global_phases(envfix, corner_phases, edge_phases, allocator) end -function fix_global_phases(env::CTMRGEnv, corner_phases, edge_phases) +function fix_global_phases(env::CTMRGEnv, corner_phases, edge_phases, allocator) cornersgfix = env.corners .* inv.(corner_phases) edgesgfix = env.edges .* inv.(edge_phases) return CTMRGEnv(cornersgfix, edgesgfix) end -function fix_global_phases(env::C4vCTMRGEnv, corner_phases, edge_phases) +function fix_global_phases(env::C4vCTMRGEnv, corner_phases, edge_phases, allocator) cornersgfix = env.corners .* inv.(corner_phases) edgesgfix = env.edges .* inv.(edge_phases) return C4vCTMRGEnv(cornersgfix, edgesgfix) diff --git a/src/algorithms/ctmrg/initialization.jl b/src/algorithms/ctmrg/initialization.jl index 42d7a6d87..c458ab000 100644 --- a/src/algorithms/ctmrg/initialization.jl +++ b/src/algorithms/ctmrg/initialization.jl @@ -39,10 +39,12 @@ function initialize_ctmrg_environment( elt::Type, n::InfiniteSquareNetwork, alg::ApplicationInitialization, - env0 = ProductStateEnv(alg.f, elt, n) + env0 = ProductStateEnv(alg.f, elt, n); + scheduler = Defaults.scheduler[], + allocator = default_allocator(n, scheduler) ) dummy_alg = SimultaneousCTMRG(trunc = (; alg = :notrunc)) - env, = ctmrg_iteration(n, CTMRGEnv(env0), dummy_alg) + env, = ctmrg_iteration(n, CTMRGEnv(env0), dummy_alg, allocator) return env end diff --git a/src/algorithms/ctmrg/projectors/fullinfinite.jl b/src/algorithms/ctmrg/projectors/fullinfinite.jl index 5fdc4855f..17ec0f745 100644 --- a/src/algorithms/ctmrg/projectors/fullinfinite.jl +++ b/src/algorithms/ctmrg/projectors/fullinfinite.jl @@ -36,7 +36,7 @@ end PROJECTOR_SYMBOLS[:FullInfiniteProjector] = FullInfiniteProjector -function compute_projector(enlarged_corners, alg::FullInfiniteProjector) +function compute_projector(enlarged_corners, alg::FullInfiniteProjector, allocator) halfinf_left = half_infinite_environment(enlarged_corners[1], enlarged_corners[2]) halfinf_right = half_infinite_environment(enlarged_corners[3], enlarged_corners[4]) @@ -56,6 +56,6 @@ function compute_projector(enlarged_corners, alg::FullInfiniteProjector) end end - P_left, P_right = contract_projectors(U, S, V, halfinf_left, halfinf_right) + P_left, P_right = contract_projectors(U, S, V, halfinf_left, halfinf_right, allocator) return (P_left, P_right), (; U, S, V, truncation_error) end diff --git a/src/algorithms/ctmrg/projectors/halfinfinite.jl b/src/algorithms/ctmrg/projectors/halfinfinite.jl index 86e4caa9f..3b647eabf 100644 --- a/src/algorithms/ctmrg/projectors/halfinfinite.jl +++ b/src/algorithms/ctmrg/projectors/halfinfinite.jl @@ -37,12 +37,12 @@ end PROJECTOR_SYMBOLS[:HalfInfiniteProjector] = HalfInfiniteProjector """ - compute_projector(enlarged_corners, alg::ProjectorAlgorithm) + compute_projector(enlarged_corners, alg::ProjectorAlgorithm, allocator) Determine left and right projectors at the bond given determined by the enlarged corners using the specified `alg`. """ -function compute_projector(enlarged_corners, alg::HalfInfiniteProjector) +function compute_projector(enlarged_corners, alg::HalfInfiniteProjector, allocator) # SVD half-infinite environment halfinf = half_infinite_environment(enlarged_corners...) svd_alg = decomposition_algorithm(alg) @@ -59,6 +59,6 @@ function compute_projector(enlarged_corners, alg::HalfInfiniteProjector) end end - P_left, P_right = contract_projectors(U, S, V, enlarged_corners...) + P_left, P_right = contract_projectors(U, S, V, enlarged_corners..., allocator) return (P_left, P_right), (; U, S, V, truncation_error) end diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index f71606cb2..ff7602df1 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -42,7 +42,7 @@ CTMRG_SYMBOLS[:SequentialCTMRG] = SequentialCTMRG Perform sequential CTMRG left move on the `col`-th column. """ -function ctmrg_leftmove(col::Int, network, env::CTMRGEnv, alg::SequentialCTMRG) +function ctmrg_leftmove(col::Int, network, env::CTMRGEnv, alg::SequentialCTMRG, allocator) #= ----> left move C1 ← T1 ← r-1 @@ -52,16 +52,16 @@ function ctmrg_leftmove(col::Int, network, env::CTMRGEnv, alg::SequentialCTMRG) C4 → T3 → r+1 c-1 c =# - projectors, info = sequential_projectors(col, network, env, alg.projector_alg) - env = renormalize_sequentially(col, projectors, network, env) + projectors, info = sequential_projectors(col, network, env, alg.projector_alg, allocator) + env = renormalize_sequentially(col, projectors, network, env, allocator) return env, info end -function ctmrg_iteration(network, env::CTMRGEnv, alg::SequentialCTMRG) +function ctmrg_iteration(network, env::CTMRGEnv, alg::SequentialCTMRG, allocator) truncation_error = zero(real(scalartype(network))) for _ in 1:4 # rotate for col in 1:size(network, 2) # left move column-wise - env, info = ctmrg_leftmove(col, network, env, alg) + env, info = ctmrg_leftmove(col, network, env, alg, allocator) ignore_derivatives() do truncation_error = max(truncation_error, info.truncation_error) end @@ -73,26 +73,26 @@ function ctmrg_iteration(network, env::CTMRGEnv, alg::SequentialCTMRG) end """ - sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm) - sequential_projectors(coordinate::NTuple{3,Int}, network::InfiniteSquareNetwork, env::CTMRGEnv, alg::ProjectorAlgorithm) + sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm, allocator) + sequential_projectors(coordinate::NTuple{3,Int}, network::InfiniteSquareNetwork, env::CTMRGEnv, alg::ProjectorAlgorithm, allocator) Compute CTMRG projectors in the `:SequentialCTMRG` scheme either for an entire column `col` or for a specific `coordinate` (where `dir=WEST` is already implied in the `:SequentialCTMRG` scheme). """ -function sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm) +function sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm, allocator) coordinates = eachcoordinate(env)[:, col] T_dst = Base.promote_op( - sequential_projectors, NTuple{3, Int}, typeof(network), typeof(env), typeof(alg) + sequential_projectors, NTuple{3, Int}, typeof(network), typeof(env), typeof(alg), typeof(allocator) ) proj_and_info = similar(coordinates, T_dst) proj_and_info′::typeof(proj_and_info) = dtmap!!(proj_and_info, coordinates) do (r, c) - proj, info = sequential_projectors((WEST, r, c), network, env, alg) + proj, info = sequential_projectors((WEST, r, c), network, env, alg, allocator) return proj, info end return _split_proj_and_info(proj_and_info′) end function sequential_projectors( - coordinate::NTuple{3, Int}, network, env::CTMRGEnv, alg::HalfInfiniteProjector + coordinate::NTuple{3, Int}, network, env::CTMRGEnv, alg::HalfInfiniteProjector, allocator ) _, r, c = coordinate r′ = r - 1 @@ -100,10 +100,10 @@ function sequential_projectors( alg´ = _set_decomposition_truncation(alg, trunc) Q1 = TensorMap(EnlargedCorner(network, env, (SOUTHWEST, r, c))) Q2 = TensorMap(EnlargedCorner(network, env, (NORTHWEST, r′, c))) - return compute_projector((Q1, Q2), alg´) + return compute_projector((Q1, Q2), alg´, allocator) end function sequential_projectors( - coordinate::NTuple{3, Int}, network, env::CTMRGEnv, alg::FullInfiniteProjector + coordinate::NTuple{3, Int}, network, env::CTMRGEnv, alg::FullInfiniteProjector, allocator ) rowsize, colsize = size(env)[2:3] coordinate_nw = _next_coordinate(coordinate, rowsize, colsize) @@ -117,15 +117,15 @@ function sequential_projectors( TensorMap(EnlargedCorner(network, env, coordinate_nw)), TensorMap(EnlargedCorner(network, env, coordinate_ne)), ) - return compute_projector(ec, alg´) + return compute_projector(ec, alg´, allocator) end """ - renormalize_sequentially(col::Int, projectors, network, env) + renormalize_sequentially(col::Int, projectors, network, env, allocator) Renormalize one column of the CTMRG environment. """ -function renormalize_sequentially(col::Int, projectors, network, env) +function renormalize_sequentially(col::Int, projectors, network, env, allocator) corners = Zygote.Buffer(env.corners) edges = Zygote.Buffer(env.edges) @@ -140,13 +140,13 @@ function renormalize_sequentially(col::Int, projectors, network, env) # Apply projectors to renormalize corners and edge for row in axes(env.corners, 2) - C_southwest = renormalize_southwest_corner((row, col), env, projectors) + C_southwest = renormalize_southwest_corner((row, col), env, projectors, allocator) corners[SOUTHWEST, row, col] = normalize(C_southwest) - C_northwest = renormalize_northwest_corner((row, col), env, projectors) + C_northwest = renormalize_northwest_corner((row, col), env, projectors, allocator) corners[NORTHWEST, row, col] = normalize(C_northwest) - E_west = renormalize_west_edge((row, col), env, projectors, network) + E_west = renormalize_west_edge((row, col), env, projectors, network, allocator) edges[WEST, row, col] = normalize(E_west) end diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index 4d4be9e89..a9c55c57e 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -36,7 +36,7 @@ end CTMRG_SYMBOLS[:SimultaneousCTMRG] = SimultaneousCTMRG -function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG) +function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG, allocator) coordinates = eachcoordinate(network, 1:4) T_corners = Base.promote_op( TensorMap ∘ EnlargedCorner, typeof(network), typeof(env), eltype(coordinates) @@ -46,8 +46,8 @@ function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG) dtmap!!(enlarged_corners′, eachcoordinate(network, 1:4)) do idx return TensorMap(EnlargedCorner(network, env, idx)) end # expand environment - projectors, info = simultaneous_projectors(enlarged_corners, env, alg.projector_alg) # compute projectors on all coordinates - env′ = renormalize_simultaneously(enlarged_corners, projectors, network, env) # renormalize enlarged corners + projectors, info = simultaneous_projectors(enlarged_corners, env, alg.projector_alg, allocator) # compute projectors on all coordinates + env′ = renormalize_simultaneously(enlarged_corners, projectors, network, env, allocator) # renormalize enlarged corners info = (; contraction_metrics = (; info.truncation_error), info.U, info.S, info.V, @@ -69,38 +69,38 @@ function _split_proj_and_info(proj_and_info) end """ - simultaneous_projectors(enlarged_corners::Array{E,3}, env::CTMRGEnv, alg::ProjectorAlgorithm) - simultaneous_projectors(coordinate, enlarged_corners::Array{E,3}, alg::ProjectorAlgorithm) + simultaneous_projectors(enlarged_corners::Array{E,3}, env::CTMRGEnv, alg::ProjectorAlgorithm, allocator) + simultaneous_projectors(coordinate, enlarged_corners::Array{E,3}, alg::ProjectorAlgorithm, allocator) Compute CTMRG projectors in the `:SimultaneousCTMRG` scheme either for all provided enlarged corners or on a specific `coordinate`. """ function simultaneous_projectors( - enlarged_corners::Array{E, 3}, env::CTMRGEnv, alg::ProjectorAlgorithm + enlarged_corners::Array{E, 3}, env::CTMRGEnv, alg::ProjectorAlgorithm, allocator ) where {E} coordinates = eachcoordinate(env, 1:4) T_dst = Base.promote_op( simultaneous_projectors, - NTuple{3, Int}, typeof(enlarged_corners), typeof(env), typeof(alg), + NTuple{3, Int}, typeof(enlarged_corners), typeof(env), typeof(alg), typeof(allocator) ) proj_and_info′ = similar(coordinates, T_dst) proj_and_info::typeof(proj_and_info′) = dtmap!!(proj_and_info′, coordinates) do coordinate - return simultaneous_projectors(coordinate, enlarged_corners, env, alg) + return simultaneous_projectors(coordinate, enlarged_corners, env, alg, allocator) end return _split_proj_and_info(proj_and_info) end function simultaneous_projectors( - coordinate, enlarged_corners::Array{E, 3}, env, alg::HalfInfiniteProjector + coordinate, enlarged_corners::Array{E, 3}, env, alg::HalfInfiniteProjector, allocator ) where {E} coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...) trunc = truncation_strategy(alg, env.edges[coordinate[1], coordinate′[2:3]...]) alg′ = _set_decomposition_truncation(alg, trunc) ec = (enlarged_corners[coordinate...], enlarged_corners[coordinate′...]) - return compute_projector(ec, alg′) + return compute_projector(ec, alg′, allocator) end function simultaneous_projectors( - coordinate, enlarged_corners::Array{E, 3}, env, alg::FullInfiniteProjector + coordinate, enlarged_corners::Array{E, 3}, env, alg::FullInfiniteProjector, allocator ) where {E} rowsize, colsize = size(env)[2:3] coordinate2 = _next_coordinate(coordinate, rowsize, colsize) @@ -114,7 +114,7 @@ function simultaneous_projectors( enlarged_corners[coordinate2...], enlarged_corners[coordinate3...], ) - return compute_projector(ec, alg′) + return compute_projector(ec, alg′, allocator) end """ @@ -122,7 +122,7 @@ $(SIGNATURES) Renormalize all enlarged corners and edges simultaneously. """ -function renormalize_simultaneously(enlarged_corners, projectors, network, env) +function renormalize_simultaneously(enlarged_corners, projectors, network, env, allocator) P_left, P_right = projectors coordinates = eachcoordinate(env, 1:4) T_CE = Tuple{cornertype(env), edgetype(env)} @@ -131,24 +131,24 @@ function renormalize_simultaneously(enlarged_corners, projectors, network, env) dtmap!!(corners_edges′, coordinates) do (dir, r, c) if dir == NORTH corner = renormalize_northwest_corner( - (r, c), enlarged_corners, P_left, P_right + (r, c), enlarged_corners, P_left, P_right, allocator ) - edge = renormalize_north_edge((r, c), env, P_left, P_right, network) + edge = renormalize_north_edge((r, c), env, P_left, P_right, network, allocator) elseif dir == EAST corner = renormalize_northeast_corner( - (r, c), enlarged_corners, P_left, P_right + (r, c), enlarged_corners, P_left, P_right, allocator ) - edge = renormalize_east_edge((r, c), env, P_left, P_right, network) + edge = renormalize_east_edge((r, c), env, P_left, P_right, network, allocator) elseif dir == SOUTH corner = renormalize_southeast_corner( - (r, c), enlarged_corners, P_left, P_right + (r, c), enlarged_corners, P_left, P_right, allocator ) - edge = renormalize_south_edge((r, c), env, P_left, P_right, network) + edge = renormalize_south_edge((r, c), env, P_left, P_right, network, allocator) elseif dir == WEST corner = renormalize_southwest_corner( - (r, c), enlarged_corners, P_left, P_right + (r, c), enlarged_corners, P_left, P_right, allocator ) - edge = renormalize_west_edge((r, c), env, P_left, P_right, network) + edge = renormalize_west_edge((r, c), env, P_left, P_right, network, allocator) end return normalize(corner), normalize(edge) end diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 59f7efea1..0df43d413 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -80,17 +80,17 @@ function TensorKit.TensorMap(Q::EnlargedCorner) end end -function renormalize_northwest_corner(ec::EnlargedCorner, P_left, P_right) - return renormalize_northwest_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A) +function renormalize_northwest_corner(ec::EnlargedCorner, P_left, P_right, allocator) + return renormalize_northwest_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A, allocator) end -function renormalize_northeast_corner(ec::EnlargedCorner, P_left, P_right) - return renormalize_northeast_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A) +function renormalize_northeast_corner(ec::EnlargedCorner, P_left, P_right, allocator) + return renormalize_northeast_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A, allocator) end -function renormalize_southeast_corner(ec::EnlargedCorner, P_left, P_right) - return renormalize_southeast_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A) +function renormalize_southeast_corner(ec::EnlargedCorner, P_left, P_right, allocator) + return renormalize_southeast_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A, allocator) end -function renormalize_southwest_corner(ec::EnlargedCorner, P_left, P_right) - return renormalize_southwest_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A) +function renormalize_southwest_corner(ec::EnlargedCorner, P_left, P_right, allocator) + return renormalize_southwest_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.A, allocator) end # Wrapper around half_infinite_environment contraction using EnlargedCorners (used in ctmrg_projectors) @@ -99,12 +99,12 @@ function half_infinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) end # Compute left and right projectors sparsely without constructing enlarged corners explicitly -function contract_projectors(U, S, V, Q::EnlargedCorner, Q_next::EnlargedCorner) +function contract_projectors(U, S, V, Q::EnlargedCorner, Q_next::EnlargedCorner, allocator) Ar = _rotate_north_localsandwich(Q.A, _prev(Q.dir, 4)) Ar_next = _rotate_north_localsandwich(Q_next.A, Q_next.dir) isqS = sdiag_pow(S, -0.5) - P_left = left_projector(Q_next.E_1, Q_next.C, Q_next.E_2, V, isqS, Ar_next) - P_right = right_projector(Q.E_1, Q.C, Q.E_2, U, isqS, Ar) + P_left = left_projector(Q_next.E_1, Q_next.C, Q_next.E_2, V, isqS, Ar_next, allocator) + P_right = right_projector(Q.E_1, Q.C, Q.E_2, U, isqS, Ar, allocator) return P_left, P_right end @@ -185,23 +185,23 @@ function (env::HalfInfiniteEnv)(x, ::Val{true}) # Adjoint linear map: env()' * ) end -function contract_projectors(U, S, V, env::HalfInfiniteEnv) +function contract_projectors(U, S, V, env::HalfInfiniteEnv, allocator) Q = EnlargedCorner(env.C_1, env.E_1, env.E_2, env.A_1, env.dir) Q_next = EnlargedCorner(env.C_2, env.E_3, env.E_4, env.A_2, _next(env.dir, 4)) - return contract_projectors(U, S, V, Q, Q_next) + return contract_projectors(U, S, V, Q, Q_next, allocator) end -function contract_projectors(U, S, V, henv::HalfInfiniteEnv, henv_next::HalfInfiniteEnv) +function contract_projectors(U, S, V, henv::HalfInfiniteEnv, henv_next::HalfInfiniteEnv, allocator) isqS = sdiag_pow(S, -0.5) P_left = left_projector( henv_next.E_1, henv_next.C_1, henv_next.E_2, henv_next.E_3, henv_next.C_2, henv_next.E_4, V, isqS, - henv_next.A_1r, henv_next.A_2r + henv_next.A_1r, henv_next.A_2r, allocator ) P_right = right_projector( henv.E_1, henv.C_1, henv.E_2, henv.E_3, henv.C_2, henv.E_4, U, isqS, - henv.A_1r, henv.A_2r + henv.A_1r, henv.A_2r, allocator ) return P_left, P_right end @@ -346,7 +346,7 @@ function full_infinite_environment( return FullInfiniteEnv(ec_1, ec_2, ec_3, ec_4) end -function contract_projectors(U, S, V, env::FullInfiniteEnv) +function contract_projectors(U, S, V, env::FullInfiniteEnv, allocator) ndir = _next(env.dir, 4) nndir = _next(ndir, 4) henv = HalfInfiniteEnv( @@ -359,7 +359,7 @@ function contract_projectors(U, S, V, env::FullInfiniteEnv) env.E_5, env.E_6, env.E_7, env.E_8, env.A_3, env.A_4, nndir, ) - return contract_projectors(U, S, V, henv, henv_next) + return contract_projectors(U, S, V, henv, henv_next, allocator) end diff --git a/src/algorithms/expectation_value/expectation_value.jl b/src/algorithms/expectation_value/expectation_value.jl index e8d25e3eb..70fd3c46a 100644 --- a/src/algorithms/expectation_value/expectation_value.jl +++ b/src/algorithms/expectation_value/expectation_value.jl @@ -11,19 +11,29 @@ In the latter case the first signature corresponds to a single layer PEPO contra the second signature yields a bilayer contraction instead. """ function MPSKit.expectation_value( - bra::S, O::LocalOperator, ket::S, env + bra::S, O::LocalOperator, ket::S, env; + scheduler = Defaults.scheduler[], + allocator = default_allocator(bra, scheduler) ) where {S <: InfiniteState} checklattice(bra, O, ket) term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly - return local_expectation_value(inds, bra, operator, ket, env) + return local_expectation_value(inds, bra, operator, ket, env, allocator) end return sum(term_vals) end -MPSKit.expectation_value(peps::InfinitePEPS, O::LocalOperator, env) = expectation_value(peps, O, peps, env) -function MPSKit.expectation_value(state::InfinitePEPO, O::LocalOperator, env) +MPSKit.expectation_value( + peps::InfinitePEPS, O::LocalOperator, env; + scheduler = Defaults.scheduler[], + allocator = default_allocator(peps, scheduler) +) = expectation_value(peps, O, peps, env; scheduler, allocator) +function MPSKit.expectation_value( + state::InfinitePEPO, O::LocalOperator, env; + scheduler = Defaults.scheduler[], + allocator = default_allocator(state, scheduler) + ) checklattice(state, O) term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly - return local_expectation_value(inds, state, operator, env) + return local_expectation_value(inds, state, operator, env, allocator) end return sum(term_vals) end @@ -33,8 +43,8 @@ end # --------------------------------------------- """ - local_expectation_value(inds, bra, operator, ket, env) - local_expectation_value(inds, state, operator, env) + local_expectation_value(inds, bra, operator, ket, env, allocator) + local_expectation_value(inds, state, operator, env, allocator) Compute the contribution of a single term of a [`LocalOperator`](@ref) to the expectation value ⟨bra|O|ket⟩ / ⟨bra|ket⟩ or tr(O * state) / tr(state), where `operator` is the local @@ -45,12 +55,12 @@ The implementation is overloaded based on the type of operator to be evaluated function local_expectation_value end # AbstractTensorMap evaluation goes through reduced density matrix -function local_expectation_value(inds, bra, operator::AbstractTensorMap, ket, env) - ρ = reduced_densitymatrix(inds, ket, bra, env) +function local_expectation_value(inds, bra, operator::AbstractTensorMap, ket, env, allocator) + ρ = reduced_densitymatrix(inds, ket, bra, env, allocator) return trmul(operator, ρ) end -function local_expectation_value(inds, state, operator::AbstractTensorMap, env) - ρ = reduced_densitymatrix(inds, state, env) +function local_expectation_value(inds, state, operator::AbstractTensorMap, env, allocator) + ρ = reduced_densitymatrix(inds, state, env, allocator) return trmul(operator, ρ) end @@ -61,18 +71,18 @@ Compute the contribution of an [`MPOTerm`](@ref), including a [`TensorProductTer For a PEPS or purified PEPO, compute ``⟨bra|O|ket⟩ / ⟨bra|ket⟩``; for a single-layer density matrix PEPO, compute ``tr(O * state) / tr(state)``. Insert the factors into the patch contraction directly and divide by the local norm of the same patch. """ -function local_expectation_value(inds, bra, operator::MPOTerm, ket, env) - return contract_local_operator(inds, operator, ket, bra, env) / - contract_local_norm(inds, ket, bra, env) +function local_expectation_value(inds, bra, operator::MPOTerm, ket, env, allocator) + return contract_local_operator(inds, operator, ket, bra, env, allocator) / + contract_local_norm(inds, ket, bra, env, allocator) end -function local_expectation_value(inds, state, operator::MPOTerm, env) - return contract_local_operator(inds, operator, state, env) / - contract_local_norm(inds, state, env) +function local_expectation_value(inds, state, operator::MPOTerm, env, allocator) + return contract_local_operator(inds, operator, state, env, allocator) / + contract_local_norm(inds, state, env, allocator) end # TODO: Implement boundary_contraction_expr for BPEnv to contract its messages directly. -local_expectation_value(inds, bra, operator::MPOTerm, ket, env::BPEnv) = - local_expectation_value(inds, bra, operator, ket, CTMRGEnv(env)) +local_expectation_value(inds, bra, operator::MPOTerm, ket, env::BPEnv, allocator) = + local_expectation_value(inds, bra, operator, ket, CTMRGEnv(env), allocator) # Expectation value of a local partition function tensor # ------------------------------------------------------ @@ -91,13 +101,17 @@ convention. function MPSKit.expectation_value( pf::InfinitePartitionFunction, op::Pair{CartesianIndex{2}, <:AbstractTensorMap{T, S, 2, 2}}, - env, + env; + scheduler = Defaults.scheduler[], + allocator = default_allocator(pf, scheduler) ) where {T, S} - return contract_local_tensor(op[1], op[2], env) / - contract_local_tensor(op[1], pf[op[1]], env) + return contract_local_tensor(op[1], op[2], env, allocator) / + contract_local_tensor(op[1], pf[op[1]], env, allocator) end function MPSKit.expectation_value( - pf::InfinitePartitionFunction, op::Pair{Tuple{Int, Int}}, env + pf::InfinitePartitionFunction, op::Pair{Tuple{Int, Int}}, env; + scheduler = Defaults.scheduler[], + allocator = default_allocator(pf, scheduler) ) - return expectation_value(pf, CartesianIndex(op[1]) => op[2], env) + return expectation_value(pf, CartesianIndex(op[1]) => op[2], env; scheduler, allocator) end diff --git a/src/algorithms/expectation_value/network_value.jl b/src/algorithms/expectation_value/network_value.jl index fa0043db5..2c8809b9d 100644 --- a/src/algorithms/expectation_value/network_value.jl +++ b/src/algorithms/expectation_value/network_value.jl @@ -7,29 +7,29 @@ Return the value (per unit cell) of a given contractible network contracted using a given CTMRG environment. """ -function network_value(network::InfiniteSquareNetwork, env::AbstractCTMRGEnv) - return prod(Iterators.product(axes(network)...)) do (r, c) - return _contract_site((r, c), network, env) * _contract_corners((r, c), env) / - _contract_vertical_edges((r, c), env) / _contract_horizontal_edges((r, c), env) +function network_value(network::InfiniteSquareNetwork, env::AbstractCTMRGEnv, allocator; scheduler = Defaults.scheduler[]) + return tmapreduce(*, Iterators.product(axes(network)...); scheduler) do (r, c) + return _contract_site((r, c), network, env, allocator) * _contract_corners((r, c), env, allocator) / + _contract_vertical_edges((r, c), env, allocator) / _contract_horizontal_edges((r, c), env, allocator) end end -network_value(state, env::AbstractCTMRGEnv) = network_value(InfiniteSquareNetwork(state), env) +network_value(state, env::AbstractCTMRGEnv, allocator; kwargs...) = network_value(InfiniteSquareNetwork(state), env, allocator; kwargs...) -function LinearAlgebra.norm(peps::InfinitePEPS, env::AbstractCTMRGEnv) - return network_value(InfiniteSquareNetwork(peps), env) +function LinearAlgebra.norm(peps::InfinitePEPS, env::AbstractCTMRGEnv; scheduler = Defaults.scheduler[], allocator = default_allocator(peps, scheduler)) + return network_value(InfiniteSquareNetwork(peps), env, allocator; scheduler) end # Local tensor insertions # ----------------------- """ - contract_local_tensor(inds, O::PFTensor, env) + contract_local_tensor(inds, O::PFTensor, env, allocator) Contract a local tensor `O` inserted into a partition function `pf` at position `inds`, using the environment `env`. """ function contract_local_tensor( - inds::Tuple{Int, Int}, O::PFTensor, env::AbstractCTMRGEnv{C, <:CTMRG_PF_EdgeTensor} + inds::Tuple{Int, Int}, O::PFTensor, env::AbstractCTMRGEnv{C, <:CTMRG_PF_EdgeTensor}, allocator ) where {C} r, c = inds return _contract_site( @@ -39,12 +39,12 @@ function contract_local_tensor( corner(env, SOUTHWEST, r + 1, c - 1), edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1), edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1), - O, + O, allocator ) end """ - contract_local_tensor(inds, O::PEPOTensor, network, env) + contract_local_tensor(inds, O::PEPOTensor, network, env, allocator) Contract a local tensor `O` inserted into the PEPO of a given `network` at position `inds`, using the environment `env`. @@ -54,6 +54,7 @@ function contract_local_tensor( O::PEPOTensor, network::InfiniteSquareNetwork{<:PEPOSandwich}, env::AbstractCTMRGEnv, + allocator, ) r, c, h = ind sandwich´ = Base.setindex(network[r, c], O, h + 2) @@ -65,9 +66,10 @@ function contract_local_tensor( edge(env, NORTH, r - 1, c), edge(env, EAST, r, c + 1), edge(env, SOUTH, r + 1, c), edge(env, WEST, r, c - 1), sandwich´, + allocator, ) end -function contract_local_tensor(inds::CartesianIndex, O::AbstractTensorMap, env::AbstractCTMRGEnv) - return contract_local_tensor(Tuple(inds), O, env) +function contract_local_tensor(inds::CartesianIndex, O::AbstractTensorMap, env::AbstractCTMRGEnv, allocator) + return contract_local_tensor(Tuple(inds), O, env, allocator) end diff --git a/src/algorithms/expectation_value/patch_contractions.jl b/src/algorithms/expectation_value/patch_contractions.jl index 77dac9859..ca363c40e 100644 --- a/src/algorithms/expectation_value/patch_contractions.jl +++ b/src/algorithms/expectation_value/patch_contractions.jl @@ -12,9 +12,9 @@ function _check_factor_count(inds, O::MPOTerm) end """ - contract_local_operator(inds, O, ket::InfinitePEPS, bra::InfinitePEPS, env) - contract_local_operator(inds, O, ket::InfinitePEPO, bra::InfinitePEPO, env) - contract_local_operator(inds, O, state::InfinitePEPO, env) + contract_local_operator(inds, O, ket::InfinitePEPS, bra::InfinitePEPS, env, allocator) + contract_local_operator(inds, O, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator) + contract_local_operator(inds, O, state::InfinitePEPO, env, allocator) Contract a local operator `O` between `ket` and `bra` states, computing `⟨bra|O|ket⟩`, where `ket` and `bra` correspond to either a PEPS or a PEPO representing a PEPS with ancillary @@ -25,27 +25,28 @@ the result is contracted over the surrounding virtual indices using the environm function contract_local_operator( inds::Vector{CartesianIndex{2}}, O, ket::InfinitePEPS, bra::InfinitePEPS, env, + allocator, ) _check_factor_count(inds, O) static_inds = Tuple(Val.(inds)) - return _contract_local_operator(static_inds, O, (ket, bra), env) + return _contract_local_operator(static_inds, O, (ket, bra), env, allocator) end function contract_local_operator( - inds::Vector{CartesianIndex{2}}, O, state::InfinitePEPO, env + inds::Vector{CartesianIndex{2}}, O, state::InfinitePEPO, env, allocator ) size(state, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) _check_factor_count(inds, O) static_inds = Tuple(Val.(inds)) - return _contract_local_operator(static_inds, O, state, env) + return _contract_local_operator(static_inds, O, state, env, allocator) end function contract_local_operator( - inds::Vector{CartesianIndex{2}}, O, ket::InfinitePEPO, bra::InfinitePEPO, env + inds::Vector{CartesianIndex{2}}, O, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator ) size(ket) == size(bra) || throw(DimensionMismatch("incompatible bra and ket dimensions")) size(ket, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) _check_factor_count(inds, O) static_inds = Tuple(Val.(inds)) - return _contract_local_operator(static_inds, O, (ket, bra), env) + return _contract_local_operator(static_inds, O, (ket, bra), env, allocator) end function contract_local_operator(inds::Vector{Tuple{Int, Int}}, O, args...) return contract_local_operator(CartesianIndex.(inds), O, args...) @@ -59,9 +60,9 @@ Base.@deprecate( ) """ - contract_local_norm(inds, ket::InfinitePEPS, bra::InfinitePEPS, env) - contract_local_norm(inds, ket::InfinitePEPO, bra::InfinitePEPO, env) - contract_local_norm(inds, state::InfinitePEPO, env) + contract_local_norm(inds, ket::InfinitePEPS, bra::InfinitePEPS, env, allocator) + contract_local_norm(inds, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator) + contract_local_norm(inds, state::InfinitePEPO, env, allocator) Contract a local norm corresponding to the overlap `ket` and `bra` states, computing a patch of `⟨bra|ket⟩`, where `ket` and `bra` correspond to either a PEPS or a PEPO representing a @@ -74,23 +75,23 @@ particular, the patch location is precisely the same as that of the patch used i [`contract_local_operator`](@ref). """ function contract_local_norm( - inds::Vector{CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env + inds::Vector{CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env, allocator ) static_inds = Tuple(Val.(inds)) - return _contract_local_norm(static_inds, (ket, bra), env) + return _contract_local_norm(static_inds, (ket, bra), env, allocator) end -function contract_local_norm(inds::Vector{CartesianIndex{2}}, state::InfinitePEPO, env) +function contract_local_norm(inds::Vector{CartesianIndex{2}}, state::InfinitePEPO, env, allocator) size(state, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) static_inds = Tuple(Val.(inds)) - return _contract_local_norm(static_inds, state, env) + return _contract_local_norm(static_inds, state, env, allocator) end function contract_local_norm( - inds::Vector{CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env + inds::Vector{CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator ) size(ket) == size(bra) || throw(DimensionMismatch("incompatible bra and ket dimensions")) size(ket, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) static_inds = Tuple(Val.(inds)) - return _contract_local_norm(static_inds, (ket, bra), env) + return _contract_local_norm(static_inds, (ket, bra), env, allocator) end function contract_local_norm(inds::Vector{Tuple{Int, Int}}, args...) return contract_local_norm(CartesianIndex.(inds), args...) diff --git a/src/algorithms/expectation_value/reduced_densitymatrix.jl b/src/algorithms/expectation_value/reduced_densitymatrix.jl index 2767c2bc0..9e1c9e8b0 100644 --- a/src/algorithms/expectation_value/reduced_densitymatrix.jl +++ b/src/algorithms/expectation_value/reduced_densitymatrix.jl @@ -2,9 +2,9 @@ # ------------------------ @doc """ - reduced_densitymatrix(inds, ket::InfinitePEPS, bra::InfinitePEPS = ket, env) - reduced_densitymatrix(inds, ket::InfinitePEPO, bra::InfinitePEPO, env) - reduced_densitymatrix(inds, state::InfinitePEPO, env) + reduced_densitymatrix(inds, ket::InfinitePEPS, bra::InfinitePEPS = ket, env, allocator) + reduced_densitymatrix(inds, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator) + reduced_densitymatrix(inds, state::InfinitePEPO, env, allocator) Construct the reduced density matrix `ρ` of `|ket⟩⟨bra|`, where both `ket` and `bra` correspond to either a PEPS or a PEPO representing a PEPS with ancillary legs. @@ -18,38 +18,38 @@ See also [`str`](@ref). # PEPS case has fast-path specializations function reduced_densitymatrix( - inds::Vector{CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env + inds::Vector{CartesianIndex{2}}, ket::InfinitePEPS, bra::InfinitePEPS, env, allocator ) - length(inds) == 1 && return reduced_densitymatrix1x1(only(inds), ket, bra, env) + length(inds) == 1 && return reduced_densitymatrix1x1(only(inds), ket, bra, env, allocator) if length(inds) == 2 if inds[2] - inds[1] == CartesianIndex(1, 0) - return reduced_densitymatrix2x1(inds[1], ket, bra, env) + return reduced_densitymatrix2x1(inds[1], ket, bra, env, allocator) elseif inds[2] - inds[1] == CartesianIndex(0, 1) - return reduced_densitymatrix1x2(inds[1], ket, bra, env) + return reduced_densitymatrix1x2(inds[1], ket, bra, env, allocator) end end static_inds = Tuple(Val.(inds)) - return _contract_densitymatrix(static_inds, (ket, bra), env) + return _contract_densitymatrix(static_inds, (ket, bra), env, allocator) end function reduced_densitymatrix( - inds::Vector{CartesianIndex{2}}, state::InfinitePEPO, env + inds::Vector{CartesianIndex{2}}, state::InfinitePEPO, env, allocator ) size(state, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) static_inds = Tuple(Val.(inds)) - return _contract_densitymatrix(static_inds, state, env) + return _contract_densitymatrix(static_inds, state, env, allocator) end function reduced_densitymatrix( - inds::Vector{CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env + inds::Vector{CartesianIndex{2}}, ket::InfinitePEPO, bra::InfinitePEPO, env, allocator ) size(ket) == size(bra) || throw(DimensionMismatch("incompatible bra and ket dimensions")) size(ket, 3) == 1 || throw(DimensionMismatch("only single-layer densitymatrices are supported")) static_inds = Tuple(Val.(inds)) - return _contract_densitymatrix(static_inds, (ket, bra), env) + return _contract_densitymatrix(static_inds, (ket, bra), env, allocator) end -reduced_densitymatrix(inds, ket::InfinitePEPS, env) = - reduced_densitymatrix(inds, ket, ket, env) +reduced_densitymatrix(inds, ket::InfinitePEPS, env, allocator) = + reduced_densitymatrix(inds, ket, ket, env, allocator) # handle deprecations of Tuple inds specifications Base.@deprecate( @@ -77,24 +77,24 @@ Base.@deprecate( # hand-optimized contraction actually pays off. """ - reduced_densitymatrix1x1(ind, ket, bra, env) + reduced_densitymatrix1x1(ind, ket, bra, env, allocator) Construct the reduced density matrix of `|ket⟩⟨bra|` on the single site `ind`, using an optimized contraction for the environment `env`. Falls back to [`_contract_densitymatrix`](@ref) for environments without such a specialization. """ -reduced_densitymatrix1x1(ind, ket, bra, env) = - _contract_densitymatrix((Val(ind),), (ket, bra), env) +reduced_densitymatrix1x1(ind, ket, bra, env, allocator) = + _contract_densitymatrix((Val(ind),), (ket, bra), env, allocator) """ - reduced_densitymatrix2x1(ind, ket, bra, env) + reduced_densitymatrix2x1(ind, ket, bra, env, allocator) Construct the reduced density matrix of `|ket⟩⟨bra|` on the vertical pair of sites starting at `ind`, using an optimized contraction for the environment `env`. Falls back to [`_contract_densitymatrix`](@ref) for environments without such a specialization. """ -reduced_densitymatrix2x1(ind, ket, bra, env) = - _contract_densitymatrix((Val(ind), Val(ind + CartesianIndex(1, 0))), (ket, bra), env) +reduced_densitymatrix2x1(ind, ket, bra, env, allocator) = + _contract_densitymatrix((Val(ind), Val(ind + CartesianIndex(1, 0))), (ket, bra), env, allocator) """ reduced_densitymatrix1x2(ind, ket, bra, env) @@ -103,5 +103,5 @@ Construct the reduced density matrix of `|ket⟩⟨bra|` on the horizontal pair at `ind`, using an optimized contraction for the environment `env`. Falls back to [`_contract_densitymatrix`](@ref) for environments without such a specialization. """ -reduced_densitymatrix1x2(ind, ket, bra, env) = - _contract_densitymatrix((Val(ind), Val(ind + CartesianIndex(0, 1))), (ket, bra), env) +reduced_densitymatrix1x2(ind, ket, bra, env, allocator) = + _contract_densitymatrix((Val(ind), Val(ind + CartesianIndex(0, 1))), (ket, bra), env, allocator) diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 111699453..fb6daa491 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -216,7 +216,9 @@ function _rrule( ::typeof(MPSKit.leading_boundary), envinit::AbstractCTMRGEnv, state, - alg::CTMRGAlgorithm, + alg::CTMRGAlgorithm; + scheduler = Defaults.scheduler[], + allocator = default_allocator(state, scheduler) ) _check_algorithm_combination(alg, gradmode) @@ -225,14 +227,14 @@ function _rrule( # prepare iterating function corresponding to a single gauge-fixed CTMRG iteration alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation alg_gauge = _scrambling_env_gauge(alg) # select appropriate gauge-fixing algorithm - env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) - signs, = compute_gauge_fix_gauge(env_conv, env, alg_gauge) + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed, allocator) + signs, = compute_gauge_fix_gauge(env_conv, env, alg_gauge, allocator) # `p` for state, `x` for environment function gauge_fixed_iteration(p, x) # fix global phases relative to the fixed point inside the iteration return fix_global_phases( - fix_relative_phases(ctmrg_iteration(InfiniteSquareNetwork(p), x, alg_fixed)[1], signs), - env, + fix_relative_phases(ctmrg_iteration(InfiniteSquareNetwork(p), x, alg_fixed, allocator)[1], signs, allocator), + env, allocator ) end # prepare its pullback @@ -468,7 +470,9 @@ function _rrule( ::typeof(MPSKit.leading_boundary), envinit::C4vCTMRGEnv, state, - alg::C4vCTMRG, + alg::C4vCTMRG; + scheduler = Defaults.scheduler[], + allocator = default_allocator(state, scheduler) ) _check_algorithm_combination(alg, gradmode) @@ -477,9 +481,9 @@ function _rrule( # prepare iterating function corresponding to a single gauge-fixed CTMRG iteration alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation alg_gauge = _scrambling_env_gauge(alg) # select appropriate gauge-fixing algorithm - env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) - signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge) - env_fixed = fix_phases(env_conv, signs, corner_phases, edge_phases) + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed, allocator) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge, allocator) + env_fixed = fix_phases(env_conv, signs, corner_phases, edge_phases, allocator) # NOTE: explicitly keeping corner non-diagonal and (possibly) complex, for use in the backwards pass C, E = env_fixed.corners[NORTHWEST, 1, 1], env_fixed.edges[NORTH, 1, 1] @@ -617,16 +621,18 @@ function PEPSKit._rrule( ::typeof(MPSKit.leading_boundary), envinit::AbstractCTMRGEnv, state, - alg::SimultaneousCTMRG{<:HalfInfiniteProjector}, + alg::SimultaneousCTMRG{<:HalfInfiniteProjector}; + scheduler = Defaults.scheduler[], + allocator = default_allocator(state, scheduler) ) env, = leading_boundary(envinit, state, alg) # gauge-fix SVD isometries alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation - env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) - signs, = compute_gauge_fix_gauge(env_conv, env, ScramblingEnvGauge()) + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed, allocator) + signs, = compute_gauge_fix_gauge(env_conv, env, ScramblingEnvGauge(), allocator) S = normalize.(info.S) - U, V = fix_relative_phases(info.U, info.V, signs) + U, V = fix_relative_phases(info.U, info.V, signs, allocator) # pretend the singular values matrices are just arbitrary complex tensors s = map(S) do s diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 808ed04f5..89a3dcb72 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -233,6 +233,8 @@ function fixedpoint( (finalize!) = OptimKit._finalize!, hasconverged = OptimKit.DefaultHasConverged(alg.optimizer_alg.gradtol), shouldstop = OptimKit.DefaultShouldStop(alg.optimizer_alg.maxiter), + scheduler = Defaults.scheduler[], + allocator = default_allocator(peps₀, scheduler) ) # validate inputs check_input(fixedpoint, peps₀, env₀, alg) @@ -275,7 +277,7 @@ function fixedpoint( alg.precondition_alg; iter = tracked_finalizer.tol_state.iter, g_global = tracked_finalizer.tol_state.gradnorm ) - g_prec = peps_precondition(x, g, tracked_finalizer.tol_state, precondition_alg) + g_prec = peps_precondition(x, g, tracked_finalizer.tol_state, precondition_alg, allocator) return peps_project_tangent(g_prec, x[1]) end end @@ -301,6 +303,7 @@ function fixedpoint( env′, info = hook_pullback( leading_boundary, env, ψ, boundary_alg; alg_rrule = gradient_alg, + allocator = allocator ) ignore_derivatives() do alg.reuse_env && update!(env, env′) diff --git a/src/algorithms/optimization/preconditioning.jl b/src/algorithms/optimization/preconditioning.jl index acf3e504f..733bd699e 100644 --- a/src/algorithms/optimization/preconditioning.jl +++ b/src/algorithms/optimization/preconditioning.jl @@ -85,8 +85,8 @@ Apply the regularized local metric ``N / ⟨ψ|ψ⟩ + δ`` to a single PEPS ten ``N`` is obtained by contracting the CTMRG environment `env` around the unit-cell site `(r, c)` and `norm_pref` is the corresponding local norm ``⟨ψ|ψ⟩``. """ -function apply_local_preconditioner(g_rc::PEPSTensor, env::AbstractCTMRGEnv, δ, (r, c), norm_pref) - @autoopt @tensor g_rc_prec[d; D_N_below D_E_below D_S_below D_W_below] := +function apply_local_preconditioner(g_rc::PEPSTensor, env::AbstractCTMRGEnv, δ, (r, c), norm_pref, allocator) + @autoopt @tensor allocator = allocator g_rc_prec[d; D_N_below D_E_below D_S_below D_W_below] := g_rc[d; D_N_above D_E_above D_S_above D_W_above] * corner(env, NORTHWEST, r - 1, c - 1)[χ_WNW; χ_NNW] * edge(env, NORTH, r - 1, c)[χ_NNW D_N_above D_N_below; χ_NNE] * @@ -101,7 +101,7 @@ function apply_local_preconditioner(g_rc::PEPSTensor, env::AbstractCTMRGEnv, δ, end """ - peps_precondition(x, g, tol_state::NamedTuple, alg) + peps_precondition(x, g, tol_state::NamedTuple, alg, allocator) Precondition the PEPS gradient `g` at the point `x = (peps, env)` according to the preconditioner algorithm `alg`. Passing `alg = nothing` disables preconditioning and returns @@ -109,15 +109,15 @@ preconditioner algorithm `alg`. Passing `alg = nothing` disables preconditioning preconditioner, and note that its regularization strength is set based on the current gradient norm and iteration count tracked in `tol_state`. """ -peps_precondition(x, g, tol_state::NamedTuple, alg::Nothing) = g +peps_precondition(x, g, tol_state::NamedTuple, alg::Nothing, allocator) = g -function peps_precondition(x, g, tol_state::NamedTuple, alg::LocalPreconditioner) +function peps_precondition(x, g, tol_state::NamedTuple, alg::LocalPreconditioner, allocator) peps, env = x δ = alg.regularization * tol_state.gradnorm^2 / max(tol_state.iter, 1) g_prec_unitcell = dtmap(eachcoordinate(g)) do (r, c) - nf = _contract_site((r, c), InfiniteSquareNetwork(peps), env) + nf = _contract_site((r, c), InfiniteSquareNetwork(peps), env, allocator) g_rc_prec, = linsolve(g[r, c], g[r, c], alg.solver_alg) do g_in - return apply_local_preconditioner(g_in, env, δ, (r, c), nf) + return apply_local_preconditioner(g_in, env, δ, (r, c), nf, allocator) end return g_rc_prec end diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 0e18e43b4..17718da0b 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -128,7 +128,8 @@ One iteration of simple update. Pass `allocator` to reuse contraction scratch; function su_iter( state::InfiniteState, circuit::LocalCircuit, alg::SimpleUpdate, env::SUWeight; - allocator = default_allocator(state, SerialScheduler()) + scheduler = Defaults.scheduler[], + allocator = default_allocator(state, scheduler) ) state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0 # Gates are applied sequentially and share the same scratch allocator. diff --git a/test/testsuite/boundarymps/vumps.jl b/test/testsuite/boundarymps/vumps.jl index be9f20ea2..96ab54363 100644 --- a/test/testsuite/boundarymps/vumps.jl +++ b/test/testsuite/boundarymps/vumps.jl @@ -1,6 +1,7 @@ using Test using Random using PEPSKit +using PEPSKit: default_allocator, SerialScheduler using TensorKit using MPSKit using LinearAlgebra @@ -99,7 +100,8 @@ function boundary_mps_fermionic_peps(AT) rng = StableRNG(29384293742893) ctm´, = leading_boundary(CTMRGEnv((args...) -> randn(rng, args...), ComplexF64, n´, χ), n´) - N_ctm´ = abs(network_value(n´, ctm´)) + allocator = default_allocator(n, SerialScheduler()) + N_ctm´ = abs(network_value(n´, ctm´, allocator)) @show N_vumps´ @test N_vumps´ ≈ N_vumps rtol = 1.0e-2 diff --git a/test/testsuite/ctmrg/c4v_contractions.jl b/test/testsuite/ctmrg/c4v_contractions.jl index cca67b0d2..2d3effbca 100644 --- a/test/testsuite/ctmrg/c4v_contractions.jl +++ b/test/testsuite/ctmrg/c4v_contractions.jl @@ -5,7 +5,7 @@ using TensorKit using Adapt using PEPSKit: EnlargedCorner, c4v_enlarge, c4v_projector! using PEPSKit: _contract_site, _contract_corners, _contract_vertical_edges, - _contract_horizontal_edges + _contract_horizontal_edges, default_allocator, SerialScheduler using PEPSKit: NORTH, EAST, SOUTH, WEST, NORTHWEST, NORTHEAST, SOUTHEAST, SOUTHWEST stype = ComplexF64 @@ -37,13 +37,14 @@ function test_c4v_contractions(n::InfiniteSquareNetwork, env::C4vCTMRGEnv) env_ref = _expand_c4v_env(env) ind = (1, 1) + allocator = PEPSKit.default_allocator(n, SerialScheduler()) @testset "Network value contractions" begin - @test _contract_site(ind, n, env) ≈ _contract_site(ind, n, env_ref) - @test _contract_corners(ind, env) ≈ _contract_corners(ind, env_ref) - @test _contract_vertical_edges(ind, env) ≈ _contract_vertical_edges(ind, env_ref) - @test _contract_horizontal_edges(ind, env) ≈ _contract_horizontal_edges(ind, env_ref) - @test network_value(n, env) ≈ network_value(n, env_ref) - @test_throws AssertionError _contract_site((1, 2), n, env) + @test _contract_site(ind, n, env, allocator) ≈ _contract_site(ind, n, env_ref, allocator) + @test _contract_corners(ind, env, allocator) ≈ _contract_corners(ind, env_ref, allocator) + @test _contract_vertical_edges(ind, env, allocator) ≈ _contract_vertical_edges(ind, env_ref, allocator) + @test _contract_horizontal_edges(ind, env, allocator) ≈ _contract_horizontal_edges(ind, env_ref, allocator) + @test network_value(n, env, allocator) ≈ network_value(n, env_ref, allocator) + @test_throws AssertionError _contract_site((1, 2), n, env, allocator) end @testset "Enlarged corners" begin @@ -58,14 +59,14 @@ function test_c4v_contractions(n::InfiniteSquareNetwork, env::C4vCTMRGEnv) end @testset "Rotations" begin - nv = network_value(n, env) + nv = network_value(n, env, allocator) for rot in (rotl90, rotr90, rot180) env_rot = rot(env) env_ref_rot = rot(env_ref) @test env_rot isa C4vCTMRGEnv @test size(env_rot.edges) == (4, 1, 1) @test all(env_rot.edges .≈ env_ref_rot.edges) - @test network_value(rot(n), env_rot) ≈ nv + @test network_value(rot(n), env_rot, allocator) ≈ nv end @test rotl90(rotr90(env)) ≈ env @test rot180(rot180(env)) ≈ env @@ -151,10 +152,11 @@ function test_c4v_ctmrg_contractions(AT, P, V, Venv) @testset "C4v environment after one CTMRG iteration" begin peps = _random_c4v_peps(AT, P, V) n = InfiniteSquareNetwork(peps) + allocator = PEPSKit.default_allocator(n, SerialScheduler()) env₀ = initialize_random_c4v_env(peps, Venv) for projector_alg in (:C4vEighProjector, :C4vQRProjector) alg = C4vCTMRG(; projector_alg) - env, = PEPSKit.ctmrg_iteration(n, env₀, alg) + env, = PEPSKit.ctmrg_iteration(n, env₀, alg, allocator) @test env isa C4vCTMRGEnv test_c4v_contractions(n, env) end diff --git a/test/testsuite/ctmrg/c4v_sign_branches.jl b/test/testsuite/ctmrg/c4v_sign_branches.jl index 396c703c2..5ad01f9b4 100644 --- a/test/testsuite/ctmrg/c4v_sign_branches.jl +++ b/test/testsuite/ctmrg/c4v_sign_branches.jl @@ -4,7 +4,7 @@ using LinearAlgebra using TensorKit using PEPSKit using Adapt -using PEPSKit: ctmrg_iteration, compute_gauge_fix_gauge, fix_relative_phases, ScramblingEnvGauge +using PEPSKit: ctmrg_iteration, compute_gauge_fix_gauge, fix_relative_phases, ScramblingEnvGauge, default_allocator, SerialScheduler # For half-integer U(1) charges, (-1)^{2Q} on the ket bonds negates the sandwich, so fixed # points come in pairs, one of which flips (C, E) to (-C, -E) every iteration. Both should converge. @@ -26,10 +26,10 @@ function charge_parity(A, W) end # sign by which one gauge-fixed iteration multiplies the edge -function iteration_sign(network, env, alg) - env′, = ctmrg_iteration(network, env, alg) - signs, = compute_gauge_fix_gauge(env′, env, ScramblingEnvGauge()) - env′ = fix_relative_phases(env′, signs) +function iteration_sign(network, env, alg, allocator) + env′, = ctmrg_iteration(network, env, alg, allocator) + signs, = compute_gauge_fix_gauge(env′, env, ScramblingEnvGauge(), allocator) + env′ = fix_relative_phases(env′, signs, allocator) return sign(real(dot(env.edges[1], env′.edges[1]))) end @@ -47,13 +47,14 @@ function ctmrg_c4v_sign_branches(AT) symmetrize!(peps, RotateReflect()) peps = adapt(AT, InfinitePEPS(map(a -> a / norm(a), peps.A))) network = InfiniteSquareNetwork(peps) + allocator = PEPSKit.default_allocator(network, SerialScheduler()) alg = C4vCTMRG(; tol, maxiter, projector_alg = :C4vEighProjector) env, info = leading_boundary(initialize_random_c4v_env(peps, Venv), peps, alg) @test info.converged env_twin = sign_twin(env) - @test iteration_sign(network, env_twin, alg) == -iteration_sign(network, env, alg) + @test iteration_sign(network, env_twin, alg, allocator) == -iteration_sign(network, env, alg, allocator) env_twin, info_twin = leading_boundary(env_twin, peps, alg) @test info_twin.converged H_AT = adapt(AT, H) diff --git a/test/testsuite/ctmrg/contractions.jl b/test/testsuite/ctmrg/contractions.jl index af39cbac7..3c3e985d9 100644 --- a/test/testsuite/ctmrg/contractions.jl +++ b/test/testsuite/ctmrg/contractions.jl @@ -3,7 +3,7 @@ using Random using PEPSKit using TensorKit using Adapt -using PEPSKit: eachcoordinate, _next_coordinate +using PEPSKit: eachcoordinate, _next_coordinate, default_allocator, SerialScheduler using PEPSKit: EnlargedCorner, HalfInfiniteEnv, FullInfiniteEnv using PEPSKit: half_infinite_environment, full_infinite_environment using PEPSKit: simultaneous_projectors, contract_projectors @@ -115,13 +115,14 @@ function test_contractions(n::InfiniteSquareNetwork, env::CTMRGEnv) ) end + allocator = default_allocator(env, SerialScheduler()) # SVD half and full infinite environments (P_left_half, P_right_half), info_half = simultaneous_projectors( - dense_enlarged_corners, env, HalfInfiniteProjector() + dense_enlarged_corners, env, HalfInfiniteProjector(), allocator ) U_half, S_half, V_half = info_half.U, info_half.S, info_half.V (P_left_full, P_right_full), info_full = simultaneous_projectors( - dense_enlarged_corners, env, FullInfiniteProjector() + dense_enlarged_corners, env, FullInfiniteProjector(), allocator ) U_full, S_full, V_full = info_full.U, info_full.S, info_full.V @@ -149,11 +150,11 @@ function test_contractions(n::InfiniteSquareNetwork, env::CTMRGEnv) # projector computation P_left_sparse, P_right_sparse = contract_projectors( - U_half[dir, r, c], S_half[dir, r, c], V_half[dir, r, c], shenv + U_half[dir, r, c], S_half[dir, r, c], V_half[dir, r, c], shenv, allocator ) P_left_dense, P_right_dense = contract_projectors( U_half[dir, r, c], S_half[dir, r, c], V_half[dir, r, c], - dense_enlarged_corners[co...], dense_enlarged_corners[co2...], + dense_enlarged_corners[co...], dense_enlarged_corners[co2...], allocator ) @test P_left_sparse ≈ P_left_dense @test P_right_sparse ≈ P_right_dense @@ -176,7 +177,7 @@ function test_contractions(n::InfiniteSquareNetwork, env::CTMRGEnv) # projector computation P_left_sparse, P_right_sparse = contract_projectors( - U_full[dir, r, c], S_full[dir, r, c], V_full[dir, r, c], sfenv + U_full[dir, r, c], S_full[dir, r, c], V_full[dir, r, c], sfenv, allocator ) P_left_dense, P_right_dense = contract_projectors( U_full[dir, r, c], S_full[dir, r, c], V_full[dir, r, c], @@ -185,7 +186,7 @@ function test_contractions(n::InfiniteSquareNetwork, env::CTMRGEnv) ), half_infinite_environment( dense_enlarged_corners[co2...], dense_enlarged_corners[co3...] - ), + ), allocator, ) @test P_left_sparse ≈ P_left_dense @test P_right_sparse ≈ P_right_dense @@ -199,10 +200,10 @@ function test_contractions(n::InfiniteSquareNetwork, env::CTMRGEnv) ## Corner renormalization C_sparse = renormalize_corner_fns[dir]( - (r, c), sparse_enlarged_corners, P_left_half, P_right_half + (r, c), sparse_enlarged_corners, P_left_half, P_right_half, allocator ) C_dense = renormalize_corner_fns[dir]( - (r, c), dense_enlarged_corners, P_left_half, P_right_half + (r, c), dense_enlarged_corners, P_left_half, P_right_half, allocator ) @test C_sparse ≈ C_dense end diff --git a/test/testsuite/ctmrg/fixed_iterscheme.jl b/test/testsuite/ctmrg/fixed_iterscheme.jl index 835b8cfb5..eabb1c23c 100644 --- a/test/testsuite/ctmrg/fixed_iterscheme.jl +++ b/test/testsuite/ctmrg/fixed_iterscheme.jl @@ -9,11 +9,13 @@ using PEPSKit using PEPSKit: ctmrg_iteration, compute_gauge_fix_gauge, + default_allocator, fix_phases, fix_relative_phases, calc_elementwise_convergence, peps_normalize, - ScramblingEnvGauge + ScramblingEnvGauge, + SerialScheduler using PEPSKit.Defaults: ctmrg_tol # initialize parameters @@ -40,17 +42,18 @@ function ctmrg_fixed_iterscheme_asymmetric(AT; svd_alg = :DivideAndConquer) env_conv1, = leading_boundary(CTMRGEnv(psi, ComplexSpace(χ)), psi, ctm_alg) + allocator = default_allocator(n, SerialScheduler()) # do extra iteration and gauge fix - env_conv2, = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg) - env_fixed = gauge_fix(env_conv2, env_conv1, ScramblingEnvGauge()) + env_conv2, = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg, allocator) + env_fixed = gauge_fix(env_conv2, env_conv1, ScramblingEnvGauge(), allocator) @test calc_elementwise_convergence(env_conv1, env_fixed) ≈ 0 atol = atol # fix gauge of single iteration signs, corner_phases, edge_phases = - compute_gauge_fix_gauge(env_conv2, env_conv1, ScramblingEnvGauge()) + compute_gauge_fix_gauge(env_conv2, env_conv1, ScramblingEnvGauge(), allocator) gauge_fixed_iteration(env::CTMRGEnv) = fix_phases( - ctmrg_iteration(n, env, ctm_alg)[1], - signs, corner_phases, edge_phases, + ctmrg_iteration(n, env, ctm_alg, allocator)[1], + signs, corner_phases, edge_phases, allocator ) # do gauge-fixed iteration @@ -84,19 +87,20 @@ function ctmrg_fixed_iterscheme_c4v(AT; eigh_alg = :QRIteration) env_conv1, info = leading_boundary(env₀, psi, ctm_alg) # do extra iteration to check gauge fixing - env_conv2, info = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg) # CHECK + allocator = default_allocator(n, SerialScheduler()) + env_conv2, info = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg, allocator) # CHECK - env_fixed = gauge_fix(env_conv2, env_conv1, ScramblingEnvGauge()) + env_fixed = gauge_fix(env_conv2, env_conv1, ScramblingEnvGauge(), allocator) env_diff = calc_elementwise_convergence(env_conv1, env_fixed) @info "Diff between iters = $(env_diff)" @test env_diff ≈ 0 atol = atol # fix gauge of single iteration signs, corner_phases, edge_phases = - compute_gauge_fix_gauge(env_conv2, env_conv1, ScramblingEnvGauge()) + compute_gauge_fix_gauge(env_conv2, env_conv1, ScramblingEnvGauge(), allocator) gauge_fixed_iteration(env::PEPSKit.AbstractCTMRGEnv) = fix_phases( - ctmrg_iteration(n, env, ctm_alg)[1], - signs, corner_phases, edge_phases, + ctmrg_iteration(n, env, ctm_alg, allocator)[1], + signs, corner_phases, edge_phases, allocator ) # do gauge-fixed iteration @@ -120,27 +124,28 @@ function ctmrg_fixed_iterscheme_divide_and_conquer(AT) env₀ = CTMRGEnv(psi, ComplexSpace(χ)) env_conv1, = leading_boundary(env₀, psi, ctm_alg_iter) + allocator = default_allocator(n, SerialScheduler()) # do extra iteration to get gauge fixing - env_conv2_iter, info_iter = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg_iter) - env_fix_iter = gauge_fix(env_conv2_iter, env_conv1, ScramblingEnvGauge()) + env_conv2_iter, info_iter = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg_iter, allocator) + env_fix_iter = gauge_fix(env_conv2_iter, env_conv1, ScramblingEnvGauge(), allocator) @test calc_elementwise_convergence(env_conv1, env_fix_iter) ≈ 0 atol = atol - env_conv2_full, info_full = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg_full) - env_fix_full = gauge_fix(env_conv2_full, env_conv1, ScramblingEnvGauge()) + env_conv2_full, info_full = @testinferred ctmrg_iteration(n, env_conv1, ctm_alg_full, allocator) + env_fix_full = gauge_fix(env_conv2_full, env_conv1, ScramblingEnvGauge(), allocator) @test calc_elementwise_convergence(env_conv1, env_fix_full) ≈ 0 atol = atol # fix gauge of single iteration signs_iter, corner_phases_iter, edge_phases_iter = - compute_gauge_fix_gauge(env_conv2_iter, env_conv1, ScramblingEnvGauge()) + compute_gauge_fix_gauge(env_conv2_iter, env_conv1, ScramblingEnvGauge(), allocator) gauge_fixed_iteration_iter(env::CTMRGEnv) = fix_phases( - ctmrg_iteration(n, env, ctm_alg_iter)[1], - signs_iter, corner_phases_iter, edge_phases_iter, + ctmrg_iteration(n, env, ctm_alg_iter, allocator)[1], + signs_iter, corner_phases_iter, edge_phases_iter, allocator ) signs_full, corner_phases_full, edge_phases_full = - compute_gauge_fix_gauge(env_conv2_full, env_conv1, ScramblingEnvGauge()) + compute_gauge_fix_gauge(env_conv2_full, env_conv1, ScramblingEnvGauge(), allocator) gauge_fixed_iteration_full(env::CTMRGEnv) = fix_phases( - ctmrg_iteration(n, env, ctm_alg_full)[1], - signs_full, corner_phases_full, edge_phases_full, + ctmrg_iteration(n, env, ctm_alg_full, allocator)[1], + signs_full, corner_phases_full, edge_phases_full, allocator ) # do gauge-fixed iteration @@ -167,8 +172,8 @@ function ctmrg_fixed_iterscheme_divide_and_conquer(AT) @test svalues_check # gauge-fix the isometries using computed relative signs - U_iter_fix, V_iter_fix = fix_relative_phases(info_iter.U, info_iter.V, signs_iter) - U_full_fix, V_full_fix = fix_relative_phases(info_full.U, info_full.V, signs_full) + U_iter_fix, V_iter_fix = fix_relative_phases(info_iter.U, info_iter.V, signs_iter, allocator) + U_full_fix, V_full_fix = fix_relative_phases(info_full.U, info_full.V, signs_full, allocator) # check normalization of U's and V's Us = [info_iter.U, U_iter_fix, info_full.U, U_full_fix] diff --git a/test/testsuite/ctmrg/gaugefix.jl b/test/testsuite/ctmrg/gaugefix.jl index 823610616..c6699ba76 100644 --- a/test/testsuite/ctmrg/gaugefix.jl +++ b/test/testsuite/ctmrg/gaugefix.jl @@ -3,7 +3,7 @@ using Random using PEPSKit using TensorKit using Adapt -using PEPSKit: ctmrg_iteration, calc_elementwise_convergence +using PEPSKit: default_allocator, SerialScheduler, ctmrg_iteration, calc_elementwise_convergence using PEPSKit: ScramblingEnvGauge using PEPSKit: peps_normalize @@ -98,9 +98,10 @@ function ctmrg_gaugefix_asymmetric(AT; minimal::Bool = false) psi = adapt(AT, psi) env_pre = adapt(AT, env_pre) n = InfiniteSquareNetwork(psi) + allocator = default_allocator(n, SerialScheduler()) env, = leading_boundary(env_pre, psi, alg) - env′, = ctmrg_iteration(n, env, alg) - env_fixed = gauge_fix(env′, env, gauge_alg) + env′, = ctmrg_iteration(n, env, alg, allocator) + env_fixed = gauge_fix(env′, env, gauge_alg, allocator) env_diff = calc_elementwise_convergence(env, env_fixed) @info "Diff between iters = $(env_diff)" @test env_diff ≈ 0 atol = atol @@ -119,9 +120,10 @@ function ctmrg_gaugefix_c4v(AT; minimal::Bool = false) psi = adapt(AT, psi) env_pre = adapt(AT, env_pre) n = InfiniteSquareNetwork(psi) + allocator = default_allocator(n, SerialScheduler()) env, = leading_boundary(env_pre, psi, alg) - env′, = ctmrg_iteration(n, env, alg) - env_fixed = gauge_fix(env′, env, gauge_alg) + env′, = ctmrg_iteration(n, env, alg, allocator) + env_fixed = gauge_fix(env′, env, gauge_alg, allocator) env_diff = calc_elementwise_convergence(env, env_fixed) @info "Diff between iters = $(env_diff)" @test env_diff ≈ 0 atol = atol diff --git a/test/testsuite/ctmrg/jacobian_real_linear.jl b/test/testsuite/ctmrg/jacobian_real_linear.jl index 530d4deeb..e5423a38e 100644 --- a/test/testsuite/ctmrg/jacobian_real_linear.jl +++ b/test/testsuite/ctmrg/jacobian_real_linear.jl @@ -5,7 +5,8 @@ using Zygote using Adapt using TensorKit, KrylovKit, PEPSKit using PEPSKit: - ctmrg_iteration, compute_gauge_fix_gauge, fix_phases, ScramblingEnvGauge + ctmrg_iteration, compute_gauge_fix_gauge, fix_phases, ScramblingEnvGauge, + default_allocator, SerialScheduler algs = [ SimultaneousCTMRG(; projector_alg = :HalfInfiniteProjector), @@ -22,12 +23,13 @@ function ctmrg_jacobian_real_linear(AT) env, = leading_boundary(CTMRGEnv(state, ComplexSpace(χenv)), state, ctm_alg) # follow code of _rrule - env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, ctm_alg) - signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge) + allocator = default_allocator(state, SerialScheduler()) + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, ctm_alg, allocator) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge, allocator) _, env_vjp = pullback(state, env_conv) do A, x - e, = ctmrg_iteration(InfiniteSquareNetwork(A), x, ctm_alg) - return fix_phases(e, signs, corner_phases, edge_phases) + e, = ctmrg_iteration(InfiniteSquareNetwork(A), x, ctm_alg, allocator) + return fix_phases(e, signs, corner_phases, edge_phases, allocator) end # get Jacobians of single iteration diff --git a/test/testsuite/ctmrg/partition_function.jl b/test/testsuite/ctmrg/partition_function.jl index 92bf02154..b6948fa94 100644 --- a/test/testsuite/ctmrg/partition_function.jl +++ b/test/testsuite/ctmrg/partition_function.jl @@ -2,6 +2,7 @@ using Test using Random using LinearAlgebra using PEPSKit +using PEPSKit: default_allocator, SerialScheduler using TensorKit using QuadGK using Test @@ -134,7 +135,8 @@ function ctmrg_partition_function(AT) env, = leading_boundary(env₀₀, Z; alg, maxiter = 300, projector_alg) # check observables - λ = network_value(Z, env) + allocator = default_allocator(Z, SerialScheduler()) + λ = network_value(Z, env, allocator) m = expectation_value(Z, (1, 1) => M, env) e = expectation_value(Z, (1, 1) => E, env) f_exact, m_exact, e_exact = classical_ising_exact(; beta) diff --git a/test/testsuite/ctmrg/pepo.jl b/test/testsuite/ctmrg/pepo.jl index 0258e943d..ab7429958 100644 --- a/test/testsuite/ctmrg/pepo.jl +++ b/test/testsuite/ctmrg/pepo.jl @@ -2,6 +2,9 @@ using Test using Random using LinearAlgebra using PEPSKit +using PEPSKit: + default_allocator, + SerialScheduler using TensorKit using KrylovKit using OptimKit @@ -117,6 +120,7 @@ function ctmrg_pepo_fixed_point(AT) psi0 = initializePEPS(T, χpeps) env2_0 = CTMRGEnv(InfiniteSquareNetwork(psi0), χenv) env3_0 = CTMRGEnv(InfiniteSquareNetwork(psi0, T), χenv) + allocator = default_allocator(psi0, SerialScheduler()) # optimize free energy per site (psi_final, env2_final, env3_final), f, = optimize( @@ -139,8 +143,8 @@ function ctmrg_pepo_fixed_point(AT) PEPSKit.update!(env2, env2′) PEPSKit.update!(env3, env3′) end - λ3 = network_value(n3, env3) - λ2 = network_value(n2, env2) + λ3 = network_value(n3, env3, allocator) + λ2 = network_value(n2, env2, allocator) return -log(real(λ3 / λ2)) end g = only(gs) @@ -149,8 +153,8 @@ function ctmrg_pepo_fixed_point(AT) # check energy n3_final = InfiniteSquareNetwork(psi_final, T) - m = PEPSKit.contract_local_tensor((1, 1, 1), M, n3_final, env3_final) - nrm3 = PEPSKit._contract_site((1, 1), n3_final, env3_final) + m = PEPSKit.contract_local_tensor((1, 1, 1), M, n3_final, env3_final, allocator) + nrm3 = PEPSKit._contract_site((1, 1), n3_final, env3_final, allocator) # compare to Monte-Carlo result from https://www.worldscientific.com/doi/abs/10.1142/S0129183101002383 @test abs(m / nrm3) ≈ 0.667162 rtol = 1.0e-2 diff --git a/test/testsuite/ctmrg/suweight.jl b/test/testsuite/ctmrg/suweight.jl index 293a6f75e..affb77363 100644 --- a/test/testsuite/ctmrg/suweight.jl +++ b/test/testsuite/ctmrg/suweight.jl @@ -2,7 +2,7 @@ using Test using Random using TensorKit using PEPSKit -using PEPSKit: str, twistdual, unitcell +using PEPSKit: str, twistdual, unitcell, default_allocator, SerialScheduler using Adapt Vps = Dict( @@ -39,6 +39,7 @@ function ctmrg_suweight(AT) Espaces = [Vv Vv Vv'; Vv Vv' Vv'] Pspaces = fill(Vp, size(Nspaces)) peps = adapt(AT, InfinitePEPS(randn, ComplexF64, Pspaces, Nspaces, Espaces)) + allocator = default_allocator(peps, SerialScheduler()) wts = SUWeight(peps) if init != :trivial rand!(wts) @@ -51,7 +52,7 @@ function ctmrg_suweight(AT) if init == :trivial @test ρ1 ≈ su_rdm_1x1(r, c, peps, nothing) end - ρ2 = reduced_densitymatrix([idx], peps, env) + ρ2 = reduced_densitymatrix([idx], peps, env, allocator) @test ρ1 ≈ ρ2 end end diff --git a/test/testsuite/ctmrg/unitcell.jl b/test/testsuite/ctmrg/unitcell.jl index fb7b24e70..0f185c5d9 100644 --- a/test/testsuite/ctmrg/unitcell.jl +++ b/test/testsuite/ctmrg/unitcell.jl @@ -1,7 +1,7 @@ using Test using Random using PEPSKit, Adapt -using PEPSKit: ctmrg_iteration, compute_gauge_fix_gauge, ScramblingEnvGauge +using PEPSKit: ctmrg_iteration, compute_gauge_fix_gauge, ScramblingEnvGauge, default_allocator, SerialScheduler using TensorKit # settings @@ -22,9 +22,10 @@ function test_unitcell( peps = adapt(AT, InfinitePEPS(randn, stype, Pspaces, Nspaces, Espaces)) env = CTMRGEnv(randn, stype, peps, chis_north, chis_east, chis_south, chis_west) + allocator = default_allocator(peps, SerialScheduler()) # apply one CTMRG iteration with fixeds - env′, info = ctmrg_iteration(InfiniteSquareNetwork(peps), env, ctm_alg) - env″, info = ctmrg_iteration(InfiniteSquareNetwork(peps), env′, ctm_alg) # another iteration to fix spaces + env′, info = ctmrg_iteration(InfiniteSquareNetwork(peps), env, ctm_alg, allocator) + env″, info = ctmrg_iteration(InfiniteSquareNetwork(peps), env′, ctm_alg, allocator) # another iteration to fix spaces # compute random expecation value to test matching bonds random_op = LocalOperator( @@ -42,7 +43,7 @@ function test_unitcell( @test expectation_value(peps, random_op, env′) isa Number # test if gauge fixing routines run through - signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env″, env′, ScramblingEnvGauge()) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env″, env′, ScramblingEnvGauge(), allocator) @test signs isa Array @test corner_phases isa Array @test edge_phases isa Array diff --git a/test/testsuite/toolbox/densitymatrices.jl b/test/testsuite/toolbox/densitymatrices.jl index 225240dcd..e31dea112 100644 --- a/test/testsuite/toolbox/densitymatrices.jl +++ b/test/testsuite/toolbox/densitymatrices.jl @@ -1,6 +1,6 @@ using TensorKit using PEPSKit -using PEPSKit: contract_local_operator, contract_local_norm, _contract_site +using PEPSKit: contract_local_operator, contract_local_norm, _contract_site, default_allocator, SerialScheduler using Test using TestExtras: @testinferred using Adapt @@ -22,6 +22,7 @@ function toolbox_single_layer_densitymatrix(AT) O = adapt(AT, rand(d, d)) @plansor O_pf[W S; N E] := O[p'; p] * ρ[1, 1, 1][p p'; N E S W] + allocator = default_allocator(ρ, SerialScheduler()) # Single site site = (1, 1) O_singlesite = LocalOperator(physicalspace(ρ), (site,) => O) @@ -30,8 +31,8 @@ function toolbox_single_layer_densitymatrix(AT) @test E1 ≈ E2 # the operator/norm route agrees with the density matrix one; a single layer is # passed to the contractions bare rather than wrapped in a tuple - val = contract_local_operator([site], O, ρ, env) - nrm = contract_local_norm([site], ρ, env) + val = contract_local_operator([site], O, ρ, env, allocator) + nrm = contract_local_norm([site], ρ, env, allocator) @test E1 ≈ val / nrm # two sites @@ -40,10 +41,10 @@ function toolbox_single_layer_densitymatrix(AT) [(2, 1), (1, 2), (2, 2), (2, 1), (3, 1)] ) O_twosite = LocalOperator(physicalspace(ρ), inds => O ⊗ O) - E3 = expectation_value(ρ, O_twosite, env) + E3 = expectation_value(ρ, O_twosite, env; allocator) # TODO: not defined for partition functions... - val = contract_local_operator(collect(inds), O ⊗ O, ρ, env) - nrm = contract_local_norm(collect(inds), ρ, env) + val = contract_local_operator(collect(inds), O ⊗ O, ρ, env, allocator) + nrm = contract_local_norm(collect(inds), ρ, env, allocator) @test E3 ≈ val / nrm end end @@ -59,9 +60,10 @@ function toolbox_double_layer_densitymatrix(AT) ρ_peps = @testinferred InfinitePEPS(ρ) env = CTMRGEnv(ρ_peps, χ) + allocator = default_allocator(ρ, SerialScheduler()) O = adapt(AT, rand(d, d)) F = adapt(AT, isomorphism(fuse(d ⊗ d'), d ⊗ d')) - @tensor O_doubled[-1; -2] := F[-1; 1 2] * O[1; 3] * twist(F', 2)[3 2; -2] + @tensor allocator = allocator O_doubled[-1; -2] := F[-1; 1 2] * O[1; 3] * twist(F', 2)[3 2; -2] # Single site site = (1, 1) @@ -70,13 +72,13 @@ function toolbox_double_layer_densitymatrix(AT) O_doubled_singlesite = LocalOperator(physicalspace(ρ_peps), (site,) => O_doubled) E2 = expectation_value(ρ_peps, O_doubled_singlesite, ρ_peps, env) @test E1 ≈ E2 - val = contract_local_operator([site], O_doubled, ρ_peps, ρ_peps, env) - nrm = contract_local_norm([site], ρ_peps, ρ_peps, env) - @test nrm ≈ _contract_site(site, InfiniteSquareNetwork(ρ_peps), env) + val = contract_local_operator([site], O_doubled, ρ_peps, ρ_peps, env, allocator) + nrm = contract_local_norm([site], ρ_peps, ρ_peps, env, allocator) + @test nrm ≈ _contract_site(site, InfiniteSquareNetwork(ρ_peps), env, allocator) @test E1 ≈ val / nrm # same through the two-layer PEPO sandwich rather than its fused PEPS view - val_pepo = contract_local_operator([site], O, ρ, ρ, env) - nrm_pepo = contract_local_norm([site], ρ, ρ, env) + val_pepo = contract_local_operator([site], O, ρ, ρ, env, allocator) + nrm_pepo = contract_local_norm([site], ρ, ρ, env, allocator) @test E1 ≈ val_pepo / nrm_pepo # two sites @@ -89,11 +91,11 @@ function toolbox_double_layer_densitymatrix(AT) O_doubled_twosite = LocalOperator(physicalspace(ρ_peps), inds => O_doubled ⊗ O_doubled) E2 = expectation_value(ρ_peps, O_doubled_twosite, ρ_peps, env) @test E1 ≈ E2 - val = contract_local_operator(collect(inds), O_doubled ⊗ O_doubled, ρ_peps, ρ_peps, env) - nrm = contract_local_norm(collect(inds), ρ_peps, ρ_peps, env) + val = contract_local_operator(collect(inds), O_doubled ⊗ O_doubled, ρ_peps, ρ_peps, env, allocator) + nrm = contract_local_norm(collect(inds), ρ_peps, ρ_peps, env, allocator) @test E1 ≈ val / nrm - val_pepo = contract_local_operator(collect(inds), O ⊗ O, ρ, ρ, env) - nrm_pepo = contract_local_norm(collect(inds), ρ, ρ, env) + val_pepo = contract_local_operator(collect(inds), O ⊗ O, ρ, ρ, env, allocator) + nrm_pepo = contract_local_norm(collect(inds), ρ, ρ, env, allocator) @test E1 ≈ val_pepo / nrm_pepo end end @@ -105,7 +107,8 @@ function toolbox_densitymatrix_too_many_layers(AT) ρ = adapt(AT, InfinitePEPO(d, D; unitcell = (2, 2, 1))) env = CTMRGEnv(InfinitePEPS(ρ), χ) inds = Tuple(Val.([CartesianIndex(1, 1)])) - @test_throws ArgumentError PEPSKit._contract_densitymatrix(inds, (ρ, ρ, ρ), env) + allocator = default_allocator(ρ, SerialScheduler()) + @test_throws ArgumentError PEPSKit._contract_densitymatrix(inds, (ρ, ρ, ρ), env, allocator) end end @@ -115,25 +118,26 @@ function toolbox_densitymatrix_generic_fallback(AT) peps = adapt(AT, InfinitePEPS(d, D; unitcell = (2, 2))) env = CTMRGEnv(peps, χ) + allocator = default_allocator(peps, SerialScheduler()) # environments without a hand-optimized contraction of a given shape fall through # to `_contract_densitymatrix`, which should agree with the specializations that # do exist for ind in CartesianIndices((2, 2)) - ρ_fast = reduced_densitymatrix([ind], peps, env) + ρ_fast = reduced_densitymatrix([ind], peps, env, allocator) ρ_generic = invoke( - PEPSKit.reduced_densitymatrix1x1, Tuple{Any, Any, Any, Any}, ind, peps, peps, env + PEPSKit.reduced_densitymatrix1x1, Tuple{Any, Any, Any, Any, Any}, ind, peps, peps, env, allocator, ) @test ρ_fast ≈ ρ_generic - ρ_fast = reduced_densitymatrix([ind, ind + CartesianIndex(1, 0)], peps, env) + ρ_fast = reduced_densitymatrix([ind, ind + CartesianIndex(1, 0)], peps, env, allocator) ρ_generic = invoke( - PEPSKit.reduced_densitymatrix2x1, Tuple{Any, Any, Any, Any}, ind, peps, peps, env + PEPSKit.reduced_densitymatrix2x1, Tuple{Any, Any, Any, Any, Any}, ind, peps, peps, env, allocator, ) @test ρ_fast ≈ ρ_generic - ρ_fast = reduced_densitymatrix([ind, ind + CartesianIndex(0, 1)], peps, env) + ρ_fast = reduced_densitymatrix([ind, ind + CartesianIndex(0, 1)], peps, env, allocator) ρ_generic = invoke( - PEPSKit.reduced_densitymatrix1x2, Tuple{Any, Any, Any, Any}, ind, peps, peps, env + PEPSKit.reduced_densitymatrix1x2, Tuple{Any, Any, Any, Any, Any}, ind, peps, peps, env, allocator, ) @test ρ_fast ≈ ρ_generic end From 53408c5689856ebc2caf059e38e1a4d07f78e858 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sat, 7 Nov 2026 12:39:43 -0500 Subject: [PATCH 2/2] Just use SerialScheduler --- src/algorithms/ctmrg/ctmrg.jl | 3 +-- src/algorithms/ctmrg/gaugefix.jl | 2 +- src/algorithms/ctmrg/initialization.jl | 3 +-- .../expectation_value/expectation_value.jl | 19 +++++++------------ .../expectation_value/network_value.jl | 10 +++++----- .../optimization/implicit_differentiation.jl | 9 +++------ .../optimization/peps_optimization.jl | 3 +-- src/algorithms/time_evolution/simpleupdate.jl | 3 +-- 8 files changed, 20 insertions(+), 32 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 8814c4eab..3037cf6e3 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -107,8 +107,7 @@ function leading_boundary(env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwo end function leading_boundary( env₀::AbstractCTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm; - scheduler = Defaults.scheduler[], - allocator = default_allocator(network, scheduler) + allocator = default_allocator(network, SerialScheduler()) ) check_input(leading_boundary, network, env₀, alg) log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index 150f6e586..ebde2b883 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -166,7 +166,7 @@ end function right_transfermatrix_fixedpoint( tops, bottoms, alg = Arnoldi(), ρ0 = initialize_right_fixedpoint(tops, bottoms), - allocator = default_allocator(ρ0, Defaults.scheduler[]) + allocator = default_allocator(ρ0, SerialScheduler()) ) _, vecs, info = eigsolve(ρ0, 1, :LM, alg) do ρ return foldr(zip(tops, bottoms); init = ρ) do (top, bottom), ρ diff --git a/src/algorithms/ctmrg/initialization.jl b/src/algorithms/ctmrg/initialization.jl index c458ab000..13e51fb57 100644 --- a/src/algorithms/ctmrg/initialization.jl +++ b/src/algorithms/ctmrg/initialization.jl @@ -40,8 +40,7 @@ function initialize_ctmrg_environment( n::InfiniteSquareNetwork, alg::ApplicationInitialization, env0 = ProductStateEnv(alg.f, elt, n); - scheduler = Defaults.scheduler[], - allocator = default_allocator(n, scheduler) + allocator = default_allocator(n, SerialScheduler()) ) dummy_alg = SimultaneousCTMRG(trunc = (; alg = :notrunc)) env, = ctmrg_iteration(n, CTMRGEnv(env0), dummy_alg, allocator) diff --git a/src/algorithms/expectation_value/expectation_value.jl b/src/algorithms/expectation_value/expectation_value.jl index 70fd3c46a..7bd6e437d 100644 --- a/src/algorithms/expectation_value/expectation_value.jl +++ b/src/algorithms/expectation_value/expectation_value.jl @@ -12,8 +12,7 @@ the second signature yields a bilayer contraction instead. """ function MPSKit.expectation_value( bra::S, O::LocalOperator, ket::S, env; - scheduler = Defaults.scheduler[], - allocator = default_allocator(bra, scheduler) + allocator = default_allocator(bra, SerialScheduler()) ) where {S <: InfiniteState} checklattice(bra, O, ket) term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly @@ -23,13 +22,11 @@ function MPSKit.expectation_value( end MPSKit.expectation_value( peps::InfinitePEPS, O::LocalOperator, env; - scheduler = Defaults.scheduler[], - allocator = default_allocator(peps, scheduler) -) = expectation_value(peps, O, peps, env; scheduler, allocator) + allocator = default_allocator(peps, SerialScheduler()) +) = expectation_value(peps, O, peps, env; allocator) function MPSKit.expectation_value( state::InfinitePEPO, O::LocalOperator, env; - scheduler = Defaults.scheduler[], - allocator = default_allocator(state, scheduler) + allocator = default_allocator(state, SerialScheduler()) ) checklattice(state, O) term_vals = dtmap(collect(O.terms)) do (inds, operator) # OhMyThreads can't iterate over O.terms directly @@ -102,16 +99,14 @@ function MPSKit.expectation_value( pf::InfinitePartitionFunction, op::Pair{CartesianIndex{2}, <:AbstractTensorMap{T, S, 2, 2}}, env; - scheduler = Defaults.scheduler[], - allocator = default_allocator(pf, scheduler) + allocator = default_allocator(pf, SerialScheduler()) ) where {T, S} return contract_local_tensor(op[1], op[2], env, allocator) / contract_local_tensor(op[1], pf[op[1]], env, allocator) end function MPSKit.expectation_value( pf::InfinitePartitionFunction, op::Pair{Tuple{Int, Int}}, env; - scheduler = Defaults.scheduler[], - allocator = default_allocator(pf, scheduler) + allocator = default_allocator(pf, SerialScheduler()) ) - return expectation_value(pf, CartesianIndex(op[1]) => op[2], env; scheduler, allocator) + return expectation_value(pf, CartesianIndex(op[1]) => op[2], env; allocator) end diff --git a/src/algorithms/expectation_value/network_value.jl b/src/algorithms/expectation_value/network_value.jl index 2c8809b9d..c32cc1624 100644 --- a/src/algorithms/expectation_value/network_value.jl +++ b/src/algorithms/expectation_value/network_value.jl @@ -7,16 +7,16 @@ Return the value (per unit cell) of a given contractible network contracted using a given CTMRG environment. """ -function network_value(network::InfiniteSquareNetwork, env::AbstractCTMRGEnv, allocator; scheduler = Defaults.scheduler[]) - return tmapreduce(*, Iterators.product(axes(network)...); scheduler) do (r, c) +function network_value(network::InfiniteSquareNetwork, env::AbstractCTMRGEnv, allocator) + return prod(Iterators.product(axes(network)...)) do (r, c) return _contract_site((r, c), network, env, allocator) * _contract_corners((r, c), env, allocator) / _contract_vertical_edges((r, c), env, allocator) / _contract_horizontal_edges((r, c), env, allocator) end end -network_value(state, env::AbstractCTMRGEnv, allocator; kwargs...) = network_value(InfiniteSquareNetwork(state), env, allocator; kwargs...) +network_value(state, env::AbstractCTMRGEnv, allocator) = network_value(InfiniteSquareNetwork(state), env, allocator) -function LinearAlgebra.norm(peps::InfinitePEPS, env::AbstractCTMRGEnv; scheduler = Defaults.scheduler[], allocator = default_allocator(peps, scheduler)) - return network_value(InfiniteSquareNetwork(peps), env, allocator; scheduler) +function LinearAlgebra.norm(peps::InfinitePEPS, env::AbstractCTMRGEnv; allocator = default_allocator(peps, SerialScheduler())) + return network_value(InfiniteSquareNetwork(peps), env, allocator) end # Local tensor insertions diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index fb6daa491..75406d4e7 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -217,8 +217,7 @@ function _rrule( envinit::AbstractCTMRGEnv, state, alg::CTMRGAlgorithm; - scheduler = Defaults.scheduler[], - allocator = default_allocator(state, scheduler) + allocator = default_allocator(state, SerialScheduler()) ) _check_algorithm_combination(alg, gradmode) @@ -471,8 +470,7 @@ function _rrule( envinit::C4vCTMRGEnv, state, alg::C4vCTMRG; - scheduler = Defaults.scheduler[], - allocator = default_allocator(state, scheduler) + allocator = default_allocator(state, SerialScheduler()) ) _check_algorithm_combination(alg, gradmode) @@ -622,8 +620,7 @@ function PEPSKit._rrule( envinit::AbstractCTMRGEnv, state, alg::SimultaneousCTMRG{<:HalfInfiniteProjector}; - scheduler = Defaults.scheduler[], - allocator = default_allocator(state, scheduler) + allocator = default_allocator(state, SerialScheduler()) ) env, = leading_boundary(envinit, state, alg) diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 89a3dcb72..4fcf958b6 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -233,8 +233,7 @@ function fixedpoint( (finalize!) = OptimKit._finalize!, hasconverged = OptimKit.DefaultHasConverged(alg.optimizer_alg.gradtol), shouldstop = OptimKit.DefaultShouldStop(alg.optimizer_alg.maxiter), - scheduler = Defaults.scheduler[], - allocator = default_allocator(peps₀, scheduler) + allocator = default_allocator(peps₀, SerialScheduler()) ) # validate inputs check_input(fixedpoint, peps₀, env₀, alg) diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 17718da0b..0e18e43b4 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -128,8 +128,7 @@ One iteration of simple update. Pass `allocator` to reuse contraction scratch; function su_iter( state::InfiniteState, circuit::LocalCircuit, alg::SimpleUpdate, env::SUWeight; - scheduler = Defaults.scheduler[], - allocator = default_allocator(state, scheduler) + allocator = default_allocator(state, SerialScheduler()) ) state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0 # Gates are applied sequentially and share the same scratch allocator.