Basis expansion method with expand

Project.toml

@@ -30,6 +30,7 @@ julia = "1.6"
 [extras]
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [targets]
 test = ["Test"]

src/ITensorTDVP.jl

@@ -1,5 +1,5 @@
 module ITensorTDVP
-export TimeDependentSum, dmrg_x, linsolve, tdvp, to_vec
+export TimeDependentSum, basis_extend, dmrg_x, linsolve, tdvp, to_vec
 include("ITensorsExtensions.jl")
 using .ITensorsExtensions: to_vec
 include("applyexp.jl")
@@ -17,8 +17,9 @@ include("contract.jl")
 include("reducedconstantterm.jl")
 include("reducedlinearproblem.jl")
 include("linsolve.jl")
+include("basis_extend.jl")
 using PackageExtensionCompat: @require_extensions
 function __init__()
   @require_extensions
 end
-end
+end

src/basis_extend.jl

@@ -0,0 +1,112 @@
+import ITensors: pause
+#
+# Possible improvements
+#  - allow a maxdim argument to be passed to `extend`
+#    and through `basis_extend`
+#  - current behavior is letting bond dimension get too
+#    big when used in imaginary time evolution
+#  - come up with better names:
+#    > should `basis_extend` be called `krylov_extend`?
+#    > should `extend` be called `basis_extend`?
+#  - Use (1-tau*H)|psi> to generate "Krylov" vectors
+#    instead of H|psi>. Needed?
+#
+
+"""
+Given an MPS psi and a collection of MPS phis,
+returns an MPS which is equal to psi
+(has fidelity 1.0 with psi) but whose MPS basis
+is extended to contain a portion of the basis of
+the phis that is orthogonal to the MPS basis of psi.
+"""
+function extend(psi::MPS, phis::Vector{MPS}; kwargs...)
+  cutoff = get(kwargs, :cutoff, 1E-14)
+  N = length(psi)
+  psi = copy(psi)
+  phis = copy(phis)
+
+  orthogonalize!(psi, N)
+  for phi in phis
+    orthogonalize!(phi, N)
+  end
+
+  s = siteinds(psi)
+
+  for j in reverse(2:N)
+    # SVD psi[j] to compute B
+    linds = (s[j - 1], linkind(psi, j - 1))
+    _, S, B = svd(psi[j], linds; righttags="bψ_$j,Link")
+    rinds = uniqueinds(B, S)
+
+    # Make projector
+    Id = ITensor(1.0)
+    for r in rinds
+      Id *= denseblocks(delta(r', dag(r)))
+    end
+    P = Id - prime(B, rinds) * dag(B)
+
+    # Sum phi density matrices
+    rho = ITensor()
+    for phi in phis
+      rho += prime(phi[j], rinds) * dag(phi[j])
+    end
+    rho /= tr(rho)
+
+    # Apply projector
+    PrhoP = apply(apply(P, rho), P)
+
+    if norm(PrhoP) > 1E-12
+      # Diagonalize projected density matrix PrhoP
+      # to compute Bphi, which spans part of right basis 
+      # of phis which is orthogonal to right basis of psi
+      D, Bphi = eigen(PrhoP; cutoff, ishermitian=true, righttags="bϕ_$j,Link")
+
+      ## Test Bphi is ortho to B
+      #O = Bphi*B
+      #if norm(O) > 1E-10
+      #  @show norm(O)
+      #  error("Non-zero overlap of extended basis with original basis")
+      #end
+
+      # Form direct sum of B and Bphi over left index
+      bψ = commonind(B, S)
+      bϕ = commonind(Bphi, D)
+      if !hasqns(bψ)
+        bx = Index(dim(bψ) + dim(bϕ), "bx_$(j-1)")
+      else
+        bx = Index(vcat(space(bψ), space(bϕ)), "bx_$(j-1)")
+      end
+      D1, D2 = ITensors.directsum_itensors(bψ, bϕ, dag(bx))
+      Bx = D1 * B + D2 * Bphi
+    else
+      Bx = B
+    end
+
+    # Shift ortho center one site left using dag(Bx)
+    # and replace tensor at site j with Bx
+    psi[j - 1] = psi[j - 1] * (psi[j] * dag(Bx))
+    psi[j] = Bx
+    for phi in phis
+      phi[j - 1] = phi[j - 1] * (phi[j] * dag(Bx))
+      phi[j] = Bx
+    end
+  end
+
+  return psi
+end
+
+function basis_extend(psi::MPS, H::MPO; kwargs...)
+  kdim = get(kwargs, :extension_krylovdim, 2)
+  maxdim = 1 + maxlinkdim(psi)
+
+  phis = Vector{MPS}(undef, kdim)
+  for k in 1:kdim
+    prev = k == 1 ? psi : phis[k - 1]
+    phis[k] = apply(H, prev; maxdim)
+    normalize!(phis[k])
+  end
+
+  cutoff = get(kwargs, :extension_cutoff, 1E-8)
+  psix = extend(psi, phis; cutoff)
+  return psix
+end

test/test_basis_extend.jl

@@ -0,0 +1,98 @@
+using ITensors
+import ITensors: pause
+using ITensorTDVP
+using Random
+using Test
+using Printf
+
+@testset "extend function" begin
+  N = 6
+  s = siteinds(2, N)
+
+  psi = randomMPS(s; linkdims=4)
+  phi = randomMPS(s; linkdims=2)
+
+  psix = ITensorTDVP.extend(psi, [phi])
+  @test inner(psix, psi) ≈ inner(psi, psi)
+  @test inner(psix, phi) ≈ inner(psi, phi)
+end
+
+@testset "extend function with QNs" begin
+  N = 6
+  s = siteinds("S=1/2", N; conserve_qns=true)
+  state = [isodd(n) ? "Up" : "Dn" for n in 1:N]
+  psi = randomMPS(s, state; linkdims=4)
+  phi = randomMPS(s, state; linkdims=2)
+
+  psix = ITensorTDVP.extend(psi, [phi])
+  @test inner(psix, psi) ≈ inner(psi, psi)
+  @test inner(psix, phi) ≈ inner(psi, phi)
+end
+
+@testset "basis_extend" begin
+  N = 10
+
+  s = siteinds("S=1/2", 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
+  H = MPO(os, s)
+
+  ψ0 = productMPS(s, n -> isodd(n) ? "Up" : "Dn")
+  ψx = basis_extend(ψ0, H)
+  @test maxlinkdim(ψx) > 1
+  @test inner(ψx, ψ0) ≈ inner(ψ0, ψ0)
+end
+
+@testset "Decoupled Ladder" begin
+  cutoff = 1E-5
+  Nx = 10
+  Ny = 2
+  N = Nx * Ny
+  s = siteinds("S=1/2", N)
+
+  hterms = OpSum()
+  for j in 1:2:(N - 2)
+    hterms += "Sz", j, "Sz", j + 2
+    hterms += 1 / 2, "S+", j, "S-", j + 2
+    hterms += 1 / 2, "S-", j, "S+", j + 2
+  end
+  for j in 2:2:(N - 2)
+    hterms += "Sz", j, "Sz", j + 2
+    hterms += 1 / 2, "S+", j, "S-", j + 2
+    hterms += 1 / 2, "S-", j, "S+", j + 2
+  end
+  H = MPO(hterms, s)
+
+  tau = -0.5
+  Nstep = 40
+
+  ψ0 = randomMPS(s; linkdims=2)
+  energy, ψg = dmrg(
+    H, ψ0; nsweeps=10, noise=1E-10, maxdim=[10, 10, 20, 20, 40, 80, 100], cutoff=1E-8
+  )
+  @show energy
+  pause()
+
+  ψ = randomMPS(s; linkdims=1)
+  @show maxlinkdim(ψ)
+  @show inner(ψ', H, ψ)
+  for n in 1:Nstep
+    @show maxlinkdim(ψ)
+    if n % 4 == 1
+      ψ = basis_extend(ψ, H)
+      @show linkdims(ψ)
+    end
+    ψ = tdvp(H, tau, ψ; cutoff)
+    normalize!(ψ)
+    @show maxlinkdim(ψ)
+    @show inner(ψ', H, ψ)
+    println()
+  end
+end
+
+nothing