TimeEvolver-owned allocator

src/algorithms/time_evolution/simpleupdate.jl

@@ -38,6 +38,10 @@ struct SUState{S <: InfiniteState, E <: SUWeight, N <: Number}
     env::E
 end
 
+# Also covers trivial simple update used by SUGauge.
+TimeEvolver(alg::SimpleUpdate, dt::Number, nstep::Int, circuit, state::SUState) =
+    TimeEvolver(alg, dt, nstep, circuit, state, default_allocator(state.psi, SerialScheduler()))
+
 """
     TimeEvolver(
         psi0::Union{InfinitePEPS, InfinitePEPO}, H::LocalOperator, dt::Number,
@@ -118,15 +122,16 @@ function _su_iter!(
 end
 
 """
-One iteration of simple update
+One iteration of simple update. Pass `allocator` to reuse contraction scratch;
+`TimeEvolver` supplies its owned allocator automatically.
 """
 function su_iter(
         state::InfiniteState, circuit::LocalCircuit,
-        alg::SimpleUpdate, env::SUWeight
+        alg::SimpleUpdate, env::SUWeight;
+        allocator = default_allocator(state, SerialScheduler())
     )
     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())
+    # Gates are applied sequentially and share the same scratch allocator.
     for (sites, gate) in circuit.gates
         if length(sites) == 1
             # 1-site gate
@@ -162,7 +167,7 @@ end
 function Base.iterate(it::TimeEvolver{<:SimpleUpdate}, state = it.state)
     iter, t = state.iter, state.t
     (iter == it.nstep) && return nothing
-    psi, env, ϵ = su_iter(state.psi, it.circuit, it.alg, state.env)
+    psi, env, ϵ = su_iter(state.psi, it.circuit, it.alg, state.env; allocator = it.allocator)
     # update internal state
     iter += 1
     t += it.dt

src/algorithms/time_evolution/time_evolve.jl

@@ -6,15 +6,15 @@ Abstract super type for time evolution algorithms of InfinitePEPS or InfinitePEP
 abstract type TimeEvolution end
 
 """
-    mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number}
+    mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number, A}
 
 Iterator for Trotter-based time evolution of InfinitePEPS or InfinitePEPO.
 
 ## Fields
 
 $(TYPEDFIELDS)
 """
-mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number}
+mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number, A}
     "Time evolution algorithm (currently supported: `SimpleUpdate`)"
     alg::TE
     "Trotter time step"
@@ -26,8 +26,15 @@ mutable struct TimeEvolver{TE <: TimeEvolution, C, S, N <: Number}
     "Internal state of the iterator, including the number of
     already performed iterations, evolved time, PEPS/PEPO and its environment"
     state::S
+    "Reusable contraction scratch retained across timesteps and state resets.
+    An evolver must not be advanced concurrently from multiple tasks."
+    allocator::A
 end
 
+# Preserve construction for algorithms that do not yet reuse contraction scratch.
+TimeEvolver(alg::TimeEvolution, dt::Number, nstep::Int, circuit, state) =
+    TimeEvolver(alg, dt, nstep, circuit, state, DefaultAllocator())
+
 Base.iterate(it::TimeEvolver) = iterate(it, it.state)
 
 """

test/testsuite/timeevol/timestep.jl

@@ -15,6 +15,7 @@ function timeevol_timestep(AT)
         dt, nstep = 1.0e-2, 50
         # manual timestep
         evolver = TimeEvolver(ψ0, H, dt, nstep, alg, env0)
+        allocator = evolver.allocator
         ψ1, env1, info1 = deepcopy(ψ0), deepcopy(env0), nothing
         for iter in 0:(nstep - 1)
             ψ1, env1, info1 = timestep(evolver, ψ1, env1)
@@ -24,12 +25,17 @@ function timeevol_timestep(AT)
         # for-loop syntax
         # Restart the exhausted evolver using the public timestep keywords.
         ψ3, env3, info3 = timestep(evolver, ψ0, env0; iter = 0, t = 0.0)
+        saved = (ψ3, env3, info3)
+        snapshot = deepcopy(saved)
+        @test (ψ3, env3, info3.ϵ) == PEPSKit.su_iter(ψ0, evolver.circuit, evolver.alg, env0)
         for state in evolver
             ψ3, env3, info3 = state
         end
         # results should be *exactly* the same
         @test ψ1 == ψ2 == ψ3
         @test env1 == env2 == env3
         @test info1 == info2 == info3
+        @test evolver.allocator === allocator
+        @test saved == snapshot # later scratch reuse must not overwrite earlier outputs
     end
 end