Refactor interfaces built around alternating update

src/ITensorNetworks.jl

@@ -106,6 +106,11 @@ include("tensornetworkoperators.jl")
 include(joinpath("ITensorsExt", "itensorutils.jl"))
 include(joinpath("Graphs", "abstractgraph.jl"))
 include(joinpath("Graphs", "abstractdatagraph.jl"))
+include(joinpath("solvers", "eigsolve.jl"))
+include(joinpath("solvers", "exponentiate.jl"))
+include(joinpath("treetensornetworks", "solvers", "applyexp.jl")) #this defines the primitive before the solver function
+include(joinpath("solvers", "applyexp.jl"))
+include(joinpath("solvers", "dmrg_x_solver.jl"))
 include(joinpath("treetensornetworks", "abstracttreetensornetwork.jl"))
 include(joinpath("treetensornetworks", "ttn.jl"))
 include(joinpath("treetensornetworks", "opsum_to_ttn.jl"))
@@ -114,7 +119,6 @@ include(joinpath("treetensornetworks", "projttns", "projttn.jl"))
 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", "update_step.jl"))
 include(joinpath("treetensornetworks", "solvers", "alternating_update.jl"))
 include(joinpath("treetensornetworks", "solvers", "tdvp.jl"))

src/solvers/applyexp.jl

@@ -0,0 +1,25 @@
+function applyexp_solver()
+  function solver(
+    init;
+    psi_ref!,
+    PH_ref!,
+    region,
+    sweep_regions,
+    sweep_step,
+    solver_krylovdim=30,
+    solver_outputlevel=0,
+    solver_tol=1E-8,
+    substep,
+    time_step,
+    normalize,
+  )
+    H = PH_ref![]
+    #applyexp tol is absolute, compute from tol_per_unit_time:
+    tol = abs(time_step) * tol_per_unit_time
+    psi, exp_info = applyexp(
+      H, time_step, init; tol, maxiter=solver_krylovdim, outputlevel=solver_outputlevel
+    )
+    return psi, (; info=exp_info)
+  end
+  return solver
+end

src/solvers/dmrg_x_solver.jl

@@ -0,0 +1,19 @@
+function dmrg_x_solver(
+  init;
+  psi_ref!,
+  PH_ref!,
+  normalize=nothing,
+  region,
+  sweep_regions,
+  sweep_step,
+  half_sweep,
+  step_kwargs...,
+)
+  H = contract(PH_ref![], ITensor(1.0))
+  D, U = eigen(H; ishermitian=true)
+  u = uniqueind(U, H)
+  max_overlap, max_ind = findmax(abs, array(dag(init) * U))
+  U_max = U * dag(onehot(u => max_ind))
+  # TODO: improve this to return the energy estimate too
+  return U_max, NamedTuple()
+end

src/solvers/eigsolve.jl

@@ -0,0 +1,38 @@
+
+function eigsolve_updater(
+  init;
+  psi_ref!,
+  PH_ref!,
+  outputlevel,
+  which_sweep,
+  region_updates,
+  which_region_update,
+  region_kwargs,
+  updater_kwargs,
+)
+  default_updater_kwargs = (;
+    solver_which_eigenvalue=:SR,
+    ishermitian=true,
+    tol=1e-14,
+    krylovdim=3,
+    maxiter=1,
+    outputlevel=0,
+    eager=false,
+  )
+  updater_kwargs = merge(default_updater_kwargs, updater_kwargs)  #last collection has precedence
+  howmany = 1
+  which = updater_kwargs[:solver_which_eigenvalue]
+  vals, vecs, info = KrylovKit.eigsolve(
+    PH_ref![],
+    init,
+    howmany,
+    which;
+    ishermitian=updater_kwargs[:ishermitian],
+    tol=updater_kwargs[:tol],
+    krylovdim=updater_kwargs[:krylovdim],
+    maxiter=updater_kwargs[:maxiter],
+    verbosity=updater_kwargs[:outputlevel],
+    eager=updater_kwargs[:eager],
+  )
+  return vecs[1], (; info, energies=vals)
+end

src/solvers/exponentiate.jl

@@ -0,0 +1,48 @@
+function exponentiate_updater(
+  init;
+  psi_ref!,
+  PH_ref!,
+  outputlevel,
+  which_sweep,
+  region_updates,
+  which_region_update,
+  region_kwargs,
+  updater_kwargs,
+)
+  default_updater_kwargs = (;
+    krylovdim=30,  #from here only solver kwargs
+    maxiter=100,
+    outputlevel=0,
+    tol=1E-12,
+    ishermitian=true,
+    issymmetric=true,
+    eager=true,
+  )
+  updater_kwargs = merge(default_updater_kwargs, updater_kwargs)  #last collection has precedence
+  #H=copy(PH_ref![])
+  H = PH_ref![] ###since we are not changing H we don't need the copy
+  # let's test whether given region and sweep regions we can find out what the previous and next region were
+  # this will be needed in subspace expansion
+  #@show step_kwargs
+  substep = get(region_kwargs, :substep, nothing)
+  time_step = get(region_kwargs, :time_step, nothing)
+  @assert !isnothing(time_step) && !isnothing(substep)
+  region_ind = which_region_update
+  next_region =
+    region_ind == length(region_updates) ? nothing : first(region_updates[region_ind + 1])
+  previous_region = region_ind == 1 ? nothing : first(region_updates[region_ind - 1])
+
+  phi, exp_info = KrylovKit.exponentiate(
+    H,
+    time_step,
+    init;
+    ishermitian=updater_kwargs[:ishermitian],
+    issymmetric=updater_kwargs[:issymmetric],
+    tol=updater_kwargs[:tol],
+    krylovdim=updater_kwargs[:krylovdim],
+    maxiter=updater_kwargs[:maxiter],
+    verbosity=updater_kwargs[:outputlevel],
+    eager=updater_kwargs[:eager],
+  )
+  return phi, (; info=exp_info)
+end

src/treetensornetworks/solvers/alternating_update.jl

@@ -26,9 +26,9 @@ function process_sweeps(
   return maxdim, mindim, cutoff, noise, kwargs
 end
 
-function sweep_printer(; outputlevel, psi, sweep, sw_time)
+function sweep_printer(; outputlevel, psi, which_sweep, sw_time)
   if outputlevel >= 1
-    print("After sweep ", sweep, ":")
+    print("After sweep ", which_sweep, ":")
     print(" maxlinkdim=", maxlinkdim(psi))
     print(" cpu_time=", round(sw_time; digits=3))
     println()
@@ -37,7 +37,7 @@ function sweep_printer(; outputlevel, psi, sweep, sw_time)
 end
 
 function alternating_update(
-  solver,
+  updater,
   PH,
   psi0::AbstractTTN;
   checkdone=(; kws...) -> false,
@@ -46,55 +46,60 @@ function alternating_update(
   (sweep_observer!)=observer(),
   sweep_printer=sweep_printer,
   write_when_maxdim_exceeds::Union{Int,Nothing}=nothing,
+  updater_kwargs,
   kwargs...,
 )
   maxdim, mindim, cutoff, noise, kwargs = process_sweeps(nsweeps; kwargs...)
 
   psi = copy(psi0)
 
-  insert_function!(sweep_observer!, "sweep_printer" => sweep_printer)
+  insert_function!(sweep_observer!, "sweep_printer" => sweep_printer) # FIX THIS
 
-  for sweep in 1:nsweeps
-    if !isnothing(write_when_maxdim_exceeds) && maxdim[sweep] > write_when_maxdim_exceeds
+  for which_sweep in 1:nsweeps
+    if !isnothing(write_when_maxdim_exceeds) &&
+      maxdim[which_sweep] > write_when_maxdim_exceeds
       if outputlevel >= 2
         println(
-          "write_when_maxdim_exceeds = $write_when_maxdim_exceeds and maxdim[sweep] = $(maxdim[sweep]), writing environment tensors to disk",
+          "write_when_maxdim_exceeds = $write_when_maxdim_exceeds and maxdim[which_sweep] = $(maxdim[which_sweep]), writing environment tensors to disk",
         )
       end
       PH = disk(PH)
     end
 
     sw_time = @elapsed begin
-      psi, PH = update_step(
-        solver,
+      psi, PH = sweep_update(
+        updater,
         PH,
         psi;
         outputlevel,
-        sweep,
-        maxdim=maxdim[sweep],
-        mindim=mindim[sweep],
-        cutoff=cutoff[sweep],
-        noise=noise[sweep],
+        which_sweep,
+        sweep_params=(;
+          maxdim=maxdim[which_sweep],
+          mindim=mindim[which_sweep],
+          cutoff=cutoff[which_sweep],
+          noise=noise[which_sweep],
+        ),
+        updater_kwargs,
         kwargs...,
       )
     end
 
-    update!(sweep_observer!; psi, sweep, sw_time, outputlevel)
+    update!(sweep_observer!; psi, which_sweep, sw_time, outputlevel)
 
-    checkdone(; psi, sweep, outputlevel, kwargs...) && break
+    checkdone(; psi, which_sweep, outputlevel, kwargs...) && break
   end
-  select!(sweep_observer!, Observers.DataFrames.Not("sweep_printer")) # remove sweep_printer
+  select!(sweep_observer!, Observers.DataFrames.Not("sweep_printer"))
   return psi
 end
 
-function alternating_update(solver, H::AbstractTTN, psi0::AbstractTTN; kwargs...)
+function alternating_update(updater, H::AbstractTTN, psi0::AbstractTTN; kwargs...)
   check_hascommoninds(siteinds, H, psi0)
   check_hascommoninds(siteinds, H, psi0')
   # Permute the indices to have a better memory layout
   # and minimize permutations
   H = ITensors.permute(H, (linkind, siteinds, linkind))
   PH = ProjTTN(H)
-  return alternating_update(solver, PH, psi0; kwargs...)
+  return alternating_update(updater, PH, psi0; kwargs...)
 end
 
 """
@@ -116,12 +121,14 @@ each step of the algorithm when optimizing the MPS.
 Returns:
 * `psi::MPS` - time-evolved MPS
 """
-function alternating_update(solver, Hs::Vector{<:AbstractTTN}, psi0::AbstractTTN; kwargs...)
+function alternating_update(
+  updater, Hs::Vector{<:AbstractTTN}, psi0::AbstractTTN; kwargs...
+)
   for H in Hs
     check_hascommoninds(siteinds, H, psi0)
     check_hascommoninds(siteinds, H, psi0')
   end
   Hs .= ITensors.permute.(Hs, Ref((linkind, siteinds, linkind)))
   PHs = ProjTTNSum(Hs)
-  return alternating_update(solver, PHs, psi0; kwargs...)
+  return alternating_update(updater, PHs, psi0; kwargs...)
 end

src/treetensornetworks/solvers/dmrg.jl

@@ -1,58 +1,34 @@
-function eigsolve_solver(;
-  solver_which_eigenvalue=:SR,
-  ishermitian=true,
-  solver_tol=1e-14,
-  solver_krylovdim=3,
-  solver_maxiter=1,
-  solver_verbosity=0,
-)
-  function solver(H, init; normalize=nothing, region=nothing, half_sweep=nothing)
-    howmany = 1
-    which = solver_which_eigenvalue
-    vals, vecs, info = eigsolve(
-      H,
-      init,
-      howmany,
-      which;
-      ishermitian,
-      tol=solver_tol,
-      krylovdim=solver_krylovdim,
-      maxiter=solver_maxiter,
-      verbosity=solver_verbosity,
-    )
-    psi = vecs[1]
-    return psi, (; solver_info=info, energies=vals)
-  end
-  return solver
-end
-
 """
 Overload of `ITensors.dmrg`.
 """
+
+function dmrg_sweep(
+  nsite::Int, graph::AbstractGraph; root_vertex=default_root_vertex(graph)
+)
+  return tdvp_sweep(2, nsite, Inf, graph; root_vertex, reverse_step=false)
+end
+
 function dmrg(
+  updater,
   H,
   init::AbstractTTN;
-  solver_which_eigenvalue=:SR,
-  ishermitian=true,
-  solver_tol=1e-14,
-  solver_krylovdim=3,
-  solver_maxiter=1,
-  solver_verbosity=0,
+  nsweeps,  #it makes sense to require this to be defined
+  nsite=2,
+  (sweep_observer!)=observer(),
+  root_vertex=default_root_vertex(init),
+  updater_kwargs=NamedTuple(;),
   kwargs...,
 )
-  return alternating_update(
-    eigsolve_solver(;
-      solver_which_eigenvalue,
-      ishermitian,
-      solver_tol,
-      solver_krylovdim,
-      solver_maxiter,
-      solver_verbosity,
-    ),
-    H,
-    init;
-    kwargs...,
+  region_updates = dmrg_sweep(nsite, init; root_vertex)
+
+  psi = alternating_update(
+    updater, H, init; nsweeps, sweep_observer!, region_updates, updater_kwargs, kwargs...
   )
+  return psi
+end
+
+function dmrg(H, init::AbstractTTN; updater=eigsolve_updater, kwargs...)
+  return dmrg(updater, H, init; kwargs...)
 end
 
 """

src/treetensornetworks/solvers/dmrg_x.jl

@@ -1,16 +1,4 @@
-function dmrg_x_solver(
-  PH, init; normalize=nothing, region=nothing, half_sweep=nothing, reverse_step=nothing
-)
-  H = contract(PH, ITensor(1.0))
-  D, U = eigen(H; ishermitian=true)
-  u = uniqueind(U, H)
-  max_overlap, max_ind = findmax(abs, array(dag(init) * U))
-  U_max = U * dag(onehot(u => max_ind))
-  # TODO: improve this to return the energy estimate too
-  return U_max, NamedTuple()
-end
-
 function dmrg_x(PH, init::AbstractTTN; kwargs...)
-  psi = alternating_update(dmrg_x_solver, PH, init; kwargs...)
+  psi = alternating_update(ITensorNetworks.dmrg_x_solver, PH, init; kwargs...)
   return psi
 end