Merge pull request #18 from tensor4all/17-improvements-for-compress-function-using-svd

17 improvements for compress function using svd

Author: shinaoka

src/abstracttensortrain.jl

@@ -269,3 +269,25 @@ Subtraction of two tensor trains. If `c = a - b`, then `c(v) ≈ a(v) - b(v)` at
 function Base.:-(lhs::AbstractTensorTrain{V}, rhs::AbstractTensorTrain{V}) where {V}
     return subtract(lhs, rhs)
 end
+
+"""
+Squared Frobenius norm of a tensor train.
+"""
+function LA.norm2(tt::AbstractTensorTrain{V})::Float64 where {V}
+    function _f(n)::Matrix{V}
+        t = sitetensor(tt, n)
+        t3 = reshape(t, size(t)[1], :, size(t)[end])
+        # (lc, s, rc) * (l, s, r) => (lc, rc, l, r)
+        tct = _contract(conj.(t3), t3, (2,), (2,))
+        tct = permutedims(tct, (1, 3, 2, 4))
+        return reshape(tct, size(tct, 1) * size(tct, 2), size(tct, 3) * size(tct, 4))
+    end
+    return real(only(reduce(*, (_f(n) for n in 1:length(tt)))))
+end
+
+"""
+Frobenius norm of a tensor train.
+"""
+function LA.norm(tt::AbstractTensorTrain{V})::Float64 where {V}
+    sqrt(LA.norm2(tt))
+end

src/tensortrain.jl

@@ -93,25 +93,44 @@ function tensortrain(tci)
 end
 
 function _factorize(
-    A::AbstractMatrix{V}, method::Symbol; tolerance::Float64, maxbonddim::Int
+    A::AbstractMatrix{V}, method::Symbol; tolerance::Float64, maxbonddim::Int, leftorthogonal::Bool=false, normalizeerror=true
 )::Tuple{Matrix{V},Matrix{V},Int} where {V}
+    reltol = 1e-14
+    abstol = 0.0
+    if normalizeerror
+        reltol = tolerance
+    else
+        abstol = tolerance
+    end
     if method === :LU
-        factorization = rrlu(A, abstol=tolerance, maxrank=maxbonddim)
+        factorization = rrlu(A, abstol=abstol, reltol=reltol, maxrank=maxbonddim, leftorthogonal=leftorthogonal)
         return left(factorization), right(factorization), npivots(factorization)
     elseif method === :CI
-        factorization = MatrixLUCI(A, abstol=tolerance, maxrank=maxbonddim)
+        factorization = MatrixLUCI(A, abstol=abstol, reltol=reltol, maxrank=maxbonddim, leftorthogonal=leftorthogonal)
         return left(factorization), right(factorization), npivots(factorization)
     elseif method === :SVD
         factorization = LinearAlgebra.svd(A)
+        err = [sum(factorization.S[n+1:end] .^ 2) for n in 1:length(factorization.S)]
+        normalized_err = err ./ sum(factorization.S .^ 2)
+
         trunci = min(
-            replacenothing(findlast(>(tolerance), factorization.S), 1),
+            replacenothing(findfirst(<(abstol^2), err), length(err)),
+            replacenothing(findfirst(<(reltol^2), normalized_err), length(normalized_err)),
             maxbonddim
         )
-        return (
-            factorization.U[:, 1:trunci],
-            Diagonal(factorization.S[1:trunci]) * factorization.Vt[1:trunci, :],
-            trunci
-        )
+        if leftorthogonal
+            return (
+                factorization.U[:, 1:trunci],
+                Diagonal(factorization.S[1:trunci]) * factorization.Vt[1:trunci, :],
+                trunci
+            )
+        else
+            return (
+                factorization.U[:, 1:trunci] * Diagonal(factorization.S[1:trunci]),
+                factorization.Vt[1:trunci, :],
+                trunci
+            )
+        end
     else
         error("Not implemented yet.")
     end
@@ -131,25 +150,28 @@ function compress!(
     tt::TensorTrain{V,N},
     method::Symbol=:LU;
     tolerance::Float64=1e-12,
-    maxbonddim::Int=typemax(Int)
+    maxbonddim::Int=typemax(Int),
+    normalizeerror::Bool=true
 ) where {V,N}
+    # From left to right
     for ell in 1:length(tt)-1
         shapel = size(tt.sitetensors[ell])
         left, right, newbonddim = _factorize(
             reshape(tt.sitetensors[ell], prod(shapel[1:end-1]), shapel[end]),
-            method; tolerance, maxbonddim
+            method; tolerance=0.0, maxbonddim=typemax(Int), leftorthogonal=true # no truncation
         )
         tt.sitetensors[ell] = reshape(left, shapel[1:end-1]..., newbonddim)
         shaper = size(tt.sitetensors[ell+1])
         nexttensor = right * reshape(tt.sitetensors[ell+1], shaper[1], prod(shaper[2:end]))
         tt.sitetensors[ell+1] = reshape(nexttensor, newbonddim, shaper[2:end]...)
     end
 
+    # From right to left
     for ell in length(tt):-1:2
         shaper = size(tt.sitetensors[ell])
         left, right, newbonddim = _factorize(
             reshape(tt.sitetensors[ell], shaper[1], prod(shaper[2:end])),
-            method; tolerance, maxbonddim
+            method; tolerance, maxbonddim, normalizeerror, leftorthogonal=false
         )
         tt.sitetensors[ell] = reshape(right, newbonddim, shaper[2:end]...)
         shapel = size(tt.sitetensors[ell-1])
@@ -212,6 +234,7 @@ function Base.reverse(tt::AbstractTensorTrain{V}) where {V}
     ]))
 end
 
+
 """
 Fitting data with a TensorTrain object.
 This may be useful when the interpolated function is noisy.
@@ -266,4 +289,4 @@ function fulltensor(obj::TensorTrain{T,N})::Array{T} where {T,N}
     end
     returnsize = collect(Iterators.flatten(sitedims_))
     return reshape(result, returnsize...)
-end
+end

test/test_blockstructure.jl

@@ -1,4 +1,5 @@
 import TensorCrossInterpolation as TCI
+import LinearAlgebra as LA
 using Random
 using Zygote
 using Optim
@@ -196,6 +197,47 @@ end
     @test ttmultileg2.(indicesmultileg) ≈ 2 .* ttmultileg.(indicesmultileg)
 end
 
+
+@testset "norm" begin
+    T = Float64
+    sitedims_ = [[2], [2], [2]]
+    N = length(sitedims_)
+    bonddims = [1, 1, 1, 1]
+
+    tt = TCI.TensorTrain([
+        ones(bonddims[n], sitedims_[n]..., bonddims[n+1]) for n in 1:N
+    ])
+
+    @test LA.norm2(tt) ≈ prod(only.(sitedims_))
+    @test LA.norm2(2 * tt) ≈ 4 * prod(only.(sitedims_))
+    @test LA.norm2(tt) ≈ LA.norm(tt)^2
+end
+
+@testset "compress! (SVD)" for T in [Float64, ComplexF64]
+    Random.seed!(1234)
+    T = Float64
+    N = 10
+    sitedims_ = [[2] for _ in 1:N]
+    χ = 10
+
+    tol = 0.1
+    bonddims = vcat(1, χ * ones(Int, N - 1), 1)
+
+    tt = TCI.TensorTrain([
+        randn(bonddims[n], sitedims_[n]..., bonddims[n+1]) for n in 1:N
+    ])
+
+    # normalizeerror=true
+    tt_compressed = deepcopy(tt)
+    TCI.compress!(tt_compressed, :SVD; tolerance=tol)
+    @test sqrt(LA.norm2(tt - tt_compressed) / LA.norm2(tt)) < sqrt(N) * tol
+
+    # normalizeerror=false
+    tt_compressed = deepcopy(tt)
+    TCI.compress!(tt_compressed, :SVD; tolerance=LA.norm(tt) * tol, normalizeerror=false)
+    @test sqrt(LA.norm2(tt - tt_compressed) / LA.norm2(tt)) < sqrt(N) * tol
+end
+
 @testset "tensor train cast" begin
     Random.seed!(10)
     localdims = [2, 2, 2]