From 0b17b861aa85511979ab2edd89fefafa2464c033 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 | 24 ++++++++++++++++++------
 2 files changed, 19 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..2c1f9f3265 100644
--- a/src/lib/ITensorMPS/src/abstractmps.jl
+++ b/src/lib/ITensorMPS/src/abstractmps.jl
@@ -1677,30 +1677,42 @@ function truncate!(M::AbstractMPS; alg="frobenius", 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`