Refactor alternating_update to use sweep_regions

src/ITensorNetworks.jl

@@ -105,8 +105,6 @@ include(joinpath("treetensornetworks", "projttns", "projttnsum.jl"))
 include(joinpath("treetensornetworks", "projttns", "projttn_apply.jl"))
 include(joinpath("treetensornetworks", "solvers", "solver_utils.jl"))
 include(joinpath("treetensornetworks", "solvers", "applyexp.jl"))
-include(joinpath("treetensornetworks", "solvers", "tdvporder.jl"))
-include(joinpath("treetensornetworks", "solvers", "tdvpinfo.jl"))
 include(joinpath("treetensornetworks", "solvers", "update_step.jl"))
 include(joinpath("treetensornetworks", "solvers", "alternating_update.jl"))
 include(joinpath("treetensornetworks", "solvers", "tdvp.jl"))

src/treetensornetworks/solvers/alternating_update.jl

@@ -3,7 +3,7 @@ function _extend_sweeps_param(param, nsweeps)
   if param isa Number
     eparam = fill(param, nsweeps)
   else
-    length(param) == nsweeps && return param
+    length(param) >= nsweeps && return param[1:nsweeps]
     eparam = Vector(undef, nsweeps)
     eparam[1:length(param)] = param
     eparam[(length(param) + 1):end] .= param[end]
@@ -23,34 +23,25 @@ function process_sweeps(
   mindim = _extend_sweeps_param(mindim, nsweeps)
   cutoff = _extend_sweeps_param(cutoff, nsweeps)
   noise = _extend_sweeps_param(noise, nsweeps)
-  return maxdim, mindim, cutoff, noise
+  return maxdim, mindim, cutoff, noise, kwargs
 end
 
 function alternating_update(
   solver,
   PH,
   psi0::AbstractTTN;
   checkdone=nothing,
-  tdvp_order=TDVPOrder(2, Base.Forward),
   outputlevel=0,
-  time_start=0.0,
-  time_step=0.0,
   nsweeps=1,
   write_when_maxdim_exceeds::Union{Int,Nothing}=nothing,
   kwargs...,
 )
-  maxdim, mindim, cutoff, noise = process_sweeps(nsweeps; kwargs...)
+  maxdim, mindim, cutoff, noise, kwargs = process_sweeps(nsweeps; kwargs...)
 
   step_observer = get(kwargs, :step_observer!, nothing)
 
   psi = copy(psi0)
 
-  # Keep track of the start of the current time step.
-  # Helpful for tracking the total time, for example
-  # when using time-dependent solvers.
-  # This will be passed as a keyword argument to the
-  # `solver`.
-  current_time = time_start
   info = nothing
   for sw in 1:nsweeps
     if !isnothing(write_when_maxdim_exceeds) && maxdim[sw] > write_when_maxdim_exceeds
@@ -64,31 +55,26 @@ function alternating_update(
 
     sw_time = @elapsed begin
       psi, PH, info = update_step(
-        tdvp_order,
         solver,
         PH,
-        time_step,
         psi;
-        kwargs...,
-        current_time,
         outputlevel,
         sweep=sw,
         maxdim=maxdim[sw],
         mindim=mindim[sw],
         cutoff=cutoff[sw],
         noise=noise[sw],
+        kwargs...,
       )
     end
 
-    current_time += time_step
-
-    update!(step_observer; psi, sweep=sw, outputlevel, current_time)
+    update!(step_observer; psi, sweep=sw, outputlevel)
 
     if outputlevel >= 1
       print("After sweep ", sw, ":")
       print(" maxlinkdim=", maxlinkdim(psi))
       @printf(" maxerr=%.2E", info.maxtruncerr)
-      print(" current_time=", round(current_time; digits=3))
+      #print(" current_time=", round(current_time; digits=3))
       print(" time=", round(sw_time; digits=3))
       println()
       flush(stdout)

src/treetensornetworks/solvers/contract.jl

@@ -41,9 +41,8 @@ function contract(
   ##   )
   ## end
 
-  reverse_step = false
   PH = ProjTTNApply(tn2, tn1)
-  psi = alternating_update(contract_solver, PH, init; nsweeps, reverse_step, kwargs...)
+  psi = alternating_update(contract_solver, PH, init; nsweeps, kwargs...)
 
   return psi
 end

src/treetensornetworks/solvers/dmrg.jl

@@ -20,14 +20,10 @@ end
 Overload of `ITensors.dmrg`.
 """
 function dmrg(H, init::AbstractTTN; kwargs...)
-  reverse_step = false
-  psi = alternating_update(eigsolve_solver(; kwargs...), H, init; reverse_step, kwargs...)
-  return psi
+  return alternating_update(eigsolve_solver(; kwargs...), H, init; kwargs...)
 end
 
 """
 Overload of `KrylovKit.eigsolve`.
 """
-function eigsolve(H, init::AbstractTTN; kwargs...)
-  return dmrg(H, init; kwargs...)
-end
+eigsolve(H, init::AbstractTTN; kwargs...) = dmrg(H, init; kwargs...)

src/treetensornetworks/solvers/dmrg_x.jl

@@ -7,7 +7,7 @@ function dmrg_x_solver(PH, init; kwargs...)
   return U_max, nothing
 end
 
-function dmrg_x(PH, init::AbstractTTN; reverse_step=false, kwargs...)
-  psi = alternating_update(dmrg_x_solver, PH, init; reverse_step, kwargs...)
+function dmrg_x(PH, init::AbstractTTN; kwargs...)
+  psi = alternating_update(dmrg_x_solver, PH, init; kwargs...)
   return psi
 end

src/treetensornetworks/solvers/linsolve.jl

@@ -52,5 +52,5 @@ function linsolve(
   # TODO: Define `itensornetwork_cache`
   # TODO: Define `linsolve_cache`
   P = linsolve_cache(itensornetwork_cache(x₀', A, x₀), itensornetwork_cache(x₀', b))
-  return alternating_update(linsolve_solver, P, x₀; reverse_step=false, kwargs...)
+  return alternating_update(linsolve_solver, P, x₀; kwargs...)
 end

src/treetensornetworks/solvers/tdvp.jl

@@ -4,6 +4,7 @@ function exponentiate_solver(; kwargs...)
     init;
     ishermitian=true,
     issymmetric=true,
+    region,
     solver_krylovdim=30,
     solver_maxiter=100,
     solver_outputlevel=0,
@@ -75,19 +76,53 @@ function _compute_nsweeps(nsteps, t, time_step)
   return nsweeps
 end
 
+function sub_time_steps(order)
+  if order == 1
+    return [1.0]
+  elseif order == 2
+    return [1 / 2, 1 / 2]
+  elseif order == 4
+    s = 1.0 / (2 - 2^(1 / 3))
+    return [s / 2, s / 2, (1 - 2 * s) / 2, (1 - 2 * s) / 2, s / 2, s / 2]
+  else
+    error("Trotter order of $order not supported")
+  end
+end
+
+function tdvp_sweep(
+  order::Int, nsite::Int, time_step::Number, graph::AbstractGraph; kwargs...
+)
+  sweep = []
+  for (substep, fac) in enumerate(sub_time_steps(order))
+    sub_time_step = time_step * fac
+    half = half_sweep(
+      direction(substep),
+      graph,
+      make_region;
+      nsite,
+      region_args=(; substep, time_step=sub_time_step),
+      reverse_args=(; substep, time_step=-sub_time_step),
+      reverse_step=true,
+    )
+    append!(sweep, half)
+  end
+  return sweep
+end
+
 function tdvp(
   solver,
   H,
   t::Number,
   init::AbstractTTN;
   time_step::Number=t,
+  nsite=2,
   nsteps=nothing,
   order::Integer=2,
   kwargs...,
 )
   nsweeps = _compute_nsweeps(nsteps, t, time_step)
-  tdvp_order = TDVPOrder(order, Base.Forward)
-  return alternating_update(solver, H, init; nsweeps, tdvp_order, time_step, kwargs...)
+  sweep_regions = tdvp_sweep(order, nsite, time_step, init; kwargs...)
+  return alternating_update(solver, H, init; nsweeps, sweep_regions, nsite, kwargs...)
 end
 
 """

src/treetensornetworks/solvers/tree_sweeping.jl

@@ -1,76 +1,65 @@
-# 
-# Sweep step
-# 
+direction(step_number) = isodd(step_number) ? Base.Forward : Base.Reverse
 
-"""
-  struct SweepStep{V}
-
-Auxiliary object specifying a single local update step in a tree sweeping algorithm.
-"""
-struct SweepStep{V} # TODO: parametrize on position type instead of vertex type
-  pos::Union{Vector{<:V},NamedEdge{V}}
-  time_direction::Int
-end
-
-# field access
-pos(st::SweepStep) = st.pos
-nsite(st::SweepStep) = isa(pos(st), AbstractEdge) ? 0 : length(pos(st))
-time_direction(st::SweepStep) = st.time_direction
-
-# utility
-current_ortho(st::SweepStep) = current_ortho(typeof(pos(st)), st)
-current_ortho(::Type{<:Vector{<:V}}, st::SweepStep{V}) where {V} = first(pos(st)) # not very clean...
-current_ortho(::Type{NamedEdge{V}}, st::SweepStep{V}) where {V} = src(pos(st))
-
-# reverse
-Base.reverse(s::SweepStep{V}) where {V} = SweepStep{V}(reverse(pos(s)), time_direction(s))
-
-function Base.:(==)(s1::SweepStep{V}, s2::SweepStep{V}) where {V}
-  return pos(s1) == pos(s2) && time_direction(s1) == time_direction(s2)
-end
-
-# 
-# Pre-defined sweeping paradigms
-# 
-
-function one_site_sweep(
-  direction::Base.ForwardOrdering,
-  graph::AbstractGraph{V},
-  root_vertex::V,
-  reverse_step;
-  kwargs...,
-) where {V}
-  edges = post_order_dfs_edges(graph, root_vertex)
-  steps = SweepStep{V}[]
-  for e in edges
-    push!(steps, SweepStep{V}([src(e)], +1))
-    reverse_step && push!(steps, SweepStep{V}(e, -1))
+function make_region(
+  edge;
+  last_edge=false,
+  nsite=1,
+  region_args=(;),
+  reverse_args=region_args,
+  reverse_step=false,
+)
+  if nsite == 1
+    site = ([src(edge)], region_args)
+    bond = (edge, reverse_args)
+    region = reverse_step ? (site, bond) : (site,)
+    if last_edge
+      return (region..., ([dst(edge)], region_args))
+    else
+      return region
+    end
+  elseif nsite == 2
+    sites_two = ([src(edge), dst(edge)], region_args)
+    sites_one = ([dst(edge)], reverse_args)
+    region = reverse_step ? (sites_two, sites_one) : (sites_two,)
+    if last_edge
+      return (sites_two,)
+    else
+      return region
+    end
+  else
+    error("nsite=$nsite not supported in alternating_update / update_step")
   end
-  push!(steps, SweepStep{V}([root_vertex], +1))
-  return steps
 end
 
-function one_site_sweep(direction::Base.ReverseOrdering, args...; kwargs...)
-  return reverse(reverse.(one_site_sweep(Base.Forward, args...; kwargs...)))
+#
+# Helper functions to take a tuple like ([1],[2])
+# and append an empty named tuple to it, giving ([1],[2],(;))
+#
+prepend_missing_namedtuple(t::Tuple) = ((;), t...)
+prepend_missing_namedtuple(t::Tuple{<:NamedTuple,Vararg}) = t
+function append_missing_namedtuple(t::Tuple)
+  return reverse(prepend_missing_namedtuple(reverse(t)))
 end
 
-function two_site_sweep(
-  direction::Base.ForwardOrdering,
-  graph::AbstractGraph{V},
-  root_vertex::V,
-  reverse_step;
+function half_sweep(
+  dir::Base.ForwardOrdering,
+  graph::AbstractGraph,
+  region_function;
+  root_vertex=default_root_vertex(graph),
   kwargs...,
-) where {V}
+)
   edges = post_order_dfs_edges(graph, root_vertex)
-  steps = SweepStep{V}[]
-  for e in edges[1:(end - 1)]
-    push!(steps, SweepStep{V}([src(e), dst(e)], +1))
-    reverse_step && push!(steps, SweepStep{V}([dst(e)], -1))
-  end
-  push!(steps, SweepStep{V}([src(edges[end]), dst(edges[end])], +1))
+  steps = collect(
+    flatten(map(e -> region_function(e; last_edge=(e == edges[end]), kwargs...), edges))
+  )
+  # Append empty namedtuple to each element if not already present
+  steps = append_missing_namedtuple.(to_tuple.(steps))
   return steps
 end
 
-function two_site_sweep(direction::Base.ReverseOrdering, args...; kwargs...)
-  return reverse(reverse.(two_site_sweep(Base.Forward, args...; kwargs...)))
+function half_sweep(dir::Base.ReverseOrdering, args...; kwargs...)
+  return map(
+    region -> (reverse(region[1]), region[2:end]...),
+    reverse(half_sweep(Base.Forward, args...; kwargs...)),
+  )
 end