From 4bb88e5991443fccec6b288953ed593fd2637a86 Mon Sep 17 00:00:00 2001
From: Benedikt Kloss <bene.kloss@gmail.com>
Date: Mon, 12 Feb 2024 10:10:17 -0500
Subject: [PATCH 1/2] Initial commit for randomized SVD.

---
 src/ITensorNetworks.jl    |   2 +
 src/linalg/rsvd_aux.jl    | 276 ++++++++++++++++++++++++++++++++++++++
 src/linalg/rsvd_linalg.jl |  87 ++++++++++++
 test/test_rsvd.jl         | 100 ++++++++++++++
 4 files changed, 465 insertions(+)
 create mode 100644 src/linalg/rsvd_aux.jl
 create mode 100644 src/linalg/rsvd_linalg.jl
 create mode 100644 test/test_rsvd.jl

diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl
index d4c993a2..e5357c23 100644
--- a/src/ITensorNetworks.jl
+++ b/src/ITensorNetworks.jl
@@ -105,6 +105,8 @@ include("tensornetworkoperators.jl")
 include(joinpath("ITensorsExt", "itensorutils.jl"))
 include(joinpath("Graphs", "abstractgraph.jl"))
 include(joinpath("Graphs", "abstractdatagraph.jl"))
+include(joinpath("linalg","rsvd_linalg.jl"))
+include(joinpath("linalg","rsvd_aux.jl"))
 include(joinpath("solvers", "eigsolve.jl"))
 include(joinpath("solvers", "exponentiate.jl"))
 include(joinpath("solvers", "dmrg_x.jl"))
diff --git a/src/linalg/rsvd_aux.jl b/src/linalg/rsvd_aux.jl
new file mode 100644
index 00000000..149e699c
--- /dev/null
+++ b/src/linalg/rsvd_aux.jl
@@ -0,0 +1,276 @@
+function get_column_space(A::Vector{<:ITensor}, lc::ITensors.QNIndex, rc::ITensors.QNIndex)
+  #gets non-zero blocks in rc by sticking in lc and contracting through
+  viable_sectors = Vector{Pair{QN,Int64}}()
+  for s in space(lc)
+    qn = first(s)
+    trial = dag(randomITensor(qn, lc))
+    adtrial = foldl(*, A; init=trial)
+    nnzblocks(adtrial) == 0 && continue
+    @assert nnzblocks(adtrial) == 1
+    i = Int(nzblocks(adtrial)[1])
+    thesector = space(only(inds(adtrial)))[i]
+    push!(viable_sectors, thesector)
+  end
+  return viable_sectors
+end
+
+function get_column_space(A::Vector{<:ITensor}, lc::Index{Int64}, rc::Index{Int64})
+  #gets non-zero blocks in rc by sticking in lc and contracting through
+  viable_sectors = Vector{Pair{Int64,Int64}}()
+  !any(isempty.(A)) && (viable_sectors = [dim(rc) => dim(rc)])
+  return viable_sectors
+end
+
+# remove since code passes through vector of ITensors
+#=
+function get_column_space(A::ITensor, lc::Index,rc::Index)
+  #gets non-zero blocks in rc by sticking in lc and contracting through
+  viable_sectors=Vector{Pair{QN,Int64}}()
+  ind_loc=only(findall(isequal(rc),inds(A)))
+  unique_qns=unique(qns(A,ind_loc))
+  for s in space(rc)
+    qn = first(s)
+    qn in unique_qns || continue 
+    push!(viable_sectors, s)
+  end
+  return viable_sectors
+end
+=#
+
+qns(t::ITensor) = qns(collect(eachnzblock(t)), t)
+qns(t::ITensor, i::Int) = qns(collect(eachnzblock(t)), t, i)
+
+function qns(bs::Vector{Block{n}}, t::ITensor) where {n}
+  return [qns(b, t) for b in bs]
+end
+
+function qns(bs::Vector{Block{n}}, t::ITensor, i::Int) where {n}
+  return [qns(b, t)[i] for b in bs]
+end
+
+function qns(b::Block{n}, t::ITensor) where {n}
+  theqns = QN[]
+  for i in 1:order(t)
+    theb = getindex(b, i)
+    theind = inds(t)[i]
+    push!(theqns, ITensors.qn(space(theind), Int(theb)))
+  end
+  return theqns
+end
+
+#QN version
+function build_guess_matrix(
+  eltype::Type{<:Number}, ind, sectors::Vector{Pair{QN,Int64}}, n::Int, p::Int
+)
+  aux_spaces = Pair{QN,Int64}[]
+  for s in sectors
+    thedim = last(s)
+    qn = first(s)
+    en = min(n + p, thedim)
+    push!(aux_spaces, Pair(qn, en))
+  end
+  aux_ind = Index(aux_spaces; dir=dir(ind))
+  M = randomITensor(eltype, dag(ind), aux_ind)    #defaults to zero flux
+  @assert nnzblocks(M) != 0
+  return M
+end
+
+#non-QN version
+function build_guess_matrix(
+  eltype::Type{<:Number}, ind, sectors::Vector{Pair{Int,Int64}}, n::Int, p::Int
+)
+  thedim = dim(ind)
+  en = min(n + p, thedim)
+  aux_ind = Index(en)
+  M = randomITensor(eltype, ind, aux_ind)
+  return M
+end
+
+#non-QN version
+function build_guess_matrix(
+  eltype::Type{<:Number}, ind, sectors::Vector{Pair{Int,Int}}, ndict::Dict;
+) ##given ndict, sectors is not necessary here
+  thedim = dim(ind)
+  en = min(ndict[space(ind)], thedim)
+  aux_ind = Index(en)
+  M = randomITensor(eltype, ind, aux_ind)
+  return M
+end
+
+#QN version
+function build_guess_matrix(
+  eltype::Type{<:Number}, ind, sectors::Vector{Pair{QN,Int64}}, ndict::Dict;
+) ##given ndict, sectors is not necessary here
+  aux_spaces = Pair{QN,Int64}[]
+  for s in sectors
+    thedim = last(s)
+    qn = first(s)
+    en = min(ndict[qn], thedim)
+    push!(aux_spaces, Pair(qn, en))
+  end
+  aux_ind = Index(aux_spaces; dir=dir(ind))
+  M = randomITensor(eltype, dag(ind), aux_ind)
+  @assert nnzblocks(M) != 0
+  return M
+end
+
+# non-QN version
+function init_guess_sizes(cind, sectors::Nothing, n::Int, rule)
+  thedim = dim(cind)
+  ndict = Dict{Int64,Int64}()
+  ndict[thedim] = min(thedim, rule(n))
+  return ndict
+end
+
+# QN version
+function init_guess_sizes(cind, sectors::Vector{Pair{QN,Int64}}, n::Int, rule)
+  @assert hasqns(cind)
+  ndict = Dict{QN,Int64}()
+  for s in sectors
+    thedim = last(s)
+    qn = first(s)
+    ndict[qn] = min(rule(n), thedim)
+  end
+  return ndict
+end
+
+function increment_guess_sizes(ndict::Dict{QN,Int64}, n_inc::Int, rule)
+  for key in keys(ndict)
+    ndict[key] = ndict[key] + rule(n_inc)
+  end
+  return ndict
+end
+
+function increment_guess_sizes(ndict::Dict{Int64,Int64}, n_inc::Int, rule)
+  for key in keys(ndict)
+    ndict[key] = ndict[key] + rule(n_inc)
+  end
+  return ndict
+end
+
+function approx_the_same(o, n; abs_eps=1e-12, rel_eps=1e-6)
+  absdev = abs.(o .- n)
+  reldev = abs.((o .- n) ./ n)
+  abs_conv = absdev .< abs_eps
+  rel_conv = reldev .< rel_eps
+  return all(abs_conv .|| rel_conv)
+end
+
+function is_converged_block(o, n; svd_kwargs...)
+  maxdim = get(svd_kwargs, :maxdim, Inf)
+  if length(o) != length(n)
+    return false
+  else
+    r = min(maxdim, length(o))
+    #ToDo: Pass kwargs?
+    return approx_the_same(o[1:r], n[1:r])
+  end
+end
+
+function is_converged!(ndict, old_fact, new_fact; n_inc=1, has_qns=true, svd_kwargs...)
+  oU, oS, oV = old_fact.U, old_fact.S, old_fact.V
+  nU, nS, nV = new_fact.U, new_fact.S, new_fact.V
+
+  # check for non-existing tensor
+  (isempty(nS) && isempty(oS)) && return true
+  (isempty(nS) && !(isempty(oS))) && warning("New randomized SVD is empty
+  while prior iteration was not. It is very unlikely that this is correct.
+  Exiting.")
+
+  theflux = flux(nS)
+  oldflux = flux(oS)
+
+  if !has_qns
+    if norm(oS) == 0.0
+      if norm(nS) == 0.0
+        return true
+      else
+        ndict[first(keys(ndict))] *= 2
+        return false
+      end
+    end
+  end
+
+  maxdim = get(svd_kwargs, :maxdim, Inf)
+  # deal with non QN tensors first and a simple case for blocksparse
+  os = sort(storage(oS); rev=true)
+  ns = sort(storage(nS); rev=true)
+  if length(os) >= maxdim && length(ns) >= maxdim
+    conv_global = approx_the_same(os[1:maxdim], ns[1:maxdim])
+  elseif length(os) != length(ns)
+    conv_global = false
+  else
+    r = length(ns)
+    conv_global = approx_the_same(os[1:r], ns[1:r]) #shouldn't this error for os shorter than ns?
+  end
+  if !hasqns(oS)
+    if conv_global == false
+      #ndict has only one key 
+      ndict[first(keys(ndict))] *= 2
+    end
+    return conv_global
+  end
+  conv_bool_total = true
+
+  # deal with QN tensors now
+  ncrind = commonind(nS, nV)
+  ocrind = commonind(oS, oV)
+  n_ind_loc = only(findall(isequal(ncrind), inds(nS)))
+  nqns = first.(space(ncrind))
+  oqns = first.(space(ocrind))
+
+  for qn in keys(ndict)
+    if !(qn in nqns) && !(qn in oqns)
+      conv_bool = true
+      continue
+      #qn
+    elseif !(qn in nqns)
+      ndict[qn] *= 2
+      conv_bool_total = false
+      warn("QN was in old factorization but not in new one, this shouldn't happen often!")
+      continue
+    elseif !(qn in oqns)
+      ndict[qn] *= 2
+      conv_bool_total = false
+      continue
+    end
+    #qn present in both old and new factorization, grab singular values to compare
+    ovals = get_qnblock_vals(qn, ocrind, oS)
+    nvals = get_qnblock_vals(qn, ncrind, nS)
+    conv_bool = is_converged_block(collect(ovals), collect(nvals); svd_kwargs...)
+    if conv_bool == false
+      ndict[qn] *= 2
+    end
+    conv_bool_total *= conv_bool
+  end
+  if conv_bool_total == true
+    @assert conv_global == true
+  else
+    if conv_global == true
+      println(
+        "Subspace expansion, rand. svd: singular vals converged globally, but may not be optimal, doing another iteration",
+      )
+    end
+  end
+  return conv_bool_total::Bool
+end
+
+"""Extracts (singular) values in block associated with `qn` in `ind`` from diagonal ITensor `svalt``."""
+function get_qnblock_vals(qn, ind, svalt)
+  s = space(ind)
+  ind_loc = only(findall(isequal(ind), inds(svalt)))
+  theblocks = eachnzblock(svalt)
+  blockdict = Int.(getindex.(theblocks, ind_loc))
+  qnindtoblock = Dict(collect(values(blockdict)) .=> collect(keys(blockdict))) #refactor
+
+  qnind = findfirst((first.(s)) .== [qn]) # refactor
+  theblock = qnindtoblock[qnind]
+  vals = diag(svalt[theblock])
+
+  return vals
+end
+
+"""Check that the factorization isn't empty. / both empty"""
+function _sanity_checks()
+  return nothing
+end
diff --git a/src/linalg/rsvd_linalg.jl b/src/linalg/rsvd_linalg.jl
new file mode 100644
index 00000000..221e5840
--- /dev/null
+++ b/src/linalg/rsvd_linalg.jl
@@ -0,0 +1,87 @@
+function prepare_rsvd(A::Vector{ITensor}, linds::Vector{<:Index})
+  if length(A) == 1
+    rinds = uniqueinds(only(A), linds)
+  else
+    rinds = uniqueinds(A[end], unioninds(A[1:(end - 1)]...))
+  end
+  CL = combiner(linds)
+  CR = combiner(rinds)
+  AC = copy(A)
+  AC[1] = CL * first(AC)
+  AC[end] = last(AC) * CR
+  cL = combinedind(CL)
+  cR = combinedind(CR)
+  return AC, (CL, cL), (CR, cR)
+end
+
+function rsvd_iterative(A::ITensor, linds::Vector{<:Index}; svd_kwargs...)
+  return rsvd_iterative([A], linds; svd_kwargs...)
+end
+
+# ToDo: not type stable for empty factorizations.
+function rsvd_iterative(A::Vector{ITensor}, linds::Vector{<:Index}; svd_kwargs...)
+  AC, (CL, cL), (CR, cR) = prepare_rsvd(A, linds)
+  scalar_type = eltype(first(AC))
+  nonzero_sectors = get_column_space(AC, cL, cR)
+  isempty(nonzero_sectors) && return nothing, nothing, nothing
+
+  n_init = 1
+  p_rule(n) = 2 * n
+  ndict = init_guess_sizes(cR, nonzero_sectors, n_init, p_rule)
+
+  M = build_guess_matrix(scalar_type, cR, nonzero_sectors, ndict)
+  fact, Q = rsvd_core(AC, M; svd_kwargs...)
+  n_inc = 1
+  ndict = increment_guess_sizes(ndict, n_inc, p_rule)
+  new_fact = nothing
+  while true
+    M = build_guess_matrix(scalar_type, cR, nonzero_sectors, ndict)
+    new_fact, Q = rsvd_core(AC, M; svd_kwargs...)
+    isnothing(Q) && return nothing, nothing, nothing
+    if is_converged!(ndict, fact, new_fact; n_inc, has_qns=any(hasqns.(AC)), svd_kwargs...)
+      break
+    else
+      fact = new_fact
+    end
+  end
+  vals = diag(array(new_fact.S))
+  (length(vals) == 1 && vals[1]^2 ≤ get(svd_kwargs, :cutoff, 0.0)) &&
+    return nothing, nothing, nothing
+  return dag(CL) * Q * new_fact.U, new_fact.S, new_fact.V * dag(CR)
+end
+
+function rsvd(A::ITensor, linds::Vector{<:Index}, n::Int, p::Int; svd_kwargs...)
+  return rsvd([A], linds, n, p; svd_kwargs...)
+end
+
+function rsvd(A::Vector{<:ITensor}, linds::Vector{<:Index}, n::Int, p::Int; svd_kwargs...)
+  AC, (CL, cL), (CR, cR) = prepare_rsvd(A, linds)
+  nonzero_sectors = get_column_space(AC, cL, cR)
+  isempty(nonzero_sectors) && return nothing, nothing, nothing
+  M = build_guess_matrix(eltype(first(AC)), cR, nonzero_sectors, n, p)
+  fact, Q = rsvd_core(AC, M; svd_kwargs...)
+  return dag(CL) * Q * fact.U, fact.S, fact.V * dag(CR)
+end
+
+function rsvd_core(AC::Vector{ITensor}, M; svd_kwargs...)
+  @assert !isnothing(commonind(last(AC), M))
+  Q = foldr(*, AC; init=M)
+  Q = ITensors.qr(Q, commoninds(Q, first(AC)))[1]
+  any(isequal(0), dims(Q)) && return nothing, nothing, nothing, nothing
+  QAC = foldl(*, AC; init=dag(Q))
+  @assert typeof(QAC) <: ITensor
+  fact = svd(QAC, commoninds(dag(Q), QAC); svd_kwargs...)
+  return fact, Q
+end
+
+# ToDo : Remove this, not needed since everything passes via Vector of ITensors.
+#=
+function rsvd_core(AC::ITensor, M; svd_kwargs...)
+  Q = AC * M
+  #@show dims(Q)
+  Q = ITensors.qr(Q, commoninds(AC, Q))[1]
+  QAC = dag(Q) * AC
+  fact = svd(QAC, uniqueind(QAC, AC); svd_kwargs...)
+  return fact, Q
+end
+=#
diff --git a/test/test_rsvd.jl b/test/test_rsvd.jl
new file mode 100644
index 00000000..f611e385
--- /dev/null
+++ b/test/test_rsvd.jl
@@ -0,0 +1,100 @@
+using ITensors
+using ITensorNetworks
+using ITensorNetworks: rsvd, rsvd_iterative
+using Test
+using Random
+
+ITensors.enable_auto_fermion()
+@testset "RSVD single shot non-fermionic no QNs" begin
+  N = 12
+  Na = div(N, 2) - 2
+  s = siteinds("S=1/2", N; conserve_qns=false)
+  linds = s[1:Na]
+  rinds = s[(Na + 1):end]
+  #target_flux=QN("Sz",4,1)
+  #T=randomITensor(target_flux,linds...,rinds...)
+  T = randomITensor(linds..., rinds...)
+  maxdim = 7
+  U, S, V = svd(T, linds...; maxdim)
+  Tt = U * S * V
+  Tt /= norm(Tt)
+  Ur, Sr, Vr = rsvd(Tt, linds, maxdim, div(maxdim, 2))
+  @test norm(Ur * Sr * Vr - Tt) ≈ 0.0 atol = 1e-5
+end
+
+@testset "RSVD single shot fermionic" begin
+  @show ITensors.using_auto_fermion()
+  N = 6
+  Na = div(N, 2)
+  s = siteinds("Fermion", N; conserve_qns=true)
+  linds = s[1:Na]
+  rinds = s[(Na + 1):end]
+  target_flux = QN("Nf", div(N, 2), -1)
+  T = randomITensor(target_flux, linds..., rinds...)
+  maxdim = 7
+  U, S, V = ITensors.svd(T, linds...; maxdim)
+  Tt = U * S * V
+  Tt /= norm(Tt)
+  Ur, Sr, Vr = rsvd(Tt, linds, maxdim, div(maxdim, 2);)
+  @test norm(Ur * Sr * Vr - Tt) ≈ 0.0 atol = 1e-5
+end
+
+@testset "RSVD vector single shot fermionic" begin
+  @show ITensors.using_auto_fermion()
+  maxdim = 8
+  N = 15
+  Na = div(N, 3)
+  s = siteinds("Fermion", N; conserve_qns=true)
+  linds = s[1:Na]
+  cinds = s[(Na + 1):(2 * Na)]
+  rinds = s[(2 * Na + 1):end]
+  target_flux = QN("Nf", div(N, 3) - 1, -1)
+  T1 = randomITensor(target_flux, linds..., cinds...)
+  uT1, sT1, vT1 = ITensors.svd(T1, linds...; maxdim)
+  target_flux = QN("Nf", div(N, 3) - 2, -1)
+  T2 = randomITensor(target_flux, dag.(cinds)..., rinds...)
+  tmap = [uT1 * sT1, vT1 * T2]
+  maxdim = 7
+  Ur, Sr, Vr = rsvd(tmap, linds, maxdim, div(maxdim, 2))
+  Ur, Sr, Vr = rsvd(tmap, linds, maxdim, div(maxdim, 2))
+  @test norm(Ur * Sr * Vr - contract(tmap)) ≈ 0.0 atol = 1e-5
+end
+
+@testset "RSVD ITensor iterative fermionic" begin
+  @show ITensors.using_auto_fermion()
+  N = 8
+  Na = div(N, 2)
+  s = siteinds("Fermion", N; conserve_qns=true)
+  linds = s[1:Na]
+  rinds = s[(Na + 1):end]
+  target_flux = QN("Nf", div(N, 2), -1)
+  T = randomITensor(target_flux, linds..., rinds...)
+  maxdim = 7
+  U, S, V = ITensors.svd(T, linds...; maxdim)
+  Tt = U * S * V
+  old_norm = norm(Tt)
+  Tt /= old_norm
+  Ur, Sr, Vr = rsvd_iterative(Tt, linds; maxdim)
+  @test norm(Ur * Sr * Vr - Tt) ≈ 0.0 atol = 1e-5
+end
+
+@testset "RSVD Vector iterative fermionic" begin
+  @show ITensors.using_auto_fermion()
+  maxdim = 8
+  N = 15
+  Na = div(N, 3)
+  s = siteinds("Fermion", N; conserve_qns=true)
+  linds = s[1:Na]
+  cinds = s[(Na + 1):(2 * Na)]
+  rinds = s[(2 * Na + 1):end]
+  target_flux = QN("Nf", div(N, 3) - 1, -1)
+  T1 = randomITensor(target_flux, linds..., cinds...)
+  uT1, sT1, vT1 = ITensors.svd(T1, linds...; maxdim)
+  target_flux = QN("Nf", div(N, 3) - 2, -1)
+  T2 = randomITensor(target_flux, dag.(cinds)..., rinds...)
+  tmap = [uT1 * sT1, vT1 * T2]
+  Ur, Sr, Vr = rsvd_iterative(tmap, linds; maxdim)
+  @test norm(Ur * Sr * Vr - contract(tmap)) ≈ 0.0 atol = 1e-5
+end
+ITensors.disable_auto_fermion()
+nothing

From 1930c8c1b4ba5f6fa252022726447af8be0dfc9b Mon Sep 17 00:00:00 2001
From: Benedikt Kloss <bene.kloss@gmail.com>
Date: Mon, 12 Feb 2024 10:29:38 -0500
Subject: [PATCH 2/2] Format

---
 src/ITensorNetworks.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl
index e5357c23..23a2ff1b 100644
--- a/src/ITensorNetworks.jl
+++ b/src/ITensorNetworks.jl
@@ -105,8 +105,8 @@ include("tensornetworkoperators.jl")
 include(joinpath("ITensorsExt", "itensorutils.jl"))
 include(joinpath("Graphs", "abstractgraph.jl"))
 include(joinpath("Graphs", "abstractdatagraph.jl"))
-include(joinpath("linalg","rsvd_linalg.jl"))
-include(joinpath("linalg","rsvd_aux.jl"))
+include(joinpath("linalg", "rsvd_linalg.jl"))
+include(joinpath("linalg", "rsvd_aux.jl"))
 include(joinpath("solvers", "eigsolve.jl"))
 include(joinpath("solvers", "exponentiate.jl"))
 include(joinpath("solvers", "dmrg_x.jl"))