Refactor interfaces built around alternating update

src/ITensorNetworks.jl

@@ -106,6 +106,11 @@
 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", "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,27 @@
+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,40 @@
+
+function eigsolve_solver(;
+    solver_which_eigenvalue=:SR,   #TODO: settle on pattern to pass solver kwargs
+    ishermitian=true,
+    solver_tol=1e-14,
+    solver_krylovdim=3,
+    solver_maxiter=1,
+    solver_verbosity=0,
+  )
+
+    function solver(
+        init;
+        psi_ref!,
+        PH_ref!,
+        normalize,
+        region,
+        sweep_regions,
+        sweep_step,
+        sweep_kwargs...
+        # slurp solver_kwargs?  #TODO: homogenize how the solver kwargs are passed
+        )
+      howmany = 1
+      which = solver_which_eigenvalue
+      vals, vecs, info = eigsolve(
+        PH_ref![],
+        init,
+        howmany,
+        which;
+        ishermitian,
+        tol=solver_tol,
+        krylovdim=solver_krylovdim,
+        maxiter=solver_maxiter,
+        verbosity=solver_verbosity,
+      )
+      phi = vecs[1]
+      return phi, (; solver_info=info, energies=vals)
+    end
+    return solver
+  end
+

src/solvers/exponentiate.jl

@@ -0,0 +1,42 @@
+function exponentiate_solver()
+    function solver(
+      init;
+      psi_ref!,
+      PH_ref!,
+      ishermitian=true,
+      issymmetric=true,
+      region,
+      sweep_regions,
+      sweep_step,
+      solver_krylovdim=30,
+      solver_maxiter=100,
+      solver_outputlevel=0,
+      solver_tol=1E-12,
+      substep,
+      normalize,
+      time_step,
+    )
+        #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
+        region_ind=sweep_step
+        next_region=region_ind==length(sweep_regions) ? nothing : first(sweep_regions[region_ind+1])
+        previous_region=region_ind==1 ? nothing : first(sweep_regions[region_ind-1])
+
+        phi, exp_info = KrylovKit.exponentiate(
+        H,
+        time_step,
+        init;
+        ishermitian,
+        issymmetric,
+        tol=solver_tol,
+        krylovdim=solver_krylovdim,
+        maxiter=solver_maxiter,
+        verbosity=solver_outputlevel,
+        eager=true,
+        )
+        return phi, (; info=exp_info)
+    end
+    return solver
+end

src/treetensornetworks/solvers/dmrg.jl

@@ -1,31 +1,3 @@
-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`.
 """
@@ -51,6 +23,7 @@ function dmrg(
     ),
     H,
     init;
+    reverse_step=false,
     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

src/treetensornetworks/solvers/tdvp.jl

@@ -1,57 +1,3 @@
-function exponentiate_solver()
-  function solver(
-    H,
-    init;
-    ishermitian=true,
-    issymmetric=true,
-    region,
-    solver_krylovdim=30,
-    solver_maxiter=100,
-    solver_outputlevel=0,
-    solver_tol=1E-12,
-    substep,
-    normalize,
-    time_step,
-  )
-    psi, exp_info = KrylovKit.exponentiate(
-      H,
-      time_step,
-      init;
-      ishermitian,
-      issymmetric,
-      tol=solver_tol,
-      krylovdim=solver_krylovdim,
-      maxiter=solver_maxiter,
-      verbosity=solver_outputlevel,
-      eager=true,
-    )
-    return psi, (; info=exp_info)
-  end
-  return solver
-end
-
-function applyexp_solver()
-  function solver(
-    H,
-    init;
-    tdvp_order,
-    solver_krylovdim=30,
-    solver_outputlevel=0,
-    solver_tol=1E-8,
-    substep,
-    time_step,
-    normalize,
-  )
-    #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
-
 function _compute_nsweeps(nsteps, t, time_step, order)
   nsweeps_per_step = order / 2
   nsweeps = 1
@@ -168,15 +114,6 @@ Optional keyword arguments:
 * `observer` - object implementing the Observer interface which can perform measurements and stop early
 * `write_when_maxdim_exceeds::Int` - when the allowed maxdim exceeds this value, begin saving tensors to disk to free memory in large calculations
 """
-function tdvp(H, t::Number, init::AbstractTTN; solver_backend="exponentiate", kwargs...)
-  if solver_backend == "exponentiate"
-    solver = exponentiate_solver
-  elseif solver_backend == "applyexp"
-    solver = applyexp_solver
-  else
-    error(
-      "solver_backend=$solver_backend not recognized (options are \"applyexp\" or \"exponentiate\")",
-    )
-  end
+function tdvp(H, t::Number, init::AbstractTTN; solver=exponentiate_solver, kwargs...)
   return tdvp(solver(), H, t, init; kwargs...)
 end

src/treetensornetworks/solvers/update_step.jl

@@ -44,7 +44,8 @@ function update_step(
   step_printer=step_printer,
   (step_observer!)=observer(),
   sweep::Int=1,
-  sweep_regions=default_sweep_regions(nsite, psi),
+  reverse_step::Bool=false,
+  sweep_regions=default_sweep_regions(nsite, psi;reverse_step),
   kwargs...,
 )
   PH = copy(PH)
@@ -182,7 +183,11 @@ function local_update(
   nsites = (region isa AbstractEdge) ? 0 : length(region)
   PH = set_nsite(PH, nsites)
   PH = position(PH, psi, region)
-  phi, info = solver(PH, phi; normalize, region, step_kwargs..., solver_kwargs...)
+  (psi_ref!) = Ref(psi) # create references, in case solver does (out-of-place) modify PH or psi
+  (PH_ref!) = Ref(PH)
+  phi, info = solver(phi;(psi_ref!),(PH_ref!), normalize, region, sweep_regions, sweep_step, step_kwargs..., solver_kwargs...)  # args passed by reference are supposed to be modified out of place
+  psi = psi_ref![] # dereference
+  PH = PH_ref![]
   if !(phi isa ITensor && info isa NamedTuple)
     println("Solver returned the following types: $(typeof(phi)), $(typeof(info))")
     error("In alternating_update, solver must return an ITensor and a NamedTuple")

test/test_treetensornetworks/test_solvers/test_tdvp.jl

@@ -4,6 +4,7 @@ using KrylovKit: exponentiate
 using Observers
 using Random
 using Test
+#exponentiate_solver=ITensorNetworks.exponentiate_solver #ToDo: how to best handle importing etc.
 
 @testset "MPS TDVP" begin
   @testset "Basic TDVP" begin
@@ -31,7 +32,7 @@ using Test
     #
     #Different backend solvers, default solver_backend = "applyexp"
     ψ1_exponentiate_backend = tdvp(
-      H, -0.1im, ψ0; nsteps=1, cutoff, nsite=1, solver_backend="exponentiate"
+      H, -0.1im, ψ0; nsteps=1, cutoff, nsite=1, solver=exponentiate_solver
     )
     @test ψ1 ≈ ψ1_exponentiate_backend rtol = 1e-7
 
@@ -48,7 +49,7 @@ using Test
 
     #Different backend solvers, default solver_backend = "applyexp"
     ψ2_exponentiate_backend = tdvp(
-      H, +0.1im, ψ1; nsteps=1, cutoff, solver_backend="exponentiate"
+      H, +0.1im, ψ1; nsteps=1, cutoff, solver=exponentiate_solver
     )
     @test ψ2 ≈ ψ2_exponentiate_backend rtol = 1e-7
 
@@ -84,7 +85,7 @@ using Test
 
     #Different backend solvers, default solver_backend = "applyexp"
     ψ1_exponentiate_backend = tdvp(
-      Hs, -0.1im, ψ0; nsteps=1, cutoff, nsite=1, solver_backend="exponentiate"
+      Hs, -0.1im, ψ0; nsteps=1, cutoff, nsite=1, solver=exponentiate_solver
     )
     @test ψ1 ≈ ψ1_exponentiate_backend rtol = 1e-7
 
@@ -101,7 +102,7 @@ using Test
 
     #Different backend solvers, default solver_backend = "applyexp"
     ψ2_exponentiate_backend = tdvp(
-      Hs, +0.1im, ψ1; nsteps=1, cutoff, solver_backend="exponentiate"
+      Hs, +0.1im, ψ1; nsteps=1, cutoff, solver=exponentiate_solver
     )
     @test ψ2 ≈ ψ2_exponentiate_backend rtol = 1e-7
 
@@ -163,10 +164,11 @@ using Test
 
     ψ0 = random_mps(s; internal_inds_space=10)
 
-    function solver(PH, psi0; time_step, kwargs...)
+    function solver(psi0; (psi_ref!),(PH_ref!), time_step, kwargs...)
       solver_kwargs = (;
         ishermitian=true, tol=1e-12, krylovdim=30, maxiter=100, verbosity=0, eager=true
       )
+      PH=PH_ref![]
       psi, exp_info = exponentiate(PH, time_step, psi0; solver_kwargs...)
       return psi, (; info=exp_info)
     end
@@ -248,7 +250,7 @@ using Test
         solver_tol=1e-12,
         solver_maxiter=500,
         solver_krylovdim=25,
-        solver_backend="exponentiate",
+        solver=exponentiate_solver,
       )
       # TODO: What should `expect` output? Right now
       # it outputs a dictionary.
@@ -384,7 +386,7 @@ using Test
         reverse_step,
         normalize=true,
         solver_krylovdim=15,
-        solver_backend="exponentiate",
+        solver=ITensorNetworks.exponentiate_solver,
       )
     end
 
@@ -551,10 +553,11 @@ end
 
     ψ0 = normalize!(random_ttn(s; link_space=10))
 
-    function solver(PH, psi0; time_step, kwargs...)
+    function solver(psi0; (psi_ref!),(PH_ref!), time_step, kwargs...)
       solver_kwargs = (;
         ishermitian=true, tol=1e-12, krylovdim=30, maxiter=100, verbosity=0, eager=true
       )
+      PH=PH_ref![]
       psi, exp_info = exponentiate(PH, time_step, psi0; solver_kwargs...)
       return psi, (; info=exp_info)
     end