Output eigenvalue from DMRG and DMRG-X (#167)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworks"
 uuid = "2919e153-833c-4bdc-8836-1ea460a35fc7"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>, Joseph Tindall <jtindall@flatironinstitute.org> and contributors"]
-version = "0.10.4"
+version = "0.11.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -63,7 +63,7 @@ IterTools = "1.4.0"
 KrylovKit = "0.6, 0.7"
 NamedGraphs = "0.6.0"
 NDTensors = "0.3"
-Observers = "0.2"
+Observers = "0.2.4"
 OMEinsumContractionOrders = "0.8.3"
 PackageExtensionCompat = "1"
 SerializedElementArrays = "0.1"

src/solvers/dmrg.jl

@@ -4,11 +4,24 @@ using KrylovKit: KrylovKit
 """
 Overload of `ITensors.ITensorMPS.dmrg`.
 """
-
 function ITensorMPS.dmrg(
-  operator, init_state; nsweeps, nsites=2, updater=eigsolve_updater, kwargs...
+  operator,
+  init_state;
+  nsweeps,
+  nsites=2,
+  updater=eigsolve_updater,
+  (region_observer!)=nothing,
+  kwargs...,
 )
-  return alternating_update(operator, init_state; nsweeps, nsites, updater, kwargs...)
+  eigvals_ref = Ref{Any}()
+  region_observer! = compose_observers(
+    region_observer!, ValuesObserver((; eigvals=eigvals_ref))
+  )
+  state = alternating_update(
+    operator, init_state; nsweeps, nsites, updater, region_observer!, kwargs...
+  )
+  eigval = only(eigvals_ref[])
+  return eigval, state
 end
 
 """

src/solvers/dmrg_x.jl

@@ -1,5 +1,19 @@
 function dmrg_x(
-  operator, init_state::AbstractTTN; nsweeps, nsites=2, updater=dmrg_x_updater, kwargs...
+  operator,
+  init_state::AbstractTTN;
+  nsweeps,
+  nsites=2,
+  updater=dmrg_x_updater,
+  (region_observer!)=nothing,
+  kwargs...,
 )
-  return alternating_update(operator, init_state; nsweeps, nsites, updater, kwargs...)
+  eigvals_ref = Ref{Any}()
+  region_observer! = compose_observers(
+    region_observer!, ValuesObserver((; eigvals=eigvals_ref))
+  )
+  state = alternating_update(
+    operator, init_state; nsweeps, nsites, updater, region_observer!, kwargs...
+  )
+  eigval = only(eigvals_ref[])
+  return eigval, state
 end

src/solvers/local_solvers/dmrg_x.jl

@@ -12,12 +12,11 @@ function dmrg_x_updater(
   which_region_update,
   internal_kwargs,
 )
-  #ToDo: Implement this via KrylovKit or similar for better scaling
   H = contract(projected_operator![], ITensor(true))
   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, (;)
+  eigvals = [((onehot(u => max_ind)' * D) * dag(onehot(u => max_ind)))[]]
+  return U_max, (; eigvals)
 end

src/update_observer.jl

@@ -9,3 +9,24 @@ end
 function update_observer!(observer::Nothing; kwargs...)
   return nothing
 end
+
+struct ComposedObservers{Observers<:Tuple}
+  observers::Observers
+end
+compose_observers(observers...) = ComposedObservers(observers)
+function update_observer!(observer::ComposedObservers; kwargs...)
+  for observerᵢ in observer.observers
+    update_observer!(observerᵢ; kwargs...)
+  end
+  return observer
+end
+
+struct ValuesObserver{Values<:NamedTuple}
+  values::Values
+end
+function update_observer!(observer::ValuesObserver; kwargs...)
+  for key in keys(observer.values)
+    observer.values[key][] = kwargs[key]
+  end
+  return observer
+end

test/test_treetensornetworks/test_solvers/test_contract.jl

@@ -44,11 +44,11 @@ using Test: @test, @test_broken, @testset
   @test inner(psi, Hpsi) ≈ inner(psi', H, psi) rtol = 1e-5
   # Test variational compression via DMRG
   Hfit = ProjOuterProdTTN(psi', H)
-  Hpsi_via_dmrg = dmrg(Hfit, psi; updater_kwargs=(; which_eigval=:LR,), nsweeps=1)
+  e, Hpsi_via_dmrg = dmrg(Hfit, psi; updater_kwargs=(; which_eigval=:LR,), nsweeps=1)
   @test abs(inner(Hpsi_via_dmrg, Hpsi / norm(Hpsi))) ≈ 1 rtol = 1e-4
   # Test whether the interface works for ProjTTNSum with factors
   Hfit = ProjTTNSum([ProjOuterProdTTN(psi', H), ProjOuterProdTTN(psi', H)], [-0.2, -0.8])
-  Hpsi_via_dmrg = dmrg(Hfit, psi; nsweeps=1, updater_kwargs=(; which_eigval=:SR,))
+  e, Hpsi_via_dmrg = dmrg(Hfit, psi; nsweeps=1, updater_kwargs=(; which_eigval=:SR,))
   @test abs(inner(Hpsi_via_dmrg, Hpsi / norm(Hpsi))) ≈ 1 rtol = 1e-4
 
   # Test basic usage for use with multiple ProjOuterProdTTN with default parameters
@@ -66,9 +66,7 @@ using Test: @test, @test_broken, @testset
   # Test the above via DMRG
   # ToDo: Investigate why this is broken
   Hfit = ProjTTNSum([ProjOuterProdTTN(psi', H), ProjOuterProdTTN(psi', identity)], [-1, 1])
-  Hpsi_normalized = ITensorNetworks.dmrg(
-    Hfit, psi; nsweeps=3, updater_kwargs=(; which_eigval=:SR)
-  )
+  e, Hpsi_normalized = dmrg(Hfit, psi; nsweeps=3, updater_kwargs=(; which_eigval=:SR))
   @test_broken abs(inner(Hpsi, (Hpsi_normalized) / norm(Hpsi))) ≈ 1 rtol = 1e-5
 
   #

test/test_treetensornetworks/test_solvers/test_dmrg.jl

@@ -53,22 +53,24 @@ ITensors.disable_auto_fermion()
   psi_mps = ITensorMPS.MPS([psi[v] for v in 1:nv(psi)])
   e2, psi2 = dmrg(H_mpo, psi_mps; nsweeps, maxdim, outputlevel=0)
 
-  psi = dmrg(
+  e, psi = dmrg(
     H, psi; nsweeps, maxdim, cutoff, nsites, updater_kwargs=(; krylovdim=3, maxiter=1)
   )
+  @test inner(psi', H, psi) ≈ e
   @test inner(psi', H, psi) ≈ inner(psi2', H_mpo, psi2)
 
   # Alias for `ITensorNetworks.dmrg`
-  psi = eigsolve(
+  e, psi = eigsolve(
     H, psi; nsweeps, maxdim, cutoff, nsites, updater_kwargs=(; krylovdim=3, maxiter=1)
   )
+  @test inner(psi', H, psi) ≈ e
   @test inner(psi', H, psi) ≈ inner(psi2', H_mpo, psi2)
 
   # Test custom sweep regions #BROKEN, ToDo: Make proper custom sweep regions for test
   #=
   orig_E = inner(psi', H, psi)
   sweep_regions = [[1], [2], [3], [3], [2], [1]]
-  psi = dmrg(H, psi; nsweeps, maxdim, cutoff, sweep_regions)
+  e, psi = dmrg(H, psi; nsweeps, maxdim, cutoff, sweep_regions)
   new_E = inner(psi', H, psi)
   @test new_E ≈ orig_E
   =#
@@ -101,13 +103,14 @@ end
   energy(; eigvals, kw...) = eigvals[1]
   region_observer! = observer(region, sweep, energy)
 
-  psi = dmrg(H, psi; nsweeps, maxdim, cutoff, sweep_observer!, region_observer!)
+  e, psi = dmrg(H, psi; nsweeps, maxdim, cutoff, sweep_observer!, region_observer!)
 
   #
   # Test out certain values
   #
   @test region_observer![9, :region] == [2, 1]
   @test region_observer![30, :energy] < -4.25
+  @test region_observer![30, :energy] ≈ e rtol = 1e-6
 end
 
 @testset "Cache to Disk" begin
@@ -126,7 +129,7 @@ end
   nsweeps = 4
   maxdim = [10, 20, 40, 80]
 
-  @test_broken psi = dmrg(
+  @test_broken e, psi = dmrg(
     H,
     psi;
     nsweeps,
@@ -160,7 +163,7 @@ end
   maxdim = [200, 250, 400, 600, 800, 1200, 2000, 2400, 2600, 3000]
   cutoff = [1e-10, 1e-10, 1e-12, 1e-12, 1e-12, 1e-12, 1e-14, 1e-14, 1e-14, 1e-14]
 
-  psi = dmrg(H, psi; nsweeps, maxdim, cutoff)
+  e, psi = dmrg(H, psi; nsweeps, maxdim, cutoff)
 end
 
 @testset "Tree DMRG" for nsites in [2]
@@ -195,7 +198,7 @@ end
     nsweeps = 10
     maxdim = [10, 20, 40, 100]
     @show use_qns
-    psi = dmrg(
+    e, psi = dmrg(
       H, psi; nsweeps, maxdim, cutoff, nsites, updater_kwargs=(; krylovdim=3, maxiter=1)
     )
 
@@ -252,7 +255,7 @@ end
   end
   states = v -> d[v]
   psi = ttn(states, s)
-  psi = dmrg(
+  e, psi = dmrg(
     H, psi; nsweeps, maxdim, cutoff, nsites, updater_kwargs=(; krylovdim=3, maxiter=1)
   )
 
@@ -280,7 +283,7 @@ end
   os = ModelHamiltonians.heisenberg(c)
   H = ttn(os, s)
   psi = random_ttn(s; link_space=5)
-  psi = dmrg(H, psi; nsweeps, maxdim, nsites)
+  e, psi = dmrg(H, psi; nsweeps, maxdim, nsites)
 
   @test all(edge_data(linkdims(psi)) .<= maxdim)
 end

test/test_treetensornetworks/test_solvers/test_dmrg_x.jl

@@ -10,6 +10,7 @@ using NamedGraphs.NamedGraphGenerators: named_comb_tree
 using Random: Random
 using Test: @test, @testset
 
+# TODO: Combine MPS and TTN tests.
 @testset "MPS DMRG-X" for conserve_qns in (false, true)
   n = 10
   s = siteinds("S=1/2", n; conserve_qns)
@@ -25,14 +26,16 @@ using Test: @test, @testset
 
   dmrg_x_kwargs = (nsweeps=20, normalize=true, maxdim=20, cutoff=1e-10, outputlevel=0)
 
-  ϕ = dmrg_x(H, ψ; nsites=2, dmrg_x_kwargs...)
+  e, ϕ = dmrg_x(H, ψ; nsites=2, dmrg_x_kwargs...)
 
+  @test inner(ϕ', H, ϕ) / inner(ϕ, ϕ) ≈ e
   @test inner(ψ', H, ψ) / inner(ψ, ψ) ≈ inner(ϕ', H, ϕ) / inner(ϕ, ϕ) rtol = 1e-1
   @test inner(H, ψ, H, ψ) ≉ inner(ψ', H, ψ)^2 rtol = 1e-7
   @test inner(H, ϕ, H, ϕ) ≈ inner(ϕ', H, ϕ)^2 rtol = 1e-7
 
-  ϕ̃ = dmrg_x(H, ϕ; nsites=1, dmrg_x_kwargs...)
+  e, ϕ̃ = dmrg_x(H, ϕ; nsites=1, dmrg_x_kwargs...)
 
+  @test inner(ϕ̃', H, ϕ̃) / inner(ϕ̃, ϕ̃) ≈ e
   @test inner(ψ', H, ψ) / inner(ψ, ψ) ≈ inner(ϕ̃', H, ϕ̃) / inner(ϕ̃, ϕ̃) rtol = 1e-1
   @test inner(H, ϕ̃, H, ϕ̃) ≈ inner(ϕ̃', H, ϕ̃)^2 rtol = 1e-3
   # Sometimes broken, sometimes not
@@ -56,14 +59,16 @@ end
 
   dmrg_x_kwargs = (nsweeps=20, normalize=true, maxdim=20, cutoff=1e-10, outputlevel=0)
 
-  ϕ = dmrg_x(H, ψ; nsites=2, dmrg_x_kwargs...)
+  e, ϕ = dmrg_x(H, ψ; nsites=2, dmrg_x_kwargs...)
 
+  @test inner(ϕ', H, ϕ) / inner(ϕ, ϕ) ≈ e
   @test inner(ψ', H, ψ) / inner(ψ, ψ) ≈ inner(ϕ', H, ϕ) / inner(ϕ, ϕ) rtol = 1e-1
   @test inner(H, ψ, H, ψ) ≉ inner(ψ', H, ψ)^2 rtol = 1e-2
   @test inner(H, ϕ, H, ϕ) ≈ inner(ϕ', H, ϕ)^2 rtol = 1e-7
 
-  ϕ̃ = dmrg_x(H, ϕ; nsites=1, dmrg_x_kwargs...)
+  e, ϕ̃ = dmrg_x(H, ϕ; nsites=1, dmrg_x_kwargs...)
 
+  @test inner(ϕ̃', H, ϕ̃) / inner(ϕ̃, ϕ̃) ≈ e
   @test inner(ψ', H, ψ) / inner(ψ, ψ) ≈ inner(ϕ̃', H, ϕ̃) / inner(ϕ̃, ϕ̃) rtol = 1e-1
   @test inner(H, ϕ̃, H, ϕ̃) ≈ inner(ϕ̃', H, ϕ̃)^2 rtol = 1e-6
   # Sometimes broken, sometimes not