Subspace expansion

src/ITensorTDVP.jl

@@ -1,14 +1,18 @@
 module ITensorTDVP
 
 using ITensors
-using KrylovKit: exponentiate, eigsolve
+using KrylovKit: exponentiate, eigsolve, svdsolve
 using Printf
+using LinearAlgebra
+
 using TimerOutputs
 using Observers
 
 using ITensors:
   AbstractMPS, @debug_check, @timeit_debug, check_hascommoninds, orthocenter, set_nsite!
-
+using ITensors.NDTensors
+using ITensors.NDTensors: eachdiagblock, blockview
+using ITensors.ITensorNetworkMaps
 # Compatibility of ITensor observer and Observers
 include("update_observer.jl")
 
@@ -25,6 +29,8 @@ include("tdvp_generic.jl")
 include("tdvp.jl")
 include("dmrg.jl")
 include("dmrg_x.jl")
+include("nullspace.jl")
+include("subspace_expansion.jl")
 
 export tdvp, dmrg_x, to_vec, TimeDependentSum
 

src/nullspace.jl

@@ -0,0 +1,131 @@
+
+ITensors.itensor(A::ITensor) = A
+
+# Insert missing diagonal blocks
+function insert_diag_blocks!(T::Tensor)
+  for b in eachdiagblock(T)
+    blockT = blockview(T, b)
+    if isnothing(blockT)
+      # Block was not found in the list, insert it
+      insertblock!(T, b)
+    end
+  end
+end
+insert_diag_blocks!(T::ITensor) = insert_diag_blocks!(tensor(T))
+
+# Reshape into an order-2 ITensor
+matricize(T::ITensor, inds::Index...) = matricize(T, inds)
+
+function matricize(T::ITensor, inds)
+  left_inds = commoninds(T, inds)
+  right_inds = uniqueinds(T, inds)
+  return matricize(T, left_inds, right_inds)
+end
+
+function matricize(T::ITensor, left_inds, right_inds)
+  CL = combiner(left_inds; dir=ITensors.Out, tags="CL")
+  CR = combiner(right_inds; dir=ITensors.In, tags="CR")
+  M = (T * CL) * CR
+  return M, CL, CR
+end
+
+function setdims(t::NTuple{N,Pair{QN,Int}}, dims::NTuple{N,Int}) where {N}
+  return first.(t) .=> dims
+end
+
+# XXX: generalize this function
+function _getindex(T::DenseTensor{ElT,N}, I1::Colon, I2::UnitRange{Int64}) where {ElT,N}
+  A = array(T)[I1, I2]
+  return tensor(Dense(vec(A)), setdims(inds(T), size(A)))
+end
+
+function getblock(i::Index, n::Integer)
+  return ITensors.space(i)[n]
+end
+
+# Make `Pair{QN,Int}` act like a regular `dim`
+NDTensors.dim(qnv::Pair{QN,Int}) = last(qnv)
+
+Base.:*(qnv::Pair{QN,Int}, d::ITensors.Arrow) = qn(qnv) * d => dim(qnv)
+
+function getblock_preserve_qns(T::Tensor, b::Block)
+  # TODO: make `T[b]` preserve QNs
+  Tb = T[b]
+  indsTb = getblock.(inds(T), Tuple(b)) .* dir.(inds(T))
+  return ITensors.setinds(Tb, indsTb)
+end
+
+function blocksparsetensor(blocks::Dict{B,TB}) where {B,TB}
+  b1, Tb1 = first(pairs(blocks))
+  N = length(b1)
+  indstypes = typeof.(inds(Tb1))
+  blocktype = eltype(Tb1)
+  indsT = getindex.(indstypes)
+  # Determine the indices from the blocks
+  for (b, Tb) in pairs(blocks)
+    indsTb = inds(Tb)
+    for n in 1:N
+      bn = b[n]
+      indsTn = indsT[n]
+      if bn > length(indsTn)
+        resize!(indsTn, bn)
+      end
+      indsTn[bn] = indsTb[n]
+    end
+  end
+  T = BlockSparseTensor(blocktype, indsT)
+  for (b, Tb) in pairs(blocks)
+    if !isempty(Tb)
+      T[b] = Tb
+    end
+  end
+  return T
+end
+
+function _nullspace_hermitian(M::Tensor; atol::Real=0.0)
+  tol = atol
+  # Insert any missing diagonal blocks
+  insert_diag_blocks!(M)
+  #D, U = eigen(Hermitian(M))
+  Dᵢₜ, Uᵢₜ = eigen(itensor(M); ishermitian=true)
+  D = tensor(Dᵢₜ)
+  U = tensor(Uᵢₜ)
+  nullspace_blocks = Dict()
+  for bU in nzblocks(U)
+    bM = Block(bU[1], bU[1])
+    bD = Block(bU[2], bU[2])
+    # Assume sorted from largest to smallest
+    indstart = sum(d -> abs(d) .> tol, storage(D[bD])) + 1
+    Ub = getblock_preserve_qns(U, bU)
+    indstop = lastindex(Ub, 2)
+    Nb = _getindex(Ub, :, indstart:indstop)
+    nullspace_blocks[bU] = Nb
+  end
+  return blocksparsetensor(nullspace_blocks)
+end
+
+function LinearAlgebra.nullspace(M::Hermitian{<:Number,<:Tensor}; kwargs...)
+  return _nullspace_hermitian(parent(M); kwargs...)
+end
+
+function LinearAlgebra.nullspace(::Order{2}, M::ITensor, left_inds, right_inds; kwargs...)
+  @assert order(M) == 2
+  M² = prime(dag(M), right_inds) * M
+  M² = permute(M², right_inds'..., right_inds...)
+  M²ₜ = tensor(M²)
+  Nₜ = nullspace(Hermitian(M²ₜ); kwargs...)
+  indsN = (Index(ind(Nₜ, 1); dir=ITensors.In), Index(ind(Nₜ, 2); dir=ITensors.In))
+  N = dag(itensor(ITensors.setinds(Nₜ, indsN)))
+  # Make the index match the input index
+  Ñ = replaceinds(N, (ind(N, 1),) => right_inds)
+  return Ñ
+end
+
+function LinearAlgebra.nullspace(T::ITensor, is...; kwargs...)
+  M, CL, CR = matricize(T, is...)
+  @assert order(M) == 2
+  cL = commoninds(M, CL)
+  cR = commoninds(M, CR)
+  N₂ = nullspace(Order(2), M, cL, cR; kwargs...)
+  return N₂ * CR
+end

src/subspace_expansion.jl

@@ -0,0 +1,206 @@
+function replaceind_indval(IV::Tuple, iĩ::Pair)
+  i, ĩ = iĩ
+  return ntuple(n -> first(IV[n]) == i ? ĩ => last(IV[n]) : IV[n], length(IV))
+end
+
+# atol controls the tolerance cutoff for determining which eigenvectors are in the null
+# space of the isometric MPS tensors. Setting to 1e-2 since we only want to keep
+# the eigenvectors corresponding to eigenvalues of approximately 1.
+
+###the most general implementation may be passing MPS, lims, and bondtensor (defaulting to the Id-matrix)
+###then we can tensorsum both left and right tensor, and just apply the bondtensor to restore input gauge
+###ideally we would also be able to apply tensorsum to the bond tensor instead of that loop below
+"""
+expand subspace (2-site like) in a sweep 
+"""
+function subspace_expansion_sweep!(
+  ψ::MPS, PH::Union{ProjMPO,ProjMPOSum}; maxdim, cutoff, atol=1e-10, kwargs...
+)
+  N = length(ψ)
+
+  if !isortho(ψ) || orthocenter(ψ) != 1
+    orthogonalize!(ψ, 1)
+  end
+
+  PH.nsite = 2
+  nsite = 2
+  position!(PH, ψ, 1)
+
+  for (b, ha) in sweepnext(N; ncenter=2)
+
+    ##TODO: figure out whether these calls should be here or inside subspace expansion, currently we do both?
+    orthogonalize!(ψ, b)
+    position!(PH, ψ, b)
+
+    if (ha == 1 && (b + nsite - 1 != N)) || (ha == 2 && b != 1)
+      b1 = (ha == 1 ? b + 1 : b)
+      Δ = (ha == 1 ? +1 : -1)
+      inds = (ha == 1 ? (b, b + Δ) : (b + Δ, b))
+      subspace_expansion!(
+        ψ, PH, (ψ.llim, ψ.rlim), inds; maxdim, cutoff=cutoff, atol=atol, kwargs...
+      )
+    end
+  end
+  return nothing
+end
+
+function subspace_expansion!(
+  ψ::MPS,
+  PH,
+  lims::Tuple{Int,Int},
+  b::Tuple{Int,Int};
+  bondtensor=nothing,
+  maxdim,
+  cutoff,
+  atol=1e-10,
+  kwargs...,
+)
+  """
+  expands in nullspace of site-tensors b
+  """
+  #atol refers to the tolerance in nullspace determination (for finite MPS can probably be set rather small)
+  #cutoff refers to largest singular value of gradient (acceleration of population gain of expansion vectors) to keep
+  #this should be related to the cutoff in two-site TDVP: \Delta_rho_i = 0.5 * lambda_y * tau **2 
+  #note that in the initial SVD there is another cutoff parameter, that we set to roughly machine precision for the time being
+  #(allows to move bond-dimension between different partial QN sectors, if no such truncation occurs distribution of bond-dimensions
+  #between different QNs locally is static once bond dimension saturates maxdim.)
+
+  ##this should only work for the case where rlim-llim > 1
+  ##not a valid MPS otherwise anyway (since bond matrix not part of MPS unless so defined like in VidalMPS struct)
+  llim, rlim = lims
+  n1, n2 = b
+  PH.nsite = 2
+  old_linkdim = dim(commonind(ψ[n1], ψ[n2]))
+  cutoff_compress = get(kwargs, :cutoff_compress, 1e-12)
+
+  ###move orthogonality center to bond, check whether there are vanishing contributions to the wavefunctions and truncate accordingly
+  if llim == n1
+    @assert rlim == n2 + 1
+    U, S, V = svd(
+      ψ[n2], uniqueinds(ψ[n2], ψ[n1]); maxdim=maxdim, cutoff=cutoff_compress, kwargs...
+    )  ##lookup svd interface again
+    ϕ_1 = ψ[n1] * V
+    ϕ_2 = U
+    old_linkdim = dim(commonind(U, S))
+    bondtensor = S
+
+  elseif rlim == n2
+    @assert llim == n1 - 1
+    U, S, V = svd(
+      ψ[n1], uniqueinds(ψ[n1], ψ[n2]); maxdim=maxdim, cutoff=cutoff_compress, kwargs...
+    )
+    ϕ_1 = U
+    ϕ_2 = ψ[n2] * V
+    old_linkdim = dim(commonind(U, S))
+    bondtensor = S
+  end
+
+  ###don't expand if we are already at maxdim
+  if old_linkdim >= maxdim
+    println("not expanding")
+    return nothing
+  end
+  position!(PH, ψ, min(n1, n2))
+
+  #orthogonalize(ψ,n1)
+  linkind_l = commonind(ϕ_1, bondtensor)
+  linkind_r = commonind(ϕ_2, bondtensor)
+
+  NL = nullspace(ϕ_1, linkind_l; atol=atol)
+  NR = nullspace(ϕ_2, linkind_r; atol=atol)
+
+  ###NOTE: This will fail for rank-1 trees with physical DOFs on leafs
+  ###NOTE: one-sided subspace expansion seems to not work well at least for trees according to Lachlan Lindoy
+  if norm(NL) == 0.0 || norm(NR) == 0.0
+    return bondtensor
+  end
+
+  ###form 2site wavefunction
+  ϕ = ϕ_1 * bondtensor * ϕ_2
+
+  ###get subspace expansion
+  newL, S, newR, success = _subspace_expand_core(
+    ϕ, PH, NL, NR, ; maxdim=maxdim - old_linkdim, cutoff, kwargs...
+  )
+  @assert success
+  #@show expanS
+  if success == false
+    return nothing
+  end
+
+  ###add expansion direction to current site tensors
+  ALⁿ¹, newl = ITensors.directsum(
+    ϕ_1, dag(newL), uniqueinds(ϕ_1, newL), uniqueinds(newL, ϕ_1); tags=("Left",)
+  )
+  ARⁿ², newr = ITensors.directsum(
+    ϕ_2, dag(newR), uniqueinds(ϕ_2, newR), uniqueinds(newR, ϕ_2); tags=("Right",)
+  )
+
+  ###TODO remove assertions regarding expansion not exceeding maxdim ?
+  @assert (dim(commonind(newL, S)) + old_linkdim) <= maxdim
+  @assert dim(commonind(newL, S)) == dim(commonind(newR, S))
+  @assert(dim(uniqueind(ϕ_1, newL)) + dim(uniqueind(newL, ϕ_1)) == dim(newl))
+  @assert(dim(uniqueind(ϕ_2, newR)) + dim(uniqueind(newR, ϕ_2)) == dim(newr))
+  @assert (old_linkdim + dim(commonind(newL, S))) <= maxdim
+  @assert (old_linkdim + dim(commonind(newR, S))) <= maxdim
+  @assert dim(newl) <= maxdim
+  @assert dim(newr) <= maxdim
+
+  ###zero-pad bond-tensor (the orthogonality center)
+  C = ITensor(dag(newl)..., dag(newr)...)
+  ψC = bondtensor
+  for I in eachindex(ψC)
+    v = ψC[I]
+    if !iszero(v)
+      C[I] = ψC[I]
+    end
+  end
+
+  ###move orthogonality center back to site (should restore input orthogonality limits)
+  if rlim == n2
+    ψ[n2] = ARⁿ²
+  elseif rlim > n2
+    ψ[n2] = noprime(ARⁿ² * C)
+  end
+
+  if llim == n1
+    ψ[n1] = ALⁿ¹
+  elseif llim < n1
+    ψ[n1] = noprime(ALⁿ¹ * C)
+  end
+
+  return nothing
+end
+
+function _subspace_expand_core(
+  centerwf::Vector{ITensor}, env, NL, NR; maxdim, cutoff, kwargs...
+)
+  ϕ = ITensor(1.0)
+  for atensor in centerwf
+    ϕ *= atensor
+  end
+  return _subspace_expand_core(ϕ, env, NL, NR; maxdim, cutoff, kwargs...)
+end
+
+function _subspace_expand_core(ϕ::ITensor, env, NL, NR; maxdim, cutoff, kwargs...)
+  ϕH = noprime(env(ϕ))   #add noprime?
+  ϕH = NL * ϕH * NR
+  if norm(ϕH) == 0.0
+    return false, false, false, false
+  end
+  U, S, V = svd(
+    ϕH,
+    commoninds(ϕH, NL);
+    maxdim,
+    cutoff,
+    use_relative_cutoff=false,
+    use_absolute_cutoff=true,
+    kwargs...,
+  )
+
+  @assert dim(commonind(U, S)) <= maxdim
+
+  NL *= dag(U)
+  NR *= dag(V)
+  return NL, S, NR, true
+end