Move ITensorTDVP.jl MPS solver code into ITensorNetworks.jl

Author: mtfishman

Project.toml

@@ -10,12 +10,16 @@ Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
+KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
+Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 Compat = "3, 4"

README.md

@@ -36,7 +36,7 @@ and 3 edge(s):
 3 => 4
 
 with vertex data:
-4-element Dictionary{Int64, Any}
+4-element Dictionaries.Dictionary{Int64, Any}
  1 │ ((dim=2|id=739|"1↔2"),)
  2 │ ((dim=2|id=739|"1↔2"), (dim=2|id=920|"2↔3"))
  3 │ ((dim=2|id=920|"2↔3"), (dim=2|id=761|"3↔4"))
@@ -88,7 +88,7 @@ and 4 edge(s):
 (1, 2) => (2, 2)
 
 with vertex data:
-4-element Dictionary{Tuple{Int64, Int64}, Any}
+4-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
  (1, 1) │ ((dim=2|id=74|"1×1↔2×1"), (dim=2|id=723|"1×1↔1×2"))
  (2, 1) │ ((dim=2|id=74|"1×1↔2×1"), (dim=2|id=823|"2×1↔2×2"))
  (1, 2) │ ((dim=2|id=723|"1×1↔1×2"), (dim=2|id=712|"1×2↔2×2"))
@@ -118,7 +118,7 @@ and 1 edge(s):
 (1, 1) => (1, 2)
 
 with vertex data:
-2-element Dictionary{Tuple{Int64, Int64}, Any}
+2-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
  (1, 1) │ ((dim=2|id=74|"1×1↔2×1"), (dim=2|id=723|"1×1↔1×2"))
  (1, 2) │ ((dim=2|id=723|"1×1↔1×2"), (dim=2|id=712|"1×2↔2×2"))
 
@@ -132,7 +132,7 @@ and 1 edge(s):
 (2, 1) => (2, 2)
 
 with vertex data:
-2-element Dictionary{Tuple{Int64, Int64}, Any}
+2-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
  (2, 1) │ ((dim=2|id=74|"1×1↔2×1"), (dim=2|id=823|"2×1↔2×2"))
  (2, 2) │ ((dim=2|id=823|"2×1↔2×2"), (dim=2|id=712|"1×2↔2×2"))
 ```
@@ -155,13 +155,13 @@ and 2 edge(s):
 2 => 3
 
 with vertex data:
-3-element Dictionary{Int64, Vector{Index}}
+3-element Dictionaries.Dictionary{Int64, Vector{Index}}
  1 │ Index[(dim=2|id=598|"S=1/2,Site,n=1")]
  2 │ Index[(dim=2|id=457|"S=1/2,Site,n=2")]
  3 │ Index[(dim=2|id=683|"S=1/2,Site,n=3")]
 
 and edge data:
-0-element Dictionary{NamedEdge{Int64}, Vector{Index}}
+0-element Dictionaries.Dictionary{NamedGraphs.NamedEdge{Int64}, Vector{Index}}
 
 julia> tn1 = ITensorNetwork(s; link_space=2)
 ITensorNetwork{Int64} with 3 vertices:
@@ -175,7 +175,7 @@ and 2 edge(s):
 2 => 3
 
 with vertex data:
-3-element Dictionary{Int64, Any}
+3-element Dictionaries.Dictionary{Int64, Any}
  1 │ ((dim=2|id=598|"S=1/2,Site,n=1"), (dim=2|id=123|"1↔2"))
  2 │ ((dim=2|id=457|"S=1/2,Site,n=2"), (dim=2|id=123|"1↔2"), (dim=2|id=656|"2↔3…
  3 │ ((dim=2|id=683|"S=1/2,Site,n=3"), (dim=2|id=656|"2↔3"))
@@ -192,7 +192,7 @@ and 2 edge(s):
 2 => 3
 
 with vertex data:
-3-element Dictionary{Int64, Any}
+3-element Dictionaries.Dictionary{Int64, Any}
  1 │ ((dim=2|id=598|"S=1/2,Site,n=1"), (dim=2|id=382|"1↔2"))
  2 │ ((dim=2|id=457|"S=1/2,Site,n=2"), (dim=2|id=382|"1↔2"), (dim=2|id=190|"2↔3…
  3 │ ((dim=2|id=683|"S=1/2,Site,n=3"), (dim=2|id=190|"2↔3"))

examples/ttns/comb_tree.jl renamed to examples/treetensornetworks/comb_tree.jl

@@ -0,0 +1,42 @@
+using ITensors
+using ITensorNetworks
+
+n = 10
+s = siteinds("S=1/2", n)
+
+function heisenberg(n)
+  os = OpSum()
+  for j in 1:(n - 1)
+    os += 0.5, "S+", j, "S-", j + 1
+    os += 0.5, "S-", j, "S+", j + 1
+    os += "Sz", j, "Sz", j + 1
+  end
+  return os
+end
+
+H = MPO(heisenberg(n), s)
+ψ = randomMPS(s, "↑"; linkdims=10)
+
+@show inner(ψ', H, ψ) / inner(ψ, ψ)
+
+ϕ = tdvp(
+  H,
+  -1.0,
+  ψ;
+  nsweeps=20,
+  reverse_step=false,
+  normalize=true,
+  maxdim=30,
+  cutoff=1e-10,
+  outputlevel=1,
+)
+
+@show inner(ϕ', H, ϕ) / inner(ϕ, ϕ)
+
+e2, ϕ2 = dmrg(H, ψ; nsweeps=10, maxdim=20, cutoff=1e-10)
+
+@show inner(ϕ2', H, ϕ2) / inner(ϕ2, ϕ2)
+
+ϕ3 = ITensorNetworks.dmrg(H, ψ; nsweeps=10, maxdim=20, cutoff=1e-10, outputlevel=1)
+
+@show inner(ϕ3', H, ϕ3) / inner(ϕ3, ϕ3)

examples/treetensornetworks/solvers/01_tdvp.jl

@@ -0,0 +1,44 @@
+using ITensors
+using ITensorNetworks
+using LinearAlgebra
+
+function heisenberg(n; h=zeros(n))
+  os = OpSum()
+  for j in 1:(n - 1)
+    os += 0.5, "S+", j, "S-", j + 1
+    os += 0.5, "S-", j, "S+", j + 1
+    os += "Sz", j, "Sz", j + 1
+  end
+  for j in 1:n
+    if h[j] ≠ 0
+      os -= h[j], "Sz", j
+    end
+  end
+  return os
+end
+
+n = 10
+s = siteinds("S=1/2", n)
+
+using Random
+Random.seed!(12)
+
+# MBL when W > 3.5-4
+W = 12
+# Random fields h ∈ [-W, W]
+h = W * (2 * rand(n) .- 1)
+H = MPO(heisenberg(n; h), s)
+
+initstate = rand(["↑", "↓"], n)
+ψ = MPS(s, initstate)
+
+dmrg_x_kwargs = (
+  nsweeps=10, reverse_step=false, normalize=true, maxdim=20, cutoff=1e-10, outputlevel=1
+)
+
+ϕ = dmrg_x(H, ψ; dmrg_x_kwargs...)
+
+@show inner(ψ', H, ψ) / inner(ψ, ψ)
+@show inner(H, ψ, H, ψ) - inner(ψ', H, ψ)^2
+@show inner(ϕ', H, ϕ) / inner(ϕ, ϕ)
+@show inner(H, ϕ, H, ϕ) - inner(ϕ', H, ϕ)^2

examples/treetensornetworks/solvers/02_dmrg-x.jl

@@ -0,0 +1,20 @@
+using ITensors
+
+function heisenberg(n; J=1.0, J2=0.0)
+  ℋ = OpSum()
+  if !iszero(J)
+    for j in 1:(n - 1)
+      ℋ += J / 2, "S+", j, "S-", j + 1
+      ℋ += J / 2, "S-", j, "S+", j + 1
+      ℋ += J, "Sz", j, "Sz", j + 1
+    end
+  end
+  if !iszero(J2)
+    for j in 1:(n - 2)
+      ℋ += J2 / 2, "S+", j, "S-", j + 2
+      ℋ += J2 / 2, "S-", j, "S+", j + 2
+      ℋ += J2, "Sz", j, "Sz", j + 2
+    end
+  end
+  return ℋ
+end

examples/treetensornetworks/solvers/03_models.jl

@@ -0,0 +1,37 @@
+using DifferentialEquations
+using ITensors
+using ITensorNetworks
+using KrylovKit: exponentiate
+
+function ode_solver(
+  H::TimeDependentSum,
+  time_step,
+  ψ₀;
+  current_time=0.0,
+  outputlevel=0,
+  solver_alg=Tsit5(),
+  kwargs...,
+)
+  if outputlevel ≥ 3
+    println("    In ODE solver, current_time = $current_time, time_step = $time_step")
+  end
+
+  time_span = (current_time, current_time + time_step)
+  u₀, ITensor_from_vec = to_vec(ψ₀)
+  f(ψ::ITensor, p, t) = H(t)(ψ)
+  f(u::Vector, p, t) = to_vec(f(ITensor_from_vec(u), p, t))[1]
+  prob = ODEProblem(f, u₀, time_span)
+  sol = solve(prob, solver_alg; kwargs...)
+  uₜ = sol.u[end]
+  return ITensor_from_vec(uₜ), nothing
+end
+
+function krylov_solver(
+  H::TimeDependentSum, time_step, ψ₀; current_time=0.0, outputlevel=0, kwargs...
+)
+  if outputlevel ≥ 3
+    println("    In Krylov solver, current_time = $current_time, time_step = $time_step")
+  end
+  ψₜ, info = exponentiate(H(current_time), time_step, ψ₀; kwargs...)
+  return ψₜ, info
+end

examples/treetensornetworks/solvers/03_solvers.jl

@@ -0,0 +1,153 @@
+using DifferentialEquations
+using ITensors
+using ITensorNetworks
+using KrylovKit
+using LinearAlgebra
+using Random
+
+Random.seed!(1234)
+
+# Define the time-independent model
+include("03_models.jl")
+
+# Define the solvers needed for TDVP
+include("03_solvers.jl")
+
+# Time dependent Hamiltonian is:
+# H(t) = H₁(t) + H₂(t) + …
+#      = f₁(t) H₁(0) + f₂(t) H₂(0) + …
+#      = cos(ω₁t) H₁(0) + cos(ω₂t) H₂(0) + …
+
+# Number of sites
+n = 6
+
+# How much information to output from TDVP
+# Set to 2 to get information about each bond/site
+# evolution, and 3 to get information about the
+# solver.
+outputlevel = 1
+
+# Frequency of time dependent terms
+ω₁ = 0.1
+ω₂ = 0.2
+
+# Nearest and next-nearest neighbor
+# Heisenberg couplings.
+J₁ = 1.0
+J₂ = 1.0
+
+time_step = 0.1
+time_stop = 1.0
+
+# nsite-update TDVP
+nsite = 2
+
+# Starting state bond/link dimension.
+# A product state starting state can
+# cause issues for TDVP without
+# subspace expansion.
+start_linkdim = 4
+
+# TDVP truncation parameters
+maxdim = 100
+cutoff = 1e-8
+
+tol = 1e-15
+
+# ODE solver parameters
+ode_alg = Tsit5()
+ode_kwargs = (; reltol=tol, abstol=tol)
+
+# Krylov solver parameters
+krylov_kwargs = (; tol=tol, eager=true)
+
+@show n
+@show ω₁, ω₂
+@show J₁, J₂
+@show maxdim, cutoff, nsite
+@show start_linkdim
+@show time_step, time_stop
+@show ode_alg
+@show ode_kwargs
+@show krylov_kwargs
+
+ω⃗ = [ω₁, ω₂]
+f⃗ = [t -> cos(ω * t) for ω in ω⃗]
+
+# H₀ = H(0) = H₁(0) + H₂(0) + …
+ℋ₁₀ = heisenberg(n; J=J₁, J2=0.0)
+ℋ₂₀ = heisenberg(n; J=0.0, J2=J₂)
+ℋ⃗₀ = [ℋ₁₀, ℋ₂₀]
+
+s = siteinds("S=1/2", n)
+
+H⃗₀ = [MPO(ℋ₀, s) for ℋ₀ in ℋ⃗₀]
+
+# Initial state, ψ₀ = ψ(0)
+# Initialize as complex since that is what DifferentialEquations.jl
+# expects.
+ψ₀ = complex.(randomMPS(s, j -> isodd(j) ? "↑" : "↓"; linkdims=start_linkdim))
+
+@show norm(ψ₀)
+
+println()
+println("#"^100)
+println("Running TDVP with ODE solver")
+println("#"^100)
+println()
+
+function ode_solver(H⃗₀, time_step, ψ₀; kwargs...)
+  return ode_solver(
+    -im * TimeDependentSum(f⃗, H⃗₀),
+    time_step,
+    ψ₀;
+    solver_alg=ode_alg,
+    ode_kwargs...,
+    kwargs...,
+  )
+end
+
+ψₜ_ode = tdvp(ode_solver, H⃗₀, time_stop, ψ₀; time_step, maxdim, cutoff, nsite, outputlevel)
+
+println()
+println("Finished running TDVP with ODE solver")
+println()
+
+println()
+println("#"^100)
+println("Running TDVP with Krylov solver")
+println("#"^100)
+println()
+
+function krylov_solver(H⃗₀, time_step, ψ₀; kwargs...)
+  return krylov_solver(
+    -im * TimeDependentSum(f⃗, H⃗₀), time_step, ψ₀; krylov_kwargs..., kwargs...
+  )
+end
+
+ψₜ_krylov = tdvp(krylov_solver, H⃗₀, time_stop, ψ₀; time_step, cutoff, nsite, outputlevel)
+
+println()
+println("Finished running TDVP with Krylov solver")
+println()
+
+println()
+println("#"^100)
+println("Running full state evolution with ODE solver")
+println("#"^100)
+println()
+
+@disable_warn_order begin
+  ψₜ_full, _ = ode_solver(prod.(H⃗₀), time_stop, prod(ψ₀); outputlevel)
+end
+
+println()
+println("Finished full state evolution with ODE solver")
+println()
+
+@show norm(ψₜ_ode)
+@show norm(ψₜ_krylov)
+@show norm(ψₜ_full)
+
+@show 1 - abs(inner(prod(ψₜ_ode), ψₜ_full))
+@show 1 - abs(inner(prod(ψₜ_krylov), ψₜ_full))