From a3b5f75e2b78bf3eaf64e87af05e825d876dda63 Mon Sep 17 00:00:00 2001
From: Bobby <bnpinkston@gmail.com>
Date: Mon, 8 Jul 2024 23:40:39 -0400
Subject: [PATCH] [ITensors] add kwarg to truncate! to return
 RankFactorization.Spectrum (#1516)

The purpose of this commit is to allow the truncation error from an
operation to be returned. This is the first step envisioned for that
functionality.

The spectrum and truncation error from calling svd() on ITensors is not
accessible from calls to truncate! and truncate. This allows the
Spectrum type (which contains the spectrum and the error) to be
optionally returned. This required importing RankFactorization.Spectrum
in ITensors (so an added import in imports.jl) and then slightly
refactoring the definition of truncate! and truncate.

This commit makes a named tuple with each element being "bond_n" where n
is `1` to `N-1` bonds and each element of the named tuple corresponds to
the Spectrum returned from the call to svd as the bonds are swept over
during the call to truncate.

All tests were run with 92979 passing and 73 broken. Total tests were
93052.
---
 src/imports.jl                        |  1 +
 src/lib/ITensorMPS/src/abstractmps.jl | 25 +++++++++++++++++++------
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/imports.jl b/src/imports.jl
index 727ab54f83..b41e28202f 100644
--- a/src/imports.jl
+++ b/src/imports.jl
@@ -182,6 +182,7 @@ import ITensors.NDTensors:
   truncate!,
   using_tblis,
   vector,
+  RankFactorization.Spectrum,
   # Deprecated
   addblock!,
   store
diff --git a/src/lib/ITensorMPS/src/abstractmps.jl b/src/lib/ITensorMPS/src/abstractmps.jl
index 629075f183..c38ba78581 100644
--- a/src/lib/ITensorMPS/src/abstractmps.jl
+++ b/src/lib/ITensorMPS/src/abstractmps.jl
@@ -1671,36 +1671,49 @@ provided as keyword arguments.
 
 Keyword arguments:
 * `site_range`=1:N - only truncate the MPS bonds between these sites
+* `return_spectrum=false` - in addition to returning the MPS, also return a named tuple containing the  of type `Spectrum` for each bond in the MPS. `Spectrum` contains a Vector of `Float64` called `eigs` that contains the singular values at each bond and the truncation error performed at that bond. For example, `_, spec = truncate!(m, maxdim=5)` and then inspect the singular values and truncation error at bond 3 by doing `spec.bond_3.eigs` or `spec.bond_3.truncerr`, respectively. 
 """
 function truncate!(M::AbstractMPS; alg="frobenius", kwargs...)
   return truncate!(Algorithm(alg), M; kwargs...)
 end
 
 function truncate!(
-  ::Algorithm"frobenius", M::AbstractMPS; site_range=1:length(M), kwargs...
+  ::Algorithm"frobenius",
+  M::AbstractMPS;
+  site_range=1:length(M),
+  return_spectrum=false,
+  kwargs...,
 )
   N = length(M)
+  Nbonds = N - 1
+  bond_names = Tuple([Symbol("bond_$(α)") for α in 1:Nbonds])
+  # TODO: Not sure how ths will work. Its a local module or "lib" as NDTensors calls it. Need to investigate this more to see how getting it here would work. Or just make it type Any?
+  spectrums = Vector{Spectrum}(undef, Nbonds)
 
   # Left-orthogonalize all tensors to make
   # truncations controlled
   orthogonalize!(M, last(site_range))
 
   # Perform truncations in a right-to-left sweep
-  for j in reverse((first(site_range) + 1):last(site_range))
+  for (idx, j) in enumerate(reverse((first(site_range) + 1):last(site_range)))
     rinds = uniqueinds(M[j], M[j - 1])
     ltags = tags(commonind(M[j], M[j - 1]))
-    U, S, V = svd(M[j], rinds; lefttags=ltags, kwargs...)
+    U, S, V, spec = svd(M[j], rinds; lefttags=ltags, kwargs...)
+    spectrums[idx] = spec
     M[j] = U
     M[j - 1] *= (S * V)
     setrightlim!(M, j)
   end
-  return M
+  if return_spectrum
+    return M, NamedTuple{bond_names}(spectrums)
+  else
+    return M
+  end
 end
 
 function truncate(ψ0::AbstractMPS; kwargs...)
   ψ = copy(ψ0)
-  truncate!(ψ; kwargs...)
-  return ψ
+  return truncate!(ψ; kwargs...)
 end
 
 # Make `*` an alias for `contract` of two `AbstractMPS`