add backend and allocator to DMRG and IDMRG

Author: AFeuerpfeil

src/algorithms/groundstate/dmrg.jl

@@ -22,6 +22,8 @@ struct DMRG{A, F} <: Algorithm
     "algorithm used for the eigenvalue solvers"
     alg_eigsolve::A
 
+    prepare::Bool
+
     trscheme::TruncationStrategy
     expscheme::TruncationStrategy
 
@@ -31,19 +33,21 @@ end
 function DMRG(;
         tol = Defaults.tol, maxiter = Defaults.maxiter, alg_eigsolve = (;),
         verbosity = Defaults.verbosity, finalize = Defaults._finalize,
-        miniter = 0, trscheme = notrunc(), expscheme = noexpand()
+        miniter = 0, trscheme = notrunc(), expscheme = noexpand(),
+        prepare=true
     )
     alg_eigsolve′ = alg_eigsolve isa NamedTuple ? Defaults.alg_eigsolve(; alg_eigsolve...) :
         alg_eigsolve
-    return DMRG(tol, maxiter, miniter, verbosity, alg_eigsolve′, trscheme, expscheme, finalize)
+    return DMRG(tol, maxiter, miniter, verbosity, alg_eigsolve′, prepare, trscheme, expscheme, finalize)
 end
 
-function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG, envs = environments(ψ, H))
+function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG, envs = environments(ψ, H), backend=DefaultBackend(), allocator=BufferAllocator())
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
     log = IterLog("DMRG")
     timeroutput = TimerOutput("DMRG")
     alg.verbosity > 3 || disable_timer!(timeroutput)
+    prepare = alg.prepare
 
     LoggingExtras.withlevel(; alg.verbosity) do
         @infov 2 loginit!(log, ϵ, expectation_value(ψ, H, envs))
@@ -59,7 +63,7 @@ function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG, envs = environm
                 for pos in [1:(length(ψ) - 1); length(ψ):-1:2]
                     @timeit timeroutput "AC_eigsolve" begin
                         pos == length(ψ) && (dir = -1)
-                        h = AC_hamiltonian(pos, ψ, H, ψ, envs)
+                        h = AC_hamiltonian(pos, ψ, H, ψ, envs; prepare, backend, allocator)
                         _, vec = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
                     end
                     ϵs[pos] = max(ϵs[pos], calc_galerkin(pos, ψ, H, ψ, envs))
@@ -125,6 +129,7 @@ struct DMRG2{A, S, F} <: Algorithm
     "algorithm used for the singular value decomposition"
     alg_svd::S
 
+    prepare::Bool
     "algorithm used for [truncation](@extref MatrixAlgebraKit.TruncationStrategy) of the two-site update"
     trscheme::TruncationStrategy
     expscheme::TruncationStrategy
@@ -136,14 +141,14 @@ end
 function DMRG2(;
         tol = Defaults.tol, maxiter = Defaults.maxiter, verbosity = Defaults.verbosity,
         miniter = 0, alg_eigsolve = (;), alg_svd = Defaults.alg_svd(), trscheme = notrunc(),
-        expscheme = noexpand(), finalize = Defaults._finalize
+        expscheme = noexpand(), finalize = Defaults._finalize, prepare=true
     )
     alg_eigsolve′ = alg_eigsolve isa NamedTuple ? Defaults.alg_eigsolve(; alg_eigsolve...) :
         alg_eigsolve
-    return DMRG2(tol, maxiter, miniter, verbosity, alg_eigsolve′, alg_svd, trscheme, expscheme, finalize)
+    return DMRG2(tol, maxiter, miniter, verbosity, alg_eigsolve′, alg_svd, prepare, trscheme, expscheme, finalize)
 end
 
-function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG2, envs = environments(ψ, H), prepare=false)
+function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG2, envs = environments(ψ, H), backend=DefaultBackend(), allocator=BufferAllocator())
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
     log = IterLog("DMRG2")
@@ -164,15 +169,15 @@ function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG2, envs = environ
                 for pos in 1:(length(ψ) - 1)
                     local ac2, newA2center, al, c, ar
                     @timeit timeroutput "AC2_eigsolve" begin
-                        @plansor ac2[-1 -2; -3 -4] := ψ.AC[pos][-1 -2; 1] * ψ.AR[pos + 1][1 -4; -3]
-                        Hac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; prepare)
+                        @plansor backend=backend allocator=allocator ac2[-1 -2; -3 -4] := ψ.AC[pos][-1 -2; 1] * ψ.AR[pos + 1][1 -4; -3]
+                        Hac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; prepare, backend, allocator)
                         _, newA2center = fixedpoint(Hac2, ac2, :SR, alg_eigsolve)
                     end
                     @timeit timeroutput "svd_trunc" begin
                         al, c, ar = svd_trunc!(newA2center; trunc = trscheme, alg = alg.alg_svd)
                         al, c = changebonds_left(al, c, expscheme; ac2 = ac2)
                 normalize!(c)
-                        v = @plansor ac2[1 2; 3 4] * conj(al[1 2; 5]) * conj(c[5; 6]) * conj(ar[6; 3 4])
+                        v = @plansor backend=backend allocator=allocator ac2[1 2; 3 4] * conj(al[1 2; 5]) * conj(c[5; 6]) * conj(ar[6; 3 4])
                         ϵs[pos] = max(ϵs[pos], abs(1 - abs(v)))
                     end
                     @timeit timeroutput "update_AC" begin
@@ -185,15 +190,15 @@ function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG2, envs = environ
                 for pos in (length(ψ) - 2):-1:1
                     local ac2, newA2center, al, c, ar
                     @timeit timeroutput "AC2_eigsolve" begin
-                        @plansor ac2[-1 -2; -3 -4] := ψ.AL[pos][-1 -2; 1] * ψ.AC[pos + 1][1 -4; -3]
-                        Hac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; prepare)
+                        @plansor backend=backend allocator=allocator ac2[-1 -2; -3 -4] := ψ.AL[pos][-1 -2; 1] * ψ.AC[pos + 1][1 -4; -3]
+                        Hac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; prepare = prepare, backend = backend, allocator = allocator)
                         _, newA2center = fixedpoint(Hac2, ac2, :SR, alg_eigsolve)
                     end
                     @timeit timeroutput "svd_trunc" begin
                         al, c, ar = svd_trunc!(newA2center; trunc = trscheme, alg = alg.alg_svd)
                         c, ar = changebonds_right(c, ar, expscheme; ac2 = ac2)
-                normalize!(c)
-                        v = @plansor ac2[1 2; 3 4] * conj(al[1 2; 5]) * conj(c[5; 6]) * conj(ar[6; 3 4])
+                        normalize!(c)
+                        v = @plansor backend=backend allocator=allocator ac2[1 2; 3 4] * conj(al[1 2; 5]) * conj(c[5; 6]) * conj(ar[6; 3 4])
                         ϵs[pos] = max(ϵs[pos], abs(1 - abs(v)))
                     end
                     @timeit timeroutput "update_AC" begin

src/algorithms/groundstate/idmrg.jl

@@ -69,20 +69,22 @@ end
 
 
 # Internal state of the IDMRG algorithm
-struct IDMRGState{S, O, E, T}
+struct IDMRGState{S, O, E, T, B, A}
     mps::S
     operator::O
     envs::E
     iter::Int
     ϵ::Float64 # TODO: Could be any <:Real
     energy::T
     timeroutput::TimerOutput
+    backend::B
+    allocator::A
 end
 function IDMRGState{T}(
         mps::S, operator::O, envs::E, iter::Int, ϵ::Float64, energy,
-        timeroutput::TimerOutput,
-    ) where {S, O, E, T}
-    return IDMRGState{S, O, E, T}(mps, operator, envs, iter, ϵ, T(energy), timeroutput)
+        timeroutput::TimerOutput, backend::B, allocator::A = BufferAllocator()
+    ) where {S, O, E, T, B, A}
+    return IDMRGState{S, O, E, T, B, A}(mps, operator, envs, iter, ϵ, T(energy), timeroutput, backend, allocator)
 end
 
 function find_groundstate!(
@@ -105,7 +107,7 @@ function find_groundstate!(
         end
     end
 
-    state = IDMRGState(mps, operator, envs, iter, ϵ, E, timeroutput)
+    state = IDMRGState(mps, operator, envs, iter, ϵ, E, timeroutput, DefaultBackend(), BufferAllocator())
     it = IterativeSolver(alg, state)
 
     return LoggingExtras.withlevel(; alg.verbosity) do
@@ -146,7 +148,7 @@ function Base.iterate(
 
     # update state
     it.state = IDMRGState{T}(
-        mps, state.operator, envs, state.iter + 1, ϵ, E_new, timeroutput,
+        mps, state.operator, envs, state.iter + 1, ϵ, E_new, timeroutput, state.backend, state.allocator
     )
 
     return (mps, envs, ϵ, ΔE), it.state
@@ -160,7 +162,7 @@ function localupdate_step!(
     trscheme = updatetruncation(it.trscheme; iter = state.iter)
     return _localupdate_sweep_idmrg!(
         state.mps, state.operator, state.envs, alg_eigsolve, state.timeroutput,
-        trscheme, expscheme
+        trscheme, expscheme, state.backend, state.allocator
     )
 end
 
@@ -173,46 +175,46 @@ function localupdate_step!(
     return _localupdate_sweep_idmrg2!(
         state.mps, state.operator, state.envs, alg_eigsolve,
         trscheme, it.alg_svd, state.timeroutput,
-        expscheme)
+        expscheme, state.backend, state.allocator)
 end
 
-function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, timeroutput::TimerOutput, alg_trscheme, expscheme)
+function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, timeroutput::TimerOutput, alg_trscheme, expscheme, backend = DefaultBackend(), allocator = BufferAllocator())
     local E
     C_old = ψ.C[0]
 
     # left to right sweep
     _idmrg_move_right!(ψ, H, envs, 1, alg_trscheme, expscheme) # We don't update the first site, as the backwards sweep will update it!
     for pos in 2:length(ψ)
         println("Left: pos = $pos")
-        h = AC_hamiltonian(pos, ψ, H, ψ, envs)
+        h = AC_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator)
         _, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
         _idmrg_move_right!(ψ, H, envs, pos, alg_trscheme, expscheme)
     end
-    _idmrg_move_left!(ψ, H, envs, length(ψ), alg_trscheme, expscheme)
+    _idmrg_move_left!(ψ, H, envs, length(ψ), alg_trscheme, expscheme, backend, allocator)
 
     # right to left sweep
     for pos in length(ψ)-1:-1:1
         println("Right: pos = $pos")
-        h = AC_hamiltonian(pos, ψ, H, ψ, envs)
+        h = AC_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator)
         E, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
 
-        _idmrg_move_left!(ψ, H, envs, pos, alg_trscheme, expscheme)
+        _idmrg_move_left!(ψ, H, envs, pos, alg_trscheme, expscheme, backend, allocator)
     end
 
 
     return ψ, envs, C_old, E
 end
 
-function _idmrg_move_right!(ψ, H,envs, pos, alg_trscheme, expscheme)
+function _idmrg_move_right!(ψ, H,envs, pos, alg_trscheme, expscheme, backend = DefaultBackend(), allocator = BufferAllocator())
     ψ.AL[pos], ψ.C[pos] = left_orth!(ψ.AC[pos]; trunc = alg_trscheme)
     ψ.AL[pos], ψ.C[pos] = changebonds_left(ψ.AL[pos], ψ.C[pos], expscheme)
     ψ.AC[pos + 1] = _mul_front(ψ.C[pos], ψ.AR[pos + 1])
     if pos == length(ψ)
         ψ.AL[pos + 1] = ψ.AC[pos + 1] / ψ.C[pos + 1]
     end
-    transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
+    transfer_leftenv!(envs, ψ, H, ψ, pos + 1, backend, allocator)
 end
-function _idmrg_move_left!(ψ, H, envs, pos, alg_trscheme, expscheme)
+function _idmrg_move_left!(ψ, H, envs, pos, alg_trscheme, expscheme, backend = DefaultBackend(), allocator = BufferAllocator())
     C, temp = right_orth!(_transpose_tail(ψ.AC[pos]); trunc = alg_trscheme)
     C, temp = changebonds_right(C, temp, expscheme)
     ψ.C[pos] = C
@@ -222,15 +224,15 @@ function _idmrg_move_left!(ψ, H, envs, pos, alg_trscheme, expscheme)
         ψ.AR[pos - 1] = ψ.AC[pos - 1] / ψ.C[pos - 2]
     end
 
-    transfer_rightenv!(envs, ψ, H, ψ, pos - 1)
+    transfer_rightenv!(envs, ψ, H, ψ, pos - 1, backend, allocator)
 end
 
-function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg_svd, expscheme, tol=1e-8)
+function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg_svd, expscheme, tol=1e-8, backend = DefaultBackend(), allocator = BufferAllocator())
     # sweep from left to right
     for pos in 1:(length(ψ) - 1)
         @timeit timeroutput "AC2_eigsolve" begin
             ac2 = AC2(ψ, pos; kind = :ACAR)
-            h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
+            h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator)
             _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
         end
         @timeit timeroutput "svd_trunc" begin
@@ -244,8 +246,8 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
             ψ.AC[pos + 1] = _transpose_front(c * ar)
         end
         @timeit timeroutput "transfer_env" begin            
-            transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
-            transfer_rightenv!(envs, ψ, H, ψ, pos)
+            transfer_leftenv!(envs, ψ, H, ψ, pos + 1, backend, allocator)
+            transfer_rightenv!(envs, ψ, H, ψ, pos, backend, allocator)
         end
     end
 
@@ -256,7 +258,7 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
     ψ.AC[1] = _mul_tail(ψ.AL[1], ψ.C[1])
     @timeit timeroutput "AC2_eigsolve" begin
         ac2 = AC2(ψ, length(ψ); kind = :ALAC)
-        h_ac2 = AC2_hamiltonian(length(ψ), ψ, H, ψ, envs)
+        h_ac2 = AC2_hamiltonian(length(ψ), ψ, H, ψ, envs; backend, allocator)
         _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
     end
     @timeit timeroutput "svd_trunc" begin
@@ -279,15 +281,15 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
 
     # update environments
     @timeit timeroutput "transfer_env" begin
-        transfer_leftenv!(envs, ψ, H, ψ, 1)
-        transfer_rightenv!(envs, ψ, H, ψ, 0)
+        transfer_leftenv!(envs, ψ, H, ψ, 1, backend, allocator)
+        transfer_rightenv!(envs, ψ, H, ψ, 0, backend, allocator)
     end
 
     # sweep from right to left
     for pos in (length(ψ) - 1):-1:1
         @timeit timeroutput "AC2_eigsolve" begin
             ac2 = AC2(ψ, pos; kind = :ALAC)
-            h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
+            h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs; backend, allocator)
             _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
         end
         @timeit timeroutput "svd_trunc" begin
@@ -302,8 +304,8 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
             ψ.AC[pos + 1] = _transpose_front(c * ar)
         end
         @timeit timeroutput "transfer_env" begin
-            transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
-            transfer_rightenv!(envs, ψ, H, ψ, pos)
+            transfer_leftenv!(envs, ψ, H, ψ, pos + 1, backend, allocator)
+            transfer_rightenv!(envs, ψ, H, ψ, pos, backend, allocator)
         end
     end
 
@@ -314,7 +316,7 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
         ψ.AR[1] = _transpose_front(pinv(ψ.C[0]; atol=tol) * _transpose_tail(ψ.AC[1]))
 
         ac2 = AC2(ψ, 0; kind = :ACAR)
-        h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs)
+        h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs; backend, allocator)
         E, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
     end
     @timeit timeroutput "svd_trunc" begin
@@ -332,8 +334,8 @@ function _localupdate_sweep_idmrg2!(ψ, H, envs, alg_eigsolve, alg_trscheme, alg
     end
 
     @timeit timeroutput "transfer_env" begin
-        transfer_leftenv!(envs, ψ, H, ψ, 1)
-        transfer_rightenv!(envs, ψ, H, ψ, 0)
+        transfer_leftenv!(envs, ψ, H, ψ, 1, backend, allocator)
+        transfer_rightenv!(envs, ψ, H, ψ, 0, backend, allocator)
     end
     return ψ, envs, C_old, E
 end