Changes to the iterator interface (#255)

Co-authored-by: Jack Dunham <[email protected]>
Co-authored-by: Matt Fishman <[email protected]>

Author: 3 people

src/solvers/abstract_problem.jl

@@ -1,4 +1,4 @@
 
 abstract type AbstractProblem end
 
-set_truncation_info(P::AbstractProblem, args...; kws...) = P
+set_truncation_info!(P::AbstractProblem, args...; kws...) = P

src/solvers/adapters.jl

@@ -1,32 +1,46 @@
+"""
+  struct PauseAfterIncrement{S<:AbstractNetworkIterator}
 
-#
-# TupleRegionIterator
-#
-# Adapts outputs to be (region, region_kwargs) tuples
-#
-# More generic design? maybe just assuming RegionIterator
-# or its outputs implement some interface function that
-# generates each tuple?
-#
-
-mutable struct TupleRegionIterator{RegionIter}
-  region_iterator::RegionIter
+Iterator wrapper whos `compute!` function simply returns itself, doing nothing in the 
+process. This allows one to manually call a custom `compute!` or insert their own code it in
+the loop body in place of `compute!`.
+"""
+struct IncrementOnly{S<:AbstractNetworkIterator} <: AbstractNetworkIterator
+  parent::S
 end
 
-region_iterator(T::TupleRegionIterator) = T.region_iterator
+islaststep(adapter::IncrementOnly) = islaststep(adapter.parent)
+state(adapter::IncrementOnly) = state(adapter.parent)
+increment!(adapter::IncrementOnly) = increment!(adapter.parent)
+compute!(adapter::IncrementOnly) = adapter
 
-function Base.iterate(T::TupleRegionIterator, which=1)
-  state = iterate(region_iterator(T), which)
-  isnothing(state) && return nothing
-  (current_region, region_kwargs) = current_region_plan(region_iterator(T))
-  return (current_region, region_kwargs), last(state)
-end
+IncrementOnly(adapter::IncrementOnly) = adapter
 
 """
-  region_tuples(R::RegionIterator)
+  struct EachRegion{SweepIterator} <: AbstractNetworkIterator
 
-The `region_tuples` adapter converts a RegionIterator into an 
-iterator which outputs a tuple of the form (current_region, current_region_kwargs)
-at each step.
+Adapter that flattens each region iterator in the parent sweep iterator into a single
+iterator.
 """
-region_tuples(R::RegionIterator) = TupleRegionIterator(R)
+struct EachRegion{SI<:SweepIterator} <: AbstractNetworkIterator
+  parent::SI
+end
+
+# In keeping with Julia convention.
+eachregion(iter::SweepIterator) = EachRegion(iter)
+
+# Essential definitions
+function islaststep(adapter::EachRegion)
+  region_iter = region_iterator(adapter.parent)
+  return islaststep(adapter.parent) && islaststep(region_iter)
+end
+function increment!(adapter::EachRegion)
+  region_iter = region_iterator(adapter.parent)
+  islaststep(region_iter) ? increment!(adapter.parent) : increment!(region_iter)
+  return adapter
+end
+function compute!(adapter::EachRegion)
+  region_iter = region_iterator(adapter.parent)
+  compute!(region_iter)
+  return adapter
+end

src/solvers/applyexp.jl

@@ -1,5 +1,4 @@
 using Printf: @printf
-using Accessors: @set
 
 @kwdef mutable struct ApplyExpProblem{State} <: AbstractProblem
   operator
@@ -11,66 +10,69 @@ operator(A::ApplyExpProblem) = A.operator
 state(A::ApplyExpProblem) = A.state
 current_exponent(A::ApplyExpProblem) = A.current_exponent
 function current_time(A::ApplyExpProblem)
-  t = im*A.current_exponent
+  t = im * A.current_exponent
   return iszero(imag(t)) ? real(t) : t
 end
 
-set_operator(A::ApplyExpProblem, operator) = (@set A.operator = operator)
-set_state(A::ApplyExpProblem, state) = (@set A.state = state)
-set_current_exponent(A::ApplyExpProblem, exponent) = (@set A.current_exponent = exponent)
-
-function region_plan(A::ApplyExpProblem; nsites, time_step, sweep_kwargs...)
-  return applyexp_regions(state(A), time_step; nsites, sweep_kwargs...)
+# Rename region_plan
+function region_plan(A::ApplyExpProblem; nsites, exponent_step, sweep_kwargs...)
+  # The `exponent_step` kwarg for the `update!` function needs some pre-processing.
+  return applyexp_regions(state(A), exponent_step; nsites, sweep_kwargs...)
 end
 
-function update(
-  prob::ApplyExpProblem,
-  local_state,
-  region_iterator;
+function update!(
+  region_iter::RegionIterator{<:ApplyExpProblem},
+  local_state;
   nsites,
   exponent_step,
   solver=runge_kutta_solver,
-  outputlevel,
-  kws...,
 )
-  iszero(abs(exponent_step)) && return prob, local_state
+  prob = problem(region_iter)
+
+  if iszero(abs(exponent_step))
+    return region_iter, local_state
+  end
 
-  local_state, info = solver(
-    x->optimal_map(operator(prob), x), exponent_step, local_state; kws...
+  solver_kwargs = region_kwargs(solver, region_iter)
+
+  local_state, _ = solver(
+    x -> optimal_map(operator(prob), x), exponent_step, local_state; solver_kwargs...
   )
-  if nsites==1
-    curr_reg = current_region(region_iterator)
-    next_reg = next_region(region_iterator)
+  if nsites == 1
+    curr_reg = current_region(region_iter)
+    next_reg = next_region(region_iter)
     if !isnothing(next_reg) && next_reg != curr_reg
       next_edge = first(edge_sequence_between_regions(state(prob), curr_reg, next_reg))
       v1, v2 = src(next_edge), dst(next_edge)
       psi = copy(state(prob))
       psi[v1], R = qr(local_state, uniqueinds(local_state, psi[v2]))
-      shifted_operator = position(operator(prob), psi, NamedEdge(v1=>v2))
-      R_t, _ = solver(x->optimal_map(shifted_operator, x), -exponent_step, R; kws...)
-      local_state = psi[v1]*R_t
+      shifted_operator = position(operator(prob), psi, NamedEdge(v1 => v2))
+      R_t, _ = solver(
+        x -> optimal_map(shifted_operator, x), -exponent_step, R; solver_kwargs...
+      )
+      local_state = psi[v1] * R_t
     end
   end
 
-  prob = set_current_exponent(prob, current_exponent(prob)+exponent_step)
+  prob.current_exponent += exponent_step
 
-  return prob, local_state
+  return region_iter, local_state
 end
 
-function sweep_callback(
-  problem::ApplyExpProblem;
+function default_sweep_callback(
+  sweep_iterator::SweepIterator{<:ApplyExpProblem};
   exponent_description="exponent",
-  outputlevel,
-  sweep,
-  nsweeps,
+  outputlevel=0,
   process_time=identity,
-  kws...,
 )
   if outputlevel >= 1
+    the_problem = problem(sweep_iterator)
     @printf(
-      "  Current %s = %s, ", exponent_description, process_time(current_exponent(problem))
+      "  Current %s = %s, ",
+      exponent_description,
+      process_time(current_exponent(the_problem))
     )
-    @printf("maxlinkdim=%d", maxlinkdim(state(problem)))
+    @printf("maxlinkdim=%d", maxlinkdim(state(the_problem)))
     println()
     flush(stdout)
   end
@@ -79,19 +81,20 @@ end
 function applyexp(
   init_prob::AbstractProblem,
   exponents;
-  extract_kwargs=(;),
-  update_kwargs=(;),
-  insert_kwargs=(;),
-  outputlevel=0,
-  nsites=1,
+  sweep_callback=default_sweep_callback,
   order=4,
-  kws...,
+  nsites=2,
+  sweep_kwargs...,
 )
   exponent_steps = diff([zero(eltype(exponents)); exponents])
-  sweep_kws = (; outputlevel, extract_kwargs, insert_kwargs, nsites, order, update_kwargs)
-  kws_array = [(; sweep_kws..., time_step=t) for t in exponent_steps]
-  sweep_iter = sweep_iterator(init_prob, kws_array)
-  converged_prob = sweep_solve(sweep_iter; outputlevel, kws...)
+
+  kws_array = [
+    (; order, nsites, sweep_kwargs..., exponent_step) for exponent_step in exponent_steps
+  ]
+  sweep_iter = SweepIterator(init_prob, kws_array)
+
+  converged_prob = problem(sweep_solve!(sweep_callback, sweep_iter))
+
   return state(converged_prob)
 end
 
@@ -111,11 +114,10 @@ function time_evolve(
   time_points,
   init_state;
   process_time=process_real_times,
-  sweep_callback=(
-    a...; k...
-  )->sweep_callback(a...; exponent_description="time", process_time, k...),
-  kws...,
+  sweep_callback=iter ->
+    default_sweep_callback(iter; exponent_description="time", process_time),
+  sweep_kwargs...,
 )
-  exponents = [-im*t for t in time_points]
-  return applyexp(operator, exponents, init_state; sweep_callback, kws...)
+  exponents = [-im * t for t in time_points]
+  return applyexp(operator, exponents, init_state; sweep_callback, sweep_kwargs...)
 end

src/solvers/eigsolve.jl

@@ -1,4 +1,3 @@
-using Accessors: @set
 using Printf: @printf
 using ITensors: truncerror
 
@@ -14,42 +13,43 @@ state(E::EigsolveProblem) = E.state
 operator(E::EigsolveProblem) = E.operator
 max_truncerror(E::EigsolveProblem) = E.max_truncerror
 
-set_operator(E::EigsolveProblem, operator) = (@set E.operator = operator)
-set_eigenvalue(E::EigsolveProblem, eigenvalue) = (@set E.eigenvalue = eigenvalue)
-set_state(E::EigsolveProblem, state) = (@set E.state = state)
-set_max_truncerror(E::EigsolveProblem, truncerror) = (@set E.max_truncerror = truncerror)
-
-function set_truncation_info(E::EigsolveProblem; spectrum=nothing)
+function set_truncation_info!(E::EigsolveProblem; spectrum=nothing)
   if !isnothing(spectrum)
-    E = set_max_truncerror(E, max(max_truncerror(E), truncerror(spectrum)))
+    E.max_truncerror = max(max_truncerror(E), truncerror(spectrum))
   end
   return E
 end
 
-function update(
-  prob::EigsolveProblem,
-  local_state,
-  region_iterator;
-  outputlevel,
+function update!(
+  region_iter::RegionIterator{<:EigsolveProblem},
+  local_state;
+  outputlevel=0,
   solver=eigsolve_solver,
-  kws...,
 )
-  eigval, local_state = solver(ψ->optimal_map(operator(prob), ψ), local_state; kws...)
-  prob = set_eigenvalue(prob, eigval)
+  prob = problem(region_iter)
+
+  eigval, local_state = solver(
+    ψ -> optimal_map(operator(prob), ψ), local_state; region_kwargs(solver, region_iter)...
+  )
+
+  prob.eigenvalue = eigval
+
   if outputlevel >= 2
-    @printf(
-      "  Region %s: energy = %.12f\n", current_region(region_iterator), eigenvalue(prob)
-    )
+    @printf("  Region %s: energy = %.12f\n", current_region(region_iter), eigenvalue(prob))
   end
-  return prob, local_state
+  return region_iter, local_state
 end
 
-function sweep_callback(problem::EigsolveProblem; outputlevel, sweep, nsweeps, kws...)
+function default_sweep_callback(
+  sweep_iterator::SweepIterator{<:EigsolveProblem}; outputlevel=0
+)
   if outputlevel >= 1
-    if nsweeps >= 10
-      @printf("After sweep %02d/%d ", sweep, nsweeps)
+    nsweeps = length(sweep_iterator)
+    current_sweep = sweep_iterator.which_sweep
+    if length(sweep_iterator) >= 10
+      @printf("After sweep %02d/%d ", current_sweep, nsweeps)
     else
-      @printf("After sweep %d/%d ", sweep, nsweeps)
+      @printf("After sweep %d/%d ", current_sweep, nsweeps)
     end
     @printf("eigenvalue=%.12f", eigenvalue(problem))
     @printf(" maxlinkdim=%d", maxlinkdim(state(problem)))
@@ -60,24 +60,22 @@ function sweep_callback(problem::EigsolveProblem; outputlevel, sweep, nsweeps, k
 end
 
 function eigsolve(
-  operator,
-  init_state;
-  nsweeps,
-  nsites=1,
-  outputlevel=0,
-  extract_kwargs=(;),
-  update_kwargs=(;),
-  insert_kwargs=(;),
-  kws...,
+  operator, init_state; nsweeps, nsites=1, outputlevel=0, factorize_kwargs, sweep_kwargs...
 )
   init_prob = EigsolveProblem(;
     state=align_indices(init_state), operator=ProjTTN(align_indices(operator))
   )
-  sweep_iter = sweep_iterator(
-    init_prob, nsweeps; nsites, outputlevel, extract_kwargs, update_kwargs, insert_kwargs
+  sweep_iter = SweepIterator(
+    init_prob,
+    nsweeps;
+    nsites,
+    outputlevel,
+    factorize_kwargs,
+    subspace_expand!_kwargs=(; eigen_kwargs=factorize_kwargs),
+    sweep_kwargs...,
   )
-  prob = sweep_solve(sweep_iter; outputlevel, kws...)
+  prob = problem(sweep_solve!(sweep_iter))
   return eigenvalue(prob), state(prob)
 end
 
-dmrg(args...; kws...) = eigsolve(args...; kws...)
+dmrg(operator, init_state; kwargs...) = eigsolve(operator, init_state; kwargs...)

src/solvers/extract.jl

@@ -1,12 +1,16 @@
-function extract(problem, region_iterator; sweep, trunc=(;), kws...)
-  trunc = truncation_parameters(sweep; trunc...)
-  region = current_region(region_iterator)
-  psi = orthogonalize(state(problem), region)
+function extract!(region_iter::RegionIterator; subspace_algorithm="nothing")
+  prob = problem(region_iter)
+  region = current_region(region_iter)
+
+  psi = orthogonalize(state(prob), region)
   local_state = prod(psi[v] for v in region)
-  problem = set_state(problem, psi)
-  problem, local_state = subspace_expand(
-    problem, local_state, region_iterator; sweep, trunc, kws...
-  )
-  shifted_operator = position(operator(problem), state(problem), region)
-  return set_operator(problem, shifted_operator), local_state
+
+  prob.state = psi
+
+  _, local_state = subspace_expand!(region_iter, local_state; subspace_algorithm)
+  shifted_operator = position(operator(prob), state(prob), region)
+
+  prob.operator = shifted_operator
+
+  return region_iter, local_state
 end