diff --git a/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl b/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl index 2d41d668e..a33ae4f92 100644 --- a/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl +++ b/benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl @@ -14,9 +14,11 @@ const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json") const SUITE = BenchmarkGroup() const MODULES = Dict{String, Symbol}( "ctmrg" => :CTMRGBenchmarks, + "timeevol" => :TimeEvolBenchmarks, ) include("ctmrg/CTMRGBenchmarks.jl") +include("timeevol/TimeEvolBenchmarks.jl") load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...) diff --git a/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl b/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl new file mode 100644 index 000000000..4b2cefd5b --- /dev/null +++ b/benchmark/PEPSKitBenchmarks/timeevol/TimeEvolBenchmarks.jl @@ -0,0 +1,76 @@ +module TimeEvolBenchmarks + +using BenchmarkTools +using TensorKit +using PEPSKit +using PEPSKit.MPSKitModels: hubbard_space + +const SUITE = BenchmarkGroup() + +""" +A simple update scenario on a 2 x 2 unit cell: the Hamiltonian `H` and the initial state `psi`, +evolved with bond dimension `D` (the truncation keeps the initial virtual spaces). +""" +struct SUSpec{H, S} + H::H + psi::S + D::Int + purified::Bool +end + +function heisenberg_peps(D::Int; T = Float64) + H = heisenberg_XYZ(InfiniteSquare(2, 2)) + psi = InfinitePEPS(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2)) + return SUSpec(H, psi, D, true) +end + +function heisenberg_u1_peps(n::Int; T = Float64) + H = heisenberg_XXZ(T, U1Irrep, InfiniteSquare(2, 2)) + V = U1Space(0 => 2n, 1 // 2 => n, -1 // 2 => n) + psi = InfinitePEPS(randn, T, physicalspace(H)[1, 1], V; unitcell = (2, 2)) + return SUSpec(H, psi, dim(V), true) +end + +function hubbard_peps(n::Int; T = Float64) + H = hubbard_model(T, Trivial, U1Irrep, InfiniteSquare(2, 2); t = 1.0, U = 6.0, mu = 3.0) + V = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2n, (1, 1 // 2) => n, (1, -1 // 2) => n) + psi = InfinitePEPS(randn, T, hubbard_space(Trivial, U1Irrep), V; unitcell = (2, 2)) + return SUSpec(H, psi, dim(V), true) +end + +function heisenberg_pepo(D::Int; T = Float64) + H = heisenberg_XYZ(InfiniteSquare(2, 2)) + psi = InfinitePEPO(randn, T, ℂ^2, ℂ^D; unitcell = (2, 2, 1)) + return SUSpec(H, psi, D, false) +end + +""" +Benchmark one simple update iteration, i.e. the application of all Trotter gates of `spec.H`, +starting from weights obtained from a few warm-up iterations. +""" +function su_iter_benchmark(spec::SUSpec; dt = 1.0e-2, nwarmup = 3) + alg = SimpleUpdate(; trunc = truncrank(spec.D), spec.purified) + circuit = PEPSKit.trotterize(spec.H, dt) + psi, env = spec.psi, SUWeight(spec.psi) + for _ in 1:nwarmup + psi, env, = PEPSKit.su_iter(psi, circuit, alg, env) + end + return @benchmarkable PEPSKit.su_iter($psi, $circuit, $alg, $env) +end + +const SCENARIOS = [ + ("Heisenberg PEPS", "Trivial", "D8") => () -> heisenberg_peps(8), + ("Heisenberg PEPS", "Trivial", "D16") => () -> heisenberg_peps(16), + ("Heisenberg PEPS", "U1", "D8") => () -> heisenberg_u1_peps(2), + ("Heisenberg PEPS", "U1", "D16") => () -> heisenberg_u1_peps(4), + ("Hubbard PEPS", "fZ2xU1", "D8") => () -> hubbard_peps(2), + ("Hubbard PEPS", "fZ2xU1", "D16") => () -> hubbard_peps(4), + ("Heisenberg PEPO", "Trivial", "D4") => () -> heisenberg_pepo(4), + ("Heisenberg PEPO", "Trivial", "D6") => () -> heisenberg_pepo(6), +] + +for (groupkeys, spec) in SCENARIOS + SUITE[("su_iter", groupkeys...)] = su_iter_benchmark(spec()) +end + +end # module diff --git a/benchmark/README.md b/benchmark/README.md index 9a9290c4a..f530291d9 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -1,9 +1,13 @@ # PEPSKit benchmarks A [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl)-based suite that exercises -performance-critical paths in PEPSKit. The suite is organised into modules; the only module -currently registered is `ctmrg`, which times a single `ctmrg_iteration` across several CTMRG -algorithms, problem scenarios, and symmetry sectors. +performance-critical paths in PEPSKit. The suite is organised into modules: + +- `ctmrg` times a single `ctmrg_iteration` across several CTMRG algorithms, problem scenarios, + and symmetry sectors. +- `timeevol` times a single simple update iteration (`su_iter`, i.e. one sweep of all Trotter + gates over a 2 x 2 unit cell) for PEPS and PEPO time evolution, across models, symmetry sectors + and bond dimensions. ## Layout @@ -16,11 +20,13 @@ bench/benchmark/ └── PEPSKitBenchmarks/ ├── PEPSKitBenchmarks.jl # module registry (MODULES), load!, loadall! ├── utils/BenchUtils.jl # TOML <-> VectorSpace conversion helpers - └── ctmrg/ - ├── CTMRGBenchmarks.jl # populates SUITE["ctmrg"] - ├── ctmrg_iteration_benchmarks.jl - ├── default.toml # Trivial + NonUniform scenarios - └── su3_hubbard.toml # SU(3) fermionic scenarios + ├── ctmrg/ + │ ├── CTMRGBenchmarks.jl # populates SUITE["ctmrg"] + │ ├── ctmrg_iteration_benchmarks.jl + │ ├── default.toml # Trivial + NonUniform scenarios + │ └── su3_hubbard.toml # SU(3) fermionic scenarios + └── timeevol/ + └── TimeEvolBenchmarks.jl # populates SUITE["timeevol"] ``` New modules are registered in the `MODULES` dict in `PEPSKitBenchmarks.jl`. @@ -54,6 +60,18 @@ julia --project=bench/benchmark bench/benchmark/benchmarks.jl --modules=ctmrg The script only populates `SUITE`; call `run(SUITE)` afterwards (e.g. from an interactive session) to execute. Available IDs are the keys of `MODULES` in `PEPSKitBenchmarks.jl`. +### Simple update + +`SUITE["timeevol"]` is indexed as `"su_iter" → model → symmetry → "D"`, e.g. + +```julia +include("bench/benchmark/benchmarks.jl") +run(SUITE["timeevol"]["su_iter"]["Hubbard PEPS"]["fZ2xU1"]["D16"]) +``` + +The scenarios are defined in the `SCENARIOS` list in `timeevol/TimeEvolBenchmarks.jl`; each +entry maps its group keys to a function that builds the Hamiltonian and initial state. + ### By group key `SUITE` is a nested `BenchmarkGroup` indexed as: diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 2aeb865b6..e3b683987 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -23,6 +23,8 @@ using KrylovKit using KrylovKit: Lanczos, BlockLanczos using TensorOperations, OptimKit +using TensorOperations: DefaultAllocator, DefaultBackend, tensoralloc_add +using TensorKit.Strided: StridedView, sreshape import OptimKit: fixedpoint using ChainRulesCore, Zygote import TupleTools @@ -32,6 +34,7 @@ using MPSKit: MPSTensor, MPOTensor, GenericMPSTensor, MPSBondTensor, ProductTran using MPSKit: InfiniteEnvironments using MPSKit: DynamicTol using MPSKit.DynamicTols: adapt_solver +using MPSKit: default_allocator using MPSKit: @log_initialization, @log_iteration, @log_convergence, @log_nonconvergence using MPSKit.IterativeLoggers: with_verbosity import MPSKit.DynamicTols: _updatetol @@ -49,7 +52,7 @@ import MPSKitModels: transverse_field_ising, heisenberg_XYZ, heisenberg_XXZ, hubbard_model, bose_hubbard_model, tj_model using FiniteDifferences -using OhMyThreads: tmap, tmap! +using OhMyThreads: tmap, tmap!, SerialScheduler using DocStringExtensions include("Defaults.jl") # Include first to allow for docstring interpolation with Defaults values diff --git a/src/algorithms/contractions/absorb_weight.jl b/src/algorithms/contractions/absorb_weight.jl index f6b764a5a..79b27c7a0 100644 --- a/src/algorithms/contractions/absorb_weight.jl +++ b/src/algorithms/contractions/absorb_weight.jl @@ -48,18 +48,58 @@ function absorb_weight( t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight, row::Int, col::Int, virt_axes::NTuple{N, Int}; inv::Bool = false ) where {N} - Np = numout(t) - vax = first(virt_axes) - wt = weight_to_absorb(weights, row, col, vax; inv) - axes, t2 = absorb_first_weight(t, wt, vax) - for vax in Base.tail(virt_axes) - axes, biperm = _permute_to_last(axes, vax + Np) - wt = weight_to_absorb(weights, row, col, vax; inv) - # use `*` to make absorption/removal twist-free - t2 = permute(t2, biperm) * wt + wts = map(ax -> weight_to_absorb(weights, row, col, ax; inv), virt_axes) + T = promote_type(scalartype(t), map(scalartype, wts)...) + return _absorb_weights!(similar(t, T), t, wts, virt_axes) +end + +""" + absorb_weight!(t, weights::SUWeight, rowcol::CartesianIndex{2}, virt_axes; inv = false) + +In-place version of [`absorb_weight`](@ref). +""" +function absorb_weight!( + t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight, + rowcol::CartesianIndex{2}, virt_axes::NTuple{N, Int}; inv::Bool = false + ) where {N} + wts = map(ax -> weight_to_absorb(weights, rowcol[1], rowcol[2], ax; inv), virt_axes) + return _absorb_weights!(t, t, wts, virt_axes) +end + +# Write `src` into `dst` with `weights[i]` absorbed into virtual axis `virtual_axes[i]`, +# where `domain(weights[i]) == domain(src)[virtual_axes[i]]`; `dst === src` is allowed. +# A diagonal weight only rescales the basis vectors of its leg, by a factor that depends on +# the sector of the leg and the index within it. This commutes with any permutation, so it is +# applied directly on each fusion-tree block of `src` in its own leg order; it is equivalent +# to `permute(src, ...) * wt` followed by the inverse permutation. +# Since the domain of `wt` is the space of the leg itself (dual legs get `transpose(wt)` from +# `weight_to_absorb`), `wt` is composed onto the leg without bending it, so even for dual legs +# no twist is involved. +function _absorb_weights!( + dst::AbstractTensorMap, src::AbstractTensorMap, + weights::NTuple{N, DiagonalTensorMap}, virtual_axes::NTuple{N, Int} + ) where {N} + Np, M = numout(src), Val(numind(src)) # static rank, so the reshaped weights are inferred + space(dst) == space(src) || + throw(SpaceMismatch(lazy"destination $(space(dst)) does not match source $(space(src))")) + for (wt, ax) in zip(weights, virtual_axes) + domain(wt)[1] == domain(src)[ax] || + throw(SpaceMismatch(lazy"weight $(space(wt)) does not match virtual axis $ax of $(space(src))")) + end + # Look up every block of the weights only once, reshaped to broadcast along their leg, and + # iterate over the subblocks of `dst` and `src` (which have the same order) with `subblocks`, + # as the per-block `block(wt, c)` and `t[f₁, f₂]` would both redo the structure lookups. + λdicts = map(weights, virtual_axes) do wt, ax + return Dict( + c => sreshape(StridedView(b.diag), ntuple(i -> i == Np + ax ? length(b.diag) : 1, M)) + for (c, b) in blocks(wt) + ) end - perm_back = invperm(axes) - return permute(t2, (perm_back[begin:numout(t)], perm_back[(numout(t) + 1):end])) + for (((f₁, f₂), d), (_, s)) in zip(subblocks(dst), subblocks(src)) + λs = map((λdict, ax) -> λdict[f₂.uncoupled[ax]], λdicts, virtual_axes) + @. d = *(s, λs...) + end + return dst end """ @@ -86,13 +126,3 @@ function weight_to_absorb( (ax == SOUTH || ax == WEST) && return transpose(wt) return wt end - -function absorb_first_weight( - t::Union{PEPSTensor, PEPOTensor}, wt::DiagonalTensorMap, vax::Int - ) - Np = numout(t) - axes, biperm = _permute_to_last(ntuple(identity, numind(t)), vax + Np) - # use `*` to make absorption/removal twist-free - t2 = permute(t, biperm) * wt - return axes, t2 -end diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl index 9b6f811aa..5ed1c4cf8 100644 --- a/src/algorithms/time_evolution/apply_gate.jl +++ b/src/algorithms/time_evolution/apply_gate.jl @@ -20,33 +20,112 @@ function _apply_sitegate( return a′ end +# --- 2-site gates on nearest neighbor bonds --------------------------------------------- +# +# The bond between `s₁` and `s₂` is updated in the native PEPS/PEPO leg order, using the +# bond direction to locate the legs. Each site tensor is split as `T = Q * R`, where the +# reduced tensor `R` holds the physical leg acted on by the gate and the bond leg. The gate is +# applied to `R₁ -- R₂` and split again by a truncated SVD, after which `Q` and the new `R` +# are contracted back into the native leg order. +# +# Temporaries allocated with `allocator` never outlive the function that allocates them: +# each function below frees its own temporaries and resets `allocator` before returning. + +""" + _orient_bond(state, sites, gate) -> (s₁, k₁), (s₂, k₂), gate + +Order the two `sites` of a nearest neighbor bond and the virtual axes `k₁`, `k₂` of their +bond legs, such that the bond leg of `s₁` is non-dual. Putting `s₁` on the codomain side +of the SVD then reproduces the original arrow on the bond, so nothing has to be flipped. +""" +function _orient_bond(state::InfiniteState, sites::Vector{CartesianIndex{2}}, gate::NNGate) + s₁, s₂ = sites + k₁ = first(_nn_bonddir(s₁, s₂)) + k₂ = mod1(k₁ + 2, 4) + T₁ = state[s₁] + # `space` of a domain leg returns its dual + isdual(space(T₁, numout(T₁) + k₁)) && return (s₁, k₁), (s₂, k₂), gate + return (s₂, k₂), (s₁, k₁), permute(gate, ((2, 1), (4, 3))) +end + """ -$(SIGNATURES) + _bond_qr(T, k, gate_ax, [allocator]) -> Q, R, p -Apply 2-site `gate` on the reduced bond tensors `a`, `b` +Split `T = Q * R`, where the reduced tensor `R : x ← p ⊗ v` holds the physical axis `gate_ax` +and the virtual axis `k` (the bond) of `T`, while `Q` holds all other axes. `p` is the +permutation that brings the product back to the leg order of `T`. For example, for a +PEPSTensor `T[p; n e s w]` with the bond `v = e` on its east axis (`k = 2`), ``` - -1← a --- 3 --- b ← -4 -2 -3 - ↓ ↓ ↓ ↓ - 1 2 |----gate---| - ↓ ↓ or ↓ ↓ - |----gate---| 1 2 - ↓ ↓ ↓ ↓ - -2 -3 -1← a --- 3 --- b ← -4 + n n + | | + w - T - v = w - Q ← x ← R - v + | ↘ | ↘ + s p s p ``` +For a PEPOTensor, the physical axis that is not `gate_ax` stays in `Q`. +The permuted copy of `T` that is factorized is a temporary allocated with `allocator`. """ -function _apply_gate(a::MPSTensor, b::MPSTensor, gate::NNGate, trunc::TruncationStrategy) - V = space(b, 1) - need_flip = isdual(V) - if isdual(space(a, 2)) - @tensor a2b2[-1 -2; -3 -4] := gate[1 2; -2 -3] * a[-1 1 3] * b[3 2 -4] +function _bond_qr( + T::Union{PEPSTensor, PEPOTensor}, k::Int, gate_ax::Int, allocator = DefaultAllocator() + ) + Np, M = numout(T), numind(T) + r = (gate_ax, Np + k) + q = TupleTools.deleteat(ntuple(identity, M), r) + cp = allocator_checkpoint!(allocator) + T′ = tensoralloc_add(scalartype(T), T, (q, r), false, Val(true), allocator) + tensoradd!(T′, T, (q, r), false, One(), Zero(), DefaultBackend(), allocator) + Q, R = left_orth!(T′; positive = true) + tensorfree!(T′, allocator) + allocator_reset!(allocator, cp) + return Q, R, _inv_mpo_perm((q..., r...), Val(Np)) +end + +""" + _undo_bond_qr(Q, X, pX, p, [allocator]) + +Inverse of [`_bond_qr`](@ref): contract `Q` with `X`, whose legs selected by `pX` are +ordered as `x ← p ⊗ v`, into the leg order specified by `p`. +""" +function _undo_bond_qr( + Q::AbstractTensorMap, X::AbstractTensorMap, pX, p, allocator = DefaultAllocator() + ) + N = numind(Q) + pQ = (ntuple(identity, N - 1), (N,)) + cp = allocator_checkpoint!(allocator) + T = tensorcontract(Q, pQ, false, X, pX, false, p, One(), DefaultBackend(), allocator) + allocator_reset!(allocator, cp) + return T +end + +""" + _apply_bond_gate(R₁, R₂, gate, trunc, [allocator]) -> U, S, Vᴴ, ϵ + +Apply the 2-site `gate` to the reduced tensors `R₁ : x₁ ← p₁ ⊗ v` and `R₂ : x₂ ← p₂ ⊗ v'` +``` + x₁ ← R₁ --- v --- R₂ → x₂ p₁ p₂ + ↘ ↘ ↓ ↓ + p₁ p₂ |----gate-----| + ↓ ↓ or ↓ ↓ + |----gate----| p₁ p₂ + ↓ ↓ ↘ ↘ + p₁ p₂ x₁ ← R₁ --- v --- R₂ → x₂ +``` +where the second form is used when the physical legs are dual (the second physical axis of +a PEPO), and split the result as `U S Vᴴ` with a truncated SVD. The square root of `S` is +absorbed into `U : x₁ ⊗ p₁ ← v′` and `Vᴴ : v′ ← x₂ ⊗ p₂`. +""" +function _apply_bond_gate( + R₁::AbstractTensorMap, R₂::AbstractTensorMap, gate::NNGate, + trunc::TruncationStrategy, allocator = DefaultAllocator() + ) + θ = if isdual(space(R₁, 2)) + @tensor allocator = allocator θ[-1 -2; -3 -4] := + R₁[-1; 1 3] * R₂[-3; 2 3] * gate[1 2; -2 -4] else - @tensor a2b2[-1 -2; -3 -4] := gate[-2 -3; 1 2] * a[-1 1 3] * b[3 2 -4] - end - a, s, b, ϵ = svd_trunc!(a2b2; trunc) - a, b = absorb_s(a, s, b) - if need_flip - a, s, b = flip(a, numind(a)), _fliptwist_s!(s), flip(b, 1) + @tensor allocator = allocator θ[-1 -2; -3 -4] := + R₁[-1; 1 3] * R₂[-3; 2 3] * gate[-2 -4; 1 2] end - b = permute(b, ((1, 2), (3,))) - return a, s, b, ϵ + U, S, Vᴴ, ϵ = svd_trunc!(θ; trunc) + absorb_s!(U, S, Vᴴ) + return U, S, Vᴴ, ϵ end diff --git a/src/algorithms/time_evolution/apply_mpo.jl b/src/algorithms/time_evolution/apply_mpo.jl index c4a61b68d..511efb47d 100644 --- a/src/algorithms/time_evolution/apply_mpo.jl +++ b/src/algorithms/time_evolution/apply_mpo.jl @@ -247,15 +247,28 @@ function _flip_virtuals!( Ms::Vector{T}, flips::Vector{Bool}; inv::Bool = false ) where {T <: GenericMPSTensor} @assert length(flips) == length(Ms) - 1 - for (n, flip) in enumerate(flips) - !flip && continue + for (n, fl) in enumerate(flips) + !fl && continue M1, M2 = Ms[n], Ms[n + 1] - Ms[n] = TensorKit.flip(M1, numind(M1); inv) - Ms[n + 1] = TensorKit.flip(M2, 1; inv) + Ms[n] = flip(M1, numind(M1); inv) + Ms[n + 1] = flip(M2, 1; inv) end return Ms end +"Apply a bond projector to the left leg of an MPS tensor." +function _apply_proj_left(Pb::MPSBondTensor, m::GenericMPSTensor{S, N}) where {S, N} + pP = ((1,), (2,)) + pM = ((1,), ntuple(i -> i + 1, N)) + pPM = (codomainind(m), domainind(m)) + return tensorcontract(Pb, pP, false, m, pM, false, pPM) +end + +"Apply a bond projector to the right leg of an MPS tensor." +function _apply_proj_right(m::GenericMPSTensor, Pa::MPSBondTensor) + return m * twistdual(Pa, 1) +end + """ Find projectors to truncate internal bonds of the cluster `Ms`. """ @@ -266,11 +279,8 @@ function _cluster_truncate!( # apply projectors # M1 -- (Pa1,wt1,Pb1) -- M2 -- (Pa2,wt2,Pb2) -- M3 for (i, (Pa, Pb)) in enumerate(zip(Pas, Pbs)) - Ms[i] = Ms[i] * twistdual(Pa, 1) - pP = ((1,), (2,)) - pM = ((1,), ntuple(i -> i + 1, numind(eltype(Ms)) - 1)) - pPM = (codomainind(Ms[i + 1]), domainind(Ms[i + 1])) - Ms[i + 1] = tensorcontract(Pb, pP, false, Ms[i + 1], pM, false, pPM) + Ms[i] = _apply_proj_right(Ms[i], Pa) + Ms[i + 1] = _apply_proj_left(Pb, Ms[i + 1]) end return wts, ϵs, Pas, Pbs end diff --git a/src/algorithms/time_evolution/get_cluster.jl b/src/algorithms/time_evolution/get_cluster.jl index 452c0cd3b..da94c9449 100644 --- a/src/algorithms/time_evolution/get_cluster.jl +++ b/src/algorithms/time_evolution/get_cluster.jl @@ -50,43 +50,28 @@ function _nn_vec_direction(nn_vec::CartesianIndex{2}) end """ -Given `site1`, `site2` connected by a nearest neighbor bond, -return the bond index and whether it is reversed from the -standard orientation (`site1` on the west/south of `site2`). +Given `site1`, `site2` connected by a nearest neighbor bond, return `(dir, r, c)` where `dir` ∈ {NORTH, EAST, SOUTH, WEST} is the direction from `site1` to `site2`. +`(r, c)` is the position of the west site for x-bonds (`dir == EAST` or `dir == WEST`) and the south site for y-bonds (`dir == NORTH` or `dir == SOUTH`). """ -function _nn_bondrev(site1::CartesianIndex{2}, site2::CartesianIndex{2}) - diff = site1 - site2 - if diff == CartesianIndex(0, -1) +function _nn_bonddir(site1::CartesianIndex{2}, site2::CartesianIndex{2}) + diff = site2 - site1 + if diff == CartesianIndex(0, 1) r, c = site1[1], site1[2] - return (1, r, c), false - elseif diff == CartesianIndex(0, 1) - r, c = site2[1], site2[2] - return (1, r, c), true - elseif diff == CartesianIndex(1, 0) - r, c = site1[1], site1[2] - return (2, r, c), false + return (EAST, r, c) elseif diff == CartesianIndex(-1, 0) + r, c = site1[1], site1[2] + return (NORTH, r, c) + elseif diff == CartesianIndex(1, 0) + r, c = site2[1], site2[2] + return (SOUTH, r, c) + elseif diff == CartesianIndex(0, -1) r, c = site2[1], site2[2] - return (2, r, c), true + return (WEST, r, c) else error("`site1` and `site2` are not nearest neighbors.") end end -function _bond_rotation(x, bonddir::Int, rev::Bool; inv::Bool = false) - return if bonddir == 1 # x-bond - rev ? rot180(x) : x - elseif bonddir == 2 # y-bond - if rev - inv ? rotr90(x) : rotl90(x) - else - inv ? rotl90(x) : rotr90(x) - end - else - error("`bonddir` must be 1 (for x-bonds) or 2 (for y-bonds).") - end -end - """ Obtain the cluster `Ms` along the (open) path `sites` in `state`. @@ -135,27 +120,3 @@ function _get_cluster( Ms = map(site -> state[site], sites) return Ms, open_vaxs, perms end - -""" -Absorb weights into the open virtual legs of -the cluster `Ms` of PEPS/PEPO tensors. -""" -function _absorb_weight!( - Ms::Vector{T}, sites::Vector{CartesianIndex{2}}, - open_vaxs::Vector{<:Tuple}, env::SUWeight; inv::Bool = false - ) where {T} - for (idx, (M, site, vaxs)) in enumerate(zip(Ms, sites, open_vaxs)) - Ms[idx] = absorb_weight(M, env, site, vaxs; inv) - end - return Ms -end - -""" -Given a vector `Ms` of `AbstractTensorMap`s, -apply permutations `perms` to each tensor. -""" -function _permute_cluster(Ms::Vector{<:AbstractTensorMap}, perms::Vector{<:Tuple}) - return map(Ms, perms) do M, perm - return permute(M, perm) - end -end diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl index 5133f428a..38a3614ea 100644 --- a/src/algorithms/time_evolution/simpleupdate.jl +++ b/src/algorithms/time_evolution/simpleupdate.jl @@ -76,43 +76,44 @@ end """ Simple update optimized for nearest neighbor gates utilizing reduced bond tensors with the physical leg. + +The tensors are updated in their native leg order, without rotating them. +Temporary tensors are allocated with `allocator`. """ function _su_iter!( state::InfiniteState, gate::NNGate, env::SUWeight, - sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate + sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate; + allocator = DefaultAllocator() ) - @assert length(sites) == 2 - trunc = only(_get_cluster_trunc(alg.trunc, sites)) - Ms, open_vaxs, = _get_cluster(state, sites) - _absorb_weight!(Ms, sites, open_vaxs, env) - # rotate - bond, rev = _nn_bondrev(sites...) - dir = first(bond) - A, B = _bond_rotation.(Ms, dir, rev; inv = false) + (s₁, k₁), (s₂, k₂), gate = _orient_bond(state, sites, gate) + trunc = only(_get_cluster_trunc(alg.trunc, [s₁, s₂])) + open₁ = TupleTools.deleteat((1, 2, 3, 4), k₁) + open₂ = TupleTools.deleteat((1, 2, 3, 4), k₂) + # absorb environment weights + T₁ = absorb_weight(state[s₁], env, s₁, open₁) + T₂ = absorb_weight(state[s₂], env, s₂, open₂) # apply gate ϵ = 0.0 - local s - for gate_ax in 1:2 - a, X = bond_tensor_first(A; gate_ax, positive = true) - b, Y = bond_tensor_last(B; gate_ax, positive = true) - a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc) + local S + for gate_ax in 1:numout(T₁) + Q₁, R₁, p₁ = _bond_qr(T₁, k₁, gate_ax, allocator) + Q₂, R₂, p₂ = _bond_qr(T₂, k₂, gate_ax, allocator) + U, S, Vᴴ, ϵ′ = _apply_bond_gate(R₁, R₂, gate, trunc, allocator) ϵ = max(ϵ, ϵ′) - A = undo_bond_tensor_first(a, X; gate_ax) - B = undo_bond_tensor_last(b, Y; gate_ax) + T₁ = _undo_bond_qr(Q₁, U, ((1,), (2, 3)), p₁, allocator) + T₂ = _undo_bond_qr(Q₂, Vᴴ, ((2,), (3, 1)), p₂, allocator) alg.purified && break # only apply gate to 1st physical leg end - # rotate back - A = _bond_rotation(A, dir, rev; inv = true) - B = _bond_rotation(B, dir, rev; inv = true) - rev && (s = transpose(s)) # remove environment weights - siteA, siteB = sites - A = absorb_weight(A, env, siteA[1], siteA[2], open_vaxs[1]; inv = true) - B = absorb_weight(B, env, siteB[1], siteB[2], open_vaxs[2]; inv = true) - # update tensor dict and weight on current bond - state[siteA] = normalize!(A, Inf) - state[siteB] = normalize!(B, Inf) - env[bond...] = normalize!(s, Inf) + absorb_weight!(T₁, env, s₁, open₁; inv = true) + absorb_weight!(T₂, env, s₂, open₂; inv = true) + # update tensor dict and weight on current bond, + # which lives on the EAST (NORTH) leg of the west (south) site + state[s₁] = normalize!(T₁, Inf) + state[s₂] = normalize!(T₂, Inf) + dir, r, c = _nn_bonddir(s₁, s₂) + d = dir in (EAST, WEST) ? 1 : 2 + env[d, r, c] = normalize!(dir in (EAST, NORTH) ? S : transpose(S), Inf) return ϵ end @@ -124,16 +125,19 @@ function su_iter( alg::SimpleUpdate, env::SUWeight ) state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0 + # the gates are applied one after another, so a single allocator serves all of them + allocator = default_allocator(state, SerialScheduler()) for (sites, gate) in circuit.gates if length(sites) == 1 # 1-site gate # TODO: special treatment for bipartite state - site = sites[1] + site = only(sites) state2[site] = _apply_sitegate(state2[site], gate; alg.purified) elseif length(sites) == 2 - (d, r, c), = _nn_bondrev(sites...) + (dir, r, c) = _nn_bonddir(sites...) + d = dir in (EAST, WEST) ? 1 : 2 alg.bipartite && iseven(r) && continue - ϵ′ = _su_iter!(state2, gate, env2, sites, alg) + ϵ′ = _su_iter!(state2, gate, env2, sites, alg; allocator) ϵ = max(ϵ, ϵ′) (!alg.bipartite) && continue if d == 1 @@ -169,17 +173,20 @@ end """ timestep( - it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight + it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight; + iter::Int = it.state.iter, t::Number = it.state.t ) -> (psi, env, info) -Given the `TimeEvolver` iterator `it`, perform one step of time evolution -on the input state `psi` and its environment `env`. +Given the `TimeEvolver` iterator `it`, perform one step of time evolution on the input state `psi` and its environment `env`. + +Use `iter` and `t` to reset the current iteration number and evolved time of `it`, respectively. """ function MPSKit.timestep( - it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight + it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight; + iter::Int = it.state.iter, t::Number = it.state.t ) _timeevol_sanity_check(psi, physicalspace(it.state.psi), it.alg) - state = SUState(it.state.iter, it.state.t, psi, env) + state = SUState(iter, t, psi, env) result = iterate(it, state) if result === nothing @warn "TimeEvolver `it` has already reached the end." diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl index 2b5b87673..1614f6e79 100644 --- a/src/algorithms/time_evolution/simpleupdate3site.jl +++ b/src/algorithms/time_evolution/simpleupdate3site.jl @@ -13,22 +13,32 @@ function _unfuse_physicalspace( return O_unfused, F end +""" +Obtain the cluster along `sites` with weights absorbed into its open virtual legs, +and permuted to MPS axis order. +""" function _get_cluster_with_weights( state::InfiniteState, sites::Vector{CartesianIndex{2}}, env::SUWeight ) Ms, open_vaxs, perms = _get_cluster(state, sites) - _absorb_weight!(Ms, sites, open_vaxs, env; inv = false) Np = (state isa InfinitePEPS) ? Val(1) : Val(2) + Ms_mps = map(Ms, sites, open_vaxs, perms) do M, site, vaxs, perm + return permute(absorb_weight(M, env, site, vaxs), perm) + end invperms = map(p -> _inv_mpo_perm(p, Np), perms) - return _permute_cluster(Ms, perms), open_vaxs, invperms + return Ms_mps, open_vaxs, invperms end """ Simple update with an N-site MPO `gate` (N ≥ 2). + +The `allocator` is accepted for compatibility with the nearest neighbor update, +but is not used yet. """ function _su_iter!( state::InfiniteState, gate::Vector{T}, env::SUWeight, - sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate + sites::Vector{CartesianIndex{2}}, alg::SimpleUpdate; + allocator = DefaultAllocator() ) where {T <: AbstractTensorMap} truncs = _get_cluster_trunc(alg.trunc, sites) Ms, open_vaxs, invperms = _get_cluster_with_weights(state, sites, env) @@ -47,18 +57,19 @@ function _su_iter!( # restore virtual arrows in `Ms` _flip_virtuals!(Ms, flips) # update env weights - bond_revs = map(sites, Iterators.drop(sites, 1)) do site1, site2 - _nn_bondrev(site1, site2) + bond_dirs = map(sites, Iterators.drop(sites, 1)) do site1, site2 + _nn_bonddir(site1, site2) end - for (wt, (bond, rev), flip) in zip(wts, bond_revs, flips) + for (wt, (dir, r, c), flip) in zip(wts, bond_dirs, flips) wt_new = flip ? _fliptwist_s!(wt) : wt - wt_new = rev ? transpose(wt_new) : wt_new - env[CartesianIndex(bond)] = normalize!(wt_new, Inf) + wt_new = (dir in (WEST, SOUTH)) ? transpose(wt_new) : wt_new + d = dir in (EAST, WEST) ? 1 : 2 + env[d, r, c] = normalize!(wt_new, Inf) end # update state tensors for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs) # restore original axes order and remove weights on open axes of the cluster - M′ = absorb_weight(permute(M, invperm), env, s, vaxs; inv = true) + M′ = absorb_weight!(permute(M, invperm), env, s, vaxs; inv = true) state[s] = normalize!(M′, Inf) end return ϵ @@ -72,9 +83,10 @@ function _get_cluster_trunc( trunc::TruncationStrategy, sites::Vector{CartesianIndex{2}} ) return map(sites, Iterators.drop(sites, 1)) do site1, site2 - (d, r, c), rev = _nn_bondrev(site1, site2) + (dir, r, c) = _nn_bonddir(site1, site2) + d = dir in (EAST, WEST) ? 1 : 2 t = truncation_strategy(trunc, d, r, c) - if rev && isa(t, TruncationSpace) + if (dir in (WEST, SOUTH)) && isa(t, TruncationSpace) t = truncspace(flip(t.space)') end return t diff --git a/src/utility/util.jl b/src/utility/util.jl index 4c450ddb4..df25d2ae5 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -66,17 +66,24 @@ end """ absorb_s(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap) + absorb_s!(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap) Given SVD result `U`, `S` and `V`, absorb singular values `S` into `U` and `V` by: ``` U -> U * sqrt(S), V -> sqrt(S) * V ``` +The in-place version `absorb_s!` overwrites `U` and `V`. """ function absorb_s(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap) @assert !isdual(space(S, 1)) sqrt_S = sdiag_pow(S, 0.5) return U * sqrt_S, sqrt_S * V end +function absorb_s!(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorMap) + @assert !isdual(space(S, 1)) + sqrt_S = sdiag_pow(S, 0.5) + return rmul!(U, sqrt_S), lmul!(sqrt_S, V) +end # avoids a conversion to a dense TensorMap by flip function _fliptwist_s!(s::DiagonalTensorMap) @@ -167,17 +174,3 @@ function random_dual!(Vs::AbstractMatrix{E}; p = 0.7) where {E <: ElementarySpac end return Vs end - -""" - _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N} - -Returns `(1, 2, ..., N)` but with `ax` moved to the end, -and the corresponding permutation for `axes` (with `ax` as the only domain index). -""" -function _permute_to_last(axes::NTuple{N, Int}, ax::Int) where {N} - codomain_axes = TupleTools.deleteat(ntuple(identity, N), ax) - q = invperm(axes) - biperm = (map(i -> q[i], codomain_axes), (q[ax],)) - new_axes = (ntuple(i -> axes[biperm[1][i]], N - 1)..., ax) - return new_axes, biperm -end diff --git a/test/testsuite/timeevol/timestep.jl b/test/testsuite/timeevol/timestep.jl index 210a7185f..7a125b435 100644 --- a/test/testsuite/timeevol/timestep.jl +++ b/test/testsuite/timeevol/timestep.jl @@ -22,9 +22,8 @@ function timeevol_timestep(AT) # time_evolve ψ2, env2, info2 = time_evolve(ψ0, H, dt, nstep, alg, env0) # for-loop syntax - ## manually reset internal state of evolver - evolver.state = PEPSKit.SUState(0, 0.0, ψ0, env0) - ψ3, env3, info3 = nothing, nothing, nothing + # Restart the exhausted evolver using the public timestep keywords. + ψ3, env3, info3 = timestep(evolver, ψ0, env0; iter = 0, t = 0.0) for state in evolver ψ3, env3, info3 = state end