From 0c252aee621a96b64b5a51bb1b60a0e511706371 Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Tue, 3 May 2022 17:07:27 -0400
Subject: [PATCH 01/19] Add extend and basis_extend functions

---
 src/ITensorTDVP.jl        |   3 +-
 src/basis_extend.jl       | 102 ++++++++++++++++++++++++++++++++++++++
 test/test_basis_extend.jl |  88 ++++++++++++++++++++++++++++++++
 3 files changed, 192 insertions(+), 1 deletion(-)
 create mode 100644 src/basis_extend.jl
 create mode 100644 test/test_basis_extend.jl

diff --git a/src/ITensorTDVP.jl b/src/ITensorTDVP.jl
index 55498ad..4c4be57 100644
--- a/src/ITensorTDVP.jl
+++ b/src/ITensorTDVP.jl
@@ -25,7 +25,8 @@ include("tdvp_generic.jl")
 include("tdvp.jl")
 include("dmrg.jl")
 include("dmrg_x.jl")
+include("basis_extend.jl")
 
-export tdvp, dmrg_x, to_vec, TimeDependentSum
+export tdvp, dmrg_x, to_vec, TimeDependentSum, basis_extend
 
 end
diff --git a/src/basis_extend.jl b/src/basis_extend.jl
new file mode 100644
index 0000000..a2295e5
--- /dev/null
+++ b/src/basis_extend.jl
@@ -0,0 +1,102 @@
+#
+# Possible improvements
+#  - allow a maxdim argument to be passed to `extend`
+#    and through `basis_extend`
+#  - 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?
+#
+
+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.)
+    for r in rinds
+      Id *= 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)
+      bx = Index(dim(bψ)+dim(bϕ),"bx_$(j-1)")
+      D1,D2 = ITensors.directsum_itensors(bψ,bϕ,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=1:kdim
+    prev = k==1 ? psi : phis[k-1]
+    phis[k] = noprime(contract(H,prev;maxdim))
+    normalize!(phis[k])
+  end
+
+  cutoff = get(kwargs,:extension_cutoff,1E-8)
+  psix = extend(psi,phis;cutoff)
+  return psix
+end
diff --git a/test/test_basis_extend.jl b/test/test_basis_extend.jl
new file mode 100644
index 0000000..aff5fd4
--- /dev/null
+++ b/test/test_basis_extend.jl
@@ -0,0 +1,88 @@
+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 "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=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=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=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

From 0c6f4d8502308806c4b0b9cf5559b02a8109ee2c Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Tue, 3 May 2022 17:09:49 -0400
Subject: [PATCH 02/19] Add docstring and notes

---
 src/basis_extend.jl | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index a2295e5..b496a88 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -2,6 +2,8 @@
 # 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`?
@@ -9,6 +11,13 @@
 #    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...)
@@ -92,7 +101,7 @@ function basis_extend(psi::MPS,
   phis = Vector{MPS}(undef,kdim)
   for k=1:kdim
     prev = k==1 ? psi : phis[k-1]
-    phis[k] = noprime(contract(H,prev;maxdim))
+    phis[k] = apply(H,prev;maxdim)
     normalize!(phis[k])
   end
 

From f72b5b238ca2b5e69a0690fb0dffd866f5b50899 Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Tue, 3 May 2022 17:18:14 -0400
Subject: [PATCH 03/19] Formatting

---
 src/basis_extend.jl       | 62 ++++++++++++++++++-------------------
 test/test_basis_extend.jl | 64 +++++++++++++++++++--------------------
 2 files changed, 60 insertions(+), 66 deletions(-)

diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index b496a88..bbdeda0 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -18,49 +18,47 @@ returns an MPS which is equal to psi
 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)
+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)
+  orthogonalize!(psi, N)
   for phi in phis
-    orthogonalize!(phi,N)
+    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)
+    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.)
+    Id = ITensor(1.0)
     for r in rinds
-      Id *= delta(r',dag(r))
+      Id *= delta(r', dag(r))
     end
-    P = Id-prime(B,rinds)*dag(B)
+    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])
+      rho += prime(phi[j], rinds) * dag(phi[j])
     end
     rho /= tr(rho)
 
     # Apply projector
-    PrhoP = apply(apply(P,rho),P)
+    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")
+      D, Bphi = eigen(PrhoP; cutoff, ishermitian=true, righttags="bϕ_$j,Link")
 
       ## Test Bphi is ortho to B
       #O = Bphi*B
@@ -70,21 +68,21 @@ function extend(psi::MPS,
       #end
 
       # Form direct sum of B and Bphi over left index
-      bψ = commonind(B,S)
-      bϕ = commonind(Bphi,D)
-      bx = Index(dim(bψ)+dim(bϕ),"bx_$(j-1)")
-      D1,D2 = ITensors.directsum_itensors(bψ,bϕ,bx)
-      Bx = D1*B + D2*Bphi
+      bψ = commonind(B, S)
+      bϕ = commonind(Bphi, D)
+      bx = Index(dim(bψ) + dim(bϕ), "bx_$(j-1)")
+      D1, D2 = ITensors.directsum_itensors(bψ, bϕ, 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 - 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 - 1] = phi[j - 1] * (phi[j] * dag(Bx))
       phi[j] = Bx
     end
   end
@@ -92,20 +90,18 @@ function extend(psi::MPS,
   return psi
 end
 
-function basis_extend(psi::MPS,
-                      H::MPO;
-                      kwargs...)
-  kdim = get(kwargs,:extension_krylovdim,2)
-  maxdim = 1+maxlinkdim(psi)
+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=1:kdim
-    prev = k==1 ? psi : phis[k-1]
-    phis[k] = apply(H,prev;maxdim)
+  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)
+  cutoff = get(kwargs, :extension_cutoff, 1E-8)
+  psix = extend(psi, phis; cutoff)
   return psix
 end
diff --git a/test/test_basis_extend.jl b/test/test_basis_extend.jl
index aff5fd4..2d9c20c 100644
--- a/test/test_basis_extend.jl
+++ b/test/test_basis_extend.jl
@@ -7,14 +7,14 @@ using Printf
 
 @testset "extend function" begin
   N = 6
-  s = siteinds(2,N)
+  s = siteinds(2, N)
 
-  psi = randomMPS(s;linkdims=4)
-  phi = randomMPS(s;linkdims=2)
+  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)
+  psix = ITensorTDVP.extend(psi, [phi])
+  @test inner(psix, psi) ≈ inner(psi, psi)
+  @test inner(psix, phi) ≈ inner(psi, phi)
 end
 
 @testset "basis_extend" begin
@@ -30,59 +30,57 @@ end
   end
   H = MPO(os, s)
 
-  ψ0 = productMPS(s,n->isodd(n) ? "Up" : "Dn")
-  ψx = basis_extend(ψ0,H)
+  ψ0 = productMPS(s, n -> isodd(n) ? "Up" : "Dn")
+  ψx = basis_extend(ψ0, H)
   @test maxlinkdim(ψx) > 1
-  @test inner(ψx,ψ0) ≈ inner(ψ0,ψ0)
+  @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)
+  N = Nx * Ny
+  s = siteinds("S=1/2", N)
 
   hterms = OpSum()
-  for j=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
+  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=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
+  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)
+  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)
+  ψ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)
+  ψ = randomMPS(s; linkdims=1)
   @show maxlinkdim(ψ)
-  @show inner(ψ',H,ψ)
-  for n=1:Nstep
+  @show inner(ψ', H, ψ)
+  for n in 1:Nstep
     @show maxlinkdim(ψ)
-    if n%4==1
-      ψ = basis_extend(ψ,H)
+    if n % 4 == 1
+      ψ = basis_extend(ψ, H)
       @show linkdims(ψ)
     end
-    ψ = tdvp(H,tau,ψ; cutoff)
+    ψ = tdvp(H, tau, ψ; cutoff)
     normalize!(ψ)
     @show maxlinkdim(ψ)
-    @show inner(ψ',H,ψ)
+    @show inner(ψ', H, ψ)
     println()
   end
-
-
 end
 
-
 nothing

From d4cfbcfe63b525e0a8656e5b297e0b8de3ea3ddb Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Tue, 3 May 2022 17:35:21 -0400
Subject: [PATCH 04/19] Update Project.toml

---
 Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Project.toml b/Project.toml
index 9d8f7c9..d6548c7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,6 +19,7 @@ julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
 [targets]
 test = ["Test"]

From 654881dbd4c256b1e4177401044dcf05e547fcf3 Mon Sep 17 00:00:00 2001
From: Miles <miles.stoudenmire@gmail.com>
Date: Fri, 6 May 2022 14:52:20 -0400
Subject: [PATCH 05/19] Fix extend for QN case

---
 src/basis_extend.jl       | 11 ++++++++---
 test/test_basis_extend.jl | 12 ++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index bbdeda0..4e4f50f 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -1,3 +1,4 @@
+import ITensors: pause
 #
 # Possible improvements
 #  - allow a maxdim argument to be passed to `extend`
@@ -40,7 +41,7 @@ function extend(psi::MPS, phis::Vector{MPS}; kwargs...)
     # Make projector
     Id = ITensor(1.0)
     for r in rinds
-      Id *= delta(r', dag(r))
+      Id *= denseblocks(delta(r', dag(r)))
     end
     P = Id - prime(B, rinds) * dag(B)
 
@@ -70,8 +71,12 @@ function extend(psi::MPS, phis::Vector{MPS}; kwargs...)
       # Form direct sum of B and Bphi over left index
       bψ = commonind(B, S)
       bϕ = commonind(Bphi, D)
-      bx = Index(dim(bψ) + dim(bϕ), "bx_$(j-1)")
-      D1, D2 = ITensors.directsum_itensors(bψ, bϕ, bx)
+      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
diff --git a/test/test_basis_extend.jl b/test/test_basis_extend.jl
index 2d9c20c..45aabb0 100644
--- a/test/test_basis_extend.jl
+++ b/test/test_basis_extend.jl
@@ -17,6 +17,18 @@ using Printf
   @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
 

From dfe56dcb3000f1d5a1adb33345b2ce3c599d529d Mon Sep 17 00:00:00 2001
From: Matt Fishman <mtfishman@users.noreply.github.com>
Date: Wed, 15 May 2024 12:09:49 -0400
Subject: [PATCH 06/19] Update basis_extend.jl

---
 src/basis_extend.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index 4e4f50f..70a7e9d 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -1,4 +1,7 @@
-import ITensors: pause
+using ITensors: ITensors, Index, ITensor, commonind, dag, delta, denseblocks, hasqns, prime, uniqueinds
+using ITensors.ITensorMPS: MPO, MPS, apply, dim, linkind, maxlinkdim, orthogonalize!
+using LinearAlgebra: normalize!, svd, tr
+
 #
 # Possible improvements
 #  - allow a maxdim argument to be passed to `extend`

From 9e3aae23ad400fcbcf26744a382fd77cefc3271b Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 12:14:30 -0400
Subject: [PATCH 07/19] Format

---
 src/ITensorTDVP.jl  | 2 +-
 src/basis_extend.jl | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/ITensorTDVP.jl b/src/ITensorTDVP.jl
index d34f80b..cf22a78 100644
--- a/src/ITensorTDVP.jl
+++ b/src/ITensorTDVP.jl
@@ -22,4 +22,4 @@ using PackageExtensionCompat: @require_extensions
 function __init__()
   @require_extensions
 end
-end
\ No newline at end of file
+end
diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index 70a7e9d..6fcd8c8 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -1,4 +1,5 @@
-using ITensors: ITensors, Index, ITensor, commonind, dag, delta, denseblocks, hasqns, prime, uniqueinds
+using ITensors:
+  ITensors, Index, ITensor, commonind, dag, delta, denseblocks, hasqns, prime, uniqueinds
 using ITensors.ITensorMPS: MPO, MPS, apply, dim, linkind, maxlinkdim, orthogonalize!
 using LinearAlgebra: normalize!, svd, tr
 

From b501b8faccae9d04c453b7c9ced64ea4a81ae9e0 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 14:07:06 -0400
Subject: [PATCH 08/19] Rewrite

---
 src/basis_extend.jl       | 181 ++++++++++++++++++--------------------
 test/test_basis_extend.jl | 149 +++++++++++++------------------
 2 files changed, 146 insertions(+), 184 deletions(-)

diff --git a/src/basis_extend.jl b/src/basis_extend.jl
index 6fcd8c8..227240f 100644
--- a/src/basis_extend.jl
+++ b/src/basis_extend.jl
@@ -1,7 +1,20 @@
 using ITensors:
-  ITensors, Index, ITensor, commonind, dag, delta, denseblocks, hasqns, prime, uniqueinds
-using ITensors.ITensorMPS: MPO, MPS, apply, dim, linkind, maxlinkdim, orthogonalize!
-using LinearAlgebra: normalize!, svd, tr
+  ITensors,
+  Algorithm,
+  Index,
+  ITensor,
+  @Algorithm_str,
+  δ,
+  commonind,
+  dag,
+  denseblocks,
+  directsum,
+  hasqns,
+  prime,
+  scalartype,
+  uniqueinds
+using ITensors.ITensorMPS: MPO, MPS, apply, dim, linkind, maxlinkdim, orthogonalize
+using LinearAlgebra: normalize, svd, tr
 
 #
 # Possible improvements
@@ -9,108 +22,82 @@ using LinearAlgebra: normalize!, svd, tr
 #    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?
+#  - Use (1-tau*operator)|state> to generate "Krylov" vectors
+#    instead of operator|state>. Needed?
 #
 
+function expand_basis(state, reference; alg, kwargs...)
+  return expand_basis(Algorithm(alg), state, reference; kwargs...)
+end
+
 """
-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.
+Given an MPS state and a collection of MPS references,
+returns an MPS which is equal to state
+(has fidelity 1.0 with state) but whose MPS basis
+is expanded to contain a portion of the basis of
+the references that is orthogonal to the MPS basis of state.
 """
-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])
+function expand_basis(
+  ::Algorithm"orthogonalize",
+  state::MPS,
+  references::Vector{MPS};
+  cutoff=10^2 * eps(real(scalartype(state))),
+)
+  n = length(state)
+  state = orthogonalize(state, n)
+  references = map(reference -> orthogonalize(reference, n), references)
+  s = siteinds(state)
+  for j in reverse(2:n)
+    # SVD state[j] to compute basisⱼ
+    linds = [s[j - 1]; linkinds(state, j - 1)]
+    _, λⱼ, basisⱼ = svd(state[j], linds; righttags="bψ_$j,Link")
+    rinds = uniqueinds(basisⱼ, λⱼ)
+    # Make projectorⱼ
+    idⱼ = prod(r -> denseblocks(δ(r', dag(r))), rinds)
+    projectorⱼ = idⱼ - prime(basisⱼ, rinds) * dag(basisⱼ)
+    # Sum reference density matrices
+    ρⱼ = sum(reference -> prime(reference[j], rinds) * dag(reference[j]), references)
+    ρⱼ /= tr(ρⱼ)
+    # Apply projectorⱼ
+    ρⱼ_projected = apply(apply(projectorⱼ, ρⱼ), projectorⱼ)
+    expanded_basisⱼ = basisⱼ
+    if norm(ρⱼ_projected) > 10^3 * eps(real(scalartype(state)))
+      # Diagonalize projected density matrix ρⱼ_projected
+      # to compute reference_basisⱼ, which spans part of right basis
+      # of references which is orthogonal to right basis of state
+      dⱼ, reference_basisⱼ = eigen(
+        ρⱼ_projected; cutoff, ishermitian=true, righttags="bϕ_$j,Link"
+      )
+      state_indⱼ = only(commoninds(basisⱼ, λⱼ))
+      reference_indⱼ = only(commoninds(reference_basisⱼ, dⱼ))
+      expanded_basisⱼ, bx = directsum(
+        basisⱼ => state_indⱼ, reference_basisⱼ => reference_indⱼ
+      )
     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
+    # Shift ortho center one site left using dag(expanded_basisⱼ)
+    # and replace tensor at site j with expanded_basisⱼ
+    state[j - 1] = state[j - 1] * (state[j] * dag(expanded_basisⱼ))
+    state[j] = expanded_basisⱼ
+    for reference in references
+      reference[j - 1] = reference[j - 1] * (reference[j] * dag(expanded_basisⱼ))
+      reference[j] = expanded_basisⱼ
     end
   end
-
-  return psi
+  return state
 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])
+function expand_basis(
+  ::Algorithm"global_krylov",
+  state::MPS,
+  operator::MPO;
+  krylovdim=2,
+  cutoff=(√(eps(scalartype(state)))),
+)
+  maxdim = maxlinkdim(state) + 1
+  references = Vector{MPS}(undef, krylovdim)
+  for k in 1:krylovdim
+    prev = k == 1 ? state : references[k - 1]
+    references[k] = normalize(apply(operator, prev; maxdim))
   end
-
-  cutoff = get(kwargs, :extension_cutoff, 1E-8)
-  psix = extend(psi, phis; cutoff)
-  return psix
+  return expand_basis(state, references; alg="orthogonalize", cutoff)
 end
diff --git a/test/test_basis_extend.jl b/test/test_basis_extend.jl
index 45aabb0..81a20a1 100644
--- a/test/test_basis_extend.jl
+++ b/test/test_basis_extend.jl
@@ -1,98 +1,73 @@
-using ITensors
-import ITensors: pause
-using ITensorTDVP
-using Random
-using Test
-using Printf
+@eval module $(gensym())
+using ITensors.ITensorMPS:
+  OpSum, MPO, MPS, dmrg, inner, linkdims, maxlinkdim, randomMPS, siteinds
+using ITensorTDVP: expand_basis, tdvp
+using LinearAlgebra: normalize
+using Test: @test, @testset
+const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
+@testset "expand_basis_orthogonalize (conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
+                                                                                    (
+    false, true
+  ),
+  elt in elts
 
-@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)
+  n = 6
+  s = siteinds("S=1/2", n; conserve_qns)
+  state = randomMPS(s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
+  reference = randomMPS(s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
+  state_expanded = expand_basis(state, [reference]; alg="orthogonalize")
+  @test inner(state_expanded, state) ≈ inner(state, state)
+  @test inner(state_expanded, reference) ≈ inner(state, reference)
 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
+  n = 10
+  s = siteinds("S=1/2", n)
+  opsum = OpSum()
+  for j in 1:(n - 1)
+    opsum += 0.5, "S+", j, "S-", j + 1
+    opsum += 0.5, "S-", j, "S+", j + 1
+    opsum += "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)
+  operator = MPO(opsum, s)
+  state = MPS(s, n -> isodd(n) ? "Up" : "Dn")
+  state_expanded = expand_basis(state, operator; alg="global_krylov")
+  @test maxlinkdim(state_expanded) > 1
+  @test inner(state_expanded, state) ≈ inner(state, state)
 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
+  nx = 10
+  ny = 2
+  n = nx * ny
+  s = siteinds("S=1/2", n)
+  opsum = OpSum()
+  for j in 1:2:(n - 2)
+    opsum += "Sz", j, "Sz", j + 2
+    opsum += 1 / 2, "S+", j, "S-", j + 2
+    opsum += 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
+  for j in 2:2:(n - 2)
+    opsum += "Sz", j, "Sz", j + 2
+    opsum += 1 / 2, "S+", j, "S-", j + 2
+    opsum += 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
+  operator = MPO(opsum, s)
+  nexpansions = 10
+  init = randomMPS(s; linkdims=2)
+  reference_energy, reference_state = dmrg(
+    operator,
+    init;
+    nsweeps=10,
+    maxdim=[10, 10, 20, 20, 40, 80, 100],
+    cutoff=1e-8,
+    noise=1e-10,
   )
-  @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()
+  state = randomMPS(s)
+  tau = 0.5
+  for step in 1:nexpansions
+    state = expand_basis(state, operator; alg="global_krylov", cutoff=1e-4)
+    state = tdvp(operator, -4tau, state; nsteps=4, cutoff=1e-5)
+    state = normalize(state)
   end
+  @test inner(state', operator, state) ≈ reference_energy rtol = 1e-3
+end
 end
-
-nothing

From cb894beb10c014f13b8bb2e01eaebc42bf130198 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 18:22:58 -0400
Subject: [PATCH 09/19] Rename

---
 src/ITensorTDVP.jl                       |  4 +-
 src/{basis_extend.jl => expand_basis.jl} |  7 +-
 src/sweep_update.jl                      | 10 ++-
 test/test_basis_extend.jl                | 73 --------------------
 test/test_expand_basis.jl                | 86 ++++++++++++++++++++++++
 test/test_exports.jl                     |  2 +-
 6 files changed, 102 insertions(+), 80 deletions(-)
 rename src/{basis_extend.jl => expand_basis.jl} (93%)
 delete mode 100644 test/test_basis_extend.jl
 create mode 100644 test/test_expand_basis.jl

diff --git a/src/ITensorTDVP.jl b/src/ITensorTDVP.jl
index cf22a78..09aecdc 100644
--- a/src/ITensorTDVP.jl
+++ b/src/ITensorTDVP.jl
@@ -1,5 +1,5 @@
 module ITensorTDVP
-export TimeDependentSum, basis_extend, dmrg_x, linsolve, tdvp, to_vec
+export TimeDependentSum, dmrg_x, expand_basis, linsolve, tdvp, to_vec
 include("ITensorsExtensions.jl")
 using .ITensorsExtensions: to_vec
 include("applyexp.jl")
@@ -17,7 +17,7 @@ include("contract.jl")
 include("reducedconstantterm.jl")
 include("reducedlinearproblem.jl")
 include("linsolve.jl")
-include("basis_extend.jl")
+include("expand_basis.jl")
 using PackageExtensionCompat: @require_extensions
 function __init__()
   @require_extensions
diff --git a/src/basis_extend.jl b/src/expand_basis.jl
similarity index 93%
rename from src/basis_extend.jl
rename to src/expand_basis.jl
index 227240f..3b988c1 100644
--- a/src/basis_extend.jl
+++ b/src/expand_basis.jl
@@ -53,11 +53,12 @@ function expand_basis(
     _, λⱼ, basisⱼ = svd(state[j], linds; righttags="bψ_$j,Link")
     rinds = uniqueinds(basisⱼ, λⱼ)
     # Make projectorⱼ
-    idⱼ = prod(r -> denseblocks(δ(r', dag(r))), rinds)
+    idⱼ = prod(r -> denseblocks(δ(scalartype(state), r', dag(r))), rinds)
     projectorⱼ = idⱼ - prime(basisⱼ, rinds) * dag(basisⱼ)
     # Sum reference density matrices
     ρⱼ = sum(reference -> prime(reference[j], rinds) * dag(reference[j]), references)
-    ρⱼ /= tr(ρⱼ)
+    # TODO: Fix bug that `tr` isn't preserving the element type.
+    ρⱼ /= scalartype(state)(tr(ρⱼ))
     # Apply projectorⱼ
     ρⱼ_projected = apply(apply(projectorⱼ, ρⱼ), projectorⱼ)
     expanded_basisⱼ = basisⱼ
@@ -91,7 +92,7 @@ function expand_basis(
   state::MPS,
   operator::MPO;
   krylovdim=2,
-  cutoff=(√(eps(scalartype(state)))),
+  cutoff=(√(eps(real(scalartype(state))))),
 )
   maxdim = maxlinkdim(state) + 1
   references = Vector{MPS}(undef, krylovdim)
diff --git a/src/sweep_update.jl b/src/sweep_update.jl
index bdaf3d3..ed97685 100644
--- a/src/sweep_update.jl
+++ b/src/sweep_update.jl
@@ -1,6 +1,14 @@
 using ITensors: ITensors, uniqueinds
 using ITensors.ITensorMPS:
-  ITensorMPS, MPS, isortho, orthocenter, orthogonalize!, position!, replacebond!, set_nsite!
+  ITensorMPS,
+  MPS,
+  isortho,
+  noiseterm,
+  orthocenter,
+  orthogonalize!,
+  position!,
+  replacebond!,
+  set_nsite!
 using LinearAlgebra: norm, normalize!, svd
 using Printf: @printf
 
diff --git a/test/test_basis_extend.jl b/test/test_basis_extend.jl
deleted file mode 100644
index 81a20a1..0000000
--- a/test/test_basis_extend.jl
+++ /dev/null
@@ -1,73 +0,0 @@
-@eval module $(gensym())
-using ITensors.ITensorMPS:
-  OpSum, MPO, MPS, dmrg, inner, linkdims, maxlinkdim, randomMPS, siteinds
-using ITensorTDVP: expand_basis, tdvp
-using LinearAlgebra: normalize
-using Test: @test, @testset
-const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
-@testset "expand_basis_orthogonalize (conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
-                                                                                    (
-    false, true
-  ),
-  elt in elts
-
-  n = 6
-  s = siteinds("S=1/2", n; conserve_qns)
-  state = randomMPS(s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
-  reference = randomMPS(s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
-  state_expanded = expand_basis(state, [reference]; alg="orthogonalize")
-  @test inner(state_expanded, state) ≈ inner(state, state)
-  @test inner(state_expanded, reference) ≈ inner(state, reference)
-end
-@testset "basis_extend" begin
-  n = 10
-  s = siteinds("S=1/2", n)
-  opsum = OpSum()
-  for j in 1:(n - 1)
-    opsum += 0.5, "S+", j, "S-", j + 1
-    opsum += 0.5, "S-", j, "S+", j + 1
-    opsum += "Sz", j, "Sz", j + 1
-  end
-  operator = MPO(opsum, s)
-  state = MPS(s, n -> isodd(n) ? "Up" : "Dn")
-  state_expanded = expand_basis(state, operator; alg="global_krylov")
-  @test maxlinkdim(state_expanded) > 1
-  @test inner(state_expanded, state) ≈ inner(state, state)
-end
-@testset "Decoupled Ladder" begin
-  nx = 10
-  ny = 2
-  n = nx * ny
-  s = siteinds("S=1/2", n)
-  opsum = OpSum()
-  for j in 1:2:(n - 2)
-    opsum += "Sz", j, "Sz", j + 2
-    opsum += 1 / 2, "S+", j, "S-", j + 2
-    opsum += 1 / 2, "S-", j, "S+", j + 2
-  end
-  for j in 2:2:(n - 2)
-    opsum += "Sz", j, "Sz", j + 2
-    opsum += 1 / 2, "S+", j, "S-", j + 2
-    opsum += 1 / 2, "S-", j, "S+", j + 2
-  end
-  operator = MPO(opsum, s)
-  nexpansions = 10
-  init = randomMPS(s; linkdims=2)
-  reference_energy, reference_state = dmrg(
-    operator,
-    init;
-    nsweeps=10,
-    maxdim=[10, 10, 20, 20, 40, 80, 100],
-    cutoff=1e-8,
-    noise=1e-10,
-  )
-  state = randomMPS(s)
-  tau = 0.5
-  for step in 1:nexpansions
-    state = expand_basis(state, operator; alg="global_krylov", cutoff=1e-4)
-    state = tdvp(operator, -4tau, state; nsteps=4, cutoff=1e-5)
-    state = normalize(state)
-  end
-  @test inner(state', operator, state) ≈ reference_energy rtol = 1e-3
-end
-end
diff --git a/test/test_expand_basis.jl b/test/test_expand_basis.jl
new file mode 100644
index 0000000..81d9ee5
--- /dev/null
+++ b/test/test_expand_basis.jl
@@ -0,0 +1,86 @@
+@eval module $(gensym())
+using ITensors: scalartype
+using ITensors.ITensorMPS: OpSum, MPO, MPS, inner, linkdims, maxlinkdim, randomMPS, siteinds
+using ITensorTDVP: dmrg, expand_basis, tdvp
+using LinearAlgebra: normalize
+using Test: @test, @testset
+const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
+@testset "expand_basis (eltype=$elt)" for elt in elts
+  @testset "expand_basis (alg=\"orthogonalize\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
+                                                                                               (
+    false, true
+  )
+    n = 6
+    s = siteinds("S=1/2", n; conserve_qns)
+    state = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
+    reference = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
+    state_expanded = expand_basis(state, [reference]; alg="orthogonalize")
+    @test scalartype(state_expanded) === elt
+    @test inner(state_expanded, state) ≈ inner(state, state)
+    @test inner(state_expanded, reference) ≈ inner(state, reference)
+  end
+  @testset "basis_extend (alg=\"global_krylov\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
+                                                                                               (
+    false, true
+  )
+    n = 10
+    s = siteinds("S=1/2", n; conserve_qns)
+    opsum = OpSum()
+    for j in 1:(n - 1)
+      opsum += 0.5, "S+", j, "S-", j + 1
+      opsum += 0.5, "S-", j, "S+", j + 1
+      opsum += "Sz", j, "Sz", j + 1
+    end
+    operator = MPO(elt, opsum, s)
+    state = MPS(elt, s, j -> isodd(j) ? "↑" : "↓")
+    state_expanded = expand_basis(state, operator; alg="global_krylov")
+    @test scalartype(state_expanded) === elt
+    @test maxlinkdim(state_expanded) > 1
+    @test inner(state_expanded, state) ≈ inner(state, state)
+  end
+  @testset "Decoupled ladder (alg=\"global_krylov\", eltype=$elt)" begin
+    nx = 10
+    ny = 2
+    n = nx * ny
+    s = siteinds("S=1/2", n)
+    opsum = OpSum()
+    for j in 1:2:(n - 2)
+      opsum += 1 / 2, "S+", j, "S-", j + 2
+      opsum += 1 / 2, "S-", j, "S+", j + 2
+      opsum += "Sz", j, "Sz", j + 2
+    end
+    for j in 2:2:(n - 2)
+      opsum += 1 / 2, "S+", j, "S-", j + 2
+      opsum += 1 / 2, "S-", j, "S+", j + 2
+      opsum += "Sz", j, "Sz", j + 2
+    end
+    operator = MPO(elt, opsum, s)
+    init = randomMPS(elt, s; linkdims=30)
+    reference_energy, reference_state = dmrg(
+      operator,
+      init;
+      nsweeps=15,
+      maxdim=[10, 10, 20, 20, 40, 80, 100],
+      cutoff=(√(eps(real(elt)))),
+      noise=(√(eps(real(elt)))),
+    )
+    state = randomMPS(elt, s)
+    nexpansions = 10
+    tau = elt(0.5)
+    for step in 1:nexpansions
+      state = expand_basis(state, operator; alg="global_krylov", cutoff=(∜(eps(real(elt)))))
+      state = tdvp(
+        operator,
+        -4tau,
+        state;
+        nsteps=4,
+        cutoff=1e-5,
+        updater_kwargs=(; tol=1e-3, krylovdim=5),
+      )
+      state = normalize(state)
+    end
+    @test scalartype(state) === elt
+    @test inner(state', operator, state) ≈ reference_energy rtol = 2 * ∜(eps(real(elt)))
+  end
+end
+end
diff --git a/test/test_exports.jl b/test/test_exports.jl
index ffc1f85..f46c40b 100644
--- a/test/test_exports.jl
+++ b/test/test_exports.jl
@@ -4,7 +4,7 @@ using Test: @test, @testset
 @testset "Test exports" begin
   @test issetequal(
     names(ITensorTDVP),
-    [:ITensorTDVP, :TimeDependentSum, :dmrg_x, :linsolve, :tdvp, :to_vec],
+    [:ITensorTDVP, :TimeDependentSum, :dmrg_x, :expand_basis, :linsolve, :tdvp, :to_vec],
   )
 end
 end

From bd038b03792a8e5d9713604c4c2f00ed601d9995 Mon Sep 17 00:00:00 2001
From: Matt Fishman <mtfishman@users.noreply.github.com>
Date: Wed, 15 May 2024 18:23:48 -0400
Subject: [PATCH 10/19] Update Project.toml

---
 Project.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 2738582..15e422a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -30,7 +30,6 @@ 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"]

From eb88efef0a3fa4b498e7d60cf901ce6b01e8243f Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 22:41:55 -0400
Subject: [PATCH 11/19] Rename expand_basis to expand, change definition of
 krylovdim

---
 src/ITensorTDVP.jl                 |  4 +--
 src/{expand_basis.jl => expand.jl} | 39 +++++++++++++++---------------
 test/test_expand_basis.jl          | 23 +++++++++++-------
 3 files changed, 35 insertions(+), 31 deletions(-)
 rename src/{expand_basis.jl => expand.jl} (77%)

diff --git a/src/ITensorTDVP.jl b/src/ITensorTDVP.jl
index 09aecdc..4457cd1 100644
--- a/src/ITensorTDVP.jl
+++ b/src/ITensorTDVP.jl
@@ -1,5 +1,5 @@
 module ITensorTDVP
-export TimeDependentSum, dmrg_x, expand_basis, linsolve, tdvp, to_vec
+export TimeDependentSum, dmrg_x, expand, linsolve, tdvp, to_vec
 include("ITensorsExtensions.jl")
 using .ITensorsExtensions: to_vec
 include("applyexp.jl")
@@ -17,7 +17,7 @@ include("contract.jl")
 include("reducedconstantterm.jl")
 include("reducedlinearproblem.jl")
 include("linsolve.jl")
-include("expand_basis.jl")
+include("expand.jl")
 using PackageExtensionCompat: @require_extensions
 function __init__()
   @require_extensions
diff --git a/src/expand_basis.jl b/src/expand.jl
similarity index 77%
rename from src/expand_basis.jl
rename to src/expand.jl
index 3b988c1..f0e63c8 100644
--- a/src/expand_basis.jl
+++ b/src/expand.jl
@@ -16,18 +16,15 @@ using ITensors:
 using ITensors.ITensorMPS: MPO, MPS, apply, dim, linkind, maxlinkdim, orthogonalize
 using LinearAlgebra: normalize, svd, tr
 
-#
-# 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
+# Possible improvements:
+#  - Allow a maxdim argument to be passed to `expand`.
+#  - Current behavior is letting bond dimension get too
+#    big when used in imaginary time evolution.
 #  - Use (1-tau*operator)|state> to generate "Krylov" vectors
-#    instead of operator|state>. Needed?
-#
+#    instead of operator|state>. Is that needed?
 
-function expand_basis(state, reference; alg, kwargs...)
-  return expand_basis(Algorithm(alg), state, reference; kwargs...)
+function expand(state, reference; alg, kwargs...)
+  return expand(Algorithm(alg), state, reference; kwargs...)
 end
 
 """
@@ -37,7 +34,7 @@ returns an MPS which is equal to state
 is expanded to contain a portion of the basis of
 the references that is orthogonal to the MPS basis of state.
 """
-function expand_basis(
+function expand(
   ::Algorithm"orthogonalize",
   state::MPS,
   references::Vector{MPS};
@@ -71,7 +68,7 @@ function expand_basis(
       )
       state_indⱼ = only(commoninds(basisⱼ, λⱼ))
       reference_indⱼ = only(commoninds(reference_basisⱼ, dⱼ))
-      expanded_basisⱼ, bx = directsum(
+      expanded_basisⱼ, expanded_indⱼ = directsum(
         basisⱼ => state_indⱼ, reference_basisⱼ => reference_indⱼ
       )
     end
@@ -87,18 +84,20 @@ function expand_basis(
   return state
 end
 
-function expand_basis(
+function expand(
   ::Algorithm"global_krylov",
   state::MPS,
   operator::MPO;
-  krylovdim=2,
+  krylovdim=3,
   cutoff=(√(eps(real(scalartype(state))))),
 )
-  maxdim = maxlinkdim(state) + 1
-  references = Vector{MPS}(undef, krylovdim)
-  for k in 1:krylovdim
-    prev = k == 1 ? state : references[k - 1]
-    references[k] = normalize(apply(operator, prev; maxdim))
+  # TODO: Try replacing this logic with `Base.accumulate`.
+  references = Vector{MPS}(undef, krylovdim - 1)
+  for k in 1:(krylovdim - 1)
+    previous_reference = get(references, k - 1, state)
+    references[k] = normalize(
+      apply(operator, previous_reference; maxdim=maxlinkdim(state) + 1)
+    )
   end
-  return expand_basis(state, references; alg="orthogonalize", cutoff)
+  return expand(state, references; alg="orthogonalize", cutoff)
 end
diff --git a/test/test_expand_basis.jl b/test/test_expand_basis.jl
index 81d9ee5..cc2ec86 100644
--- a/test/test_expand_basis.jl
+++ b/test/test_expand_basis.jl
@@ -1,25 +1,25 @@
 @eval module $(gensym())
 using ITensors: scalartype
 using ITensors.ITensorMPS: OpSum, MPO, MPS, inner, linkdims, maxlinkdim, randomMPS, siteinds
-using ITensorTDVP: dmrg, expand_basis, tdvp
+using ITensorTDVP: dmrg, expand, tdvp
 using LinearAlgebra: normalize
 using Test: @test, @testset
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
-@testset "expand_basis (eltype=$elt)" for elt in elts
-  @testset "expand_basis (alg=\"orthogonalize\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
-                                                                                               (
+@testset "expand (eltype=$elt)" for elt in elts
+  @testset "expand (alg=\"orthogonalize\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
+                                                                                         (
     false, true
   )
     n = 6
     s = siteinds("S=1/2", n; conserve_qns)
     state = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
     reference = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
-    state_expanded = expand_basis(state, [reference]; alg="orthogonalize")
+    state_expanded = expand(state, [reference]; alg="orthogonalize")
     @test scalartype(state_expanded) === elt
     @test inner(state_expanded, state) ≈ inner(state, state)
     @test inner(state_expanded, reference) ≈ inner(state, reference)
   end
-  @testset "basis_extend (alg=\"global_krylov\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
+  @testset "expand (alg=\"global_krylov\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
                                                                                                (
     false, true
   )
@@ -33,7 +33,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     end
     operator = MPO(elt, opsum, s)
     state = MPS(elt, s, j -> isodd(j) ? "↑" : "↓")
-    state_expanded = expand_basis(state, operator; alg="global_krylov")
+    state_expanded = expand(state, operator; alg="global_krylov")
     @test scalartype(state_expanded) === elt
     @test maxlinkdim(state_expanded) > 1
     @test inner(state_expanded, state) ≈ inner(state, state)
@@ -68,7 +68,10 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     nexpansions = 10
     tau = elt(0.5)
     for step in 1:nexpansions
-      state = expand_basis(state, operator; alg="global_krylov", cutoff=(∜(eps(real(elt)))))
+      # TODO: Use `∜` instead of `fourthroot` in Julia 1.10 and above.
+      state = expand(
+        state, operator; alg="global_krylov", krylovdim=3, cutoff=fourthroot(eps(real(elt)))
+      )
       state = tdvp(
         operator,
         -4tau,
@@ -80,7 +83,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       state = normalize(state)
     end
     @test scalartype(state) === elt
-    @test inner(state', operator, state) ≈ reference_energy rtol = 2 * ∜(eps(real(elt)))
+    # TODO: Use `∜` instead of `fourthroot` in Julia 1.10 and above.
+    @test inner(state', operator, state) ≈ reference_energy rtol =
+      2 * fourthroot(eps(real(elt)))
   end
 end
 end

From 11f37ea64c3c5f66cc47227775c40dc9972996ed Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 22:45:57 -0400
Subject: [PATCH 12/19] Format, rename test file

---
 test/{test_expand_basis.jl => test_expand.jl} | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)
 rename test/{test_expand_basis.jl => test_expand.jl} (98%)

diff --git a/test/test_expand_basis.jl b/test/test_expand.jl
similarity index 98%
rename from test/test_expand_basis.jl
rename to test/test_expand.jl
index cc2ec86..30bc1db 100644
--- a/test/test_expand_basis.jl
+++ b/test/test_expand.jl
@@ -20,7 +20,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test inner(state_expanded, reference) ≈ inner(state, reference)
   end
   @testset "expand (alg=\"global_krylov\", conserve_qns=$conserve_qns, eltype=$elt)" for conserve_qns in
-                                                                                               (
+                                                                                         (
     false, true
   )
     n = 10

From b9ac39ea984044ab7b80ae4777f0cb6b1f01a32f Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Wed, 15 May 2024 22:59:51 -0400
Subject: [PATCH 13/19] Change back definition of krylovdim

---
 src/expand.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/expand.jl b/src/expand.jl
index f0e63c8..e93ac85 100644
--- a/src/expand.jl
+++ b/src/expand.jl
@@ -88,12 +88,12 @@ function expand(
   ::Algorithm"global_krylov",
   state::MPS,
   operator::MPO;
-  krylovdim=3,
+  krylovdim=2,
   cutoff=(√(eps(real(scalartype(state))))),
 )
   # TODO: Try replacing this logic with `Base.accumulate`.
-  references = Vector{MPS}(undef, krylovdim - 1)
-  for k in 1:(krylovdim - 1)
+  references = Vector{MPS}(undef, krylovdim)
+  for k in 1:krylovdim
     previous_reference = get(references, k - 1, state)
     references[k] = normalize(
       apply(operator, previous_reference; maxdim=maxlinkdim(state) + 1)

From 38bcb44601746c1100c031ddbb08fc27f806e20e Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 May 2024 07:43:24 -0400
Subject: [PATCH 14/19] Fix exports test

---
 test/test_exports.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/test_exports.jl b/test/test_exports.jl
index f46c40b..ab457cd 100644
--- a/test/test_exports.jl
+++ b/test/test_exports.jl
@@ -4,7 +4,7 @@ using Test: @test, @testset
 @testset "Test exports" begin
   @test issetequal(
     names(ITensorTDVP),
-    [:ITensorTDVP, :TimeDependentSum, :dmrg_x, :expand_basis, :linsolve, :tdvp, :to_vec],
+    [:ITensorTDVP, :TimeDependentSum, :dmrg_x, :expand, :linsolve, :tdvp, :to_vec],
   )
 end
 end

From 17ed2c9224de7700badfe73d1d5b5b92244be842 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 May 2024 08:12:40 -0400
Subject: [PATCH 15/19] Set RNG seeds in tests

---
 test/test_expand.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/test/test_expand.jl b/test/test_expand.jl
index 30bc1db..d14391b 100644
--- a/test/test_expand.jl
+++ b/test/test_expand.jl
@@ -3,6 +3,7 @@ using ITensors: scalartype
 using ITensors.ITensorMPS: OpSum, MPO, MPS, inner, linkdims, maxlinkdim, randomMPS, siteinds
 using ITensorTDVP: dmrg, expand, tdvp
 using LinearAlgebra: normalize
+using Random: Random
 using Test: @test, @testset
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "expand (eltype=$elt)" for elt in elts
@@ -12,6 +13,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
   )
     n = 6
     s = siteinds("S=1/2", n; conserve_qns)
+    Random.seed!(1234)
     state = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
     reference = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
     state_expanded = expand(state, [reference]; alg="orthogonalize")
@@ -55,6 +57,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       opsum += "Sz", j, "Sz", j + 2
     end
     operator = MPO(elt, opsum, s)
+    Random.seed!(1234)
     init = randomMPS(elt, s; linkdims=30)
     reference_energy, reference_state = dmrg(
       operator,
@@ -64,6 +67,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       cutoff=(√(eps(real(elt)))),
       noise=(√(eps(real(elt)))),
     )
+    Random.seed!(1234)
     state = randomMPS(elt, s)
     nexpansions = 10
     tau = elt(0.5)

From d2484ab54851d4920cc2426dfc895231ce76c2c7 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 May 2024 09:26:56 -0400
Subject: [PATCH 16/19] Add docstrings and News in README

---
 Project.toml        |  2 +-
 README.md           | 17 +++++++++++
 src/expand.jl       | 74 +++++++++++++++++++++++++++++++++++++++------
 test/test_expand.jl |  2 +-
 4 files changed, 84 insertions(+), 11 deletions(-)

diff --git a/Project.toml b/Project.toml
index 15e422a..2fd21d7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ITensorTDVP"
 uuid = "25707e16-a4db-4a07-99d9-4d67b7af0342"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org> and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
diff --git a/README.md b/README.md
index 3130aa6..5638ae7 100644
--- a/README.md
+++ b/README.md
@@ -22,6 +22,23 @@ However, as noted above we now recommend installing and loading `ITensorMPS` ins
 
 ## News
 
+### ITensorTDVP.jl v0.4.1 release notes
+
+#### Breaking changes
+
+#### New features
+
+- A new (experimental) `expand` function has been introduced for performing global Krylov expansion based on [arXiv:2005.06104](https://arxiv.org/abs/2005.06104), which can help with the accuracy of TDVP evolution in certain cases. See the docstrings of `expand` for more details:
+```julia
+julia> using ITensorTDVP
+
+julia> ?
+
+help?> expand
+# ...
+```
+Users are not given many customization options just yet as we gain more experience on the right balance between efficacy of the expansion and performance, and default values and keyword argument option are subject to change as we learn more about how to best use the method.
+
 ### ITensorTDVP.jl v0.4 release notes
 
 #### Breaking changes
diff --git a/src/expand.jl b/src/expand.jl
index e93ac85..0d3c62a 100644
--- a/src/expand.jl
+++ b/src/expand.jl
@@ -20,25 +20,61 @@ using LinearAlgebra: normalize, svd, tr
 #  - Allow a maxdim argument to be passed to `expand`.
 #  - Current behavior is letting bond dimension get too
 #    big when used in imaginary time evolution.
-#  - Use (1-tau*operator)|state> to generate "Krylov" vectors
+#  - Consider switching the default to variational/fit apply
+#    when building Krylov vectors.
+#  - Use (1-tau*operator)|state> to generate Krylov vectors
 #    instead of operator|state>. Is that needed?
 
 function expand(state, reference; alg, kwargs...)
   return expand(Algorithm(alg), state, reference; kwargs...)
 end
 
+function expand_cutoff_doctring()
+  return """
+  The cutoff is used to control the truncation of the expanded
+  basis and defaults to half the precision of the scalar type
+  of the input state, i.e. ~1e-8 for `Float64`.
+  """
+end
+
+function expand_warning_doctring()
+  return """
+  !!! warning
+      Users are not given many customization options just yet as we
+      gain more experience on the right balance between efficacy of the
+      expansion and performance in various scenarios, and default values
+      and keyword arguments are subject to change as we learn more about
+      how to best use the method.
+  """
+end
+
+function expand_citation_docstring()
+  return """
+  [^global_expansion]: Time Dependent Variational Principle with Ancillary Krylov Subspace. Mingru Yang, Steven R. White, [arXiv:2005.06104](https://arxiv.org/abs/2005.06104)
+  """
+end
+
 """
-Given an MPS state and a collection of MPS references,
-returns an MPS which is equal to state
-(has fidelity 1.0 with state) but whose MPS basis
+    expand(state::MPS, references::Vector{MPS}; alg="orthogonalize", cutoff)
+
+Given an MPS `state` and a collection of MPS `references`,
+returns an MPS which is equal to `state`
+(has fidelity 1 with `state`) but whose MPS basis
 is expanded to contain a portion of the basis of
-the references that is orthogonal to the MPS basis of state.
+the `references` that is orthogonal to the MPS basis of `state`.
+See [^global_expansion] for more details.
+
+$(expand_cutoff_doctring())
+
+$(expand_warning_doctring())
+
+$(expand_citation_docstring())
 """
 function expand(
   ::Algorithm"orthogonalize",
   state::MPS,
   references::Vector{MPS};
-  cutoff=10^2 * eps(real(scalartype(state))),
+  cutoff=(√(eps(real(scalartype(state))))),
 )
   n = length(state)
   state = orthogonalize(state, n)
@@ -84,20 +120,40 @@ function expand(
   return state
 end
 
+"""
+    expand(state::MPS, reference::MPO; alg="global_krylov", krylovdim=2, cutoff)
+
+Given an MPS `state` and an MPO `reference`,
+returns an MPS which is equal to `state`
+(has fidelity 1 with state) but whose MPS basis
+is expanded to contain a portion of the basis of
+the Krylov space built by repeated applications of
+`reference` to `state` that is orthogonal
+to the MPS basis of `state`.
+The `reference` operator is applied to `state` `krylovdim`
+number of times, with a default of 2 which should give
+a good balance between reliability and performance.
+See [^global_expansion] for more details.
+
+$(expand_cutoff_doctring())
+
+$(expand_warning_doctring())
+
+$(expand_citation_docstring())
+"""
 function expand(
   ::Algorithm"global_krylov",
   state::MPS,
   operator::MPO;
   krylovdim=2,
   cutoff=(√(eps(real(scalartype(state))))),
+  apply_kwargs=(; maxdim=maxlinkdim(state) + 1),
 )
   # TODO: Try replacing this logic with `Base.accumulate`.
   references = Vector{MPS}(undef, krylovdim)
   for k in 1:krylovdim
     previous_reference = get(references, k - 1, state)
-    references[k] = normalize(
-      apply(operator, previous_reference; maxdim=maxlinkdim(state) + 1)
-    )
+    references[k] = normalize(apply(operator, previous_reference; apply_kwargs...))
   end
   return expand(state, references; alg="orthogonalize", cutoff)
 end
diff --git a/test/test_expand.jl b/test/test_expand.jl
index d14391b..03b8035 100644
--- a/test/test_expand.jl
+++ b/test/test_expand.jl
@@ -89,7 +89,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test scalartype(state) === elt
     # TODO: Use `∜` instead of `fourthroot` in Julia 1.10 and above.
     @test inner(state', operator, state) ≈ reference_energy rtol =
-      2 * fourthroot(eps(real(elt)))
+      5 * fourthroot(eps(real(elt)))
   end
 end
 end

From e9fa36ab1c0f4e93d791bd8d3b36f63369b7909e Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 May 2024 10:08:01 -0400
Subject: [PATCH 17/19] Julia 1.6 compatibility

---
 test/test_expand.jl | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/test/test_expand.jl b/test/test_expand.jl
index 03b8035..622f84f 100644
--- a/test/test_expand.jl
+++ b/test/test_expand.jl
@@ -72,9 +72,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     nexpansions = 10
     tau = elt(0.5)
     for step in 1:nexpansions
-      # TODO: Use `∜` instead of `fourthroot` in Julia 1.10 and above.
+      # TODO: Use `fourthroot`/`∜` in Julia 1.10 and above.
       state = expand(
-        state, operator; alg="global_krylov", krylovdim=3, cutoff=fourthroot(eps(real(elt)))
+        state, operator; alg="global_krylov", krylovdim=3, cutoff=eps(real(elt))^(1//4)
       )
       state = tdvp(
         operator,
@@ -87,9 +87,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       state = normalize(state)
     end
     @test scalartype(state) === elt
-    # TODO: Use `∜` instead of `fourthroot` in Julia 1.10 and above.
-    @test inner(state', operator, state) ≈ reference_energy rtol =
-      5 * fourthroot(eps(real(elt)))
+    # TODO: Use `fourthroot`/`∜` in Julia 1.10 and above.
+    @test inner(state', operator, state) ≈ reference_energy rtol = 5 * eps(real(elt))^(1//4)
   end
 end
 end

From 9df76fadc907679e304eede49a7d1c9223e32f41 Mon Sep 17 00:00:00 2001
From: Matt Fishman <mtfishman@users.noreply.github.com>
Date: Thu, 16 May 2024 10:18:06 -0400
Subject: [PATCH 18/19] Update README.md

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 5638ae7..78c9263 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ julia> ?
 help?> expand
 # ...
 ```
-Users are not given many customization options just yet as we gain more experience on the right balance between efficacy of the expansion and performance, and default values and keyword argument option are subject to change as we learn more about how to best use the method.
+Users are not given many customization options just yet as we gain more experience on the right balance between efficacy of the expansion and performance, and default values and keyword arguments are subject to change as we learn more about how to best use the method.
 
 ### ITensorTDVP.jl v0.4 release notes
 

From 8d1b295da922292dc72cf72f0294e09dd3048e92 Mon Sep 17 00:00:00 2001
From: mtfishman <mfishman@flatironinstitute.org>
Date: Thu, 16 May 2024 11:14:34 -0400
Subject: [PATCH 19/19] Use StableRNGs in tests

---
 test/Project.toml   |  1 +
 test/test_expand.jl | 16 ++++++++--------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/test/Project.toml b/test/Project.toml
index c03e08e..7e322e9 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -8,5 +8,6 @@ Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/test_expand.jl b/test/test_expand.jl
index 622f84f..34d8072 100644
--- a/test/test_expand.jl
+++ b/test/test_expand.jl
@@ -3,7 +3,7 @@ using ITensors: scalartype
 using ITensors.ITensorMPS: OpSum, MPO, MPS, inner, linkdims, maxlinkdim, randomMPS, siteinds
 using ITensorTDVP: dmrg, expand, tdvp
 using LinearAlgebra: normalize
-using Random: Random
+using StableRNGs: StableRNG
 using Test: @test, @testset
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "expand (eltype=$elt)" for elt in elts
@@ -13,9 +13,9 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
   )
     n = 6
     s = siteinds("S=1/2", n; conserve_qns)
-    Random.seed!(1234)
-    state = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
-    reference = randomMPS(elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
+    rng = StableRNG(1234)
+    state = randomMPS(rng, elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=4)
+    reference = randomMPS(rng, elt, s, j -> isodd(j) ? "↑" : "↓"; linkdims=2)
     state_expanded = expand(state, [reference]; alg="orthogonalize")
     @test scalartype(state_expanded) === elt
     @test inner(state_expanded, state) ≈ inner(state, state)
@@ -57,8 +57,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       opsum += "Sz", j, "Sz", j + 2
     end
     operator = MPO(elt, opsum, s)
-    Random.seed!(1234)
-    init = randomMPS(elt, s; linkdims=30)
+    rng = StableRNG(1234)
+    init = randomMPS(rng, elt, s; linkdims=30)
     reference_energy, reference_state = dmrg(
       operator,
       init;
@@ -67,8 +67,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       cutoff=(√(eps(real(elt)))),
       noise=(√(eps(real(elt)))),
     )
-    Random.seed!(1234)
-    state = randomMPS(elt, s)
+    rng = StableRNG(1234)
+    state = randomMPS(rng, elt, s)
     nexpansions = 10
     tau = elt(0.5)
     for step in 1:nexpansions