Changes to the iterator interface

src/ITensorNetworks.jl

@@ -66,6 +66,7 @@ include("solvers/abstract_problem.jl")
 include("solvers/eigsolve.jl")
 include("solvers/applyexp.jl")
 include("solvers/fitting.jl")
+include("solvers/default_kwargs.jl")
 
 include("apply.jl")
 include("inner.jl")

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,39 @@
+"""
+  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 NoComputeStep{S<:AbstractNetworkIterator} <: AbstractNetworkIterator
+  parent::S
 end
 
-region_iterator(T::TupleRegionIterator) = T.region_iterator
+laststep(adapter::NoComputeStep) = laststep(adapter.parent)
+state(adapter::NoComputeStep) = state(adapter.parent)
+increment!(adapter::NoComputeStep) = increment!(adapter.parent)
+compute!(adapter::NoComputeStep) = 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
+NoComputeStep(adapter::NoComputeStep) = adapter
 
 """
-  region_tuples(R::RegionIterator)
+  struct EachRegion{RegionIterator} <: 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.
+Wapper adapter that returns a tuple (region, kwargs) at each step rather than the iterator
+itself.
 """
-region_tuples(R::RegionIterator) = TupleRegionIterator(R)
+struct EachRegion{R<:RegionIterator} <: AbstractNetworkIterator
+  parent::R
+end
+
+# Essential definitions
+Base.length(adapter::EachRegion) = length(adapter.parent)
+state(adapter::EachRegion) = state(adapter.parent)
+increment!(adapter::EachRegion) = state(adapter.parent)
+
+function compute!(adapter::EachRegion)
+  # Do the usual compute! for RegionIterator
+  compute!(adapter.parent)
+  # But now lets return something useful
+  return current_region_plan(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,64 @@ 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_iterator::RegionIterator{<:ApplyExpProblem},
+  local_state;
   nsites,
   exponent_step,
   solver=runge_kutta_solver,
-  outputlevel,
   kws...,
 )
-  iszero(abs(exponent_step)) && return prob, local_state
+  prob = problem(region_iterator)
+
+  iszero(abs(exponent_step)) && return local_state
 
-  local_state, info = solver(
-    x->optimal_map(operator(prob), x), exponent_step, local_state; kws...
+  local_state, _ = solver(
+    x -> optimal_map(operator(prob), x), exponent_step, local_state; kws...
   )
-  if nsites==1
+  if nsites == 1
     curr_reg = current_region(region_iterator)
     next_reg = next_region(region_iterator)
     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; kws...)
+      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 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 +76,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 = sweep_solve(sweep_callback, sweep_iter)
+
   return state(converged_prob)
 end
 
@@ -111,11 +109,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/default_kwargs.jl

@@ -0,0 +1,53 @@
+"""
+  default_kwargs(f, [obj = Any])
+
+Return the default keyword arguments for the function `f`. These defaults may be
+derived from the contents or type of the second arugment `obj`.
+
+## Interface
+
+Given a function `f`, one can optionally set the default keyword arguments for this
+function by specializing either of the following two-argument methods:
+```
+ITensorNetworks.default_kwargs(::typeof(f), prob::AbstractProblem)
+ITensorNetworks.default_kwargs(::typeof(f), ::Type{<:AbstractProblem})
+```
+If one does not require the contents of `prob::Prob` to generate the defaults then it is
+recommended to dispatch on `Type{<:Prob}` directly (second method) so the defaults
+can be accessed without constructing an instance of a `Prob`.
+
+The return value of `default_kwargs` should be a `NamedTuple`, and will overwrite any
+default values set in the function signature.
+"""
+default_kwargs(f) = default_kwargs(f, Any)
+default_kwargs(f, obj) = _default_kwargs_fallback(f, obj)
+
+# To avoid annoying potential method ambiguities.
+function _default_kwargs_fallback(f, iter::RegionIterator)
+  return default_kwargs(f, problem(iter))
+end
+function _default_kwargs_fallback(f, problem::AbstractProblem)
+  return default_kwargs(f, typeof(problem))
+end
+
+# Eventually we reach this if nothing is specialized.
+_default_kwargs_fallback(::Any, ::DataType) = (;)
+
+"""
+  current_kwargs(f, iter::RegionIterator)
+
+Return the keyword arguments to be passed to the function `f` for the current region
+defined by the stateful iterator `iter`.
+"""
+function current_kwargs(f::Function, iter::RegionIterator)
+  region_kwargs = get(current_region_kwargs(iter), Symbol(f, :_kwargs), (;))
+  rv = merge(default_kwargs(f, iter), region_kwargs)
+  return rv
+end
+
+# Generic
+
+# I think these should be set independent of a function, but for now:
+function default_kwargs(::typeof(factorize), ::Any)
+  return (; maxdim=typemax(Int), cutoff=0.0, mindim=1)
+end

src/solvers/eigsolve.jl

@@ -1,4 +1,3 @@
-using Accessors: @set
 using Printf: @printf
 using ITensors: truncerror
 
@@ -14,42 +13,48 @@ 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,
-  solver=eigsolve_solver,
-  kws...,
+function update!(
+  region_iterator::RegionIterator{<:EigsolveProblem}, local_state; outputlevel, solver
 )
-  eigval, local_state = solver(ψ->optimal_map(operator(prob), ψ), local_state; kws...)
-  prob = set_eigenvalue(prob, eigval)
+  prob = problem(region_iterator)
+
+  eigval, local_state = solver(
+    ψ -> optimal_map(operator(prob), ψ),
+    local_state;
+    current_kwargs(solver, region_iterator)...,
+  )
+
+  prob.eigenvalue = eigval
+
   if outputlevel >= 2
     @printf(
       "  Region %s: energy = %.12f\n", current_region(region_iterator), eigenvalue(prob)
     )
   end
-  return prob, local_state
+  return local_state
 end
 
-function sweep_callback(problem::EigsolveProblem; outputlevel, sweep, nsweeps, kws...)
+function default_kwargs(::typeof(update!), ::Type{<:EigsolveProblem})
+  return (; outputlevel=0, solver=eigsolve_solver)
+end
+
+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)))
@@ -59,25 +64,13 @@ function sweep_callback(problem::EigsolveProblem; outputlevel, sweep, nsweeps, k
   end
 end
 
-function eigsolve(
-  operator,
-  init_state;
-  nsweeps,
-  nsites=1,
-  outputlevel=0,
-  extract_kwargs=(;),
-  update_kwargs=(;),
-  insert_kwargs=(;),
-  kws...,
-)
+function eigsolve(operator, init_state; nsweeps, nsites=1, outputlevel=0, 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
-  )
-  prob = sweep_solve(sweep_iter; outputlevel, kws...)
+  sweep_iter = SweepIterator(init_prob, nsweeps; nsites, outputlevel, sweep_kwargs...)
+  prob = 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...)