diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 2aeb865b6..5bda89392 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -57,23 +57,32 @@ include("Defaults.jl") # Include first to allow for docstring interpolation wit include("utility/util.jl") include("utility/contraction_labels.jl") include("utility/tensor_traces.jl") +include("utility/contraction_labels.jl") +include("utility/tensor_traces.jl") include("utility/indexing.jl") include("utility/diffable_threads.jl") include("utility/twistdual.jl") +include("utility/twistdual.jl") include("utility/eigh.jl") include("utility/svd.jl") include("utility/orth.jl") +include("utility/orth.jl") include("utility/rotations.jl") include("utility/hook_pullback.jl") include("utility/autoopt.jl") include("utility/retractions.jl") include("utility/initialization.jl") +include("utility/initialization.jl") include("networks/tensors.jl") include("networks/local_sandwich.jl") include("networks/infinitesquarenetwork.jl") +include("networks/tensors_triangular.jl") +include("networks/local_triangular_sandwich.jl") +include("networks/infinitetriangularnetwork.jl") include("states/infinitepeps.jl") +include("states/infinitepepstriangular.jl") include("states/infinitepartitionfunction.jl") include("utility/symmetrization.jl") @@ -84,15 +93,20 @@ include("operators/localoperator.jl") include("operators/localoperator/dense.jl") include("operators/localoperator/tensorproduct.jl") include("operators/localoperator/mpo.jl") +include("operators/localoperator/dense.jl") +include("operators/localoperator/tensorproduct.jl") +include("operators/localoperator/mpo.jl") include("operators/localcircuit.jl") include("operators/lattices/squarelattice.jl") include("operators/models.jl") include("environments/ctmrg_environments.jl") +include("environments/ctmrg_environments_triangular.jl") include("environments/vumps_environments.jl") include("environments/suweight.jl") include("environments/bp_environments.jl") include("environments/product_state_environments.jl") +include("environments/product_state_environments.jl") include("algorithms/contractions/ctmrg/types.jl") include("algorithms/contractions/ctmrg/expr.jl") @@ -104,9 +118,12 @@ 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/network_value.jl") include("algorithms/contractions/ctmrg/gaugefix.jl") include("algorithms/contractions/ctmrg/characteristic_equations.jl") +include("algorithms/contractions/ctmrg/characteristic_equations.jl") +include("algorithms/contractions/absorb.jl") include("algorithms/contractions/absorb.jl") include("algorithms/contractions/absorb_weight.jl") include("algorithms/contractions/transfer.jl") @@ -118,12 +135,20 @@ include("algorithms/contractions/local_patch/patch_contractions.jl") include("algorithms/contractions/local_patch/densitymatrix/generic.jl") include("algorithms/contractions/local_patch/densitymatrix/ctmrg.jl") include("algorithms/contractions/local_patch/densitymatrix/bp.jl") +include("algorithms/contractions/bp_messages.jl") +include("algorithms/contractions/local_patch/expr_utils.jl") +include("algorithms/contractions/local_patch/network_expr.jl") +include("algorithms/contractions/local_patch/patch_contractions.jl") +include("algorithms/contractions/local_patch/densitymatrix/generic.jl") +include("algorithms/contractions/local_patch/densitymatrix/ctmrg.jl") +include("algorithms/contractions/local_patch/densitymatrix/bp.jl") include("algorithms/contractions/bondenv/benv_tools.jl") include("algorithms/contractions/bondenv/gaugefix.jl") include("algorithms/contractions/bondenv/als_solve.jl") include("algorithms/contractions/bondenv/benv_ctm.jl") include("algorithms/contractions/correlator/peps.jl") include("algorithms/contractions/correlator/pepo_purified.jl") +include("algorithms/contractions/correlator/pepo_purified.jl") include("algorithms/contractions/correlator/pepo_1layer.jl") include("algorithms/ctmrg/sparse_environments.jl") @@ -133,11 +158,28 @@ include("algorithms/ctmrg/projectors/halfinfinite.jl") include("algorithms/ctmrg/projectors/fullinfinite.jl") include("algorithms/ctmrg/projectors/c4v_eigh.jl") include("algorithms/ctmrg/projectors/c4v_qr.jl") +include("algorithms/ctmrg/projectors/projectors.jl") +include("algorithms/ctmrg/projectors/halfinfinite.jl") +include("algorithms/ctmrg/projectors/fullinfinite.jl") +include("algorithms/ctmrg/projectors/c4v_eigh.jl") +include("algorithms/ctmrg/projectors/c4v_qr.jl") include("algorithms/ctmrg/simultaneous.jl") include("algorithms/ctmrg/sequential.jl") include("algorithms/ctmrg/gaugefix.jl") include("algorithms/ctmrg/c4v.jl") include("algorithms/ctmrg/initialization.jl") +include("algorithms/ctmrg/initialization.jl") + +include("algorithms/truncation/truncationschemes.jl") +include("algorithms/truncation/fullenv_truncation.jl") +include("algorithms/truncation/bond_truncation.jl") + +include("algorithms/compress/local.jl") + +include("algorithms/contractions/ctmrg_contractions_triangular.jl") +include("algorithms/ctmrg_triangular/ctmrg.jl") +include("algorithms/ctmrg_triangular/utility.jl") +include("algorithms/ctmrg_triangular/simultaneous.jl") include("algorithms/truncation/truncationschemes.jl") include("algorithms/truncation/fullenv_truncation.jl") @@ -165,8 +207,17 @@ include("algorithms/expectation_value/patch_contractions.jl") include("algorithms/expectation_value/network_value.jl") include("algorithms/expectation_value/correlator_adapters.jl") include("algorithms/expectation_value/correlators.jl") +include("algorithms/expectation_value/expectation_value.jl") +include("algorithms/expectation_value/reduced_densitymatrix.jl") +include("algorithms/expectation_value/patch_contractions.jl") +include("algorithms/expectation_value/network_value.jl") +include("algorithms/expectation_value/correlator_adapters.jl") +include("algorithms/expectation_value/correlators.jl") include("algorithms/toolbox.jl") +include("algorithms/toolbox_triangular.jl") +include("algorithms/optimization/implicit_differentiation.jl") +include("algorithms/optimization/preconditioning.jl") include("algorithms/optimization/implicit_differentiation.jl") include("algorithms/optimization/preconditioning.jl") include("algorithms/optimization/peps_optimization.jl") @@ -176,7 +227,9 @@ include("algorithms/select_algorithm.jl") using .Defaults: set_scheduler! export set_scheduler! export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, OrthAdjoint, QRAdjoint +export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, OrthAdjoint, QRAdjoint export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG +export CTMRGEnvTriangular, SimultaneousCTMRGTriangular # CTMRGTriaEnv, SimultaneousCTMRGTria export initialize_ctmrg_environment, RandomInitialization, ProductStateInitialization, ApplicationInitialization, IdentityInitialization export corner, edge, setcorner!, setedge! @@ -195,14 +248,18 @@ export fixedpoint export LocalTruncation export compress +export LocalTruncation +export compress + export absorb_weight export ALSTruncation, FullEnvTruncation export SimpleUpdate export TimeEvolver, timestep, time_evolve -export InfiniteSquareNetwork +export InfiniteSquareNetwork, InfiniteTriangularNetwork export InfinitePartitionFunction export InfinitePEPS, InfiniteTransferPEPS +export InfinitePEPSTriangular export SUWeight export InfinitePEPO, InfiniteTransferPEPO diff --git a/src/algorithms/contractions/ctmrg_contractions_triangular.jl b/src/algorithms/contractions/ctmrg_contractions_triangular.jl new file mode 100644 index 000000000..42b777923 --- /dev/null +++ b/src/algorithms/contractions/ctmrg_contractions_triangular.jl @@ -0,0 +1,104 @@ +function build_double_corner_matrix_triangular(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, dir::Int, r::Int, c::Int) where {P <: PFTensorTriangular} + @tensor opt = true mat[-1 -2; -3 -4] := + env.C[_coordinates(dir, 0, r, c, size(network))...][6 5; 1] * + env.C[_coordinates(dir + 1, 0, r, c, size(network))...][1 3; 2] * + env.Ea[_coordinates(dir - 1, 0, r, c, size(network))...][-1 7; 6] * + env.Eb[_coordinates(dir + 1, 0, r, c, size(network))...][2 4; -3] * + rotl60(network[r, c], mod(dir - 1, 6))[7 -2 -4; 5 3 4] + return mat +end + +function build_double_corner_matrix_triangular(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, dir::Int, r::Int, c::Int) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + @tensor opt = true mat[-1 -2 -3; -4 -5 -6] := + env.C[_coordinates(dir, 0, r, c, size(network))...][6 5 8; 1] * + env.C[_coordinates(dir + 1, 0, r, c, size(network))...][1 3 9; 2] * + env.Ea[_coordinates(dir - 1, 0, r, c, size(network))...][-1 7 10; 6] * + env.Eb[_coordinates(dir + 1, 0, r, c, size(network))...][2 4 11; -4] * + rotl60(ket(O_r_c), mod(dir - 1, 6))[12; 5 3 4 -5 -2 7] * + conj(rotl60(bra(O_r_c), mod(dir - 1, 6))[12; 8 9 11 -6 -3 10]) + return mat +end + +function semi_renormalize_edge(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, Pas, Pbs, dir, r, c) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + @tensor opt = true mat[-1 -2 -3; -4 -5 -6] := + Pas[dir, r, c][-1; 1 2 8] * Pbs[dir, r, c][6 7 9; -4] * + env.Ea[_coordinates(dir, 0, r, c, size(network))...][1 3 10; 4] * + env.Eb[_coordinates(dir, 1, r, c, size(network))...][4 5 11; 6] * + rotl60(ket(O_r_c), mod(dir - 1, 6))[12; 3 5 7 -5 -2 2] * + conj(rotl60(bra(O_r_c), mod(dir - 1, 6))[12; 10 11 9 -6 -3 8]) + return mat / norm(mat) +end + +function semi_renormalize_edge(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, Pas, Pbs, dir, r, c) where {P <: PFTensorTriangular} + @tensor opt = true mat[-1 -2; -3 -4] := + Pas[dir, r, c][-1; 1 2] * + Pbs[dir, r, c][6 7; -3] * + env.Ea[_coordinates(dir, 0, r, c, size(network))...][1 3; 4] * + env.Eb[_coordinates(dir, 1, r, c, size(network))...][4 5; 6] * + rotl60(network[r, c], mod(dir - 1, 6))[2 -2 -4; 3 5 7] + return mat / norm(mat) +end + +function build_halfinfinite_projectors(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, Ẽas, Ẽbs, Ẽastr, Ẽbstr, dir, r, c) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + O_rp2_c = network[_coordinates(dir + 2, 0, r, c, size(network))[2:3]...] + @tensor opt = true σL[-1 -2 -3; -4] := + env.C[_coordinates(dir, 0, r, c, size(network))...][1 2 10; 8] * + env.C[_coordinates(dir - 1, 0, r, c, size(network))...][4 3 11; 1] * + env.C[_coordinates(dir - 2, 0, r, c, size(network))...][6 5 12; 4] * + Ẽbs[_coordinates(dir, 1, r, c, size(network))...][8 9 13; -4] * + Ẽastr[_coordinates(dir - 3, 0, r, c, size(network))...][-1 7 14; 6] * + rotl60(ket(O_r_c), mod(dir - 1, 6))[15 2 9 -2 7 5 3] * + conj(rotl60(bra(O_r_c), mod(dir - 1, 6))[15; 10 13 -3 14 12 11]) + @tensor opt = true σR[-1; -2 -3 -4] := + env.C[_coordinates(dir + 1, 0, r, c, size(network))...][8 2 10; 1] * + env.C[_coordinates(dir + 2, 0, r, c, size(network))...][1 3 11; 4] * + env.C[_coordinates(dir + 3, 0, r, c, size(network))...][4 5 12; 6] * + Ẽbstr[_coordinates(dir + 3, -1, r, c, size(network))...][6 7 13; -2] * + Ẽas[_coordinates(dir, 0, r, c, size(network))...][-1 9 14; 8] * + rotl60(ket(O_rp2_c), mod(dir - 1, 6))[15; 9 2 3 5 7 -3] * + conj(rotl60(bra(O_rp2_c), mod(dir - 1, 6))[15; 14 10 11 12 13 -4]) + return σL, σR +end + +function build_halfinfinite_projectors(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, Ẽas, Ẽbs, Ẽastr, Ẽbstr, dir, r, c) where {P <: PFTensorTriangular} + @tensor opt = true σL[-1 -2; -3] := + env.C[_coordinates(dir, 0, r, c, size(network))...][1 2; 8] * + env.C[_coordinates(dir - 1, 0, r, c, size(network))...][4 3; 1] * + env.C[_coordinates(dir - 2, 0, r, c, size(network))...][6 5; 4] * + Ẽbs[_coordinates(dir, 1, r, c, size(network))...][8 9; -3] * + Ẽastr[_coordinates(dir - 3, 0, r, c, size(network))...][-1 7; 6] * + rotl60(network[r, c], mod(dir - 1, 6))[3 5 7; 2 9 -2] + @tensor opt = true σR[-1; -2 -3] := + env.C[_coordinates(dir + 1, 0, r, c, size(network))...][8 2; 1] * + env.C[_coordinates(dir + 2, 0, r, c, size(network))...][1 3; 4] * + env.C[_coordinates(dir + 3, 0, r, c, size(network))...][4 5; 6] * + Ẽbstr[_coordinates(dir + 3, -1, r, c, size(network))...][6 7; -2] * + Ẽas[_coordinates(dir, 0, r, c, size(network))...][-1 9; 8] * + rotl60(network[_coordinates(dir + 2, 0, r, c, size(network))[2:3]...], mod(dir - 1, 6))[-3 7 5; 9 2 3] + return σL, σR +end + +function renormalize_corner_triangular((dir, r, c), network::InfiniteTriangularNetwork{P}, env, Pas, Pbs) where {P <: PFTensorTriangular} + @tensor opt = true env_C_new[-1 -2; -3] := + env.C[_coordinates(dir, 0, r, c, size(network))...][1 3; 6] * + env.Ea[_coordinates(dir - 1, 0, r, c, size(network))...][4 2; 1] * + env.Eb[_coordinates(dir, 1, r, c, size(network))...][6 7; 8] * + rotl60(network[r, c], mod(dir - 1, 6))[2 5 -2; 3 7 9] * + Pas[mod1(dir - 1, 6), r, c][-1; 4 5] * Pbs[dir, r, c][8 9; -3] + return env_C_new +end + +function renormalize_corner_triangular((dir, r, c), network::InfiniteTriangularNetwork{P}, env, Pas, Pbs) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + @tensor opt = true env_C_new[-1 -2 -3; -4] := + env.C[_coordinates(dir, 0, r, c, size(network))...][1 3 10; 6] * + env.Ea[_coordinates(dir - 1, 0, r, c, size(network))...][4 2 11; 1] * + env.Eb[_coordinates(dir, 1, r, c, size(network))...][6 7 12; 8] * + rotl60(ket(O_r_c), mod(dir - 1, 6))[15; 3 7 9 -2 5 2] * + conj(rotl60(bra(O_r_c), mod(dir - 1, 6))[15; 10 12 14 -3 13 11]) * + Pas[mod1(dir - 1, 6), r, c][-1; 4 5 13] * Pbs[dir, r, c][8 9 14; -4] + return env_C_new +end diff --git a/src/algorithms/ctmrg_triangular/ctmrg.jl b/src/algorithms/ctmrg_triangular/ctmrg.jl new file mode 100644 index 000000000..3e1eff763 --- /dev/null +++ b/src/algorithms/ctmrg_triangular/ctmrg.jl @@ -0,0 +1,50 @@ +abstract type CTMRGAlgorithmTriangular end + +function leading_boundary(env₀::CTMRGEnvTriangular, network::InfiniteTriangularNetwork; kwargs...) + alg = select_algorithm(leading_boundary, env₀; kwargs...) + return leading_boundary(env₀, network, alg) +end +function leading_boundary( + env₀::CTMRGEnvTriangular, network::InfiniteTriangularNetwork, alg::CTMRGAlgorithmTriangular + ) + log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) + return with_verbosity(; alg.verbosity) do + env = deepcopy(env₀) + S_old = DiagonalTensorMap.(id.(domain.(env₀.C))) + η = one(real(scalartype(network))) + ctmrg_loginit!(log, η, network, env₀) + local info + for iter in 1:(alg.maxiter) + env, S = ctmrg_iteration(network, env, alg) # Grow and renormalize in all 4 directions + η = calc_convergence(S, S_old) + S_old = copy(S) + + if η ≤ alg.tol && iter ≥ alg.miniter + ctmrg_logfinish!(log, iter, η, network, env) + break + end + if iter == alg.maxiter + ctmrg_logcancel!(log, iter, η, network, env) + else + ctmrg_logiter!(log, iter, η, network, env) + end + info = (; convergence_metric = η) + end + return env, info + end +end +function leading_boundary(env₀::CTMRGEnvTriangular, state, args...; kwargs...) + return leading_boundary(env₀, InfiniteTriangularNetwork(state), args...; kwargs...) +end + +function calc_convergence(S_new::Array{T, 3}, S_old::Array{T, 3}) where {E, S, T <: DiagonalTensorMap{E, S}} + ε = Inf + function dist(S1, S2) + if space(S1) == space(S2) + return norm(S1^4 - S2^4) + else + return Inf + end + end + return maximum(dist.(S_new, S_old)) +end diff --git a/src/algorithms/ctmrg_triangular/simultaneous.jl b/src/algorithms/ctmrg_triangular/simultaneous.jl new file mode 100644 index 000000000..90e1321bb --- /dev/null +++ b/src/algorithms/ctmrg_triangular/simultaneous.jl @@ -0,0 +1,237 @@ +struct SimultaneousCTMRGTriangular <: CTMRGAlgorithmTriangular + tol::Float64 + maxiter::Int + miniter::Int + verbosity::Int + conditioning::Bool + projector_alg::Symbol + trunctype::Symbol +end +function SimultaneousCTMRGTriangular(; + tol = Defaults.ctmrg_tol, + maxiter = Defaults.ctmrg_maxiter, + miniter = Defaults.ctmrg_miniter, + verbosity = Defaults.ctmrg_verbosity, + conditioning = true, + projector_alg = :twothirds, + trunctype = :truncrank + ) + return SimultaneousCTMRGTriangular(tol, maxiter, miniter, verbosity, conditioning, projector_alg, trunctype) +end + +# Based on +# https://arxiv.org/pdf/2510.04907 + +function ctmrg_iteration(network::InfiniteTriangularNetwork, env::CTMRGEnvTriangular, alg::SimultaneousCTMRGTriangular) + trunc_type = getfield(@__MODULE__, alg.trunctype) + if trunc_type == FixedSpaceTruncation + trunc = FixedSpaceTruncation() + else + trunc = trunc_type(dim(domain(env.C[1, 1, 1])[1])) + end + Pas, Pbs, S = calculate_projectors(network, env, trunc, alg.projector_alg) + + env = renormalize_corners!(network, env, Pas, Pbs) + env = normalize_corners(env) + + Ẽas, Ẽbs, Ẽastr, Ẽbstr = semi_renormalize(network, env, Pas, Pbs, trunc) + Qas, Qbs = build_matrix_second_projectors(network, env, Ẽas, Ẽbs, Ẽastr, Ẽbstr, trunc; alg.conditioning) + + env = renormalize_edges(env, Ẽas, Ẽbs, Qas, Qbs) + env = normalize_edges(env) + return env, S +end + +function calculate_projectors(network, env, trunc, projector_alg) + if projector_alg == :full + return calculate_full_projectors(network, env, trunc) + elseif projector_alg == :twothirds + return calculate_twothirds_projectors(network, env, trunc) + else + @error "projector_alg = $projector_alg not defined" + end +end + +# TODO: remove once this once everything is plumbed through a proper projector algorithm +_truncation_strategy(trunc::TruncationStrategy, _) = trunc +function _truncation_strategy(::FixedSpaceTruncation, edge) + tspace = space(edge, 1) + return isdual(tspace) ? truncspace(flip(tspace)) : truncspace(tspace) +end + +function calculate_twothirds_projectors(network::InfiniteTriangularNetwork{P}, env, trunc) where {P} + projectors = dtmap(eachcoordinate(network, 1:6)) do (dir, r, c) + ρL = build_double_corner_matrix_triangular(network, env, _prev(dir, 6), r, c) + ρR = build_double_corner_matrix_triangular(network, env, _next(dir, 6), r, c) + ρρ = ρL * ρR + ρρ /= norm(ρρ) + + trunc = _truncation_strategy(trunc, env.Ea[dir, r, c]) + U, S, V = svd_trunc(ρρ; trunc) + iqsrtS = sdiag_pow(S, -0.5) + + Pb = ρR * V' * iqsrtS + Pa = iqsrtS * U' * ρL + return Pa, Pb, S + end + return getindex.(projectors, 1), getindex.(projectors, 2), getindex.(projectors, 3) +end + +function calculate_full_projectors(network::InfiniteTriangularNetwork{P}, env, trunc) where {P} + projectors = dtmap(eachcoordinate(network, 1:6)) do (dir, r, c) + ρL = build_double_corner_matrix_triangular(network, env, mod1(dir - 1, 6), r, c) + ρR = build_double_corner_matrix_triangular(network, env, mod1(dir + 1, 6), r, c) + ρ̄ = build_double_corner_matrix_triangular(network, env, mod1(dir + 3, 6), r, c) + ρ̄ /= norm(ρ̄) + Ū, S̄, V̄ᴴ = svd_full(ρ̄) + sqrtS̄ = sdiag_pow(S̄, 0.5) + ρ̄ᴿ = Ū * sqrtS̄ + ρ̄ᴸ = sqrtS̄ * V̄ᴴ + ρρ = ρ̄ᴸ * ρL * ρR * ρ̄ᴿ + ρρ /= norm(ρρ) + + trunc = _truncation_strategy(trunc, env.Ea[dir, r, c]) + U, S, Vᴴ = svd_trunc(ρρ; trunc) + isqrtS = sdiag_pow(S, -0.5) + Pb = ρR * ρ̄ᴿ * Vᴴ' * isqrtS + Pa = isqrtS * U' * ρ̄ᴸ * ρL + + return Pa, Pb, S + end + return getindex.(projectors, 1), getindex.(projectors, 2), getindex.(projectors, 3) +end + +function renormalize_corners!(network::InfiniteTriangularNetwork{P}, env, Pas, Pbs) where {P <: PEPSSandwichTriangular} + coordinates = eachcoordinate(network, 1:6) + new_corners′ = similar(env.C) + new_corners = dtmap!!(new_corners′, coordinates) do (dir, r, c) + return renormalize_corner_triangular((dir, r, c), network, env, Pas, Pbs) + end + return CTMRGEnvTriangular(new_corners, env.Ea, env.Eb) +end + +function renormalize_corners!(network::InfiniteTriangularNetwork{P}, env, Pas, Pbs) where {P <: PFTensorTriangular} + coordinates = eachcoordinate(network, 1:6) + new_corners′ = similar(env.C) + new_corners = dtmap!!(new_corners′, coordinates) do (dir, r, c) + return renormalize_corner_triangular((dir, r, c), network, env, Pas, Pbs) + end + return CTMRGEnvTriangular(new_corners, env.Ea, env.Eb) +end + +function _permute_edge(t::Union{T1, T2}) where {E, S, T1 <: AbstractTensorMap{E, S, 2, 1}, T2 <: AbstractTensorMap{E, S, 1, 2}} + return permute(t, ((1, 3), (2,))) +end + +function _permute_edge(t::Union{T1, T2}) where {E, S, T1 <: AbstractTensorMap{E, S, 3, 1}, T2 <: AbstractTensorMap{E, S, 1, 3}} + return permute(t, ((1, 3, 4), (2,))) +end + +function semi_renormalize(network::InfiniteTriangularNetwork, env::CTMRGEnvTriangular, Pas, Pbs, trunc) + projectors = dtmap(eachcoordinate(network, 1:6)) do (dir, r, c) + mat = semi_renormalize_edge(network, env, Pas, Pbs, dir, r, c) + + U, S, V = svd_full(mat) + + sqrtS = sdiag_pow(S, 0.5) + + Ẽb = U * sqrtS + Ẽa = _permute_edge(sqrtS * V) + + trunc = _truncation_strategy(trunc, env.Ea[dir, r, c]) + Utr, Str, Vtr = svd_trunc(mat; trunc) + sqrt_Str = sdiag_pow(Str, 0.5) + Ẽbtr = Utr * sqrt_Str + Ẽatr = _permute_edge(sqrt_Str * Vtr) + + return Ẽa, Ẽb, Ẽatr, Ẽbtr + end + return getindex.(projectors, 1), getindex.(projectors, 2), getindex.(projectors, 3), getindex.(projectors, 4) +end + +function build_matrix_second_projectors(network::InfiniteTriangularNetwork, env::CTMRGEnvTriangular, Ẽas, Ẽbs, Ẽastr, Ẽbstr, trunc; conditioning = true) + projectors = dtmap(eachcoordinate(network, 1:6)) do (dir, r, c) + trunc = _truncation_strategy(trunc, env.Ea[dir, r, c]) + σL, σR = build_halfinfinite_projectors(network, env, Ẽas, Ẽbs, Ẽastr, Ẽbstr, dir, r, c) + if conditioning + σL /= norm(σL) + σR /= norm(σR) + _, SL, VLᴴ = svd_full(σL) + UR, SR, _ = svd_full(σR) + + sqrtSL = sdiag_pow(SL, 0.5) + sqrtSR = sdiag_pow(SR, 0.5) + + FLU = sqrtSL * VLᴴ + FRU = UR * sqrtSR + + mat = FLU * FRU + mat /= norm(mat) + WU, SU, QUᴴ = svd_trunc(mat; trunc) + isqrtSU = sdiag_pow(SU, -0.5) + + Qa = isqrtSU * WU' * FLU + Qb = FRU * QUᴴ' * isqrtSU + else + mat = σL * σR + mat /= norm(mat) + U, S, V = svd_trunc(mat; trunc) + isqrtS = sdiag_pow(S, -0.5) + Qa = isqrtS * U' * σL + Qb = σR * V' * isqrtS + end + return Qa, Qb + end + PL = map(x -> x[1], projectors) + PR = map(x -> x[2], projectors) + return PL, PR +end + +function renormalize_edges(env::CTMRGEnvTriangular, Ẽas::Array{T1, 3}, Ẽbs::Array{T1, 3}, Qas::Array{T2, 3}, Qbs::Array{T2, 3}) where {E, S, T1 <: AbstractTensorMap{E, S, 3, 1}, T2 <: AbstractTensorMap{E, S, 1, 1}} + new_edges = dtmap(eachcoordinate(env, 1:6)) do (dir, r, c) + Eb_new = Ẽbs[dir, r, c] * Qbs[dir, r, c] + Ea_new = permute(Qas[dir, r, c] * permute(Ẽas[dir, r, c], ((1,), (2, 3, 4))), ((1, 2, 3), (4,))) + return Ea_new, Eb_new + end + Ea = map(x -> x[1], new_edges) + Eb = map(x -> x[2], new_edges) + return CTMRGEnvTriangular(env.C, Ea, Eb) +end + +function renormalize_edges(env::CTMRGEnvTriangular, Ẽas::Array{T1, 3}, Ẽbs::Array{T1, 3}, Qas::Array{T2, 3}, Qbs::Array{T2, 3}) where {T1 <: CTMRGEdgeTensor, T2 <: CTMRGCornerTensor} + new_edges = dtmap(eachcoordinate(env, 1:6)) do (dir, r, c) + Eb_new = Ẽbs[dir, r, c] * Qbs[dir, r, c] + Ea_new = permute(Qas[dir, r, c] * permute(Ẽas[dir, r, c], ((1,), (2, 3))), ((1, 2), (3,))) + return Ea_new, Eb_new + end + return CTMRGEnvTriangular(env.C, getindex.(new_edges, 1), getindex.(new_edges, 2)) +end + +function normalize_corners(env::CTMRGEnvTriangular) + C_normalized = map(env.C) do C + return C / norm(C) + end + return CTMRGEnvTriangular(C_normalized, env.Ea, env.Eb) +end + +function normalize_edges(env::CTMRGEnvTriangular) + Ea_normalized = map(env.Ea) do e + return e / norm(e) + end + Eb_normalized = map(env.Eb) do e + return e / norm(e) + end + return CTMRGEnvTriangular(env.C, Ea_normalized, Eb_normalized) +end + +function calculate_error(Ss, Ss_prev) + ε = Inf + for (S, S_prev) in zip(Ss, Ss_prev) + if space(S) == space(S_prev) + ε = norm(S^4 - S_prev^4) + else + return Inf + end + end + return ε +end diff --git a/src/algorithms/ctmrg_triangular/utility.jl b/src/algorithms/ctmrg_triangular/utility.jl new file mode 100644 index 000000000..faf1e7066 --- /dev/null +++ b/src/algorithms/ctmrg_triangular/utility.jl @@ -0,0 +1,60 @@ +function _coordinates(dir::Int, rot::Int, r::Int, c::Int, unitcell::Tuple{Int, Int}) + rows, cols = unitcell + if mod1(dir + rot, 6) == 1 + return (mod1(dir, 6), mod1(r - 1, rows), mod1(c, cols)) + elseif mod1(dir + rot, 6) == 2 + return (mod1(dir, 6), mod1(r - 1, rows), mod1(c + 1, cols)) + elseif mod1(dir + rot, 6) == 3 + return (mod1(dir, 6), mod1(r, rows), mod1(c + 1, cols)) + elseif mod1(dir + rot, 6) == 4 + return (mod1(dir, 6), mod1(r + 1, rows), mod1(c, cols)) + elseif mod1(dir + rot, 6) == 5 + return (mod1(dir, 6), mod1(r + 1, rows), mod1(c - 1, cols)) + elseif mod1(dir + rot, 6) == 6 + return (mod1(dir, 6), mod1(r, rows), mod1(c - 1, cols)) + end +end + +function _coordinates(dir::Int, r::Int, c::Int, unitcell::Tuple{Int, Int}) + return _coordinates(dir, 0, r, c; unitcell)[2:3] +end + +function _truncway(trunctype::Symbol, D::E) where {E <: ElementarySpace} + if trunctype == :truncdim + return truncdim(dim(D)) + else + return notrunc() + end +end + +function rotr60(A::T, dir::Int) where {E, S, T <: AbstractTensorMap{E, S, 1, 6}} + if dir == 0 + return A + else + return rotr60(permute(A, ((1,), (7, 2, 3, 4, 5, 6))), dir - 1) + end +end + +function rotr60(A::T, dir::Int) where {E, S, T <: AbstractTensorMap{E, S, 3, 3}} + if dir == 0 + return A + else + return rotr60(permute(A, ((2, 3, 6), (1, 4, 5))), dir - 1) + end +end + +function rotl60(A::T, dir::Int) where {E, S, T <: AbstractTensorMap{E, S, 1, 6}} + if dir == 0 + return A + else + return rotl60(permute(A, ((1,), (3, 4, 5, 6, 7, 2))), dir - 1) + end +end + +function rotl60(A::T, dir::Int) where {E, S, T <: AbstractTensorMap{E, S, 3, 3}} + if dir == 0 + return A + else + return rotl60(permute(A, ((4, 1, 2), (5, 6, 3))), dir - 1) + end +end diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/fixed_point_differentiation.jl new file mode 100644 index 000000000..c8005f7fd --- /dev/null +++ b/src/algorithms/optimization/fixed_point_differentiation.jl @@ -0,0 +1,527 @@ +abstract type GradMode{F} end + +const GRADIENT_MODE_SYMBOLS = IdDict{Symbol, Type{<:GradMode}}() +const LINSOLVER_SOLVER_SYMBOLS = IdDict{Symbol, Type{<:KrylovKit.LinearSolver}}( + :gmres => GMRES, :bicgstab => BiCGStab +) +const EIGSOLVER_SOLVER_SYMBOLS = IdDict{Symbol, Type{<:KrylovKit.KrylovAlgorithm}}( + :arnoldi => Arnoldi +) + +""" + GradMode(; kwargs...) + +Keyword argument parser returning the appropriate `GradMode` algorithm struct. +""" +function GradMode(; + alg = Defaults.gradient_alg, + tol = Defaults.gradient_tol, + maxiter = Defaults.gradient_maxiter, + verbosity = Defaults.gradient_verbosity, + iterscheme = Defaults.gradient_iterscheme, + solver_alg = (;), + ) + # replace symbol with GradMode alg type + haskey(GRADIENT_MODE_SYMBOLS, alg) || + throw(ArgumentError("unknown GradMode algorithm: $alg")) + alg_type = GRADIENT_MODE_SYMBOLS[alg] + + # parse GradMode algorithm + gradient_algorithm = if alg_type <: Union{GeomSum, ManualIter} + alg_type{iterscheme}(tol, maxiter, verbosity) + elseif alg_type <: Union{<:LinSolver, <:EigSolver} + solver = if solver_alg isa NamedTuple # determine linear/eigen solver algorithm + solver_kwargs = (; tol, maxiter, verbosity, solver_alg...) + + solver_type = if alg_type <: LinSolver # replace symbol with solver alg type + solver_kwargs = (; alg = Defaults.gradient_linsolver, solver_kwargs...) + haskey(LINSOLVER_SOLVER_SYMBOLS, solver_kwargs.alg) || throw( + ArgumentError("unknown LinSolver solver: $(solver_kwargs.alg)"), + ) + LINSOLVER_SOLVER_SYMBOLS[solver_kwargs.alg] + elseif alg_type <: EigSolver + solver_kwargs = (; + alg = Defaults.gradient_eigsolver, + eager = Defaults.gradient_eigsolver_eager, + solver_kwargs..., + ) + haskey(EIGSOLVER_SOLVER_SYMBOLS, solver_kwargs.alg) || throw( + ArgumentError("unknown EigSolver solver: $(solver_kwargs.alg)"), + ) + EIGSOLVER_SOLVER_SYMBOLS[solver_kwargs.alg] + end + + solver_kwargs = Base.structdiff(solver_kwargs, (; alg = nothing)) # remove `alg` keyword argument + solver_type(; solver_kwargs...) + else + solver_alg + end + + alg_type{iterscheme}(solver) + else + throw(ArgumentError("unknown gradient algorithm: $alg")) + end + + return gradient_algorithm +end + +iterscheme(::GradMode{F}) where {F} = F + +""" +$(TYPEDEF) + +Gradient mode for CTMRG using explicit evaluation of the geometric sum. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + GeomSum(; kwargs...) + +Construct the `GeomSum` algorithm struct based on the following keyword arguments: + +* `tol::Real=$(Defaults.gradient_tol)` : Convergence tolerance for the difference of norms of two consecutive summands in the geometric sum. +* `maxiter::Int=$(Defaults.gradient_maxiter)` : Maximal number of gradient iterations. +* `verbosity::Int=$(Defaults.gradient_verbosity)` : Output information verbosity that can be one of the following: + 0. Suppress output information + 1. Print convergence warnings + 2. Information at each gradient iteration +* `iterscheme::Symbol=:$(Defaults.gradient_iterscheme)` : Style of CTMRG iteration which is being differentiated, which can be: + - `:fixed` : the differentiated CTMRG iteration uses a pre-computed SVD with a fixed set of gauges + - `:diffgauge` : the differentiated iteration consists of a CTMRG iteration and a subsequent gauge-fixing step such that the gauge-fixing procedure is differentiated as well +""" +struct GeomSum{F} <: GradMode{F} + tol::Real + maxiter::Int + verbosity::Int +end +GeomSum(; kwargs...) = GradMode(; alg = :geomsum, kwargs...) + +GRADIENT_MODE_SYMBOLS[:geomsum] = GeomSum + +""" +$(TYPEDEF) + +Gradient mode for CTMRG using manual iteration to solve the linear problem. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + ManualIter(; kwargs...) + +Construct the `ManualIter` algorithm struct based on the following keyword arguments: + +* `tol::Real=$(Defaults.gradient_tol)` : Convergence tolerance for the norm difference of two consecutive `dx` contributions. +* `maxiter::Int=$(Defaults.gradient_maxiter)` : Maximal number of gradient iterations. +* `verbosity::Int=$(Defaults.gradient_verbosity)` : Output information verbosity that can be one of the following: + 0. Suppress output information + 1. Print convergence warnings + 2. Information at each gradient iteration +* `iterscheme::Symbol=:$(Defaults.gradient_iterscheme)` : Style of CTMRG iteration which is being differentiated, which can be: + - `:fixed` : the differentiated CTMRG iteration uses a pre-computed SVD with a fixed set of gauges + - `:diffgauge` : the differentiated iteration consists of a CTMRG iteration and a subsequent gauge-fixing step such that the gauge-fixing procedure is differentiated as well +""" +struct ManualIter{F} <: GradMode{F} + tol::Real + maxiter::Int + verbosity::Int +end +ManualIter(; kwargs...) = GradMode(; alg = :manualiter, kwargs...) + +GRADIENT_MODE_SYMBOLS[:manualiter] = ManualIter + +""" +$(TYPEDEF) + +Gradient mode wrapper around `KrylovKit.LinearSolver` for solving the gradient linear +problem using iterative solvers. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + LinSolver(; kwargs...) + +Construct the `LinSolver` algorithm struct based on the following keyword arguments: + +* `tol::Real=$(Defaults.gradient_tol)` : Convergence tolerance of the linear solver. +* `maxiter::Int=$(Defaults.gradient_maxiter)` : Maximal number of solver iterations. +* `verbosity::Int=$(Defaults.gradient_verbosity)` : Output information verbosity of the linear solver. +* `iterscheme::Symbol=:$(Defaults.gradient_iterscheme)` : Style of CTMRG iteration which is being differentiated, which can be: + - `:fixed` : the differentiated CTMRG iteration uses a pre-computed SVD with a fixed set of gauges + - `:diffgauge` : the differentiated iteration consists of a CTMRG iteration and a subsequent gauge-fixing step such that the gauge-fixing procedure is differentiated as well +* `solver_alg::Union{KrylovKit.LinearSolver,NamedTuple}=(; alg::Symbol=:$(Defaults.gradient_linsolver)` : Linear solver algorithm which, if supplied directly as a `KrylovKit.LinearSolver` overrides the above specified `tol`, `maxiter` and `verbosity`. Alternatively, it can be supplied via a `NamedTuple` where `alg` can be one of the following: + - `:gmres` : GMRES iterative linear solver, see [`KrylovKit.GMRES`](@extref) for details + - `:bicgstab` : BiCGStab iterative linear solver, see [`KrylovKit.BiCGStab`](@extref) for details +""" +struct LinSolver{F} <: GradMode{F} + solver_alg::KrylovKit.LinearSolver +end +LinSolver(; kwargs...) = GradMode(; alg = :linsolver, kwargs...) + +GRADIENT_MODE_SYMBOLS[:linsolver] = LinSolver + +""" +$(TYPEDEF) + +Gradient mode wrapper around `KrylovKit.KrylovAlgorithm` for solving the gradient linear +problem as an eigenvalue problem. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + EigSolver(; kwargs...) + +Construct the `EigSolver` algorithm struct based on the following keyword arguments: + +* `tol::Real=$(Defaults.gradient_tol)` : Convergence tolerance of the eigen solver. +* `maxiter::Int=$(Defaults.gradient_maxiter)` : Maximal number of solver iterations. +* `verbosity::Int=$(Defaults.gradient_verbosity)` : Output information verbosity of the linear solver. +* `iterscheme::Symbol=:$(Defaults.gradient_iterscheme)` : Style of CTMRG iteration which is being differentiated, which can be: + - `:fixed` : the differentiated CTMRG iteration uses a pre-computed SVD with a fixed set of gauges + - `:diffgauge` : the differentiated iteration consists of a CTMRG iteration and a subsequent gauge-fixing step such that the gauge-fixing procedure is differentiated as well +* `solver_alg::Union{KrylovKit.KrylovAlgorithm,NamedTuple}=(; alg=:$(Defaults.gradient_eigsolver)` : Eigen solver algorithm which, if supplied directly as a `KrylovKit.KrylovAlgorithm` overrides the above specified `tol`, `maxiter` and `verbosity`. Alternatively, it can be supplied via a `NamedTuple` where `alg` can be one of the following: + - `:arnoldi` : Arnoldi Krylov algorithm, see [`KrylovKit.Arnoldi`](@extref) for details +""" +struct EigSolver{F} <: GradMode{F} + solver_alg::KrylovKit.KrylovAlgorithm +end +EigSolver(; kwargs...) = GradMode(; alg = :eigsolver, kwargs...) + +GRADIENT_MODE_SYMBOLS[:eigsolver] = EigSolver + +""" + _check_algorithm_combination(boundary_alg, gradient_alg_or_symmetrization) + _check_algorithm_combination(boundary_alg, gradient_alg, symmetrization) + +Check for allowed combinations of gradient algorithm and boundary algorithm to be used for +computing the gradient of a `leading_boundary` call. Throws an error containing a +recommended fix if the combination is not allowed or broken. +""" +function _check_algorithm_combination(boundary_alg, gradient_alg_or_symmetrization) end +function _check_algorithm_combination(boundary_alg, gradient_alg, symmetrization) + _check_algorithm_combination(boundary_alg, gradient_alg) + _check_algorithm_combination(boundary_alg, symmetrization) + return nothing +end +function _check_algorithm_combination(::SequentialCTMRG, ::GradMode{:fixed}) + msg = "`:fixed` mode is not compatible with `SequentialCTMRG` since the sequential \ + application of SVDs does not allow to differentiate through a fixed set of \ + gauges; select SimultaneousCTMRG instead to use :fixed mode" + throw(ArgumentError(msg)) +end +function _check_algorithm_combination(::C4vCTMRG{<:C4vEighProjector}, ::GradMode{:diffgauge}) + msg = "`:diffgauge` mode is currently not compatible with eigh-based C4v CTMRG; \ + either switch to a different projector algorithm (e.g. `c4v_qr`), or use :fixed \ + mode for differentiation instead." + throw(ArgumentError(msg)) +end +function _check_algorithm_combination(::C4vCTMRG, symm::Union{Nothing, <:SymmetrizationStyle}) + if !(symm isa RotateReflect) + msg = "C4vCTMRG optimization is compatible only with RotateReflect symmetrization. \ + Make sure to set `symmetrization = RotateReflect()`." + throw(ArgumentError(msg)) + end + return nothing +end + +#= +Evaluating the gradient of the cost function for CTMRG: +- The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum. +- With AD, the gradient is computed by differentiating the cost function with respect to the PEPS tensors, including computing the environment tensors. +- With explicit evaluation of the geometric sum, the gradient is computed by differentiating the cost function with the environment kept fixed, and then manually adding the gradient contributions from the environments. +=# + +_scrambling_env_gauge(::CTMRGAlgorithm) = ScramblingEnvGauge() +_scrambling_env_gauge(::C4vCTMRG) = ScramblingEnvGaugeC4v() + +function _set_fixed_truncation(alg::CTMRGAlgorithm) + alg_fixed = @set alg.projector_alg = _set_truncation(alg.projector_alg, FixedSpaceTruncation()) + return alg_fixed +end + +function _rrule( + gradmode::GradMode{:diffgauge}, + config::RuleConfig, + ::typeof(leading_boundary), + envinit, + state, + alg::CTMRGAlgorithm, + ) + _check_algorithm_combination(alg, gradmode) + + env, info = leading_boundary(envinit, state, alg) + alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation + alg_gauge = _scrambling_env_gauge(alg) # TODO: make this a field in GradMode? + + # prepare iterating function corresponding to a single gauge-fixed CTMRG iteration + function f(A, x) + return gauge_fix(ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1], x, alg_gauge)[1] + end + # compute its pullback + _, env_vjp = rrule_via_ad(config, f, state, env) + # split off state and environment parts + ∂f∂A(x)::typeof(state) = env_vjp(x)[2] + ∂f∂x(x)::typeof(env) = env_vjp(x)[3] + + function leading_boundary_diffgauge_pullback((Δenv′, Δinfo)) + Δenv = unthunk(Δenv′) + + # evaluate the geometric sum + ∂F∂env = fpgrad(Δenv, ∂f∂x, ∂f∂A, Δenv, gradmode) + + return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() + end + + return (env, info), leading_boundary_diffgauge_pullback +end + +function _rrule( + gradmode::GradMode, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit, + state, + alg::CTMRGAlgorithmTriangular, + ) + msg = "Fixed-point differentiation is not available for `CTMRGAlgorithmTriangular`, \ + since gauge fixing for triangular CTMRG is currently not implemented." + throw(ArgumentError(msg)) +end +# Here f is differentiated from an pre-computed SVD with fixed U, S and V +function _rrule( + gradmode::GradMode{:fixed}, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit, + state, + alg::CTMRGAlgorithm, + ) + _check_algorithm_combination(alg, gradmode) + + env, = leading_boundary(envinit, state, alg) + alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation + alg_gauge = _scrambling_env_gauge(alg) # TODO: make this a field in GradMode? + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) + _, signs = gauge_fix(env_conv, env, alg_gauge) + + # fix decomposition + alg_fixed = gauge_fix(alg, signs, info) + + # prepare iterating function corresponding to a single CTMRG iteration with a + # gauge-fixed projector + function f(A, x) + return fix_global_phases( + ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1], x, + ) + end + # prepare its pullback + _, env_vjp = rrule_via_ad(config, f, state, env) + # split off state and environment parts + ∂f∂A(x)::typeof(state) = env_vjp(x)[2] + ∂f∂x(x)::typeof(env) = env_vjp(x)[3] + + function leading_boundary_fixed_pullback((Δenv′, Δinfo)) + Δenv = unthunk(Δenv′) + + # evaluate the geometric sum + ∂F∂env = fpgrad(Δenv, ∂f∂x, ∂f∂A, Δenv, gradmode) + + return NoTangent(), ZeroTangent(), ∂F∂env, NoTangent() + end + + return (env, info), leading_boundary_fixed_pullback +end + +function gauge_fix(alg::SVDAdjoint, signs, info) + # embed gauge signs in larger space to fix gauge of full U and V on truncated subspace + rowsize, colsize = size(signs, 2), size(signs, 3) + inds = info.truncation_indices + signs_full = map(Iterators.product(1:4, 1:rowsize, 1:colsize)) do (dir, row, col) + σ = signs[dir, row, col] + row_sign, col_sign = if dir == NORTH # take unit cell interdependency of signs into account + row, _prev(col, colsize) + elseif dir == EAST + _prev(row, rowsize), col + elseif dir == SOUTH + row, _next(col, colsize) + elseif dir == WEST + _next(row, rowsize), col + end + + ind = inds[dir, row_sign, col_sign] + extended_σ = id(scalartype(σ), domain(info.S_full[dir, row_sign, col_sign])) + for (c, b) in blocks(σ) + I = get(ind, c, nothing) + @assert !isnothing(I) + block(extended_σ, c)[I, I] = b + end + return extended_σ + end + + # fix kept and full U and V + U_fixed, V_fixed = fix_relative_phases(info.U, info.V, signs) + U_full_fixed, V_full_fixed = fix_relative_phases(info.U_full, info.V_full, signs_full) + return SVDAdjoint(; + fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, U_full_fixed, info.S_full, V_full_fixed, inds), + rrule_alg = alg.rrule_alg, + ) +end +function gauge_fix(alg::SVDAdjoint{F}, signs, info) where {F <: IterSVD} + # fix kept U and V only since iterative SVD doesn't have access to full spectrum + U_fixed, V_fixed = fix_relative_phases(info.U, info.V, signs) + return SVDAdjoint(; + fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, nothing, nothing, nothing, nothing), + rrule_alg = alg.rrule_alg, + ) +end +function gauge_fix(alg::EighAdjoint, signs, info) + σ = signs[1] + inds = info.truncation_indices + + # embed gauge signs in larger space to fix gauge of full V on truncated subspace + extended_σ = id(scalartype(σ), domain(info.D_full)) + for (c, b) in blocks(σ) + I = get(inds, c, nothing) + @assert !isnothing(I) + block(extended_σ, c)[I, I] = b + end + + # fix kept and full V + V_fixed = info.V * σ' + V_full_fixed = info.V_full * extended_σ' + return EighAdjoint(; + fwd_alg = FixedEig(info.D, V_fixed, info.D_full, V_full_fixed, inds), + rrule_alg = alg.rrule_alg, + ) +end +function gauge_fix(alg::EighAdjoint{F}, signs, info) where {F <: IterEigh} + # fix kept V only since iterative decomposition doesn't have access to full spectrum + V_fixed = info.V * signs[1]' + return EighAdjoint(; + fwd_alg = FixedEig(info.D, V_fixed, nothing, nothing, nothing), + rrule_alg = alg.rrule_alg, + ) +end +function gauge_fix(alg::QRAdjoint, signs, info) + Q_fixed = info.Q * signs[1]' + R_fixed = signs[1] * info.R + return QRAdjoint(; + fwd_alg = FixedQR(Q_fixed, R_fixed), + rrule_alg = alg.rrule_alg, + ) +end + +@doc """ + fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y0, alg) + +Compute the gradient of the CTMRG fixed point by solving the following equation: + +dx = ∑ₙ (∂f∂x)ⁿ ∂f∂A dA = (1 - ∂f∂x)⁻¹ ∂f∂A dA + +where `∂F∂x` is the gradient of the cost function with respect to the PEPS tensors, `∂f∂x` +is the partial gradient of the CTMRG iteration with respect to the environment tensors, +`∂f∂A` is the partial gradient of the CTMRG iteration with respect to the PEPS tensors, and +`y0` is the initial guess for the fixed-point iteration. The function returns the gradient +`dx` of the fixed-point iteration. +""" +fpgrad + +# TODO: can we construct an implementation that does not need to evaluate the vjp +# twice if both ∂f∂A and ∂f∂x are needed? +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, _, alg::GeomSum) + g = ∂F∂x + dx = ∂f∂A(g) # n = 0 term: ∂F∂x ∂f∂A + ϵ = 2 * alg.tol + for i in 1:(alg.maxiter) + g = ∂f∂x(g) + Σₙ = ∂f∂A(g) + dx += Σₙ + ϵnew = norm(Σₙ) # TODO: normalize this error? + Δϵ = ϵ - ϵnew + alg.verbosity > 1 && + @printf("Gradient iter: %3d ‖Σₙ‖: %.2e Δ‖Σₙ‖: %.2e\n", i, ϵnew, Δϵ) + ϵ = ϵnew + + ϵ < alg.tol && break + if alg.verbosity > 0 && i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σₙ‖ = $ϵ" + end + end + return dx +end + +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::ManualIter) + y = deepcopy(y₀) # Do not mutate y₀ + dx = ∂f∂A(y) + ϵ = 1.0 + for i in 1:(alg.maxiter) + y′ = ∂F∂x + ∂f∂x(y) + + dxnew = ∂f∂A(y′) + ϵnew = norm(dxnew - dx) + Δϵ = ϵ - ϵnew + alg.verbosity > 1 && @printf( + "Gradient iter: %3d ‖Cᵢ₊₁-Cᵢ‖/N: %.2e Δ‖Cᵢ₊₁-Cᵢ‖/N: %.2e\n", i, ϵnew, Δϵ + ) + y = y′ + dx = dxnew + ϵ = ϵnew + + ϵ < alg.tol && break + if alg.verbosity > 0 && i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Cᵢ₊₁-Cᵢ‖ = $ϵ" + end + end + return dx +end + +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::LinSolver) + y, info = reallinsolve(∂f∂x, ∂F∂x, y₀, alg.solver_alg, 1, -1) + if alg.solver_alg.verbosity > 0 && info.converged != 1 + @warn("gradient fixed-point iteration reached maximal number of iterations:", info) + end + + return ∂f∂A(y) +end + +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, x₀, alg::EigSolver) + function f(X) + y = ∂f∂x(X[1]) + return (VI.add!!(y, ∂F∂x, X[2]), X[2]) + end + X₀ = (x₀, one(scalartype(x₀))) + _, vecs, info = realeigsolve(f, X₀, 1, :LM, alg.solver_alg) + if alg.solver_alg.verbosity > 0 && info.converged < 1 + @warn("gradient fixed-point iteration reached maximal number of iterations:", info) + end + if norm(vecs[1][2]) < 1.0e-2 * alg.solver_alg.tol + @warn "Fixed-point gradient computation using Arnoldi failed:\n\tauxiliary component should be finite but was $(vecs[1][2])\n\tpossibly the Jacobian does not have a unique eigenvalue 1" + @info "Falling back to linear solver for fixed-point gradient computation." + backup_ls_alg = _alg_or_nt( + GradMode, + (; + alg = :linsolver, + tol = alg.solver_alg.tol, + maxiter = alg.solver_alg.maxiter * alg.solver_alg.krylovdim, + verbosity = alg.solver_alg.verbosity, + iterscheme = iterscheme(alg), + ) + ) + return fpgrad(∂F∂x, ∂f∂x, ∂f∂A, x₀, backup_ls_alg) + else + y = VI.scale!!(vecs[1][1], inv(vecs[1][2])) + end + + return ∂f∂A(y) +end diff --git a/src/algorithms/toolbox_triangular.jl b/src/algorithms/toolbox_triangular.jl new file mode 100644 index 000000000..99da822fc --- /dev/null +++ b/src/algorithms/toolbox_triangular.jl @@ -0,0 +1,153 @@ +function network_value(network::InfiniteTriangularNetwork, env::CTMRGEnvTriangular) + return prod(Iterators.product(axes(network)...)) do (r, c) + nw_corners = complex(_contract_corners((r, c), network, env)) + nw_full = complex(_contract_site_large((r, c), network, env)) + nw_0 = complex(_contract_edges_0((r, c), network, env)) + nw_60 = complex(_contract_edges_60((r, c), network, env)) + nw_120 = complex(_contract_edges_120((r, c), network, env)) + return (nw_full * nw_corners^2 / (nw_0 * nw_60 * nw_120))^(1 / 3) + end +end + +function _contract_edges_0((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PFTensorTriangular} + return @tensor opt = true network[r, c][DL180 DL240 DL300; DL120 DL60 DL0] * network[r, _next(c, end)][DL0 DR240 DR300; DR120 DR60 DR0] * + env.C[1, _prev(r, end), c][χNW DL120; χNa] * env.C[2, _prev(r, end), _next(c + 1, end)][χNb DR60; χNE] * env.C[3, r, _next(c + 1, end)][χNE DR0; χSE] * + env.C[4, _next(r, end), _next(c, end)][χSE DR300; χSa] * env.C[5, _next(r, end), _prev(c, end)][χSb DL240; χSW] * env.C[6, r, _prev(c, end)][χSW DL180; χNW] * + env.Eb[1, _prev(r, end), _next(c, end)][χNa DL60; χNC] * env.Ea[1, _prev(r, end), _next(c, end)][χNC DR120; χNb] * + env.Eb[4, _next(r, end), c][χSa DR240; χSC] * env.Ea[4, _next(r, end), c][χSC DL300; χSb] +end + +function _contract_edges_60((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PFTensorTriangular} + return @tensor opt = true network[r, c][DTR180 DBL60 DTR300; DTR120 DTR60 DTR0] * network[_next(r, end), _prev(c, end)][DBL180 DBL240 DBL300; DBL120 DBL60 DBL0] * + env.C[1, _prev(r, end), c][χNWb DTR120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN DTR60; χNE] * env.C[3, r, _next(c, end)][χNE DTR0; χSEa] * + env.C[4, _next(r + 1, end), _prev(c, end)][χSEb DBL300; χS] * env.C[5, _next(r + 1, end), _prev(c - 1, end)][χS DBL240; χSW] * env.C[6, _next(r, end), _prev(c - 1, end)][χSW DBL180; χNWa] * + env.Eb[3, _next(r, end), c][χSEa DTR300; χSEC] * env.Ea[3, _next(r, end), c][χSEC DBL0; χSEb] * + env.Eb[6, r, _prev(c, end)][χNWa DBL120; χNWC] * env.Ea[6, r, _prev(c, end)][χNWC DTR180; χNWb] +end + +function _contract_edges_120((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PFTensorTriangular} + return @tensor opt = true network[r, c][DTL180 DTL240 DTL300; DTL120 DTL60 DTL0] * network[_next(r, end), c][DBR180 DBR240 DBR300; DTL300 DBR60 DBR0] * + env.C[1, _prev(r, end), c][χNW DTL120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN DTL60; χNEa] * env.C[3, _next(r, end), _next(c, end)][χNEb DBR0; χSE] * + env.C[4, _next(r + 1, end), c][χSE DBR300; χS] * env.C[5, _next(r + 1, end), _prev(c, end)][χS DBR240; χSWa] * env.C[6, r, _prev(c, end)][χSWb DTL180; χNW] * + env.Eb[2, r, _next(c, end)][χNEa DTL0; χNEC] * env.Ea[2, r, _next(c, end)][χNEC DBR60; χNEb] * + env.Eb[5, _next(r, end), _prev(c, end)][χSWa DBR180; χSWC] * env.Ea[5, _next(r, end), _prev(c, end)][χSWC DTL240; χSWb] +end + +function _contract_site_large((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PFTensorTriangular} + return @tensor opt = true network[_prev(r, end), c][DNW180 DW60 DNW300; DNW120 DNW60 DNW0] * network[_prev(r, end), _next(c, end)][DNW0 DNE240 DNE300; DNE120 DNE60 DNE0] * + network[r, _next(c, end)][DE180 DE240 DE300; DNE300 DE60 DE0] * network[_next(r, end), c][DSE180 DSE240 DSE300; DSE120 DE240 DSE0] * network[_next(r, end), _prev(c, end)][DSW180 DSW240 DSW300; DSW120 DSW60 DSE180] * + network[r, _prev(c, end)][DW180 DW240 DSW120; DW120 DW60 DW0] * network[r, c][DW0 DSW60 DSE120; DNW300 DNE240 DE180] * + env.C[1, _prev(r - 1, end), c][χNWa DNW120; χNb] * env.Eb[1, _prev(r - 1, end), _next(c, end)][χNb DNW60; χNC] * env.Ea[1, _prev(r - 1, end), _next(c, end)][χNC DNE120; χNa] * + env.C[2, _prev(r - 1, end), _next(c + 1, end)][χNa DNE60; χNEb] * env.Eb[2, _prev(r, end), _next(c + 1, end)][χNEb DNE0; χNEC] * env.Ea[2, _prev(r, end), _next(c + 1, end)][χNEC DE60; χNEa] * + env.C[3, r, _next(c + 1, end)][χNEa DE0; χSEb] * env.Eb[3, _next(r, end), _next(c, end)][χSEb DE300; χSEC] * env.Ea[3, _next(r, end), _next(c, end)][χSEC DSE0; χSEa] * + env.C[4, _next(r + 1, end), c][χSEa DSE300; χSb] * env.Eb[4, _next(r + 1, end), _prev(c, end)][χSb DSE240; χSC] * env.Ea[4, _next(r + 1, end), _prev(c, end)][χSC DSW300; χSa] * + env.C[5, _next(r + 1, end), _prev(c - 1, end)][χSa DSW240; χSWb] * env.Eb[5, _next(r, end), _prev(c - 1, end)][χSWb DSW180; χSWC] * env.Ea[5, _next(r, end), _prev(c - 1, end)][χSWC DW240; χSWa] * + env.C[6, r, _prev(c - 1, end)][χSWa DW180; χNWb] * env.Eb[6, _prev(r, end), _prev(c, end)][χNWb DW120; χNWC] * env.Ea[6, _prev(r, end), _prev(c, end)][χNWC DNW180; χNWa] +end + +function _contract_corners((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PFTensorTriangular} + return @tensor opt = true env.C[1, _prev(r, end), c][χNW D120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN D60; χNE] * env.C[3, r, _next(c, end)][χNE D0; χSE] * + env.C[4, _next(r, end), c][χSE D300; χS] * env.C[5, _next(r, end), _prev(c, end)][χS D240; χSW] * env.C[6, r, _prev(c, end)][χSW D180; χNW] * + network[r, c][D180 D240 D300; D120 D60 D0] +end + +### For + +function _contract_edges_0((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + O_r_cn = network[r, _next(c, end)] + return @tensor opt = true ket(O_r_c)[dL; DLt120 DLt60 DLt0 DLt300 DLt240 DLt180] * ket(O_r_cn)[dR; DRt120 DRt60 DRt0 DRt300 DRt240 DLt0] * + conj(bra(O_r_c)[dL; DLb120 DLb60 DLb0 DLb300 DLb240 DLb180]) * conj(bra(O_r_cn)[dR; DRb120 DRb60 DRb0 DRb300 DRb240 DLb0]) * + env.C[1, _prev(r, end), c][χNW DLt120 DLb120; χNa] * env.C[2, _prev(r, end), _next(c + 1, end)][χNb DRt60 DRb60; χNE] * env.C[3, r, _next(c + 1, end)][χNE DRt0 DRb0; χSE] * + env.C[4, _next(r, end), _next(c, end)][χSE DRt300 DRb300; χSa] * env.C[5, _next(r, end), _prev(c, end)][χSb DLt240 DLb240; χSW] * env.C[6, r, _prev(c, end)][χSW DLt180 DLb180; χNW] * + env.Eb[1, _prev(r, end), _next(c, end)][χNa DLt60 DLb60; χNC] * env.Ea[1, _prev(r, end), _next(c, end)][χNC DRt120 DRb120; χNb] * + env.Eb[4, _next(r, end), c][χSa DRt240 DRb240; χSC] * env.Ea[4, _next(r, end), c][χSC DLt300 DLb300; χSb] +end + +function _contract_edges_0((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, op::AbstractTensorMap{E, S, 2, 2}) where {P <: PEPSSandwichTriangular, E, S} + O_r_c = network[r, c] + O_r_cn = network[r, _next(c, end)] + return @tensor opt = true ket(O_r_c)[dLt; DLt120 DLt60 DLt0 DLt300 DLt240 DLt180] * ket(O_r_cn)[dRt; DRt120 DRt60 DRt0 DRt300 DRt240 DLt0] * + conj(bra(O_r_c)[dLb; DLb120 DLb60 DLb0 DLb300 DLb240 DLb180]) * conj(bra(O_r_cn)[dRb; DRb120 DRb60 DRb0 DRb300 DRb240 DLb0]) * + env.C[1, _prev(r, end), c][χNW DLt120 DLb120; χNa] * env.C[2, _prev(r, end), _next(c + 1, end)][χNb DRt60 DRb60; χNE] * env.C[3, r, _next(c + 1, end)][χNE DRt0 DRb0; χSE] * + env.C[4, _next(r, end), _next(c, end)][χSE DRt300 DRb300; χSa] * env.C[5, _next(r, end), _prev(c, end)][χSb DLt240 DLb240; χSW] * env.C[6, r, _prev(c, end)][χSW DLt180 DLb180; χNW] * + env.Eb[1, _prev(r, end), _next(c, end)][χNa DLt60 DLb60; χNC] * env.Ea[1, _prev(r, end), _next(c, end)][χNC DRt120 DRb120; χNb] * + env.Eb[4, _next(r, end), c][χSa DRt240 DRb240; χSC] * env.Ea[4, _next(r, end), c][χSC DLt300 DLb300; χSb] * + op[dLb dRb; dLt dRt] +end + +function _contract_edges_60((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + O_rn_cn = network[_next(r, end), _next(c, end)] + return @tensor opt = true ket(O_r_c)[dTR; DTRt120 DTRt60 DTRt0 DTRt300 DBLt60 DTRt180] * ket(O_rn_cn)[dBL; DBLt120 DBLt60 DBLt0 DBLt300 DBLt240 DBLt180] * + conj(bra(O_r_c)[dTR; DTRb120 DTRb60 DTRb0 DTRb300 DBLb60 DTRb180]) * conj(bra(O_rn_cn)[dBL; DBLb120 DBLb60 DBLb0 DBLb300 DBLb240 DBLb180]) * + env.C[1, _prev(r, end), c][χNWb DTRt120 DTRb120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN DTRt60 DTRb60; χNE] * env.C[3, r, _next(c, end)][χNE DTRt0 DTRb0; χSEa] * + env.C[4, _next(r + 1, end), _prev(c, end)][χSEb DBLt300 DBLb300; χS] * env.C[5, _next(r + 1, end), _prev(c - 1, end)][χS DBLt240 DBLb240; χSW] * env.C[6, _next(r, end), _prev(c - 1, end)][χSW DBLt180 DBLb180; χNWa] * + env.Eb[3, _next(r, end), c][χSEa DTRt300 DTRb300; χSEC] * env.Ea[3, _next(r, end), c][χSEC DBLt0 DBLb0; χSEb] * + env.Eb[6, r, _prev(c, end)][χNWa DBLt120 DBLb120; χNWC] * env.Ea[6, r, _prev(c, end)][χNWC DTRt180 DTRb180; χNWb] +end + +function _contract_edges_120((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + O_rn_c = network[_next(r, end), c] + return @tensor opt = true ket(O_r_c)[dTL; DTLt120 DTLt60 DTLt0 DTLt300 DTLt240 DTLt180] * ket(O_rn_c)[dBR; DTLt300 DBRt60 DBRt0 DBRt300 DBRt240 DBRt180] * + conj(bra(O_r_c)[dTL; DTLb120 DTLb60 DTLb0 DTLb300 DTLb240 DTLb180]) * conj(bra(O_rn_c)[dBR; DTLb300 DBRb60 DBRb0 DBRb300 DBRb240 DBRb180]) * + env.C[1, _prev(r, end), c][χNW DTLt120 DTLb120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN DTLt60 DTLb60; χNEa] * env.C[3, _next(r, end), _next(c, end)][χNEb DBRt0 DBRb0; χSE] * + env.C[4, _next(r + 1, end), c][χSE DBRt300 DBRb300; χS] * env.C[5, _next(r + 1, end), _prev(c, end)][χS DBRt240 DBRb240; χSWa] * env.C[6, r, _prev(c, end)][χSWb DTLt180 DTLb180; χNW] * + env.Eb[2, r, _next(c, end)][χNEa DTLt0 DTLb0; χNEC] * env.Ea[2, r, _next(c, end)][χNEC DBRt60 DBRb60; χNEb] * + env.Eb[5, _next(r, end), _prev(c, end)][χSWa DBRt180 DBRb180; χSWC] * env.Ea[5, _next(r, end), _prev(c, end)][χSWC DTLt240 DTLb240; χSWb] +end + +function _contract_site_large((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + O_rp_c = network[_prev(r, end), c] + O_rp_cn = network[_prev(r, end), _next(c, end)] + O_rn_cp = network[_next(r, end), _prev(c, end)] + O_rn_c = network[_next(r, end), c] + O_r_cn = network[r, _next(c, end)] + O_r_cp = network[r, _prev(c, end)] + return @tensor opt = true ket(O_rp_c)[dNW; DNWt120 DNWt60 DNWt0 DNWt300 DWt60 DNWt180] * ket(O_rp_cn)[dNE; DNEt120 DNEt60 DNEt0 DNEt300 DNEt240 DNWt0] * + ket(O_r_cn)[dE; DNEt300 DEt60 DEt0 DEt300 DEt240 DEt180] * ket(O_rn_c)[dSE; DSEt120 DEt240 DSEt0 DSEt300 DSEt240 DSEt180] * ket(O_rn_cp)[dSW; DSWt120 DSWt60 DSEt180 DSWt300 DSWt240 DSWt180] * + ket(O_r_cp)[dW; DWt120 DWt60 DWt0 DSWt120 DWt240 DWt180] * ket(O_r_c)[dCenter; DNWt300 DNEt240 DEt180 DSEt120 DSWt60 DWt0] * + conj(bra(O_rp_c)[dNW; DNWb120 DNWb60 DNWb0 DNWb300 DWb60 DNWb180]) * conj(bra(O_rp_cn)[dNE; DNEb120 DNEb60 DNEb0 DNEb300 DNEb240 DNWb0]) * + conj(bra(O_r_cn)[dE; DNEb300 DEb60 DEb0 DEb300 DEb240 DEb180]) * conj(bra(O_rn_c)[dSE; DSEb120 DEb240 DSEb0 DSEb300 DSEb240 DSEb180]) * conj(bra(O_rn_cp)[dSW; DSWb120 DSWb60 DSEb180 DSWb300 DSWb240 DSWb180]) * + conj(bra(O_r_cp)[dW; DWb120 DWb60 DWb0 DSWb120 DWb240 DWb180]) * conj(bra(O_r_c)[dCenter; DNWb300 DNEb240 DEb180 DSEb120 DSWb60 DWb0]) * + env.C[1, _prev(r - 1, end), c][χNWa DNWt120 DNWb120; χNb] * env.Eb[1, _prev(r - 1, end), _next(c, end)][χNb DNWt60 DNWb60; χNC] * env.Ea[1, _prev(r - 1, end), _next(c, end)][χNC DNEt120 DNEb120; χNa] * + env.C[2, _prev(r - 1, end), _next(c + 1, end)][χNa DNEt60 DNEb60; χNEb] * env.Eb[2, _prev(r, end), _next(c + 1, end)][χNEb DNEt0 DNEb0; χNEC] * env.Ea[2, _prev(r, end), _next(c + 1, end)][χNEC DEt60 DEb60; χNEa] * + env.C[3, r, _next(c + 1, end)][χNEa DEt0 DEb0; χSEb] * env.Eb[3, _next(r, end), _next(c, end)][χSEb DEt300 DEb300; χSEC] * env.Ea[3, _next(r, end), _next(c, end)][χSEC DSEt0 DSEb0; χSEa] * + env.C[4, _next(r + 1, end), c][χSEa DSEt300 DSEb300; χSb] * env.Eb[4, _next(r + 1, end), _prev(c, end)][χSb DSEt240 DSEb240; χSC] * env.Ea[4, _next(r + 1, end), _prev(c, end)][χSC DSWt300 DSWb300; χSa] * + env.C[5, _next(r + 1, end), _prev(c - 1, end)][χSa DSWt240 DSWb240; χSWb] * env.Eb[5, _next(r, end), _prev(c - 1, end)][χSWb DSWt180 DSWb180; χSWC] * env.Ea[5, _next(r, end), _prev(c - 1, end)][χSWC DWt240 DWb240; χSWa] * + env.C[6, r, _prev(c - 1, end)][χSWa DWt180 DWb180; χNWb] * env.Eb[6, _prev(r, end), _prev(c, end)][χNWb DWt120 DWb120; χNWC] * env.Ea[6, _prev(r, end), _prev(c, end)][χNWC DNWt180 DNWb180; χNWa] +end + +function _contract_corners((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular) where {P <: PEPSSandwichTriangular} + O_r_c = network[r, c] + return @tensor opt = true env.C[1, _prev(r, end), c][χNW Dt120 Db120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN Dt60 Db60; χNE] * env.C[3, r, _next(c, end)][χNE Dt0 Db0; χSE] * + env.C[4, _next(r, end), c][χSE Dt300 Db300; χS] * env.C[5, _next(r, end), _prev(c, end)][χS Dt240 Db240; χSW] * env.C[6, r, _prev(c, end)][χSW Dt180 Db180; χNW] * + ket(O_r_c)[d; Dt120 Dt60 Dt0 Dt300 Dt240 Dt180] * conj(bra(O_r_c)[d; Db120 Db60 Db0 Db300 Db240 Db180]) +end + +function _contract_corners((r, c)::Tuple{Int, Int}, network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, op::AbstractTensorMap{E, S, 1, 1}) where {P <: PEPSSandwichTriangular, E, S} + O_r_c = network[r, c] + return @tensor opt = true env.C[1, _prev(r, end), c][χNW Dt120 Db120; χN] * env.C[2, _prev(r, end), _next(c, end)][χN Dt60 Db60; χNE] * env.C[3, r, _next(c, end)][χNE Dt0 Db0; χSE] * + env.C[4, _next(r, end), c][χSE Dt300 Db300; χS] * env.C[5, _next(r, end), _prev(c, end)][χS Dt240 Db240; χSW] * env.C[6, r, _prev(c, end)][χSW Dt180 Db180; χNW] * + ket(O_r_c)[dt; Dt120 Dt60 Dt0 Dt300 Dt240 Dt180] * conj(bra(O_r_c)[db; Db120 Db60 Db0 Db300 Db240 Db180]) * op[db; dt] +end + +function energy(network::InfiniteTriangularNetwork{P}, env::CTMRGEnvTriangular, onesite_op, twosite_op) where {P <: PEPSSandwichTriangular} + E = sum(Iterators.product(axes(network)...)) do (r, c) + expval_onesite = _contract_corners((r, c), network, env, onesite_op) / _contract_corners((r, c), network, env) + expval_twosite = _contract_edges_0((r, c), network, env, twosite_op) / _contract_edges_0((r, c), network, env) + return expval_onesite + expval_twosite + end + ignore_derivatives() do + isapprox(imag(E), 0; atol = sqrt(eps(real(E)))) || + @warn "Expectation value is not real: $E." + end + return real(E) +end +function energy(state::InfinitePEPSTriangular, env::CTMRGEnvTriangular, args...) + return energy(InfiniteTriangularNetwork(state), env, args...) +end + +# TODO: write a proper cost_function that works on an InfinitePEPSTriangular and a LocalOperator, and use it in the optimization loop. For now, we just compute the energy using the network contraction functions above. diff --git a/src/environments/ctmrg_environments_triangular.jl b/src/environments/ctmrg_environments_triangular.jl new file mode 100644 index 000000000..c2fea0a7c --- /dev/null +++ b/src/environments/ctmrg_environments_triangular.jl @@ -0,0 +1,119 @@ +#TODO: Add docs and figure +struct CTMRGEnvTriangular{TC, TA, TB} + "6 x rows x cols array of corner C tensors, where the first dimension specifies the spatial direction" + C::Array{TC, 3} + "6 x rows x cols array of edge Ta tensors, where the first dimension specifies the spatial direction" + Ea::Array{TA, 3} + "6 x rows x cols array of edge Tb tensors, where the first dimension specifies the spatial direction" + Eb::Array{TB, 3} +end + +function get_Ds(D::A) where {A <: ProductSpace} + return [dir > 3 ? _elementwise_dual(D) : D for dir in 1:6] +end + +function get_Ds(D::A) where {A <: ElementarySpace} + return [dir > 3 ? ProductSpace(D') : ProductSpace(D) for dir in 1:6] +end + +function CTMRGEnvTriangular( + f, T, D::Union{A, B}, chis::B; unitcell::Tuple{Int, Int} = (1, 1) + ) where {A <: ProductSpace, B <: ElementarySpace} + Ds = get_Ds(D) + N = length(Ds[1]) + st = spacetype(Ds[1]) + + T_type = tensormaptype(st, N + 1, 1, T) + + Cs = Array{T_type}(undef, 6, unitcell...) + Eas = Array{T_type}(undef, 6, unitcell...) + Ebs = Array{T_type}(undef, 6, unitcell...) + + for dir in 1:6, r in 1:unitcell[1], c in 1:unitcell[2] + C = _edge_tensor(f, T, chis, Ds[dir], chis) + Ea = _edge_tensor(f, T, chis, Ds[dir], chis) + Eb = _edge_tensor(f, T, chis, Ds[mod1(dir + 1, 6)], chis) + + C /= norm(C) + Ea /= norm(Ea) + Eb /= norm(Eb) + + Cs[dir, r, c] = C + Eas[dir, r, c] = Ea + Ebs[dir, r, c] = Eb + end + return CTMRGEnvTriangular(Cs, Eas, Ebs) +end + +Base.size(env::CTMRGEnvTriangular, args...) = size(env.C, args...) +Base.axes(x::CTMRGEnvTriangular, args...) = axes(x.C, args...) +function eachcoordinate(x::CTMRGEnvTriangular) + return collect(Iterators.product(axes(x, 2), axes(x, 3))) +end +function eachcoordinate(x::CTMRGEnvTriangular, dirs) + return collect(Iterators.product(dirs, axes(x, 2), axes(x, 3))) +end +Base.real(env::CTMRGEnvTriangular) = CTMRGEnvTriangular(real.(env.C), real.(env.Ea), real.(env.Eb)) +Base.complex(env::CTMRGEnvTriangular) = CTMRGEnvTriangular(complex.(env.C), complex.(env.Ea), complex.(env.Eb)) + +cornertype(env::CTMRGEnvTriangular) = cornertype(typeof(env)) +cornertype(::Type{CTMRGEnvTriangular{T}}) where {T} = T +edgetype(env::CTMRGEnvTriangular) = edgetype(typeof(env)) +edgetype(::Type{CTMRGEnvTriangular{T}}) where {T} = T + +TensorKit.spacetype(::Type{E}) where {E <: CTMRGEnvTriangular} = spacetype(cornertype(E)) + +# In-place update of environment +function update!(env::CTMRGEnvTriangular{T}, env´::CTMRGEnvTriangular{T}) where {T} + env.C .= env´.C + env.Ea .= env´.Ea + env.Eb .= env´.Eb + return env +end + +# Custom adjoint for CTMRGEnvTriangular constructor, needed for fixed-point differentiation +function ChainRulesCore.rrule( + ::Type{CTMRGEnvTriangular}, C::Array{T, 3}, Ea::Array{T, 3}, Eb::Array{T, 3} + ) where {T} + function triangularenv_pullback(Δenv_) + Δenv = unthunk(Δenv_) + return NoTangent(), Δenv.C, Δenv.Ea, Δenv.Eb + end + return CTMRGEnvTriangular(C, Ea, Eb), triangularenv_pullback +end + +# math basics used for tangent accumulation in backwards pass +function Base.:+(e₁::CTMRGEnvTriangular, e₂::CTMRGEnvTriangular) + return CTMRGEnvTriangular(e₁.C + e₂.C, e₁.Ea + e₂.Ea, e₁.Eb + e₂.Eb) +end +function Base.:-(e₁::CTMRGEnvTriangular, e₂::CTMRGEnvTriangular) + return CTMRGEnvTriangular(e₁.C - e₂.C, e₁.Ea - e₂.Ea, e₁.Eb - e₂.Eb) +end +Base.:*(α::Number, e::CTMRGEnvTriangular) = CTMRGEnvTriangular(α * e.C, α * e.Ea, α * e.Eb) + +# Custom adjoint for CTMRGEnvTriangular getproperty, to avoid creating named tuples in backward pass +function ChainRulesCore.rrule(::typeof(getproperty), e::CTMRGEnvTriangular, name::Symbol) + result = getproperty(e, name) + if name === :C + function C_pullback(ΔC_) + ΔC = unthunk(ΔC_) + return NoTangent(), CTMRGEnvTriangular(ΔC, zerovector(e.Ea), zerovector(e.Eb)) + end + return result, C_pullback + elseif name === :Ea + function Ea_pullback(ΔEa_) + ΔEa = unthunk(ΔEa_) + return NoTangent(), CTMRGEnvTriangular(zerovector(e.C), ΔEa, zerovector(e.Eb)) + end + return result, Ea_pullback + elseif name === :Eb + function Eb_pullback(ΔEb_) + ΔEb = unthunk(ΔEb_) + return NoTangent(), CTMRGEnvTriangular(zerovector(e.C), zerovector(e.Ea), ΔEb) + end + return result, Eb_pullback + else + # this should never happen because already errored in forwards pass + throw(ArgumentError("No rrule for getproperty of $name")) + end +end diff --git a/src/networks/infinitetriangularnetwork.jl b/src/networks/infinitetriangularnetwork.jl new file mode 100644 index 000000000..d8eebdd34 --- /dev/null +++ b/src/networks/infinitetriangularnetwork.jl @@ -0,0 +1,156 @@ +const PEPSTensorTriangular{S <: ElementarySpace} = AbstractTensorMap{<:Any, S, 1, 6} + +""" +$(TYPEDEF) + +Contractible square network. Wraps a matrix of 'rank-6-tensor-like' objects. + +## Fields + +$(TYPEDFIELDS) +""" +struct InfiniteTriangularNetwork{O} + A::Matrix{O} + InfiniteTriangularNetwork{O}(A::Matrix{O}) where {O} = new{O}(A) + function InfiniteTriangularNetwork(A::Matrix) + # for I in eachindex(IndexCartesian(), A) + # d, w = Tuple(I) + # northwest_virtualspace(A[d, w]) == + # _elementwise_dual(southeast_virtualspace(A[_prev(d, end), w])) || throw( + # SpaceMismatch("North virtual space at site $((d, w)) does not match.") + # ) + # northeast_virtualspace(A[d, w]) == + # _elementwise_dual(southwest_virtualspace(A[d, _next(w, end)])) || + # throw(SpaceMismatch("East virtual space at site $((d, w)) does not match.")) + # east_virtualspace(A[d, w]) == + # _elementwise_dual(west_virtualspace(A[d, _next(w, end)])) || + # throw(SpaceMismatch("East virtual space at site $((d, w)) does not match.")) + # end + return InfiniteTriangularNetwork{eltype(A)}(A) + end +end +InfiniteTriangularNetwork(n::InfiniteTriangularNetwork) = n + +## Unit cell interface + +unitcell(n::InfiniteTriangularNetwork) = n.A +Base.size(n::InfiniteTriangularNetwork, args...) = size(unitcell(n), args...) +Base.length(n::InfiniteTriangularNetwork) = length(unitcell(n)) +Base.eltype(n::InfiniteTriangularNetwork) = eltype(typeof(n)) +Base.eltype(::Type{InfiniteTriangularNetwork{O}}) where {O} = O + +Base.copy(n::InfiniteTriangularNetwork) = InfiniteTriangularNetwork(copy(unitcell(n))) +function Base.similar(n::InfiniteTriangularNetwork, T::Type{TorA} = scalartype(n)) where {TorA} + return InfiniteTriangularNetwork(map(t -> similar(t, T), unitcell(n))) +end +function Base.repeat(n::InfiniteTriangularNetwork, counts...) + return InfiniteTriangularNetwork(repeat(unitcell(n), counts...)) +end + +## Indexing +Base.getindex(n::InfiniteTriangularNetwork, args...) = Base.getindex(unitcell(n), args...) +function Base.setindex!(n::InfiniteTriangularNetwork, args...) + return (Base.setindex!(unitcell(n), args...); n) +end +Base.axes(n::InfiniteTriangularNetwork, args...) = axes(unitcell(n), args...) +eachcoordinate(n::InfiniteTriangularNetwork) = collect(Iterators.product(axes(n)...)) +function eachcoordinate(n::InfiniteTriangularNetwork, dirs) + return collect(Iterators.product(dirs, axes(n, 1), axes(n, 2))) +end + +## Spaces + +TensorKit.spacetype(::Type{T}) where {T <: InfiniteTriangularNetwork} = spacetype(eltype(T)) +function virtualspace(n::InfiniteTriangularNetwork, r::Int, c::Int, dir) + Nr, Nc = size(n) + return virtualspace(n[mod1(r, Nr), mod1(c, Nc)], dir) +end + +## Vector interface + +function VectorInterface.scalartype(::Type{T}) where {T <: InfiniteTriangularNetwork} + return scalartype(eltype(T)) +end +function VectorInterface.zerovector(A::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(zerovector(unitcell(A))) +end + +## Math (for Zygote accumulation) + +function Base.:+(A₁::InfiniteTriangularNetwork, A₂::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(_add_localsandwich.(unitcell(A₁), unitcell(A₂))) +end +function Base.:-(A₁::InfiniteTriangularNetwork, A₂::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(_subtract_localsandwich.(unitcell(A₁), unitcell(A₂))) +end +function Base.:*(α::Number, A::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(_mul_localsandwich.(α, unitcell(A))) +end +Base.:*(A::InfiniteTriangularNetwork, α::Number) = α * A +function Base.:/(A::InfiniteTriangularNetwork, α::Number) + return A * inv(α) +end +function LinearAlgebra.dot(A₁::InfiniteTriangularNetwork, A₂::InfiniteTriangularNetwork) + return dot(unitcell(A₁), unitcell(A₂)) +end +LinearAlgebra.norm(A::InfiniteTriangularNetwork) = norm(unitcell(A)) + +## (Approximate) equality + +function Base.:(==)(A₁::InfiniteTriangularNetwork, A₂::InfiniteTriangularNetwork) + return all(zip(unitcell(A₁), unitcell(A₂))) do (p₁, p₂) + return p₁ == p₂ + end +end +function Base.isapprox(A₁::InfiniteTriangularNetwork, A₂::InfiniteTriangularNetwork; kwargs...) + return all(zip(unitcell(A₁), unitcell(A₂))) do (p₁, p₂) + return _isapprox_localsandwich(p₁, p₂; kwargs...) + end +end + +## Rotations + +function rotl60(n::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(rotl60(_rotl60_localsandwich.(unitcell(n)))) +end +function rotr60(n::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(rotr60(_rotr60_localsandwich.(unitcell(n)))) +end +function Base.rot180(n::InfiniteTriangularNetwork) + return InfiniteTriangularNetwork(rot180(_rot180_localsandwich.(unitcell(n)))) +end + +## Chainrules + +# generic implementation +function ChainRulesCore.rrule( + ::typeof(Base.getindex), network::InfiniteTriangularNetwork, r::Int, c::Int + ) + O = network[r, c] + + function getindex_pullback(ΔO_) + ΔO = map(unthunk, ΔO_) + if ΔO isa Tangent + ΔO = ChainRulesCore.construct(typeof(O), ChainRulesCore.backing(ΔO)) + end + Δnetwork = zerovector(network) + Δnetwork[r, c] = ΔO + return NoTangent(), Δnetwork, NoTangent(), NoTangent() + end + return O, getindex_pullback +end + +# specialized PFTensor implementation +function ChainRulesCore.rrule( + ::typeof(Base.getindex), network::InfiniteTriangularNetwork{<:PFTensor}, r::Int, c::Int + ) + O = network[r, c] + + function getindex_pullback(ΔO_) + ΔO = unthunk(ΔO_) + Δnetwork = zerovector(network) + Δnetwork[r, c] = ΔO + return NoTangent(), Δnetwork, NoTangent(), NoTangent() + end + return O, getindex_pullback +end diff --git a/src/networks/local_triangular_sandwich.jl b/src/networks/local_triangular_sandwich.jl new file mode 100644 index 000000000..a5fe0ab99 --- /dev/null +++ b/src/networks/local_triangular_sandwich.jl @@ -0,0 +1,43 @@ +# +# Interface for local effective-rank-6 tensor sandwiches +# + +# route all virtualspace getters through a single method for convenience +northwest_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[1] +northeast_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[2] +east_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[3] +southeast_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[4] +southwest_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[5] +west_virtualspace(O::AbstractTensorMap{E, S, 1, 6}, args...) where {E, S} = domain(O)[6] + +## Rotations + +# generic local interface +_rotl60_localsandwich(O) = rotl60.(O) +_rotr60_localsandwich(O) = rotr60.(O) + +## PartitionFunction + +# specialized local rotation interface +_rotl60_localsandwich(O::PFTensorTriangular) = rotl60(O) +_rotr60_localsandwich(O::PFTensorTriangular) = rotr60(O) +_rot180_localsandwich(O::PFTensorTriangular) = rot180(O) + +# specialized local math interface +_add_localsandwich(O1::PFTensorTriangular, O2::PFTensorTriangular) = O1 + O2 +_subtract_localsandwich(O1::PFTensorTriangular, O2::PFTensorTriangular) = O1 - O2 +_mul_localsandwich(α::Number, O::PFTensorTriangular) = α * O +_isapprox_localsandwich(O1::PFTensorTriangular, O2::PFTensorTriangular; kwargs...) = isapprox(O1, O2; kwargs...) + +## PEPS + +const PEPSSandwichTriangular{T <: PEPSTensorTriangular} = Tuple{T, T} + +ket(O::PEPSSandwichTriangular) = O[1] +bra(O::PEPSSandwichTriangular) = O[2] + +function virtualspace(O::PEPSSandwichTriangular, dir) + return virtualspace(ket(O), dir) ⊗ virtualspace(bra(O), dir)' +end + +TensorKit.spacetype(::Type{P}) where {P <: PEPSSandwichTriangular} = spacetype(eltype(P)) diff --git a/src/networks/tensors_triangular.jl b/src/networks/tensors_triangular.jl new file mode 100644 index 000000000..00c60de23 --- /dev/null +++ b/src/networks/tensors_triangular.jl @@ -0,0 +1,101 @@ +# +# Partition function +# + +""" + const PartitionFunctionTensorTriangular{S} + +Default type for partition function tensors with 6 virtual indices, conventionally ordered +as: ``T : W ⊗ SW ⊗ SE ← NW ⊗ NE ⊗ E``. Here ``NW``, ``NE``, ``E``, ``SE``, ``SW`` and ``W`` denote the +north-west (120°), north-east (60°), east (0°), south-east (300°), south-west (240°) and west (180°) spaces, respectively, +where the angles denote the directions of the legs with respect to the positive horizontal axis. + +``` + NW NE + ╲ ╱ + ╲ ╱ + ╲ ╱ + W ----- T ----- E + ╱ ╲ + ╱ ╲ + ╱ ╲ + SW P SE +``` +""" +const PartitionFunctionTensorTriangular{S <: ElementarySpace} = AbstractTensorMap{<:Any, S, 3, 3} +const PFTensorTriangular = PartitionFunctionTensorTriangular + +""" + PartitionFunctionTensorTriangular(f, ::Type{T}, NWspace::S, + [NEspace::S], [Espace::S], [SEspace::S], + [SWspace::S], [Wspace::S]) where {T,S<:Union{Int,ElementarySpace}} + +Construct a PartitionFunctionTensorTriangular tensor based on the +north-west, north-east, east, south-east, south-west and west spaces +The tensor elements are generated based on `f` and the element type is specified in `T`. +""" +function PartitionFunctionTensorTriangular( + f, ::Type{T}, + NWspace::S, NEspace::S = NWspace, Espace::S = NWspace, SEspace::S = NWspace, + SWspace::S = NWspace, Wspace::S = NWspace, + ) where {T, S <: ElementarySpace} + return f(T, Wspace ⊗ SWspace ⊗ SEspace ← NWspace ⊗ NEspace ⊗ Espace) +end + +function virtualspace(t::PFTensorTriangular, dir) + invp = (4, 5, 6, 3, 2, 1) # internally, virtual directions are ordered as NW, NE, E, SE, SW... + return space(t, invp[dir]) +end + +# +# PEPS +# + +""" + const PEPSTensorTriangular{S} + +Default type for PEPS tensors with a single physical index, and 6 virtual indices, +conventionally ordered as: ``T : P ← NW ⊗ NE ⊗ E ⊗ SE ⊗ SW ⊗ W``. Here, ``P`` denotes the physical space +and ``NW``, ``NE``, ``E``, ``SE``, ``SW`` and ``W`` denote the north-west (120°), north-east (60°), +east (0°), south-east (300°), south-west (240°) and west (180°) spaces, respectively, +where the angles denote the directions of the legs with respect to the positive horizontal axis. + +``` +``` + NW NE + ╲ ╱ + ╲ ╱ + ╲ ╱ + W ----- T ----- E + ╱|╲ + ╱ | ╲ + ╱ | ╲ + SW P SE +``` +``` +""" +const PEPSTensorTriangular{S <: ElementarySpace} = AbstractTensorMap{<:Any, S, 1, 6} + +""" + PEPSTensor(f, ::Type{T}, Pspace::S, NWspace::S, + [NEspace::S], [Espace::S], [SEspace::S], + [SWspace::S], [Wspace::S]) where {T,S<:Union{Int,ElementarySpace}} + +Construct a PEPS tensor based on the physical, north-west, north-east, east, south-east, south-west and west spaces. +The tensor elements are generated based on `f` and the element type is specified in `T`. +""" +function PEPSTensorTriangular( + f, ::Type{T}, + Pspace::S, + NWspace::S, NEspace::S = NWspace, Espace::S = NWspace, SEspace::S = NWspace, + SWspace::S = NWspace, Wspace::S = NWspace, + ) where {T, S <: ElementarySpace} + return f(T, Pspace ← NWspace ⊗ NEspace ⊗ Espace ⊗ SEspace ⊗ SWspace ⊗ Wspace) +end + +rotl60(t::PEPSTensorTriangular) = permute(t, ((1,), (3, 4, 5, 6, 7, 2))) +rotr60(t::PEPSTensorTriangular) = permute(t, ((1,), (7, 2, 3, 4, 5, 6))) +Base.rot180(t::PEPSTensorTriangular) = permute(t, ((1,), (5, 6, 7, 2, 3, 4))) + +physicalspace(t::PEPSTensorTriangular) = space(t, 1) +virtualspace(t::PEPSTensorTriangular, dir) = space(t, dir + 1) diff --git a/src/operators/transfermatrix.jl b/src/operators/transfermatrix.jl index fd94825e1..b1162fb6a 100644 --- a/src/operators/transfermatrix.jl +++ b/src/operators/transfermatrix.jl @@ -153,10 +153,10 @@ function initialize_mps( return InfiniteMPS( [ f( - T, - virtualspaces[_prev(i, end)] * _elementwise_dual(north_virtualspace(O, i)), - virtualspaces[mod1(i, end)], - ) for i in 1:length(O) + T, + virtualspaces[_prev(i, end)] * _elementwise_dual(north_virtualspace(O, i)), + virtualspaces[mod1(i, end)], + ) for i in 1:length(O) ]; kwargs... ) end diff --git a/src/states/infinitepepstriangular.jl b/src/states/infinitepepstriangular.jl new file mode 100644 index 000000000..a8660e982 --- /dev/null +++ b/src/states/infinitepepstriangular.jl @@ -0,0 +1,231 @@ +""" + struct InfinitePEPSTriangular{T<:PEPSTensorTriangularTriangular} + +Represents an infinite projected entangled-pair state on a 2D triangular lattice. + +## Fields + +$(TYPEDFIELDS) +""" +struct InfinitePEPSTriangular{T <: PEPSTensorTriangular} + A::Matrix{T} + InfinitePEPSTriangular{T}(A::Matrix{T}) where {T <: PEPSTensorTriangular} = new{T}(A) + function InfinitePEPSTriangular(A::Array{T, 2}) where {T <: PEPSTensorTriangular} + bosonic_braiding = BraidingStyle(sectortype(T)) === Bosonic() + # TODO: check braiding and spaces + return new{T}(A) + end +end + +## Constructors + +""" + InfinitePEPSTriangular(A::AbstractMatrix{T}) + +Create an `InfinitePEPSTriangular` by specifying a matrix containing the PEPS tensors at each site in +the unit cell. +""" +function InfinitePEPSTriangular(A::AbstractMatrix{<:PEPSTensorTriangular}) + return InfinitePEPSTriangular(Array(deepcopy(A))) # TODO: find better way to copy +end + +""" + InfinitePEPSTriangular([f=randn, T=ComplexF64,] Pspaces::A, Nspaces::A, [Espaces::A]) where {A<:AbstractMatrix{ElementarySpace}} + +Create an `InfinitePEPSTriangular` by specifying the physical, north virtual and east virtual spaces +of the PEPS tensor at each site in the unit cell as a matrix. +""" +function InfinitePEPSTriangular( + f, T::Type{<:Number}, Pspaces::M, NWspaces::M, NEspaces::M = NWspaces, Espaces::M = NWspaces + ) where {M <: AbstractMatrix{<:ElementarySpace}} + size(Pspaces) == size(NEspaces) == size(NEspaces) == size(Espaces) || + throw(ArgumentError("Input spaces should have equal sizes.")) + + SEspaces = adjoint.(circshift(NWspaces, (-1, 0))) + SWspaces = adjoint.(NEspaces) + Wspaces = adjoint.(circshift(Espaces, (0, 1))) + + A = map(Pspaces, NWspaces, NEspaces, Espaces, SEspaces, SWspaces, Wspaces) do P, NW, NE, E, SE, SW, W + return PEPSTensorTriangular(f, T, P, NW, NE, E, SE, SW, W) + end + + return InfinitePEPSTriangular(A) +end +function InfinitePEPSTriangular( + Pspaces::A, virtual_spaces...; kwargs... + ) where {A <: Union{AbstractMatrix{<:ElementarySpace}, ElementarySpace}} + return InfinitePEPSTriangular(randn, ComplexF64, Pspaces, virtual_spaces...; kwargs...) +end + +""" + InfinitePEPSTriangular(A::PEPSTensorTriangular; unitcell=(1, 1)) + +Create an `InfinitePEPSTriangular` by specifying a tensor and unit cell. + +The unit cell is labeled as a matrix which means that any tensor in the unit cell, +regardless if PEPS tensor or environment tensor, is obtained by shifting the row +and column index `[r, c]` by one, respectively: +``` +TODO: diagram +``` +The unit cell has periodic boundary conditions, so `[r, c]` is indexed modulo the +size of the unit cell. +""" +function InfinitePEPSTriangular(A::T; unitcell::Tuple{Int, Int} = (1, 1)) where {T <: PEPSTensorTriangular} + return InfinitePEPSTriangular(fill(A, unitcell)) +end + +# expand PEPS spaces to unit cell size +function _fill_state_virtual_spaces_triangular( + NWspace::S, NEspace::S = NWspace, Espace::S = NWspace; unitcell::Tuple{Int, Int} = (1, 1) + ) where {S <: ElementarySpace} + return fill(NWspace, unitcell), fill(NEspace, unitcell), fill(Espace, unitcell) +end + +""" + InfinitePEPSTriangular([f=randn, T=ComplexF64,] Pspace, Nspace, [Espace]; unitcell=(1,1)) + +Create an InfinitePEPSTriangular by specifying its physical, north and east spaces and unit cell. +""" +function InfinitePEPSTriangular( + f, T::Type{<:Number}, Pspace::S, vspaces...; unitcell::Tuple{Int, Int} = (1, 1) + ) where {S <: ElementarySpace} + return InfinitePEPSTriangular( + f, T, + _fill_state_physical_spaces(Pspace; unitcell), + _fill_state_virtual_spaces_triangular(vspaces...; unitcell)..., + ) +end + +## Unit cell interface + +unitcell(t::InfinitePEPSTriangular) = t.A +Base.size(A::InfinitePEPSTriangular, args...) = size(unitcell(A), args...) +Base.length(A::InfinitePEPSTriangular) = length(unitcell(A)) +Base.eltype(::Type{InfinitePEPSTriangular{T}}) where {T} = T +Base.eltype(A::InfinitePEPSTriangular) = eltype(typeof(A)) + +Base.copy(A::InfinitePEPSTriangular) = InfinitePEPSTriangular(copy(unitcell(A))) +function Base.similar(A::InfinitePEPSTriangular, T::Type{TorA} = scalartype(A)) where {TorA} + return InfinitePEPSTriangular(map(t -> similar(t, T), unitcell(A))) +end +Base.repeat(A::InfinitePEPSTriangular, counts...) = InfinitePEPSTriangular(repeat(unitcell(A), counts...)) + +Base.getindex(A::InfinitePEPSTriangular, args...) = Base.getindex(unitcell(A), args...) +Base.setindex!(A::InfinitePEPSTriangular, args...) = (Base.setindex!(unitcell(A), args...); A) +Base.axes(A::InfinitePEPSTriangular, args...) = axes(unitcell(A), args...) +eachcoordinate(A::InfinitePEPSTriangular) = collect(Iterators.product(axes(A)...)) +function eachcoordinate(A::InfinitePEPSTriangular, dirs) + return collect(Iterators.product(dirs, axes(A, 1), axes(A, 2))) +end + +## Spaces + +TensorKit.spacetype(::Type{T}) where {T <: InfinitePEPSTriangular} = spacetype(eltype(T)) +virtualspace(n::InfinitePEPSTriangular, dir) = virtualspace.(unitcell(n), dir) +function virtualspace(n::InfinitePEPSTriangular, r::Int, c::Int, dir) + Nr, Nc = size(n) + return virtualspace(n[mod1(r, Nr), mod1(c, Nc)], dir) +end +physicalspace(n::InfinitePEPSTriangular) = physicalspace.(unitcell(n)) +function physicalspace(n::InfinitePEPSTriangular, r::Int, c::Int) + Nr, Nc = size(n) + return physicalspace(n[mod1(r, Nr), mod1(c, Nc)]) +end + +## InfiniteTriangularNetwork interface + +function InfiniteTriangularNetwork(top::InfinitePEPSTriangular, bot::InfinitePEPSTriangular = top) + size(top) == size(bot) || throw( + ArgumentError("Top PEPS, bottom PEPS and PEPO rows should have the same length") + ) + return InfiniteTriangularNetwork(map(tuple, unitcell(top), unitcell(bot))) +end + +## Vector interface + +VI.scalartype(::Type{NT}) where {NT <: InfinitePEPSTriangular} = scalartype(eltype(NT)) +VI.zerovector(A::InfinitePEPSTriangular) = InfinitePEPSTriangular(zerovector(unitcell(A))) + +function VI.scale(ψ::InfinitePEPSTriangular, α::Number) + _scale = Base.Fix2(scale, α) + return InfinitePEPSTriangular(map(_scale, unitcell(ψ))) +end +function VI.scale!(ψ::InfinitePEPSTriangular, α::Number) + _scale! = Base.Fix2(scale!, α) + foreach(_scale!, unitcell(ψ)) + return ψ +end +function VI.scale!(ψ₁::InfinitePEPSTriangular, ψ₂::InfinitePEPSTriangular, α::Number) + _scale!(x, y) = scale!(x, y, α) + foreach(_scale!, unitcell(ψ₁), unitcell(ψ₂)) + return ψ₁ +end +VI.scale!!(ψ::InfinitePEPSTriangular, α::Number) = scale!(ψ, α) +VI.scale!!(ψ₁::InfinitePEPSTriangular, ψ₂::InfinitePEPSTriangular, α::Number) = scale!(ψ₁, ψ₂, α) + +function VI.add(ψ₁::InfinitePEPSTriangular, ψ₂::InfinitePEPSTriangular, α::Number, β::Number) + _add(x, y) = add(x, y, α, β) + return InfinitePEPSTriangular(map(_add, unitcell(ψ₁), unitcell(ψ₂))) +end +function VI.add!(ψ₁::InfinitePEPSTriangular, ψ₂::InfinitePEPSTriangular, α::Number, β::Number) + _add!(x, y) = add!(x, y, α, β) + foreach(_add!, unitcell(ψ₁), unitcell(ψ₂)) + return ψ₁ +end +VI.add!!(ψ₁::InfinitePEPSTriangular, ψ₂::InfinitePEPSTriangular, α::Number, β::Number) = add!(ψ₁, ψ₂, α, β) + +## Math + +function Base.:+(A₁::InfinitePEPSTriangular, A₂::InfinitePEPSTriangular) + return InfinitePEPSTriangular(unitcell(A₁) + unitcell(A₂)) +end +function Base.:-(A₁::InfinitePEPSTriangular, A₂::InfinitePEPSTriangular) + return InfinitePEPSTriangular(unitcell(A₁) - unitcell(A₂)) +end +Base.:*(α::Number, A::InfinitePEPSTriangular) = InfinitePEPSTriangular(α * unitcell(A)) +Base.:*(A::InfinitePEPSTriangular, α::Number) = α * A +Base.:/(A::InfinitePEPSTriangular, α::Number) = InfinitePEPSTriangular(unitcell(A) / α) +LinearAlgebra.dot(A₁::InfinitePEPSTriangular, A₂::InfinitePEPSTriangular) = dot(unitcell(A₁), unitcell(A₂)) +LinearAlgebra.norm(A::InfinitePEPSTriangular) = norm(unitcell(A)) + +## (Approximate) equality +function Base.:(==)(A₁::InfinitePEPSTriangular, A₂::InfinitePEPSTriangular) + return all(zip(unitcell(A₁), unitcell(A₂))) do (p₁, p₂) + return p₁ == p₂ + end +end +function Base.isapprox(A₁::InfinitePEPSTriangular, A₂::InfinitePEPSTriangular; kwargs...) + return all(zip(unitcell(A₁), unitcell(A₂))) do (p₁, p₂) + return isapprox(p₁, p₂; kwargs...) + end +end + + +## Chainrules + +function ChainRulesCore.rrule(::typeof(Base.getindex), network::InfinitePEPSTriangular, args...) + tensor = network[args...] + + function getindex_pullback(Δtensor_) + Δtensor = unthunk(Δtensor_) + Δnetwork = zerovector(network) + Δnetwork[args...] = Δtensor + return NoTangent(), Δnetwork, NoTangent(), NoTangent() + end + return tensor, getindex_pullback +end + +function ChainRulesCore.rrule( + ::Type{InfiniteTriangularNetwork}, top::InfinitePEPSTriangular, bot::InfinitePEPSTriangular + ) + network = InfiniteTriangularNetwork(top, bot) + + function InfiniteTriangularNetwork_pullback(Δnetwork_) + Δnetwork = unthunk(Δnetwork_) + Δtop = InfinitePEPSTriangular(map(ket, unitcell(Δnetwork))) + Δbot = InfinitePEPSTriangular(map(bra, unitcell(Δnetwork))) + return NoTangent(), Δtop, Δbot + end + return network, InfiniteTriangularNetwork_pullback +end diff --git a/test/ctmrg/triangular.jl b/test/ctmrg/triangular.jl new file mode 100644 index 000000000..554a4a793 --- /dev/null +++ b/test/ctmrg/triangular.jl @@ -0,0 +1,202 @@ +using Test +using Random +using TensorKit +using MPSKit +using PEPSKit +using OptimKit +using Zygote +using LinearAlgebra +using KrylovKit + +using PEPSKit: unitcell, norm_preserving_retract + +sd = 123456 +Random.seed!(sd) + +D = 2 +χ = 7 +T = ComplexF64 + +ctmrg_tol = 1.0e-8 +ctmrg_maxiter = 300 +ctmrg_verbosity = 2 + +const ising_βc_triangular = BigFloat(BigFloat(asinh(BigFloat(sqrt(BigFloat(1.0) / BigFloat(3.0))))) / BigFloat(2.0)) +const f_onsager_triangular::BigFloat = -3.20253248660790791834355252025862951439 + +function classical_ising_triangular(β) + t = Float64[exp(β) exp(-β); exp(-β) exp(β)] + + r = eigen(t) + nt = r.vectors * sqrt(LinearAlgebra.Diagonal(r.values)) * r.vectors + + O = zeros(2, 2, 2, 2, 2, 2) + O[1, 1, 1, 1, 1, 1] = 1 + O[2, 2, 2, 2, 2, 2] = 1 + + H = [1 1; 1 -1] / sqrt(2) + + @tensor o[-1 -2 -3; -4 -5 -6] := O[1 2 3; 4 5 6] * nt[-1; 1] * nt[-2; 2] * nt[-3; 3] * nt[-4; 4] * nt[-5; 5] * nt[-6; 6] + @tensor o2[-1 -2 -3; -4 -5 -6] := o[1 2 3; 4 5 6] * H[-1; 1] * H[-2; 2] * H[-3; 3] * H[-4; 4] * H[-5; 5] * H[-6; 6] + return TensorMap(o2, ℂ^2 * ℂ^2 * ℂ^2, ℂ^2 * ℂ^2 * ℂ^2) +end + +@testset "CTM_triangular - Random tensor" begin + for conditioning in [true false] + for projector_alg in [:twothirds :full] + Random.seed!(79413165445) + alg = SimultaneousCTMRGTriangular(; + tol = ctmrg_tol, maxiter = ctmrg_maxiter, verbosity = ctmrg_verbosity, + conditioning, projector_alg + ) + pspace = ComplexSpace(2) + vspace = ComplexSpace(D) + envspace = ComplexSpace(χ) + + ket = randn(T, pspace, vspace ⊗ vspace ⊗ vspace ⊗ vspace' ⊗ vspace' ⊗ vspace') + bra = copy(ket) + pf = randn(T, vspace ⊗ vspace ⊗ vspace, vspace ⊗ vspace ⊗ vspace) + sandwiches = [pf, (ket, bra)] + # sandwiches = [(ket, bra), pf] + unitcell = (2, 2) + + for (sandwich, V) in zip(sandwiches, [vspace, vspace ⊗ vspace']) + # for (sandwich, vspace) in zip(sandwiches, [vspace ⊗ vspace', vspace]) + network = InfiniteTriangularNetwork(fill(sandwich, unitcell)) + env₀ = CTMRGEnvTriangular(randn, T, V, envspace; unitcell) + env, info = leading_boundary(env₀, network, alg) + @test info.convergence_metric < 1.0 # this is not much of a test + end + end + end +end + +@testset "CTM_triangular - Classical Ising" begin + χ_local = 20 + T_local = Float64 + + for conditioning in [true false] + for projector_alg in [:twothirds :full] + Random.seed!(156484561351) + + alg = SimultaneousCTMRGTriangular(; + tol = ctmrg_tol, maxiter = ctmrg_maxiter, verbosity = ctmrg_verbosity, + conditioning, projector_alg + ) + sz = (1, 1) + T = classical_ising_triangular(ising_βc_triangular) + pf = InfiniteTriangularNetwork(fill(T, sz)) + + vspace = codomain(T)[1] + envspace = ComplexSpace(χ_local) + env₀ = CTMRGEnvTriangular(randn, T_local, vspace, envspace; unitcell = sz) + env, info = leading_boundary(env₀, pf, alg) + + nw_value = network_value(pf, env) + lz = real(log(nw_value)) + fs = lz * -1 / ising_βc_triangular + @test fs ≈ f_onsager_triangular rtol = 1.0e-4 + end + end +end + + +## Gradient tests for triangular CTMRG +# ------------------------------------ + +Pspaces = [ComplexSpace(2)] +Vspaces = [ComplexSpace(D)] +Espaces = [ComplexSpace(χ)] +# TODO: actually add triangular lattice models... +models = [transverse_field_ising(InfiniteSquare(); J = 1.0, g = 1.0)] +names = ["Heisenberg"] + +gradtol = 1.0e-4 +ctmrg_verbosity = 0 +ctmrg_algs = [[:triangular]] +projector_algs = [[:twothirds, :full]] +conditionings = [[true, false]] +# svd_rrule_algs = nothing # TODO: allow this +gradient_algs = [[nothing]] +# gradient_iterschemes = [[:fixed, :diffgauge]] # TODO: support this +steps = -0.01:0.005:0.01 + +# have to use a custom retraction, since the state type of hardcoded... +function peps_triangular_retract(x, η, α) + peps = x[1] + env = deepcopy(x[2]) + + retractions = norm_preserving_retract.(unitcell(peps), unitcell(η), α) + peps´ = typeof(peps)(map(first, retractions)) + ξ = typeof(peps)(map(last, retractions)) + + return (peps´, env), ξ +end + + +@testset "Triangular CTMRG energy gradient test for $(names[i]) model" verbose = true for i in + eachindex( + models + ) + Pspace = Pspaces[i] + Vspace = Vspaces[i] + Espace = Espaces[i] + # calgs = ctmrg_algs[i] + palgs = projector_algs[i] + conds = conditionings[i] + # salgs = svd_rrule_algs[i] + galgs = gradient_algs[i] + # gischemes = gradient_iterschemes[i] + + # TODO: allow passing actual LocalOperator + H = models[i] + O1 = last(first(filter(t -> length(first(t)) == 1, H.terms))) + O2 = last(first(filter(t -> length(first(t)) == 2, H.terms))) + + @testset "ctmrg_alg=:simultaneous_triangular, projector_alg=:$projector_alg, conditioning=$conditioning, gradient_alg=:$gradient_alg" for ( + projector_alg, conditioning, gradient_alg, + ) in Iterators.product(palgs, conds, galgs) + + @info "optimtest of ctmrg_alg=:simultaneous_triangular, projector_alg=:$projector_alg, conditioning=$conditioning, gradient_alg=:$gradient_alg on $(names[i])" + Random.seed!(42039482030) + dir = InfinitePEPSTriangular(Pspace, Vspace) + psi = InfinitePEPSTriangular(Pspace, Vspace) + + # TODO: actually go through CTMRGAlgorithm selection + contrete_ctmrg_alg = SimultaneousCTMRGTriangular(; + projector_alg, conditioning, verbosity = ctmrg_verbosity, + trunctype = :FixedSpaceTruncation + ) + # instantiate because hook_pullback doesn't go through the keyword selector... + concrete_gradient_alg = if isnothing(gradient_alg) + nothing # TODO: add this to the PEPSKit.GradMode selector? + else + PEPSKit.GradMode(; alg = gradient_alg, tol = gradtol, iterscheme = gradient_iterscheme) + end + env, = leading_boundary( + CTMRGEnvTriangular(randn, ComplexF64, Vspace ⊗ Vspace', Espace), + psi, contrete_ctmrg_alg + ) + alphas, fs, dfs1, dfs2 = OptimKit.optimtest( + (psi, env), + dir; + alpha = steps, + retract = peps_triangular_retract, + inner = PEPSKit.real_inner, + ) do (peps, env) + E, g = Zygote.withgradient(peps) do psi + env2, = PEPSKit.hook_pullback( + leading_boundary, + env, + psi, + contrete_ctmrg_alg; + alg_rrule = concrete_gradient_alg, + ) + return PEPSKit.energy(psi, env2, O1, O2) + end + + return E, only(g) + end + @test dfs1 ≈ dfs2 atol = 1.0e-2 + end +end diff --git a/test/testsuite/ctmrg/unitcell.jl b/test/testsuite/ctmrg/unitcell.jl index fb7b24e70..abcd74ebe 100644 --- a/test/testsuite/ctmrg/unitcell.jl +++ b/test/testsuite/ctmrg/unitcell.jl @@ -31,11 +31,11 @@ function test_unitcell( Pspaces, [ (c,) => adapt( - AT, randn( - scalartype(peps), - Pspaces[c], Pspaces[c], - ) - ) for c in CartesianIndices(unitcell) + AT, randn( + scalartype(peps), + Pspaces[c], Pspaces[c], + ) + ) for c in CartesianIndices(unitcell) ]..., ) @test expectation_value(peps, random_op, env) isa Number