From 50dab6d9a9c39e2eb94c6a86016628887c955ac9 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 13:13:07 +0800 Subject: [PATCH 1/8] Add gauge transformations for PEPS, PEPO, and CTMRG environments --- src/PEPSKit.jl | 6 +- src/algorithms/bp/gaugefix.jl | 16 +- src/algorithms/contractions/ctmrg/gaugefix.jl | 227 ------------------ src/algorithms/ctmrg/gaugefix.jl | 123 +++------- src/algorithms/ctmrg/gaugetrans.jl | 130 ++++++++++ .../optimization/implicit_differentiation.jl | 10 +- src/networks/gaugetrans.jl | 188 +++++++++++++++ src/operators/infinitepepo.jl | 24 ++ src/states/infinitepeps.jl | 18 ++ test/bp/gaugefix.jl | 1 + test/ctmrg/gaugetrans.jl | 13 + test/testsuite/TestSuite.jl | 8 +- test/testsuite/bp/gaugefix.jl | 22 +- test/testsuite/ctmrg/gaugetrans.jl | 103 ++++++++ test/testsuite/ctmrg/unitcell.jl | 2 +- 15 files changed, 553 insertions(+), 338 deletions(-) delete mode 100644 src/algorithms/contractions/ctmrg/gaugefix.jl create mode 100644 src/algorithms/ctmrg/gaugetrans.jl create mode 100644 src/networks/gaugetrans.jl create mode 100644 test/ctmrg/gaugetrans.jl create mode 100644 test/testsuite/ctmrg/gaugetrans.jl diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index b715dc2ec..22d895c88 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -75,6 +75,7 @@ include("utility/initialization.jl") include("networks/tensors.jl") include("networks/local_sandwich.jl") include("networks/infinitesquarenetwork.jl") +include("networks/gaugetrans.jl") include("states/infinitepeps.jl") include("states/infinitepartitionfunction.jl") @@ -108,7 +109,6 @@ include("algorithms/contractions/ctmrg/fullinf_env.jl") include("algorithms/contractions/ctmrg/renormalize_corner.jl") include("algorithms/contractions/ctmrg/renormalize_edge.jl") include("algorithms/contractions/ctmrg/network_value.jl") -include("algorithms/contractions/ctmrg/gaugefix.jl") include("algorithms/contractions/ctmrg/characteristic_equations.jl") include("algorithms/contractions/absorb.jl") @@ -142,6 +142,7 @@ include("algorithms/ctmrg/sequential.jl") include("algorithms/ctmrg/gaugefix.jl") include("algorithms/ctmrg/c4v.jl") include("algorithms/ctmrg/initialization.jl") +include("algorithms/ctmrg/gaugetrans.jl") include("algorithms/truncation/truncationschemes.jl") include("algorithms/truncation/fullenv_truncation.jl") @@ -212,7 +213,8 @@ export InfinitePEPO, InfiniteTransferPEPO export BPEnv, BeliefPropagation export BPGauge, SUGauge -export gauge_fix +export gauge_fix, gauge_transform +export VirtualGaugeTransform, CTMRGEnvGaugeTransform export initialize_mps, initializePEPS export ReflectDepth, ReflectWidth, Rotate, RotateReflect diff --git a/src/algorithms/bp/gaugefix.jl b/src/algorithms/bp/gaugefix.jl index a86599e04..c4c9686fa 100644 --- a/src/algorithms/bp/gaugefix.jl +++ b/src/algorithms/bp/gaugefix.jl @@ -29,7 +29,7 @@ here. The available algorithms are: end """ - gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv) + gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv) -> psi′, gauge Fix the gauge of `psi` (which can be an [`InfinitePEPS`](@ref), or an [`InfinitePEPO`](@ref) interpreted as purified state with two physical legs) @@ -49,7 +49,7 @@ function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv) psi′[r + 1, 2] = copy(psi′[r, 1]) end end - return psi′, XXinv + return psi′, VirtualGaugeTransform(XXinv) end function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv) # convert to iPEPS, and store physical leg fusers @@ -61,13 +61,17 @@ function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv) psi_Fs = reshape(psi_Fs, (Nr, Nc)) psi′ = map(Base.Fix2(getindex, 1), psi_Fs) Fs = map(Base.Fix2(getindex, 2), psi_Fs) - psi′, XXinv = gauge_fix(InfinitePEPS(psi′), alg, env) + psi′, gauge = gauge_fix(InfinitePEPS(psi′), alg, env) # convert back to iPEPO psi′ = map(psi′.A, Fs) do t, F return F' * t end psi′ = reshape(psi′, (Nr, Nc, 1)) - return InfinitePEPO(psi′), XXinv + gauge_dims = (2, Nr, Nc, 1) + gauge = VirtualGaugeTransform( + reshape(gauge.matrices, gauge_dims), reshape(gauge.inverses, gauge_dims) + ) + return InfinitePEPO(psi′), gauge end function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv) @@ -94,8 +98,8 @@ along the canonical direction of the PEPS arrows (`SOUTH ← NORTH` or `WEST ← = X X⁻¹ ``` -Which are then used to update the gauge of `psi`. Thus, by convention `X` is attached to the `SOUTH`/`WEST` directions -and `X⁻¹` is attached to the `NORTH`/`EAST` directions. +These matrices update the gauge of `psi` using the [`VirtualGaugeTransform`](@ref) convention. +`X` is attached to the `NORTH`/`EAST` leg at the indexed site, and `X⁻¹` to the neighboring `SOUTH`/`WEST` leg. """ function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv, alg::BPGauge) dir, row, col = Tuple(I) diff --git a/src/algorithms/contractions/ctmrg/gaugefix.jl b/src/algorithms/contractions/ctmrg/gaugefix.jl deleted file mode 100644 index 41f9222eb..000000000 --- a/src/algorithms/contractions/ctmrg/gaugefix.jl +++ /dev/null @@ -1,227 +0,0 @@ -# Gauge fixing contractions -# ------------------------- - -# corners - -""" -$(SIGNATURES) - -Multiply corner tensor with incoming and outgoing gauge signs. - -``` - corner -- σ_in --in - | - σ_out - | - out -``` -""" -function fix_gauge_corner( - corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor - ) - return @autoopt @tensor corner_fix[χ_out; χ_in] := - σ_out[χ_out; χ1] * corner[χ1; χ2] * conj(σ_in[χ_in; χ2]) -end - -""" -$(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) - return fix_gauge_corner( - corner(env, NORTHWEST, row, col), - signs[WEST, row, col], - signs[NORTH, row, _next(col, end)], - ) -end - -""" -$(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) - return fix_gauge_corner( - corner(env, NORTHEAST, row, col), - signs[NORTH, row, col], - signs[EAST, _next(row, end), col], - ) -end - -""" -$(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) - return fix_gauge_corner( - corner(env, SOUTHEAST, row, col), - signs[EAST, row, col], - signs[SOUTH, row, _prev(col, end)], - ) -end - -""" -$(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) - return fix_gauge_corner( - corner(env, SOUTHWEST, row, col), - signs[SOUTH, row, col], - signs[WEST, _prev(row, end), col], - ) -end - -# edges - -""" -$(SIGNATURES) - -Multiply edge tensor with incoming and outgoing gauge signs. - -``` - out-- σ_out -- edge -- σ_in --in - | -``` -""" -@generated function fix_gauge_edge( - edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor - ) where {T, S, N} - edge_fix_e = tensorexpr( - :edge_fix, - (envlabel(:out), ntuple(i -> virtuallabel(i), N - 1)...), - (envlabel(:in),), - ) - edge_e = tensorexpr( - :edge, (envlabel(1), ntuple(i -> virtuallabel(i), N - 1)...), (envlabel(2),) - ) - σ_out_e = tensorexpr(:σ_out, (envlabel(:out),), (envlabel(1),)) - σ_in_e = tensorexpr(:σ_in, (envlabel(:in),), (envlabel(2),)) - return macroexpand( - @__MODULE__, - :(return @autoopt @tensor $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)), - ) -end - -""" -$(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) - return fix_gauge_edge( - edge(env, NORTH, row, col), signs[NORTH, row, col], signs[NORTH, row, _next(col, end)], - ) -end - -""" -$(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) - return fix_gauge_edge( - edge(env, EAST, row, col), signs[EAST, row, col], signs[EAST, _next(row, end), col] - ) -end - -""" -$(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) - return fix_gauge_edge( - edge(env, SOUTH, row, col), signs[SOUTH, row, col], signs[SOUTH, row, _prev(col, end)], - ) -end - -""" -$(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) - return fix_gauge_edge( - edge(env, WEST, row, col), signs[WEST, row, col], signs[WEST, _prev(row, end), col] - ) -end - -# left singular vectors - -""" -$(SIGNATURES) - -Multiply north left singular vectors with gauge signs from the right. -""" -function fix_gauge_north_left_vecs((row, col), U, signs) - return U[NORTH, row, col] * signs[NORTH, row, _next(col, end)]' -end - -""" -$(SIGNATURES) - -Multiply east left singular vectors with gauge signs from the right. -""" -function fix_gauge_east_left_vecs((row, col), U, signs) - return U[EAST, row, col] * signs[EAST, _next(row, end), col]' -end - -""" -$(SIGNATURES) - -Multiply south left singular vectors with gauge signs from the right. -""" -function fix_gauge_south_left_vecs((row, col), U, signs) - return U[SOUTH, row, col] * signs[SOUTH, row, _prev(col, end)]' -end - -""" -$(SIGNATURES) - -Multiply west left singular vectors with gauge signs from the right. -""" -function fix_gauge_west_left_vecs((row, col), U, signs) - return U[WEST, row, col] * signs[WEST, _prev(row, end), col]' -end - -# right singular vectors - -""" -$(SIGNATURES) - -Multiply north right singular vectors with gauge signs from the left. -""" -function fix_gauge_north_right_vecs((row, col), V, signs) - return signs[NORTH, row, _next(col, end)] * V[NORTH, row, col] -end - -""" -$(SIGNATURES) - -Multiply east right singular vectors with gauge signs from the left. -""" -function fix_gauge_east_right_vecs((row, col), V, signs) - return signs[EAST, _next(row, end), col] * V[EAST, row, col] -end - -""" -$(SIGNATURES) - -Multiply south right singular vectors with gauge signs from the left. -""" -function fix_gauge_south_right_vecs((row, col), V, signs) - return signs[SOUTH, row, _prev(col, end)] * V[SOUTH, row, col] -end - -""" -$(SIGNATURES) - -Multiply west right singular vectors with gauge signs from the left. -""" -function fix_gauge_west_right_vecs((row, col), V, signs) - return signs[WEST, _prev(row, end), col] * V[WEST, row, col] -end diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index f6bf3463f..ce37a578f 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -40,7 +40,7 @@ function compute_gauge_fix_gauge( # find additional global phases corner_phases, edge_phases = compute_global_phases( - fix_relative_phases(envfinal, signs), envprev + gauge_transform(envfinal, signs), envprev ) return signs, corner_phases, edge_phases @@ -67,7 +67,7 @@ function compute_gauge_fix_gauge( return signs, corner_phases, edge_phases end -function fix_phases(env::AbstractCTMRGEnv, signs, corner_phases, edge_phases) +function fix_phases(env::AbstractCTMRGEnv, signs::CTMRGEnvGaugeTransform, 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) return env_fixed_global @@ -98,19 +98,10 @@ function compute_relative_phases( randn(storagetype(T), codomain(t) ← domain(t)) end - # Find right fixed points of mixed transfer matrices - eigsolve_alg = Arnoldi() - ρprev = right_transfermatrix_fixedpoint(Tsprev, M, eigsolve_alg) - ρfinal = right_transfermatrix_fixedpoint(Tsfinal, M, eigsolve_alg) - - # Decompose and multiply - Qprev, = left_orth!(ρprev; positive = true) - Qfinal, = left_orth!(ρfinal; positive = true) - - return Qprev * Qfinal' + return _scrambling_gauge_matrix(Tsprev, Tsfinal, M, Arnoldi()) end - return signs + return CTMRGEnvGaugeTransform(signs, adjoint.(signs)) end function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, ::ScramblingEnvGauge) where {C, T} @@ -122,20 +113,17 @@ function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGE # (make Hermitian such that T-M transfer matrix has real eigenvalues) M = _project_hermitian(randn(storagetype(Tfinal), space(Tfinal))) - # Find right fixed points of mixed transfer matrices - eigsolve_alg = Lanczos() # real eigenvalues - ρprev = right_transfermatrix_fixedpoint([Tprev], [M], eigsolve_alg) - ρfinal = right_transfermatrix_fixedpoint([Tfinal], [M], eigsolve_alg) + σ = _scrambling_gauge_matrix([Tprev], [Tfinal], [M], Lanczos()) + return CTMRGEnvGaugeTransform(fill(σ, (4, 1, 1)), fill(σ', (4, 1, 1))) +end - # Decompose and multiply +"""Extract the relative unitary gauge from mixed-transfer fixed points with a shared random MPS.""" +function _scrambling_gauge_matrix(Tsprev, Tsfinal, M, eigsolve_alg) + ρprev = right_transfermatrix_fixedpoint(Tsprev, M, eigsolve_alg) + ρfinal = right_transfermatrix_fixedpoint(Tsfinal, M, eigsolve_alg) Qprev, = left_orth!(ρprev; positive = true) Qfinal, = left_orth!(ρfinal; positive = true) - - σ = Qprev * Qfinal' - - signs = fill(σ, (4, 1, 1)) - - return signs + return Qprev * Qfinal' end # this is a bit of a hack to get the fixed point of the mixed transfer matrix @@ -177,86 +165,33 @@ function right_transfermatrix_fixedpoint( return first(vecs) end -# Explicit fixing of relative phases (doing this compactly in a loop is annoying) -function fix_relative_phases(envfinal::CTMRGEnv, signs) - corners_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - Cf = if dir == NORTHWEST - fix_gauge_northwest_corner((r, c), envfinal, signs) - elseif dir == NORTHEAST - fix_gauge_northeast_corner((r, c), envfinal, signs) - elseif dir == SOUTHEAST - fix_gauge_southeast_corner((r, c), envfinal, signs) - else # dir == SOUTHWEST - fix_gauge_southwest_corner((r, c), envfinal, signs) - end - return Cf - end +fix_relative_phases(env::CTMRGEnv, signs::CTMRGEnvGaugeTransform) = gauge_transform(env, signs) - edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - Ef = if dir == NORTHWEST - fix_gauge_north_edge((r, c), envfinal, signs) - elseif dir == NORTHEAST - fix_gauge_east_edge((r, c), envfinal, signs) - elseif dir == SOUTHEAST - fix_gauge_south_edge((r, c), envfinal, signs) - else # dir == SOUTHWEST - fix_gauge_west_edge((r, c), envfinal, signs) - end - return Ef +function fix_relative_phases(env::C4vCTMRGEnv, signs::CTMRGEnvGaugeTransform) + X, Xinv = signs.matrices, signs.inverses + corners = map(env.corners) do C + return fix_gauge_corner(C, X[WEST, 1, 1], Xinv[NORTH, 1, 1]') end - - return CTMRGEnv(corners_fixed, edges_fixed) -end -function fix_relative_phases(envfinal::C4vCTMRGEnv, signs) - # all corners are the same - corner_fixed = map(eachcoordinate(envfinal, 1:1)) do (dir, r, c) - fix_gauge_northwest_corner((r, c), envfinal, signs) - end - edges_fixed = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - Ef = if dir == NORTHWEST - fix_gauge_north_edge((r, c), envfinal, signs) - elseif dir == NORTHEAST - fix_gauge_east_edge((r, c), envfinal, signs) - elseif dir == SOUTHEAST - fix_gauge_south_edge((r, c), envfinal, signs) - else # dir == SOUTHWEST - fix_gauge_west_edge((r, c), envfinal, signs) - end - return Ef + edges = map(eachcoordinate(env, 1:4)) do (dir, r, c) + return fix_gauge_edge(env.edges[dir, r, c], X[dir, r, c], Xinv[dir, r, c]') end - return C4vCTMRGEnv(corner_fixed, edges_fixed) + return C4vCTMRGEnv(corners, edges) end + function fix_relative_phases( - U::Array{Ut, 3}, V::Array{Vt, 3}, signs + U::Array{Ut, 3}, V::Array{Vt, 3}, signs::CTMRGEnvGaugeTransform ) where {Ut <: AbstractTensorMap, Vt <: AbstractTensorMap} + nr, nc = size(U)[2:3] U_fixed = map(eachindex(IndexCartesian(), U)) do I - dir, r, c = Tuple(I) - Uf = if dir == NORTHWEST - fix_gauge_north_left_vecs((r, c), U, signs) - elseif dir == NORTHEAST - fix_gauge_east_left_vecs((r, c), U, signs) - elseif dir == SOUTHEAST - fix_gauge_south_left_vecs((r, c), U, signs) - else # dir == SOUTHWEST - fix_gauge_west_left_vecs((r, c), U, signs) - end - return Uf + d, r, c = Tuple(I) + ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + return U[I] * signs.inverses[d, ri, ci] end - V_fixed = map(eachindex(IndexCartesian(), V)) do I - dir, r, c = Tuple(I) - Vf = if dir == NORTHWEST - fix_gauge_north_right_vecs((r, c), V, signs) - elseif dir == NORTHEAST - fix_gauge_east_right_vecs((r, c), V, signs) - elseif dir == SOUTHEAST - fix_gauge_south_right_vecs((r, c), V, signs) - else # dir == SOUTHWEST - fix_gauge_west_right_vecs((r, c), V, signs) - end - return Vf + d, r, c = Tuple(I) + ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + return signs.matrices[d, ri, ci] * V[I] end - return U_fixed, V_fixed end diff --git a/src/algorithms/ctmrg/gaugetrans.jl b/src/algorithms/ctmrg/gaugetrans.jl new file mode 100644 index 000000000..e818f42e7 --- /dev/null +++ b/src/algorithms/ctmrg/gaugetrans.jl @@ -0,0 +1,130 @@ +""" + gauge_transform(env::CTMRGEnv, gauge::VirtualGaugeTransform) + +Transform the network-facing virtual legs of `env` to compensate for a virtual gauge transformation of its state. +A gauge of shape `(2, rows, cols)` acts on a single-layer PEPO environment or a double-layer PEPS/PEPO environment, as determined by the number of virtual legs on each edge. +In a double layer, the same gauge acts on the ket and its complex conjugate on the bra. +A PEPO gauge with an explicit layer dimension is accepted when it contains one layer. +The corner tensors and boundary legs are unchanged. +""" +function gauge_transform(env::CTMRGEnv, gauge::VirtualGaugeTransform{3}) + _check_gauge_size(gauge, size(env)[2:3]) + X, Xinv = gauge.matrices, gauge.inverses + edges = map(eachcoordinate(env, 1:4)) do (d, r, c) + factor = if d == NORTH + Xinv[NORTH, _next(r, size(X, 2)), c] + elseif d == EAST + Xinv[EAST, r, _prev(c, size(X, 3))] + elseif d == SOUTH + X[NORTH, r, c] + else # d == WEST + X[EAST, r, c] + end + return _gauge_transform_virtual_edge(env.edges[d, r, c], factor, d == NORTH || d == EAST) + end + return CTMRGEnv(env.corners, edges) +end +function gauge_transform(env::CTMRGEnv, gauge::VirtualGaugeTransform{4}) + size(gauge.matrices, 4) == 1 || throw(ArgumentError("Virtual environment transformations require a single-layer gauge")) + return gauge_transform( + env, VirtualGaugeTransform(dropdims(gauge.matrices; dims = 4), dropdims(gauge.inverses; dims = 4)) + ) +end + +"""Apply a compensating virtual gauge to a single-layer or double-layer edge tensor.""" +function _gauge_transform_virtual_edge(edge::CTMRGEdgeTensor, X::MPSBondTensor, incoming::Bool) + if numout(edge) == 2 + if incoming + return @tensor t[a p; b] := edge[a q; b] * X[p; q] + else + return @tensor t[a p; b] := edge[a q; b] * X[q; p] + end + elseif numout(edge) == 3 + if incoming + return @tensor t[a p q; b] := edge[a p′ q′; b] * X[p; p′] * conj(X[q; q′]) + else + return @tensor t[a p q; b] := edge[a p′ q′; b] * X[p′; p] * conj(X[q′; q]) + end + end + throw(ArgumentError("Virtual environment transformations require one or two network-facing legs per edge")) +end + +""" + gauge_transform(env::CTMRGEnv, gauge::CTMRGEnvGaugeTransform) + +Transform the boundary legs of the corners and edges of `env`, leaving its network-facing legs unchanged. +Each boundary bond receives a matrix and its inverse, preserving contracted observables for general invertible gauges. +This also supports environments with more than two network-facing legs per edge. +""" +function gauge_transform(env::CTMRGEnv, gauge::CTMRGEnvGaugeTransform) + _check_gauge_size(gauge, size(env)[2:3]) + X, Xinv = gauge.matrices, gauge.inverses + nr, nc = size(env)[2:3] + corners = map(eachcoordinate(env, 1:4)) do (d, r, c) + outdir, indir = _prev(d, 4), d + ri, ci = _environment_gauge_neighbor(indir, r, c, nr, nc) + return fix_gauge_corner(env.corners[d, r, c], X[outdir, r, c], Xinv[indir, ri, ci]') + end + edges = map(eachcoordinate(env, 1:4)) do (d, r, c) + ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + return fix_gauge_edge(env.edges[d, r, c], X[d, r, c], Xinv[d, ri, ci]') + end + return CTMRGEnv(corners, edges) +end + +"""Find the next boundary bond along an edge, respecting the periodic unit cell.""" +function _environment_gauge_neighbor(d::Int, r::Int, c::Int, nr::Int, nc::Int) + d == NORTH && return r, _next(c, nc) + d == EAST && return _next(r, nr), c + d == SOUTH && return r, _prev(c, nc) + return _prev(r, nr), c +end + +""" +$(SIGNATURES) + +Multiply corner tensor with incoming and outgoing gauge signs. + +``` + corner -- σ_in --in + | + σ_out + | + out +``` +""" +function fix_gauge_corner( + corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor + ) + return @autoopt @tensor corner_fix[χ_out; χ_in] := + σ_out[χ_out; χ1] * corner[χ1; χ2] * conj(σ_in[χ_in; χ2]) +end + +""" +$(SIGNATURES) + +Multiply edge tensor with incoming and outgoing gauge signs. + +``` + out-- σ_out -- edge -- σ_in --in + | +``` +""" +@generated function fix_gauge_edge( + edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor + ) where {T, S, N} + edge_fix_e = tensorexpr( + :edge_fix, + (envlabel(:out), ntuple(i -> virtuallabel(i), N - 1)...), + (envlabel(:in),), + ) + edge_e = tensorexpr( + :edge, (envlabel(1), ntuple(i -> virtuallabel(i), N - 1)...), (envlabel(2),) + ) + σ_out_e = tensorexpr(:σ_out, (envlabel(:out),), (envlabel(1),)) + σ_in_e = tensorexpr(:σ_in, (envlabel(:in),), (envlabel(2),)) + return macroexpand( + @__MODULE__, + :(return @autoopt @tensor $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)), + ) +end diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index 111699453..cea465d62 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -401,11 +401,11 @@ _default_solver_alg(::Type{<:ImplicitGradient}) = Defaults.gradient_implicit_sol _select_solver_alg_symbol(::Type{<:ImplicitGradient}, solver_alg) = IMPLICIT_SOLVER_SYMBOLS[solver_alg] -function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vEighProjector}, signs, info) - return info.V * signs[1]' +function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vEighProjector}, signs::CTMRGEnvGaugeTransform, info) + return info.V * signs.inverses[1] end -function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vQRProjector}, signs, info) - return info.Q * signs[1]' +function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vQRProjector}, signs::CTMRGEnvGaugeTransform, info) + return info.Q * signs.inverses[1] end function _check_algorithm_combination(::SequentialCTMRG, ::ImplicitGradient) @@ -493,7 +493,7 @@ function _rrule( # instantiate the differentiable variables corresponding to the intermediate projector of the contraction algorithm u = zeros(storagetype(U), space(UL, numind(UL))' ← space(U, numind(U))') - # prepare pullback of C4v CTMRG environment constructor (artefact of reusing asymmetric environment type for C4v symmetric contraction) + # prepare pullback of the C4v CTMRG environment constructor _, c4v_env_vjp = rrule_via_ad(config, C4vCTMRGEnv, C, E) # Two tapes: the linear solve only uses the environment cotangents, so reading the diff --git a/src/networks/gaugetrans.jl b/src/networks/gaugetrans.jl new file mode 100644 index 000000000..aa133c0f2 --- /dev/null +++ b/src/networks/gaugetrans.jl @@ -0,0 +1,188 @@ +""" + VirtualGaugeTransform(matrices[, inverses]) + VirtualGaugeTransform(pairs) + +Gauge transformations on the virtual bonds of an [`InfinitePEPS`](@ref) or [`InfinitePEPO`](@ref). +`matrices` contains one invertible `MPSBondTensor` per north/east bond, with shape `(2, rows, cols)` for a PEPS or `(2, rows, cols, layers)` for a PEPO. +`inverses` contains the corresponding inverse maps and is computed when omitted; alternatively, supply an array of `(X, X⁻¹)` pairs. +Explicit inverse maps must have compatible spaces and are assumed to be the inverses of `matrices`. + +At `(r, c)`, the north and east legs receive `matrices[NORTH, r, c]` and `matrices[EAST, r, c]`. +The south and west legs receive the inverses belonging to the north bond at `(r + 1, c)` and the east bond at `(r, c - 1)`, respectively, with periodic indexing. +For a PEPO, the same convention applies independently to each layer, leaving both physical legs unchanged. +Use [`CTMRGEnvGaugeTransform`](@ref) for transformations of CTMRG boundary legs. +""" +struct VirtualGaugeTransform{N, A <: AbstractArray{<:MPSBondTensor, N}, B <: AbstractArray{<:MPSBondTensor, N}} + matrices::A + inverses::B + function VirtualGaugeTransform( + matrices::AbstractArray{<:MPSBondTensor, N}, + inverses::AbstractArray{<:MPSBondTensor, N}, + ) where {N} + N in (3, 4) || throw(ArgumentError("Virtual gauges require direction, row, column, and optionally layer indices")) + _check_gauge_matrices(matrices, inverses, 2) + return new{N, typeof(matrices), typeof(inverses)}(matrices, inverses) + end +end + +""" + CTMRGEnvGaugeTransform(matrices[, inverses]) + CTMRGEnvGaugeTransform(pairs) + +Gauge transformations on the boundary legs of a [`CTMRGEnv`](@ref), with shape `(4, rows, cols)` ordered north, east, south, west. +Each matrix is an `MPSBondTensor` acting on the first (outgoing) boundary leg of the edge at the same coordinate. +The inverse at the neighboring bond acts on its last (incoming) boundary leg, and the corner tensors are transformed consistently. +The neighbor is east, south, west, or north for a north, east, south, or west edge, respectively. +The network-facing virtual legs are unchanged. + +The matrices need not be unitary. +Inverse maps are computed when omitted, or may be supplied explicitly with compatible spaces; alternatively, supply `(X, X⁻¹)` pairs. +Explicit inverse maps are assumed to be the inverses of `matrices`. +This follows the bond indexing used by CTMRG environment gauge fixing. +""" +struct CTMRGEnvGaugeTransform{A <: AbstractArray{<:MPSBondTensor, 3}, B <: AbstractArray{<:MPSBondTensor, 3}} + matrices::A + inverses::B + function CTMRGEnvGaugeTransform( + matrices::AbstractArray{<:MPSBondTensor, 3}, + inverses::AbstractArray{<:MPSBondTensor, 3}, + ) + _check_gauge_matrices(matrices, inverses, 4) + return new{typeof(matrices), typeof(inverses)}(matrices, inverses) + end +end + +"""Validate the layout and tensor-map spaces of a collection of gauge pairs.""" +function _check_gauge_matrices( + matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor}, ndirs::Int + ) + Base.require_one_based_indexing(matrices, inverses) + axes(matrices) == axes(inverses) || throw(DimensionMismatch("Gauge matrices and inverses must have matching axes")) + size(matrices, 1) == ndirs || throw(DimensionMismatch("Expected $ndirs gauge directions")) + all(>(0), size(matrices)) || throw(ArgumentError("Gauge arrays must be nonempty")) + for (X, Xinv) in zip(matrices, inverses) + domain(X) == codomain(Xinv) && codomain(X) == domain(Xinv) || + throw(SpaceMismatch("Gauge matrices and inverses must have opposite domain and codomain spaces")) + end + return nothing +end + +VirtualGaugeTransform(matrices::AbstractArray{<:MPSBondTensor}) = + VirtualGaugeTransform(matrices, inv.(matrices)) +CTMRGEnvGaugeTransform(matrices::AbstractArray{<:MPSBondTensor, 3}) = + CTMRGEnvGaugeTransform(matrices, inv.(matrices)) +VirtualGaugeTransform(pairs::AbstractArray{<:Tuple{MPSBondTensor, MPSBondTensor}}) = + VirtualGaugeTransform(first.(pairs), last.(pairs)) +CTMRGEnvGaugeTransform(pairs::AbstractArray{<:Tuple{MPSBondTensor, MPSBondTensor}, 3}) = + CTMRGEnvGaugeTransform(first.(pairs), last.(pairs)) + +Base.inv(gauge::VirtualGaugeTransform) = VirtualGaugeTransform(gauge.inverses, gauge.matrices) +Base.inv(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform(gauge.inverses, gauge.matrices) + +"""Rotate virtual bonds counterclockwise, reversing the old north bonds into new east bonds.""" +function _rotl90_virtual_gauge( + matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} + ) + nr, nc = size(matrices)[2:3] + return map(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I + d, r, c = Tuple(I)[1:3] + layer = Tuple(I)[4:end] + # Match the leg permutation in state rotations, including fermionic signs. + return d == NORTH ? matrices[EAST, c, nc + 1 - r, layer...] : + permute(inverses[NORTH, _next(c, nr), nc + 1 - r, layer...], ((2,), (1,))) + end +end + +"""Rotate virtual bonds clockwise, reversing the old east bonds into new north bonds.""" +function _rotr90_virtual_gauge( + matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} + ) + nr, nc = size(matrices)[2:3] + return map(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I + d, r, c = Tuple(I)[1:3] + layer = Tuple(I)[4:end] + return d == NORTH ? + permute(inverses[EAST, nr + 1 - c, _prev(r, nc), layer...], ((2,), (1,))) : + matrices[NORTH, nr + 1 - c, r, layer...] + end +end + +"""Rotate virtual bonds by a half turn, reversing both bond directions using inverse maps.""" +function _rot180_virtual_gauge(inverses::AbstractArray{<:MPSBondTensor}) + nr, nc = size(inverses)[2:3] + return map(CartesianIndices(inverses)) do I + d, r, c = Tuple(I)[1:3] + layer = Tuple(I)[4:end] + r, c = nr + 1 - r, nc + 1 - c + Xinv = d == NORTH ? inverses[NORTH, _next(r, nr), c, layer...] : + inverses[EAST, r, _prev(c, nc), layer...] + return permute(Xinv, ((2,), (1,))) + end +end + +"""Rotate boundary gauge directions and unit-cell coordinates counterclockwise.""" +function _rotl90_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) + return map(CartesianIndices((4, size(matrices, 3), size(matrices, 2)))) do I + d, r, c = Tuple(I) + return matrices[_next(d, 4), c, size(matrices, 3) + 1 - r] + end +end + +"""Rotate boundary gauge directions and unit-cell coordinates clockwise.""" +function _rotr90_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) + return map(CartesianIndices((4, size(matrices, 3), size(matrices, 2)))) do I + d, r, c = Tuple(I) + return matrices[_prev(d, 4), size(matrices, 2) + 1 - c, r] + end +end + +"""Rotate boundary gauge directions and unit-cell coordinates by a half turn.""" +function _rot180_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) + return map(CartesianIndices(matrices)) do I + d, r, c = Tuple(I) + return matrices[mod1(d + 2, 4), size(matrices, 2) + 1 - r, size(matrices, 3) + 1 - c] + end +end + +""" + rotl90(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) + rotr90(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) + rot180(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) + +Rotate a gauge consistently with its state or CTMRG environment, leaving PEPO layers in place. +For each rotation `R`, `R(gauge_transform(state, gauge)) ≈ gauge_transform(R(state), R(gauge))`. +""" +Base.rotl90(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( + _rotl90_virtual_gauge(gauge.matrices, gauge.inverses), + _rotl90_virtual_gauge(gauge.inverses, gauge.matrices) +) +Base.rotl90(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( + _rotl90_boundary_gauge(gauge.matrices), _rotl90_boundary_gauge(gauge.inverses) +) +Base.rotr90(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( + _rotr90_virtual_gauge(gauge.matrices, gauge.inverses), + _rotr90_virtual_gauge(gauge.inverses, gauge.matrices) +) +Base.rotr90(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( + _rotr90_boundary_gauge(gauge.matrices), _rotr90_boundary_gauge(gauge.inverses) +) +Base.rot180(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( + _rot180_virtual_gauge(gauge.inverses), _rot180_virtual_gauge(gauge.matrices) +) +Base.rot180(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( + _rot180_boundary_gauge(gauge.matrices), _rot180_boundary_gauge(gauge.inverses) +) + +"""Check that a gauge array has the expected unit cell and layer count.""" +function _check_gauge_size(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}, dims::Tuple) + size(gauge.matrices)[2:end] == dims || throw(DimensionMismatch("Gauge unit cell does not match the object being transformed")) + return nothing +end + +"""Retrieve the north, east, south, and west maps acting on a site's virtual legs.""" +function _virtual_gauge_factors(gauge::VirtualGaugeTransform, r::Int, c::Int, layer::Int...) + X, Xinv = gauge.matrices, gauge.inverses + return X[NORTH, r, c, layer...], X[EAST, r, c, layer...], + Xinv[NORTH, _next(r, size(X, 2)), c, layer...], + Xinv[EAST, r, _prev(c, size(X, 3)), layer...] +end diff --git a/src/operators/infinitepepo.jl b/src/operators/infinitepepo.jl index 1f68a9161..7992f0a89 100644 --- a/src/operators/infinitepepo.jl +++ b/src/operators/infinitepepo.jl @@ -287,3 +287,27 @@ function _stack_tuples(A::Matrix{NTuple{N, T}}) where {N, T} end return out end + +## Gauge transform + +""" + gauge_transform(ρ::InfinitePEPO, gauge::VirtualGaugeTransform) + +Transform the virtual bonds of each layer of `ρ`, leaving its physical legs unchanged. +The gauge must have shape `(2, rows, cols, layers)`. +A gauge of shape `(2, rows, cols)` is also accepted for a single-layer PEPO. +""" +function gauge_transform(ρ::InfinitePEPO, gauge::VirtualGaugeTransform{4}) + _check_gauge_size(gauge, size(ρ)) + A = map(eachcoordinate(ρ)) do (r, c, h) + Xn, Xe, Xs, Xw = _virtual_gauge_factors(gauge, r, c, h) + return @tensor t[p q; n e s w] := ρ[r, c, h][p q; n′ e′ s′ w′] * + Xn[n′; n] * Xe[e′; e] * Xs[s; s′] * Xw[w; w′] + end + return InfinitePEPO(A) +end +function gauge_transform(ρ::InfinitePEPO, gauge::VirtualGaugeTransform{3}) + size(ρ, 3) == 1 || throw(DimensionMismatch("A multilayer PEPO requires a gauge for every layer")) + dims = (size(gauge.matrices)..., 1) + return gauge_transform(ρ, VirtualGaugeTransform(reshape(gauge.matrices, dims), reshape(gauge.inverses, dims))) +end diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index 7f6a38eac..591d86ba7 100644 --- a/src/states/infinitepeps.jl +++ b/src/states/infinitepeps.jl @@ -272,3 +272,21 @@ function ChainRulesCore.rrule( end return network, InfiniteSquareNetwork_pullback end + +## Gauge transform + +""" + gauge_transform(ψ::InfinitePEPS, gauge::VirtualGaugeTransform) + +Transform the virtual bonds of `ψ` using a gauge with shape `(2, rows, cols)`. +See [`VirtualGaugeTransform`](@ref) for the bond convention. +""" +function gauge_transform(ψ::InfinitePEPS, gauge::VirtualGaugeTransform{3}) + _check_gauge_size(gauge, size(ψ)) + A2 = map(eachcoordinate(ψ)) do (r, c) + Xn, Xe, Xs, Xw = _virtual_gauge_factors(gauge, r, c) + return @tensor t[p; n e s w] := ψ[r, c][p; n′ e′ s′ w′] * + Xn[n′; n] * Xe[e′; e] * Xs[s; s′] * Xw[w; w′] + end + return InfinitePEPS(A2) +end diff --git a/test/bp/gaugefix.jl b/test/bp/gaugefix.jl index 251792d15..b3c92788d 100644 --- a/test/bp/gaugefix.jl +++ b/test/bp/gaugefix.jl @@ -7,6 +7,7 @@ is_buildkite = get(ENV, "BUILDKITE", "false") == "true" if !is_buildkite TestSuite.bp_gaugefix_bp_vs_su(Vector) + TestSuite.bp_gaugefix_pepo(Vector) end if CUDA.functional() diff --git a/test/ctmrg/gaugetrans.jl b/test/ctmrg/gaugetrans.jl new file mode 100644 index 000000000..41cc0db6b --- /dev/null +++ b/test/ctmrg/gaugetrans.jl @@ -0,0 +1,13 @@ +using Test +using PEPSKit + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +if !is_buildkite + TestSuite.ctmrg_gaugetrans(Vector) + TestSuite.ctmrg_gaugetrans_pepo(Vector) + TestSuite.ctmrg_gaugetrans_stacked_env(Vector) +end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 3e576386d..6a80a408a 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -56,7 +56,7 @@ using .BPExpVals module BPGaugeFix include("bp/gaugefix.jl") - export bp_gaugefix_bp_vs_su + export bp_gaugefix_bp_vs_su, bp_gaugefix_pepo end using .BPGaugeFix @@ -140,6 +140,12 @@ module CTMRGGaugeFix end using .CTMRGGaugeFix +module CTMRGGaugeTrans + include("ctmrg/gaugetrans.jl") + export ctmrg_gaugetrans, ctmrg_gaugetrans_pepo, ctmrg_gaugetrans_stacked_env +end +using .CTMRGGaugeTrans + module CTMRGPartitionFunction include("ctmrg/partition_function.jl") export ctmrg_partition_function, ctmrg_partition_function_spaces diff --git a/test/testsuite/bp/gaugefix.jl b/test/testsuite/bp/gaugefix.jl index c0a4df34f..f5f7ab125 100644 --- a/test/testsuite/bp/gaugefix.jl +++ b/test/testsuite/bp/gaugefix.jl @@ -78,14 +78,15 @@ function bp_gaugefix_bp_vs_su(AT; posdef_msgs = [true, false]) @test compare_weights(wts1, wts2) < 1.0e-9 bpg_alg = BPGauge() - peps2, XXinv = @testinferred gauge_fix(peps1, bpg_alg, env) + peps2, gauge = @testinferred gauge_fix(peps1, bpg_alg, env) + @test gauge_transform(peps1, gauge) ≈ peps2 if bipartite @test _is_bipartite(peps2) end for (a1, a2) in zip(peps1.A, peps2.A) @test space(a1) == space(a2) end - for (X, Xinv) in XXinv + for (X, Xinv) in zip(gauge.matrices, gauge.inverses) # X, Xinv should contract to identity @tensor tmp[-1; -2] := X[-1; 1] * Xinv[1; -2] @test tmp ≈ twistdual(adapt(AT, TensorKit.id(space(X, 1))), 1) @@ -94,3 +95,20 @@ function bp_gaugefix_bp_vs_su(AT; posdef_msgs = [true, false]) end end end + +"""Test that the BP gauge reproduces a PEPO and preserves its CTMRG density matrices.""" +function bp_gaugefix_pepo(AT) + return @testset "BP gauge of PEPO ($AT) ($S)" for S in (U1Irrep, FermionParity) + P = Vect[S](0 => 1, 1 => 1) + V = S == U1Irrep ? Vect[S](-1 => 1, 0 => 2, 1 => 1) : Vect[S](0 => 2, 1 => 1) + N, E = random_dual!(fill(V, 2, 3)), random_dual!(fill(V, 2, 3)) + ρ = adapt(AT, InfinitePEPO(fill(P, 2, 3), N, E)) + bp_env = adapt(AT, BPEnv(randn, ComplexF64, ρ)) + ρg, gauge = gauge_fix(ρ, BPGauge(), bp_env) + @test gauge_transform(ρ, gauge) ≈ ρg + env = adapt(AT, CTMRGEnv(InfinitePEPS(ρ), V)) + envg = gauge_transform(env, gauge) + inds = [CartesianIndex(1, 1)] + @test reduced_densitymatrix(inds, ρ, ρ, env) ≈ reduced_densitymatrix(inds, ρg, ρg, envg) + end +end diff --git a/test/testsuite/ctmrg/gaugetrans.jl b/test/testsuite/ctmrg/gaugetrans.jl new file mode 100644 index 000000000..eaf213950 --- /dev/null +++ b/test/testsuite/ctmrg/gaugetrans.jl @@ -0,0 +1,103 @@ +using Test +using Random +using TensorKit +using PEPSKit +using Adapt +using PEPSKit: random_dual!, eachcoordinate, NORTH +using PEPSKit: _contract_site + +ds = Dict( + U1Irrep => U1Space(i => d for (i, d) in zip(-1:1, (1, 1, 2))), + FermionParity => Vect[FermionParity](0 => 2, 1 => 1) +) +Ds = Dict( + U1Irrep => U1Space(i => D for (i, D) in zip(-1:1, (1, 3, 2))), + FermionParity => Vect[FermionParity](0 => 3, 1 => 2) +) +χs = Dict( + U1Irrep => U1Space(i => D for (i, D) in zip(-1:1, (2, 4, 2))), + FermionParity => Vect[FermionParity](0 => 4, 1 => 4) +) + +"""Generate a non-unitary gauge near the identity to avoid ill-conditioned inverses.""" +function _random_gauge_matrix(AT, V::ElementarySpace) + return adapt(AT, id(ComplexF64, V) + 0.1 * randn(ComplexF64, V ← V)) +end + +"""Test virtual and boundary gauge invariance.""" +function ctmrg_gaugetrans(AT) + return @testset "CTMRGEnv of InfinitePEPS ($S) ($AT)" for S in keys(ds) + d, D, χ, uc = ds[S], Ds[S], χs[S], (2, 3) + N, E = random_dual!(fill(D, uc)), random_dual!(fill(D, uc)) + ψ = adapt(AT, InfinitePEPS(fill(d, uc), N, E)) + env = adapt(AT, CTMRGEnv(ψ, χ)) + gauge = VirtualGaugeTransform( + map(eachcoordinate(ψ, 1:2)) do (dir, r, c) + return _random_gauge_matrix(AT, dir == NORTH ? N[r, c] : E[r, c]) + end + ) + boundary_gauge = CTMRGEnvGaugeTransform( + map(env.edges) do edge + return _random_gauge_matrix(AT, space(edge, 1)) + end + ) + ψg, envg = gauge_transform(ψ, gauge), gauge_transform(env, gauge) + envb = gauge_transform(env, boundary_gauge) + for R in (rotl90, rotr90, rot180), (state, g, transformed) in + ((ψ, gauge, ψg), (env, gauge, envg), (env, boundary_gauge, envb)) + @test R(transformed) ≈ gauge_transform(R(state), R(g)) + end + for (r, c) in eachcoordinate(ψ) + inds = [CartesianIndex(r, c)] + ρ = reduced_densitymatrix(inds, ψ, env) + @test reduced_densitymatrix(inds, ψg, envg) ≈ ρ + @test reduced_densitymatrix(inds, ψ, envb) ≈ ρ + end + end +end + +"""Test independent PEPO layer gauges through single-layer density matrices.""" +function ctmrg_gaugetrans_pepo(AT) + return @testset "PEPO virtual gauges ($S) ($AT)" for S in keys(ds) + d, D, χ, uc = ds[S], Ds[S], χs[S], (2, 3, 2) + P, N, E = fill(d, uc), fill(D, uc), fill(D, uc) + foreach(random_dual!, eachslice(N; dims = 3)) + foreach(random_dual!, eachslice(E; dims = 3)) + ρ = adapt(AT, InfinitePEPO(P, N, E)) + X = map(Iterators.product(1:2, axes(P)...)) do (dir, r, c, h) + return _random_gauge_matrix(AT, dir == NORTH ? N[r, c, h] : E[r, c, h]) + end + gauge = VirtualGaugeTransform(X) + ρg = gauge_transform(ρ, gauge) + for R in (rotl90, rotr90, rot180) + @test R(ρg) ≈ gauge_transform(R(ρ), R(gauge)) + end + for h in axes(P, 3) + layer, layerg = InfinitePEPO(ρ.A[:, :, h:h]), InfinitePEPO(ρg.A[:, :, h:h]) + env = adapt(AT, CTMRGEnv(InfinitePartitionFunction(layer), χ)) + envg = gauge_transform(env, VirtualGaugeTransform(X[:, :, :, h:h])) + for (r, c) in Iterators.product(axes(P, 1), axes(P, 2)) + inds = [CartesianIndex(r, c)] + @test reduced_densitymatrix(inds, layerg, envg) ≈ reduced_densitymatrix(inds, layer, env) + end + end + end +end + +"""Test boundary gauge invariance with more than two network-facing edge legs.""" +function ctmrg_gaugetrans_stacked_env(AT) + return @testset "Stacked-network boundary gauges ($AT)" begin + V = ℂ^2 + ψ = adapt(AT, InfinitePEPS(V, V)) + ρ = adapt(AT, InfinitePEPO(V, V; unitcell = (1, 1, 2))) + network = InfiniteSquareNetwork(ψ, ρ, ψ) + env = adapt(AT, CTMRGEnv(network, V)) + gauge = CTMRGEnvGaugeTransform( + map(env.edges) do edge + return _random_gauge_matrix(AT, space(edge, 1)) + end + ) + @test _contract_site((1, 1), network, gauge_transform(env, gauge)) ≈ + _contract_site((1, 1), network, env) + end +end diff --git a/test/testsuite/ctmrg/unitcell.jl b/test/testsuite/ctmrg/unitcell.jl index fb7b24e70..5f2600b48 100644 --- a/test/testsuite/ctmrg/unitcell.jl +++ b/test/testsuite/ctmrg/unitcell.jl @@ -43,7 +43,7 @@ function test_unitcell( # test if gauge fixing routines run through signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env″, env′, ScramblingEnvGauge()) - @test signs isa Array + @test signs isa CTMRGEnvGaugeTransform @test corner_phases isa Array @test edge_phases isa Array return nothing From 39d9a3a5f4dcc9a17226127abcf8fa65ed62c83f Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 13:55:56 +0800 Subject: [PATCH 2/8] Deduplication --- src/algorithms/bp/gaugefix.jl | 27 +++----- src/algorithms/contractions/bp_messages.jl | 9 --- src/algorithms/ctmrg/gaugefix.jl | 4 +- src/algorithms/ctmrg/gaugetrans.jl | 75 ++++------------------ test/testsuite/ctmrg/gaugetrans.jl | 13 +++- 5 files changed, 34 insertions(+), 94 deletions(-) diff --git a/src/algorithms/bp/gaugefix.jl b/src/algorithms/bp/gaugefix.jl index c4c9686fa..1527d8059 100644 --- a/src/algorithms/bp/gaugefix.jl +++ b/src/algorithms/bp/gaugefix.jl @@ -37,19 +37,18 @@ using fixed point environment `env` of belief propagation. """ function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv) bipartite = _is_bipartite(psi) && _is_bipartite(env) - psi′ = copy(psi) XXinv = map(eachcoordinate(psi, 1:2)) do I - _, X, Xinv = _bp_gauge_fix!(CartesianIndex(I), psi′, env, alg) - return X, Xinv + return _bp_gauge(CartesianIndex(I), env, alg) end + gauge = VirtualGaugeTransform(XXinv) + psi′ = gauge_transform(psi, gauge) if bipartite - # copy 1st column to 2nd column to eliminate differences - # caused by order of applying gauge transformations + # Enforce exact equality between equivalent sites. for r in 1:2 psi′[r + 1, 2] = copy(psi′[r, 1]) end end - return psi′, VirtualGaugeTransform(XXinv) + return psi′, gauge end function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv) # convert to iPEPS, and store physical leg fusers @@ -85,7 +84,7 @@ function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv) end """ - _bp_gauge_fix!(I, psi::InfinitePEPS, env::BPEnv) -> psi, X, X⁻¹ + _bp_gauge(I, env::BPEnv, alg::BPGauge) -> X, X⁻¹ For the bond at direction `I[1]` (which can be `NORTH` or `EAST`) from site `I[2], I[3]`, we identify the following gauge matrices, @@ -101,10 +100,7 @@ along the canonical direction of the PEPS arrows (`SOUTH ← NORTH` or `WEST ← These matrices update the gauge of `psi` using the [`VirtualGaugeTransform`](@ref) convention. `X` is attached to the `NORTH`/`EAST` leg at the indexed site, and `X⁻¹` to the neighboring `SOUTH`/`WEST` leg. """ -function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv, alg::BPGauge) - dir, row, col = Tuple(I) - @assert dir == NORTH || dir == EAST - +function _bp_gauge(I::CartesianIndex{3}, env::BPEnv, alg::BPGauge) sqrtM12, isqrtM12, sqrtM21, isqrtM21 = _sqrt_bp_messages(I, env) U, Λ, Vᴴ = svd_compact!(sqrtM12 * sqrtM21; alg = alg.svd_alg) sqrtΛ = sdiag_pow(Λ, 1 / 2) @@ -113,14 +109,7 @@ function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv, alg if isdual(space(sqrtM12, 1)) X, invX = twist(flip(X, 2), 1), flip(invX, 1) end - if dir == NORTH - psi[row, col] = absorb_north_message(psi[row, col], X) - psi[row - 1, col] = absorb_south_message(psi[row - 1, col], invX) - elseif dir == EAST - psi[row, col] = absorb_east_message(psi[row, col], X) - psi[row, col + 1] = absorb_west_message(psi[row, col + 1], invX) - end - return psi, X, invX + return X, invX end """ diff --git a/src/algorithms/contractions/bp_messages.jl b/src/algorithms/contractions/bp_messages.jl index 55fdc51b9..04f5e43e6 100644 --- a/src/algorithms/contractions/bp_messages.jl +++ b/src/algorithms/contractions/bp_messages.jl @@ -38,12 +38,3 @@ function contract_west_message( M_south[DSb; DSt] * M_west[DWb; DWt] * M_north[DNt; DNb] end end - -absorb_north_message(A::PEPSTensor, M::PEPSMessage) = - @tensor A′[d; N' E S W] := A[d; N E S W] * M[N; N'] -absorb_east_message(A::PEPSTensor, M::PEPSMessage) = - @tensor A′[d; N E' S W] := A[d; N E S W] * M[E; E'] -absorb_south_message(A::PEPSTensor, M::PEPSMessage) = - @tensor A′[d; N E S' W] := A[d; N E S W] * M[S'; S] -absorb_west_message(A::PEPSTensor, M::PEPSMessage) = - @tensor A′[d; N E S W'] := A[d; N E S W] * M[W'; W] diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index ce37a578f..c573e6df0 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -170,10 +170,10 @@ fix_relative_phases(env::CTMRGEnv, signs::CTMRGEnvGaugeTransform) = gauge_transf function fix_relative_phases(env::C4vCTMRGEnv, signs::CTMRGEnvGaugeTransform) X, Xinv = signs.matrices, signs.inverses corners = map(env.corners) do C - return fix_gauge_corner(C, X[WEST, 1, 1], Xinv[NORTH, 1, 1]') + return _gauge_transform_corner(C, X[WEST, 1, 1], Xinv[NORTH, 1, 1]) end edges = map(eachcoordinate(env, 1:4)) do (dir, r, c) - return fix_gauge_edge(env.edges[dir, r, c], X[dir, r, c], Xinv[dir, r, c]') + return absorb_left_right(env.edges[dir, r, c], X[dir, r, c], Xinv[dir, r, c]) end return C4vCTMRGEnv(corners, edges) end diff --git a/src/algorithms/ctmrg/gaugetrans.jl b/src/algorithms/ctmrg/gaugetrans.jl index e818f42e7..5adeae637 100644 --- a/src/algorithms/ctmrg/gaugetrans.jl +++ b/src/algorithms/ctmrg/gaugetrans.jl @@ -9,17 +9,9 @@ The corner tensors and boundary legs are unchanged. """ function gauge_transform(env::CTMRGEnv, gauge::VirtualGaugeTransform{3}) _check_gauge_size(gauge, size(env)[2:3]) - X, Xinv = gauge.matrices, gauge.inverses edges = map(eachcoordinate(env, 1:4)) do (d, r, c) - factor = if d == NORTH - Xinv[NORTH, _next(r, size(X, 2)), c] - elseif d == EAST - Xinv[EAST, r, _prev(c, size(X, 3))] - elseif d == SOUTH - X[NORTH, r, c] - else # d == WEST - X[EAST, r, c] - end + # At the edge's coordinate, use the factor for the opposite virtual leg. + factor = _virtual_gauge_factors(gauge, r, c)[mod1(d + 2, 4)] return _gauge_transform_virtual_edge(env.edges[d, r, c], factor, d == NORTH || d == EAST) end return CTMRGEnv(env.corners, edges) @@ -61,17 +53,23 @@ function gauge_transform(env::CTMRGEnv, gauge::CTMRGEnvGaugeTransform) X, Xinv = gauge.matrices, gauge.inverses nr, nc = size(env)[2:3] corners = map(eachcoordinate(env, 1:4)) do (d, r, c) - outdir, indir = _prev(d, 4), d - ri, ci = _environment_gauge_neighbor(indir, r, c, nr, nc) - return fix_gauge_corner(env.corners[d, r, c], X[outdir, r, c], Xinv[indir, ri, ci]') + ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + return _gauge_transform_corner(env.corners[d, r, c], X[_prev(d, 4), r, c], Xinv[d, ri, ci]) end edges = map(eachcoordinate(env, 1:4)) do (d, r, c) ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) - return fix_gauge_edge(env.edges[d, r, c], X[d, r, c], Xinv[d, ri, ci]') + return absorb_left_right(env.edges[d, r, c], X[d, r, c], Xinv[d, ri, ci]) end return CTMRGEnv(corners, edges) end +"""Apply the outgoing gauge and incoming inverse with the corner's leg orientations.""" +function _gauge_transform_corner(C::CTMRGCornerTensor, Xout::MPSBondTensor, Xin::MPSBondTensor) + # Keep the explicit contraction for the fermionic signs of mixed-dual corners. + return @autoopt @tensor Cg[χ_out; χ_in] := + Xout[χ_out; χ1] * C[χ1; χ2] * Xin[χ2; χ_in] +end + """Find the next boundary bond along an edge, respecting the periodic unit cell.""" function _environment_gauge_neighbor(d::Int, r::Int, c::Int, nr::Int, nc::Int) d == NORTH && return r, _next(c, nc) @@ -79,52 +77,3 @@ function _environment_gauge_neighbor(d::Int, r::Int, c::Int, nr::Int, nc::Int) d == SOUTH && return r, _prev(c, nc) return _prev(r, nr), c end - -""" -$(SIGNATURES) - -Multiply corner tensor with incoming and outgoing gauge signs. - -``` - corner -- σ_in --in - | - σ_out - | - out -``` -""" -function fix_gauge_corner( - corner::CTMRGCornerTensor, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor - ) - return @autoopt @tensor corner_fix[χ_out; χ_in] := - σ_out[χ_out; χ1] * corner[χ1; χ2] * conj(σ_in[χ_in; χ2]) -end - -""" -$(SIGNATURES) - -Multiply edge tensor with incoming and outgoing gauge signs. - -``` - out-- σ_out -- edge -- σ_in --in - | -``` -""" -@generated function fix_gauge_edge( - edge::CTMRGEdgeTensor{T, S, N}, σ_out::CTMRGCornerTensor, σ_in::CTMRGCornerTensor - ) where {T, S, N} - edge_fix_e = tensorexpr( - :edge_fix, - (envlabel(:out), ntuple(i -> virtuallabel(i), N - 1)...), - (envlabel(:in),), - ) - edge_e = tensorexpr( - :edge, (envlabel(1), ntuple(i -> virtuallabel(i), N - 1)...), (envlabel(2),) - ) - σ_out_e = tensorexpr(:σ_out, (envlabel(:out),), (envlabel(1),)) - σ_in_e = tensorexpr(:σ_in, (envlabel(:in),), (envlabel(2),)) - return macroexpand( - @__MODULE__, - :(return @autoopt @tensor $edge_fix_e := $σ_out_e * $edge_e * conj($σ_in_e)), - ) -end diff --git a/test/testsuite/ctmrg/gaugetrans.jl b/test/testsuite/ctmrg/gaugetrans.jl index eaf213950..da8b24661 100644 --- a/test/testsuite/ctmrg/gaugetrans.jl +++ b/test/testsuite/ctmrg/gaugetrans.jl @@ -4,7 +4,7 @@ using TensorKit using PEPSKit using Adapt using PEPSKit: random_dual!, eachcoordinate, NORTH -using PEPSKit: _contract_site +using PEPSKit: @autoopt, _contract_site, _gauge_transform_corner ds = Dict( U1Irrep => U1Space(i => d for (i, d) in zip(-1:1, (1, 1, 2))), @@ -53,6 +53,17 @@ function ctmrg_gaugetrans(AT) @test reduced_densitymatrix(inds, ψg, envg) ≈ ρ @test reduced_densitymatrix(inds, ψ, envb) ≈ ρ end + if S == FermionParity + for (left, right) in ((D, D'), (D', D)) + corner = adapt(AT, randn(ComplexF64, left ← right)) + Xout = _random_gauge_matrix(AT, left) + Xin = _random_gauge_matrix(AT, right) + σin = Xin' + @autoopt @tensor expected[χ_out; χ_in] := + Xout[χ_out; χ1] * corner[χ1; χ2] * conj(σin[χ_in; χ2]) + @test _gauge_transform_corner(corner, Xout, Xin) ≈ expected + end + end end end From 0c8600f4bc65f845719d5421be50772d95beb981 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 14:43:59 +0800 Subject: [PATCH 3/8] Type stability improvements --- src/algorithms/bp/gaugefix.jl | 4 ++-- src/algorithms/ctmrg/gaugefix.jl | 5 +++-- src/networks/gaugetrans.jl | 15 ++++++++++++--- test/testsuite/bp/gaugefix.jl | 2 +- test/testsuite/ctmrg/gaugefix.jl | 4 ++-- 5 files changed, 20 insertions(+), 10 deletions(-) diff --git a/src/algorithms/bp/gaugefix.jl b/src/algorithms/bp/gaugefix.jl index 1527d8059..15356960e 100644 --- a/src/algorithms/bp/gaugefix.jl +++ b/src/algorithms/bp/gaugefix.jl @@ -58,8 +58,8 @@ function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv) return fuse_physicalspaces(t) end psi_Fs = reshape(psi_Fs, (Nr, Nc)) - psi′ = map(Base.Fix2(getindex, 1), psi_Fs) - Fs = map(Base.Fix2(getindex, 2), psi_Fs) + psi′ = first.(psi_Fs) + Fs = last.(psi_Fs) psi′, gauge = gauge_fix(InfinitePEPS(psi′), alg, env) # convert back to iPEPO psi′ = map(psi′.A, Fs) do t, F diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index c573e6df0..519e1b1e1 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -77,7 +77,7 @@ function compute_relative_phases( envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge ) where {C, T} - signs = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) + signs = _map_gauge_matrices(eachcoordinate(envfinal, 1:4)) do (dir, r, c) # Gather edge tensors and pretend they're InfiniteMPSs if dir == NORTH Tsprev = circshift(envprev.edges[dir, r, :], 1 - c) @@ -146,7 +146,8 @@ end function initialize_right_fixedpoint(tops, bottoms) ρ0 = randn( - TensorKit.promote_storagetype(tops...), space(tops[end], numind(tops[end]))' ← space(bottoms[end], numind(bottoms[end]))' + mapreduce(storagetype, TensorKit.promote_storagetype, tops), + space(tops[end], numind(tops[end]))' ← space(bottoms[end], numind(bottoms[end]))' ) return ρ0 end diff --git a/src/networks/gaugetrans.jl b/src/networks/gaugetrans.jl index aa133c0f2..d370c1078 100644 --- a/src/networks/gaugetrans.jl +++ b/src/networks/gaugetrans.jl @@ -79,12 +79,21 @@ CTMRGEnvGaugeTransform(pairs::AbstractArray{<:Tuple{MPSBondTensor, MPSBondTensor Base.inv(gauge::VirtualGaugeTransform) = VirtualGaugeTransform(gauge.inverses, gauge.matrices) Base.inv(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform(gauge.inverses, gauge.matrices) +# Preserve inferred unions of tensor maps, which ordinary map widens. +function _map_gauge_matrices(f::F, coords::AbstractArray{I, N}) where {F, I, N} + T = Base.promote_op(f, I) + if T <: MPSBondTensor && (isconcretetype(T) || T isa Union) + return map!(f, Array{T, N}(undef, size(coords)), coords) + end + return map(f, coords) +end + """Rotate virtual bonds counterclockwise, reversing the old north bonds into new east bonds.""" function _rotl90_virtual_gauge( matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} ) nr, nc = size(matrices)[2:3] - return map(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I + return _map_gauge_matrices(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I d, r, c = Tuple(I)[1:3] layer = Tuple(I)[4:end] # Match the leg permutation in state rotations, including fermionic signs. @@ -98,7 +107,7 @@ function _rotr90_virtual_gauge( matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} ) nr, nc = size(matrices)[2:3] - return map(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I + return _map_gauge_matrices(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I d, r, c = Tuple(I)[1:3] layer = Tuple(I)[4:end] return d == NORTH ? @@ -110,7 +119,7 @@ end """Rotate virtual bonds by a half turn, reversing both bond directions using inverse maps.""" function _rot180_virtual_gauge(inverses::AbstractArray{<:MPSBondTensor}) nr, nc = size(inverses)[2:3] - return map(CartesianIndices(inverses)) do I + return _map_gauge_matrices(CartesianIndices(inverses)) do I d, r, c = Tuple(I)[1:3] layer = Tuple(I)[4:end] r, c = nr + 1 - r, nc + 1 - c diff --git a/test/testsuite/bp/gaugefix.jl b/test/testsuite/bp/gaugefix.jl index f5f7ab125..f07e9ebb9 100644 --- a/test/testsuite/bp/gaugefix.jl +++ b/test/testsuite/bp/gaugefix.jl @@ -104,7 +104,7 @@ function bp_gaugefix_pepo(AT) N, E = random_dual!(fill(V, 2, 3)), random_dual!(fill(V, 2, 3)) ρ = adapt(AT, InfinitePEPO(fill(P, 2, 3), N, E)) bp_env = adapt(AT, BPEnv(randn, ComplexF64, ρ)) - ρg, gauge = gauge_fix(ρ, BPGauge(), bp_env) + ρg, gauge = @inferred gauge_fix(ρ, BPGauge(), bp_env) @test gauge_transform(ρ, gauge) ≈ ρg env = adapt(AT, CTMRGEnv(InfinitePEPS(ρ), V)) envg = gauge_transform(env, gauge) diff --git a/test/testsuite/ctmrg/gaugefix.jl b/test/testsuite/ctmrg/gaugefix.jl index 823610616..39f5cfffc 100644 --- a/test/testsuite/ctmrg/gaugefix.jl +++ b/test/testsuite/ctmrg/gaugefix.jl @@ -100,7 +100,7 @@ function ctmrg_gaugefix_asymmetric(AT; minimal::Bool = false) n = InfiniteSquareNetwork(psi) env, = leading_boundary(env_pre, psi, alg) env′, = ctmrg_iteration(n, env, alg) - env_fixed = gauge_fix(env′, env, gauge_alg) + env_fixed = @inferred gauge_fix(env′, env, gauge_alg) env_diff = calc_elementwise_convergence(env, env_fixed) @info "Diff between iters = $(env_diff)" @test env_diff ≈ 0 atol = atol @@ -121,7 +121,7 @@ function ctmrg_gaugefix_c4v(AT; minimal::Bool = false) n = InfiniteSquareNetwork(psi) env, = leading_boundary(env_pre, psi, alg) env′, = ctmrg_iteration(n, env, alg) - env_fixed = gauge_fix(env′, env, gauge_alg) + env_fixed = @inferred gauge_fix(env′, env, gauge_alg) env_diff = calc_elementwise_convergence(env, env_fixed) @info "Diff between iters = $(env_diff)" @test env_diff ≈ 0 atol = atol From 196b02f40e702c7b859c6fa8dec4ffe1fb26bf48 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 16:59:41 +0800 Subject: [PATCH 4/8] `adapt` for gauge transformations --- ext/PEPSKitAdaptExt.jl | 15 +++++++++ test/ctmrg/gaugetrans.jl | 14 ++++++++- test/testsuite/ctmrg/gaugetrans.jl | 50 +++++++++++++----------------- test/types/adapt.jl | 17 ++++++++++ 4 files changed, 67 insertions(+), 29 deletions(-) diff --git a/ext/PEPSKitAdaptExt.jl b/ext/PEPSKitAdaptExt.jl index 069a8a598..0c7a4783c 100644 --- a/ext/PEPSKitAdaptExt.jl +++ b/ext/PEPSKitAdaptExt.jl @@ -1,6 +1,7 @@ module PEPSKitAdaptExt using PEPSKit +using PEPSKit: _map_gauge_matrices using Adapt function Adapt.adapt_structure(to, x::PEPSKit.LocalOperator{T, S}) where {T, S} @@ -35,4 +36,18 @@ function Adapt.adapt_structure(to, x::PEPSKit.C4vCTMRGEnv{C, T}) where {C, T} return C4vCTMRGEnv{eltype(C′), eltype(T′)}(C′, T′) end +function Adapt.adapt_structure(to, gauge::VirtualGaugeTransform) + adapt_tensor = Base.Fix1(adapt, to) + return VirtualGaugeTransform( + _map_gauge_matrices(adapt_tensor, gauge.matrices), _map_gauge_matrices(adapt_tensor, gauge.inverses) + ) +end + +function Adapt.adapt_structure(to, gauge::CTMRGEnvGaugeTransform) + adapt_tensor = Base.Fix1(adapt, to) + return CTMRGEnvGaugeTransform( + _map_gauge_matrices(adapt_tensor, gauge.matrices), _map_gauge_matrices(adapt_tensor, gauge.inverses) + ) +end + end diff --git a/test/ctmrg/gaugetrans.jl b/test/ctmrg/gaugetrans.jl index 41cc0db6b..4bede8238 100644 --- a/test/ctmrg/gaugetrans.jl +++ b/test/ctmrg/gaugetrans.jl @@ -1,5 +1,5 @@ using Test -using PEPSKit +using PEPSKit, CUDA, AMDGPU @isdefined(TestSuite) || include("../testsuite/TestSuite.jl") using .TestSuite @@ -11,3 +11,15 @@ if !is_buildkite TestSuite.ctmrg_gaugetrans_pepo(Vector) TestSuite.ctmrg_gaugetrans_stacked_env(Vector) end + +if CUDA.functional() + TestSuite.ctmrg_gaugetrans(CuArray) + TestSuite.ctmrg_gaugetrans_pepo(CuArray) + TestSuite.ctmrg_gaugetrans_stacked_env(CuArray) +end + +if AMDGPU.functional() + TestSuite.ctmrg_gaugetrans(ROCArray) + TestSuite.ctmrg_gaugetrans_pepo(ROCArray) + TestSuite.ctmrg_gaugetrans_stacked_env(ROCArray) +end diff --git a/test/testsuite/ctmrg/gaugetrans.jl b/test/testsuite/ctmrg/gaugetrans.jl index da8b24661..9da1fcbbd 100644 --- a/test/testsuite/ctmrg/gaugetrans.jl +++ b/test/testsuite/ctmrg/gaugetrans.jl @@ -19,28 +19,24 @@ Ds = Dict( FermionParity => Vect[FermionParity](0 => 4, 1 => 4) ) -"""Generate a non-unitary gauge near the identity to avoid ill-conditioned inverses.""" -function _random_gauge_matrix(AT, V::ElementarySpace) - return adapt(AT, id(ComplexF64, V) + 0.1 * randn(ComplexF64, V ← V)) +"""Generate a non-unitary matrix near the identity.""" +function _random_gauge_matrix(V::ElementarySpace) + return id(ComplexF64, V) + 0.1 * randn(ComplexF64, V ← V) end -"""Test virtual and boundary gauge invariance.""" +"""Test invariance of reduced density matrix under gauge transformations.""" function ctmrg_gaugetrans(AT) return @testset "CTMRGEnv of InfinitePEPS ($S) ($AT)" for S in keys(ds) d, D, χ, uc = ds[S], Ds[S], χs[S], (2, 3) N, E = random_dual!(fill(D, uc)), random_dual!(fill(D, uc)) ψ = adapt(AT, InfinitePEPS(fill(d, uc), N, E)) env = adapt(AT, CTMRGEnv(ψ, χ)) - gauge = VirtualGaugeTransform( - map(eachcoordinate(ψ, 1:2)) do (dir, r, c) - return _random_gauge_matrix(AT, dir == NORTH ? N[r, c] : E[r, c]) - end - ) - boundary_gauge = CTMRGEnvGaugeTransform( - map(env.edges) do edge - return _random_gauge_matrix(AT, space(edge, 1)) - end - ) + virtual = map(eachcoordinate(ψ, 1:2)) do (dir, r, c) + return _random_gauge_matrix(dir == NORTH ? N[r, c] : E[r, c]) + end + boundary = map(edge -> _random_gauge_matrix(space(edge, 1)), env.edges) + gauge = adapt(AT, VirtualGaugeTransform(virtual)) + boundary_gauge = adapt(AT, CTMRGEnvGaugeTransform(boundary)) ψg, envg = gauge_transform(ψ, gauge), gauge_transform(env, gauge) envb = gauge_transform(env, boundary_gauge) for R in (rotl90, rotr90, rot180), (state, g, transformed) in @@ -56,8 +52,8 @@ function ctmrg_gaugetrans(AT) if S == FermionParity for (left, right) in ((D, D'), (D', D)) corner = adapt(AT, randn(ComplexF64, left ← right)) - Xout = _random_gauge_matrix(AT, left) - Xin = _random_gauge_matrix(AT, right) + Xout = adapt(AT, _random_gauge_matrix(left)) + Xin = adapt(AT, _random_gauge_matrix(right)) σin = Xin' @autoopt @tensor expected[χ_out; χ_in] := Xout[χ_out; χ1] * corner[χ1; χ2] * conj(σin[χ_in; χ2]) @@ -75,10 +71,10 @@ function ctmrg_gaugetrans_pepo(AT) foreach(random_dual!, eachslice(N; dims = 3)) foreach(random_dual!, eachslice(E; dims = 3)) ρ = adapt(AT, InfinitePEPO(P, N, E)) - X = map(Iterators.product(1:2, axes(P)...)) do (dir, r, c, h) - return _random_gauge_matrix(AT, dir == NORTH ? N[r, c, h] : E[r, c, h]) + matrices = map(Iterators.product(1:2, axes(P)...)) do (dir, r, c, h) + return _random_gauge_matrix(dir == NORTH ? N[r, c, h] : E[r, c, h]) end - gauge = VirtualGaugeTransform(X) + gauge = adapt(AT, VirtualGaugeTransform(matrices)) ρg = gauge_transform(ρ, gauge) for R in (rotl90, rotr90, rot180) @test R(ρg) ≈ gauge_transform(R(ρ), R(gauge)) @@ -86,7 +82,8 @@ function ctmrg_gaugetrans_pepo(AT) for h in axes(P, 3) layer, layerg = InfinitePEPO(ρ.A[:, :, h:h]), InfinitePEPO(ρg.A[:, :, h:h]) env = adapt(AT, CTMRGEnv(InfinitePartitionFunction(layer), χ)) - envg = gauge_transform(env, VirtualGaugeTransform(X[:, :, :, h:h])) + layer_gauge = VirtualGaugeTransform(gauge.matrices[:, :, :, h:h], gauge.inverses[:, :, :, h:h]) + envg = gauge_transform(env, layer_gauge) for (r, c) in Iterators.product(axes(P, 1), axes(P, 2)) inds = [CartesianIndex(r, c)] @test reduced_densitymatrix(inds, layerg, envg) ≈ reduced_densitymatrix(inds, layer, env) @@ -102,13 +99,10 @@ function ctmrg_gaugetrans_stacked_env(AT) ψ = adapt(AT, InfinitePEPS(V, V)) ρ = adapt(AT, InfinitePEPO(V, V; unitcell = (1, 1, 2))) network = InfiniteSquareNetwork(ψ, ρ, ψ) - env = adapt(AT, CTMRGEnv(network, V)) - gauge = CTMRGEnvGaugeTransform( - map(env.edges) do edge - return _random_gauge_matrix(AT, space(edge, 1)) - end - ) - @test _contract_site((1, 1), network, gauge_transform(env, gauge)) ≈ - _contract_site((1, 1), network, env) + env = CTMRGEnv(network, V) + matrices = map(edge -> _random_gauge_matrix(space(edge, 1)), env.edges) + gauge = adapt(AT, CTMRGEnvGaugeTransform(matrices)) + envg = gauge_transform(env, gauge) + @test _contract_site((1, 1), network, envg) ≈ _contract_site((1, 1), network, env) end end diff --git a/test/types/adapt.jl b/test/types/adapt.jl index ccef55ba6..b3fe89f36 100644 --- a/test/types/adapt.jl +++ b/test/types/adapt.jl @@ -64,6 +64,23 @@ if !is_buildkite @test TensorKit.storagetype(jl_env) <: JLArray end end + @testset "Gauge transforms" begin + X = id(ComplexF64, ℂ^2) + for (make_gauge, dims) in ( + (VirtualGaugeTransform, (2, 2, 1)), + (VirtualGaugeTransform, (2, 2, 1, 2)), + (CTMRGEnvGaugeTransform, (4, 2, 1)), + ) + gauge = make_gauge(fill(X, dims), fill(X', dims)) + jl_gauge = adapt(JLArray, gauge) + @test all(t -> storagetype(t) <: JLArray, jl_gauge.matrices) + @test all(t -> storagetype(t) <: JLArray, jl_gauge.inverses) + @test first(jl_gauge.inverses) isa TensorKit.AdjointTensorMap + host_gauge = adapt(Array, jl_gauge) + @test host_gauge.matrices == gauge.matrices + @test host_gauge.inverses == gauge.inverses + end + end @testset "C4vCTMRGEnv" begin Vphy = Z2Space(0 => 1, 1 => 1) V = Z2Space(0 => 1, 1 => 2) From fe6d1e28f38ad75852a0a2d1aee8606330838ac3 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 17:33:35 +0800 Subject: [PATCH 5/8] Test cleanup --- test/testsuite/ctmrg/gaugetrans.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/testsuite/ctmrg/gaugetrans.jl b/test/testsuite/ctmrg/gaugetrans.jl index 9da1fcbbd..c223d7ad9 100644 --- a/test/testsuite/ctmrg/gaugetrans.jl +++ b/test/testsuite/ctmrg/gaugetrans.jl @@ -30,7 +30,7 @@ function ctmrg_gaugetrans(AT) d, D, χ, uc = ds[S], Ds[S], χs[S], (2, 3) N, E = random_dual!(fill(D, uc)), random_dual!(fill(D, uc)) ψ = adapt(AT, InfinitePEPS(fill(d, uc), N, E)) - env = adapt(AT, CTMRGEnv(ψ, χ)) + env = CTMRGEnv(ψ, χ) virtual = map(eachcoordinate(ψ, 1:2)) do (dir, r, c) return _random_gauge_matrix(dir == NORTH ? N[r, c] : E[r, c]) end @@ -81,7 +81,7 @@ function ctmrg_gaugetrans_pepo(AT) end for h in axes(P, 3) layer, layerg = InfinitePEPO(ρ.A[:, :, h:h]), InfinitePEPO(ρg.A[:, :, h:h]) - env = adapt(AT, CTMRGEnv(InfinitePartitionFunction(layer), χ)) + env = CTMRGEnv(InfinitePartitionFunction(layer), χ) layer_gauge = VirtualGaugeTransform(gauge.matrices[:, :, :, h:h], gauge.inverses[:, :, :, h:h]) envg = gauge_transform(env, layer_gauge) for (r, c) in Iterators.product(axes(P, 1), axes(P, 2)) From 7f314d4a40130a05ff7814c57c1dd410c009a79b Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 20:52:54 +0800 Subject: [PATCH 6/8] Remove rotations of gauge transformations --- ext/PEPSKitAdaptExt.jl | 5 +- src/algorithms/ctmrg/gaugefix.jl | 31 ++++----- src/networks/gaugetrans.jl | 103 ----------------------------- test/testsuite/ctmrg/gaugetrans.jl | 7 -- test/types/adapt.jl | 4 +- 5 files changed, 16 insertions(+), 134 deletions(-) diff --git a/ext/PEPSKitAdaptExt.jl b/ext/PEPSKitAdaptExt.jl index 0c7a4783c..b2523f76d 100644 --- a/ext/PEPSKitAdaptExt.jl +++ b/ext/PEPSKitAdaptExt.jl @@ -1,7 +1,6 @@ module PEPSKitAdaptExt using PEPSKit -using PEPSKit: _map_gauge_matrices using Adapt function Adapt.adapt_structure(to, x::PEPSKit.LocalOperator{T, S}) where {T, S} @@ -39,14 +38,14 @@ end function Adapt.adapt_structure(to, gauge::VirtualGaugeTransform) adapt_tensor = Base.Fix1(adapt, to) return VirtualGaugeTransform( - _map_gauge_matrices(adapt_tensor, gauge.matrices), _map_gauge_matrices(adapt_tensor, gauge.inverses) + map(adapt_tensor, gauge.matrices), map(adapt_tensor, gauge.inverses) ) end function Adapt.adapt_structure(to, gauge::CTMRGEnvGaugeTransform) adapt_tensor = Base.Fix1(adapt, to) return CTMRGEnvGaugeTransform( - _map_gauge_matrices(adapt_tensor, gauge.matrices), _map_gauge_matrices(adapt_tensor, gauge.inverses) + map(adapt_tensor, gauge.matrices), map(adapt_tensor, gauge.inverses) ) end diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index 519e1b1e1..805fff144 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -77,33 +77,26 @@ function compute_relative_phases( envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge ) where {C, T} - signs = _map_gauge_matrices(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - # Gather edge tensors and pretend they're InfiniteMPSs - if dir == NORTH - Tsprev = circshift(envprev.edges[dir, r, :], 1 - c) - Tsfinal = circshift(envfinal.edges[dir, r, :], 1 - c) - elseif dir == EAST - Tsprev = circshift(envprev.edges[dir, :, c], 1 - r) - Tsfinal = circshift(envfinal.edges[dir, :, c], 1 - r) - elseif dir == SOUTH - Tsprev = circshift(reverse(envprev.edges[dir, r, :]), c) - Tsfinal = circshift(reverse(envfinal.edges[dir, r, :]), c) - elseif dir == WEST - Tsprev = circshift(reverse(envprev.edges[dir, :, c]), r) - Tsfinal = circshift(reverse(envfinal.edges[dir, :, c]), r) - end - - # Random MPS of same bond dimension + signs = map(eachcoordinate(envfinal, 1:4)) do coordinate + Tsprev = _gauge_transfer_edges(envprev.edges, coordinate) + Tsfinal = _gauge_transfer_edges(envfinal.edges, coordinate) M = map(Tsfinal) do t - randn(storagetype(T), codomain(t) ← domain(t)) + return randn(storagetype(t), codomain(t) ← domain(t)) end - return _scrambling_gauge_matrix(Tsprev, Tsfinal, M, Arnoldi()) end return CTMRGEnvGaugeTransform(signs, adjoint.(signs)) end +function _gauge_transfer_edges(edges, (dir, r, c)) + dir == NORTH && return circshift(edges[dir, r, :], 1 - c) + dir == EAST && return circshift(edges[dir, :, c], 1 - r) + dir == SOUTH && return circshift(reverse(edges[dir, r, :]), c) + dir == WEST && return circshift(reverse(edges[dir, :, c]), r) + throw(ArgumentError(lazy"invalid direction $dir")) +end + function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, ::ScramblingEnvGauge) where {C, T} # "general" algorithm from https://arxiv.org/abs/2311.11894 Tprev = envprev.edges[1, 1, 1] diff --git a/src/networks/gaugetrans.jl b/src/networks/gaugetrans.jl index d370c1078..30d921798 100644 --- a/src/networks/gaugetrans.jl +++ b/src/networks/gaugetrans.jl @@ -79,109 +79,6 @@ CTMRGEnvGaugeTransform(pairs::AbstractArray{<:Tuple{MPSBondTensor, MPSBondTensor Base.inv(gauge::VirtualGaugeTransform) = VirtualGaugeTransform(gauge.inverses, gauge.matrices) Base.inv(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform(gauge.inverses, gauge.matrices) -# Preserve inferred unions of tensor maps, which ordinary map widens. -function _map_gauge_matrices(f::F, coords::AbstractArray{I, N}) where {F, I, N} - T = Base.promote_op(f, I) - if T <: MPSBondTensor && (isconcretetype(T) || T isa Union) - return map!(f, Array{T, N}(undef, size(coords)), coords) - end - return map(f, coords) -end - -"""Rotate virtual bonds counterclockwise, reversing the old north bonds into new east bonds.""" -function _rotl90_virtual_gauge( - matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} - ) - nr, nc = size(matrices)[2:3] - return _map_gauge_matrices(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I - d, r, c = Tuple(I)[1:3] - layer = Tuple(I)[4:end] - # Match the leg permutation in state rotations, including fermionic signs. - return d == NORTH ? matrices[EAST, c, nc + 1 - r, layer...] : - permute(inverses[NORTH, _next(c, nr), nc + 1 - r, layer...], ((2,), (1,))) - end -end - -"""Rotate virtual bonds clockwise, reversing the old east bonds into new north bonds.""" -function _rotr90_virtual_gauge( - matrices::AbstractArray{<:MPSBondTensor}, inverses::AbstractArray{<:MPSBondTensor} - ) - nr, nc = size(matrices)[2:3] - return _map_gauge_matrices(CartesianIndices((2, nc, nr, size(matrices)[4:end]...))) do I - d, r, c = Tuple(I)[1:3] - layer = Tuple(I)[4:end] - return d == NORTH ? - permute(inverses[EAST, nr + 1 - c, _prev(r, nc), layer...], ((2,), (1,))) : - matrices[NORTH, nr + 1 - c, r, layer...] - end -end - -"""Rotate virtual bonds by a half turn, reversing both bond directions using inverse maps.""" -function _rot180_virtual_gauge(inverses::AbstractArray{<:MPSBondTensor}) - nr, nc = size(inverses)[2:3] - return _map_gauge_matrices(CartesianIndices(inverses)) do I - d, r, c = Tuple(I)[1:3] - layer = Tuple(I)[4:end] - r, c = nr + 1 - r, nc + 1 - c - Xinv = d == NORTH ? inverses[NORTH, _next(r, nr), c, layer...] : - inverses[EAST, r, _prev(c, nc), layer...] - return permute(Xinv, ((2,), (1,))) - end -end - -"""Rotate boundary gauge directions and unit-cell coordinates counterclockwise.""" -function _rotl90_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) - return map(CartesianIndices((4, size(matrices, 3), size(matrices, 2)))) do I - d, r, c = Tuple(I) - return matrices[_next(d, 4), c, size(matrices, 3) + 1 - r] - end -end - -"""Rotate boundary gauge directions and unit-cell coordinates clockwise.""" -function _rotr90_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) - return map(CartesianIndices((4, size(matrices, 3), size(matrices, 2)))) do I - d, r, c = Tuple(I) - return matrices[_prev(d, 4), size(matrices, 2) + 1 - c, r] - end -end - -"""Rotate boundary gauge directions and unit-cell coordinates by a half turn.""" -function _rot180_boundary_gauge(matrices::AbstractArray{<:MPSBondTensor, 3}) - return map(CartesianIndices(matrices)) do I - d, r, c = Tuple(I) - return matrices[mod1(d + 2, 4), size(matrices, 2) + 1 - r, size(matrices, 3) + 1 - c] - end -end - -""" - rotl90(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) - rotr90(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) - rot180(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}) - -Rotate a gauge consistently with its state or CTMRG environment, leaving PEPO layers in place. -For each rotation `R`, `R(gauge_transform(state, gauge)) ≈ gauge_transform(R(state), R(gauge))`. -""" -Base.rotl90(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( - _rotl90_virtual_gauge(gauge.matrices, gauge.inverses), - _rotl90_virtual_gauge(gauge.inverses, gauge.matrices) -) -Base.rotl90(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( - _rotl90_boundary_gauge(gauge.matrices), _rotl90_boundary_gauge(gauge.inverses) -) -Base.rotr90(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( - _rotr90_virtual_gauge(gauge.matrices, gauge.inverses), - _rotr90_virtual_gauge(gauge.inverses, gauge.matrices) -) -Base.rotr90(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( - _rotr90_boundary_gauge(gauge.matrices), _rotr90_boundary_gauge(gauge.inverses) -) -Base.rot180(gauge::VirtualGaugeTransform) = VirtualGaugeTransform( - _rot180_virtual_gauge(gauge.inverses), _rot180_virtual_gauge(gauge.matrices) -) -Base.rot180(gauge::CTMRGEnvGaugeTransform) = CTMRGEnvGaugeTransform( - _rot180_boundary_gauge(gauge.matrices), _rot180_boundary_gauge(gauge.inverses) -) - """Check that a gauge array has the expected unit cell and layer count.""" function _check_gauge_size(gauge::Union{VirtualGaugeTransform, CTMRGEnvGaugeTransform}, dims::Tuple) size(gauge.matrices)[2:end] == dims || throw(DimensionMismatch("Gauge unit cell does not match the object being transformed")) diff --git a/test/testsuite/ctmrg/gaugetrans.jl b/test/testsuite/ctmrg/gaugetrans.jl index c223d7ad9..01e0b66d4 100644 --- a/test/testsuite/ctmrg/gaugetrans.jl +++ b/test/testsuite/ctmrg/gaugetrans.jl @@ -39,10 +39,6 @@ function ctmrg_gaugetrans(AT) boundary_gauge = adapt(AT, CTMRGEnvGaugeTransform(boundary)) ψg, envg = gauge_transform(ψ, gauge), gauge_transform(env, gauge) envb = gauge_transform(env, boundary_gauge) - for R in (rotl90, rotr90, rot180), (state, g, transformed) in - ((ψ, gauge, ψg), (env, gauge, envg), (env, boundary_gauge, envb)) - @test R(transformed) ≈ gauge_transform(R(state), R(g)) - end for (r, c) in eachcoordinate(ψ) inds = [CartesianIndex(r, c)] ρ = reduced_densitymatrix(inds, ψ, env) @@ -76,9 +72,6 @@ function ctmrg_gaugetrans_pepo(AT) end gauge = adapt(AT, VirtualGaugeTransform(matrices)) ρg = gauge_transform(ρ, gauge) - for R in (rotl90, rotr90, rot180) - @test R(ρg) ≈ gauge_transform(R(ρ), R(gauge)) - end for h in axes(P, 3) layer, layerg = InfinitePEPO(ρ.A[:, :, h:h]), InfinitePEPO(ρg.A[:, :, h:h]) env = CTMRGEnv(InfinitePartitionFunction(layer), χ) diff --git a/test/types/adapt.jl b/test/types/adapt.jl index b3fe89f36..5b05ab07c 100644 --- a/test/types/adapt.jl +++ b/test/types/adapt.jl @@ -72,11 +72,11 @@ if !is_buildkite (CTMRGEnvGaugeTransform, (4, 2, 1)), ) gauge = make_gauge(fill(X, dims), fill(X', dims)) - jl_gauge = adapt(JLArray, gauge) + jl_gauge = @inferred adapt(JLArray, gauge) @test all(t -> storagetype(t) <: JLArray, jl_gauge.matrices) @test all(t -> storagetype(t) <: JLArray, jl_gauge.inverses) @test first(jl_gauge.inverses) isa TensorKit.AdjointTensorMap - host_gauge = adapt(Array, jl_gauge) + host_gauge = @inferred adapt(Array, jl_gauge) @test host_gauge.matrices == gauge.matrices @test host_gauge.inverses == gauge.inverses end From ef8c50ec47bbfe4b3ab8c3fed8eea17c2d31e292 Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 21:17:45 +0800 Subject: [PATCH 7/8] Deduplication on CTMRG gauge fixing --- src/algorithms/ctmrg/gaugefix.jl | 51 ++++++------------- src/algorithms/ctmrg/gaugetrans.jl | 12 +---- .../optimization/implicit_differentiation.jl | 8 +-- 3 files changed, 22 insertions(+), 49 deletions(-) diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index 805fff144..29e9eebeb 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -21,49 +21,29 @@ element-wise converged to `envprev`. function gauge_fix( envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, alg::ScramblingEnvGauge ) 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 = compute_relative_phases(envfinal, envprev, alg) + return fix_global_phases(fix_relative_phases(envfinal, signs), envprev) end -function compute_gauge_fix_gauge( - envfinal::CTMRGEnv{C, T}, envprev::CTMRGEnv{C, T}, alg::ScramblingEnvGauge - ) where {C, T} - # Check if spaces in envprev and envfinal are the same - same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) && - space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c]) +function _check_gauge_fix_spaces(envfinal::AbstractCTMRGEnv, envprev::AbstractCTMRGEnv) + (axes(envfinal.corners) == axes(envprev.corners) && axes(envfinal.edges) == axes(envprev.edges)) || + throw(DimensionMismatch("Unit cells of envfinal and envprev do not match")) + for (final, prev) in ((envfinal.corners, envprev.corners), (envfinal.edges, envprev.edges)) + same_spaces = all(zip(final, prev)) do (a, b) + return space(a) == space(b) + end + same_spaces || throw(ArgumentError("Spaces of envfinal and envprev are not the same")) end - all(same_spaces) || throw(ArgumentError("Spaces of envfinal and envprev are not the same")) - - # find relative phases - signs = compute_relative_phases(envfinal, envprev, alg) - - # find additional global phases - corner_phases, edge_phases = compute_global_phases( - gauge_transform(envfinal, signs), envprev - ) - - return signs, corner_phases, edge_phases + return nothing end function compute_gauge_fix_gauge( - envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, alg::ScramblingEnvGauge + envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, alg::ScramblingEnvGauge ) where {C, T} - # Check if spaces in envprev and envfinal are the same - same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c) - space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) && - space(envfinal.corners[1, r, c]) == space(envprev.corners[1, r, c]) - end - all(same_spaces) || throw(ArgumentError("Spaces of envfinal and envprev are not the same")) - - # find relative phases signs = compute_relative_phases(envfinal, envprev, alg) - - # find additional global phases corner_phases, edge_phases = compute_global_phases( fix_relative_phases(envfinal, signs), envprev ) - return signs, corner_phases, edge_phases end @@ -76,7 +56,7 @@ end function compute_relative_phases( envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge ) where {C, T} - + _check_gauge_fix_spaces(envfinal, envprev) signs = map(eachcoordinate(envfinal, 1:4)) do coordinate Tsprev = _gauge_transfer_edges(envprev.edges, coordinate) Tsfinal = _gauge_transfer_edges(envfinal.edges, coordinate) @@ -98,6 +78,7 @@ function _gauge_transfer_edges(edges, (dir, r, c)) end function compute_relative_phases(envfinal::C4vCTMRGEnv{C, T}, envprev::C4vCTMRGEnv{C, T}, ::ScramblingEnvGauge) where {C, T} + _check_gauge_fix_spaces(envfinal, envprev) # "general" algorithm from https://arxiv.org/abs/2311.11894 Tprev = envprev.edges[1, 1, 1] Tfinal = envfinal.edges[1, 1, 1] @@ -178,12 +159,12 @@ function fix_relative_phases( nr, nc = size(U)[2:3] U_fixed = map(eachindex(IndexCartesian(), U)) do I d, r, c = Tuple(I) - ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + _, ri, ci = _next_coordinate((d, r, c), nr, nc) return U[I] * signs.inverses[d, ri, ci] end V_fixed = map(eachindex(IndexCartesian(), V)) do I d, r, c = Tuple(I) - ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + _, ri, ci = _next_coordinate((d, r, c), nr, nc) return signs.matrices[d, ri, ci] * V[I] end return U_fixed, V_fixed diff --git a/src/algorithms/ctmrg/gaugetrans.jl b/src/algorithms/ctmrg/gaugetrans.jl index 5adeae637..4ea1c204b 100644 --- a/src/algorithms/ctmrg/gaugetrans.jl +++ b/src/algorithms/ctmrg/gaugetrans.jl @@ -53,11 +53,11 @@ function gauge_transform(env::CTMRGEnv, gauge::CTMRGEnvGaugeTransform) X, Xinv = gauge.matrices, gauge.inverses nr, nc = size(env)[2:3] corners = map(eachcoordinate(env, 1:4)) do (d, r, c) - ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + _, ri, ci = _next_coordinate((d, r, c), nr, nc) return _gauge_transform_corner(env.corners[d, r, c], X[_prev(d, 4), r, c], Xinv[d, ri, ci]) end edges = map(eachcoordinate(env, 1:4)) do (d, r, c) - ri, ci = _environment_gauge_neighbor(d, r, c, nr, nc) + _, ri, ci = _next_coordinate((d, r, c), nr, nc) return absorb_left_right(env.edges[d, r, c], X[d, r, c], Xinv[d, ri, ci]) end return CTMRGEnv(corners, edges) @@ -69,11 +69,3 @@ function _gauge_transform_corner(C::CTMRGCornerTensor, Xout::MPSBondTensor, Xin: return @autoopt @tensor Cg[χ_out; χ_in] := Xout[χ_out; χ1] * C[χ1; χ2] * Xin[χ2; χ_in] end - -"""Find the next boundary bond along an edge, respecting the periodic unit cell.""" -function _environment_gauge_neighbor(d::Int, r::Int, c::Int, nr::Int, nc::Int) - d == NORTH && return r, _next(c, nc) - d == EAST && return _next(r, nr), c - d == SOUTH && return r, _prev(c, nc) - return _prev(r, nr), c -end diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index cea465d62..73ada285f 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -226,7 +226,7 @@ function _rrule( 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) + signs = compute_relative_phases(env_conv, env, alg_gauge) # `p` for state, `x` for environment function gauge_fixed_iteration(p, x) # fix global phases relative to the fixed point inside the iteration @@ -478,8 +478,8 @@ function _rrule( 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) + signs = compute_relative_phases(env_conv, env, alg_gauge) + env_fixed = fix_global_phases(fix_relative_phases(env_conv, signs), env) # 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] @@ -624,7 +624,7 @@ function PEPSKit._rrule( # 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()) + signs = compute_relative_phases(env_conv, env, ScramblingEnvGauge()) S = normalize.(info.S) U, V = fix_relative_phases(info.U, info.V, signs) From b407c4e3a419722b5e5bed7c354347849da9b44b Mon Sep 17 00:00:00 2001 From: Yue Zhengyuan Date: Sun, 8 Nov 2026 23:25:19 +0800 Subject: [PATCH 8/8] Inference fix for Julia 1.10 --- src/algorithms/ctmrg/gaugefix.jl | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl index 29e9eebeb..f9885894c 100644 --- a/src/algorithms/ctmrg/gaugefix.jl +++ b/src/algorithms/ctmrg/gaugefix.jl @@ -57,14 +57,15 @@ function compute_relative_phases( envfinal::AbstractCTMRGEnv{C, T}, envprev::AbstractCTMRGEnv{C, T}, ::ScramblingEnvGauge ) where {C, T} _check_gauge_fix_spaces(envfinal, envprev) - signs = map(eachcoordinate(envfinal, 1:4)) do coordinate - Tsprev = _gauge_transfer_edges(envprev.edges, coordinate) - Tsfinal = _gauge_transfer_edges(envfinal.edges, coordinate) - M = map(Tsfinal) do t - return randn(storagetype(t), codomain(t) ← domain(t)) - end - return _scrambling_gauge_matrix(Tsprev, Tsfinal, M, Arnoldi()) - end + # Julia 1.10 infers the gauge array with a comprehension. + signs = [ + begin + Tsprev = _gauge_transfer_edges(envprev.edges, coordinate) + Tsfinal = _gauge_transfer_edges(envfinal.edges, coordinate) + M = map(t -> randn(storagetype(t), space(t)), Tsfinal) + _scrambling_gauge_matrix(Tsprev, Tsfinal, M, Arnoldi()) + end for coordinate in eachcoordinate(envfinal, 1:4) + ] return CTMRGEnvGaugeTransform(signs, adjoint.(signs)) end