Merge branch '37-recompression-interface' into 'main'

Add recompression interface

See merge request tensors4fields/TensorCrossInterpolation.jl!70

Author: Marc Ritter

src/tensortrain.jl

@@ -56,6 +56,64 @@ function tensortrain(tci)
     return TensorTrain(tci)
 end
 
+function _factorize(
+    A::Matrix{V}, method::Symbol; tolerance::Float64, maxbonddim::Int
+) where {V}
+    if method === :LU
+        factorization = rrlu(A, abstol=tolerance, maxrank=maxbonddim)
+        return left(factorization), right(factorization), npivots(factorization)
+    elseif method === :CI
+        factorization = MatrixLUCI(A, abstol=tolerance, maxrank=maxbonddim)
+        return left(factorization), right(factorization), npivots(factorization)
+    elseif method === :SVD
+        factorization = LinearAlgebra.svd(A)
+        trunci = min(
+            replacenothing(findlast(>(tolerance), factorization.S), 1),
+            maxbonddim
+        )
+        return (
+            factorization.U[:, 1:trunci],
+            Diagonal(factorization.S[1:trunci]) * factorization.Vt[1:trunci, :],
+            trunci
+        )
+    else
+        error("Not implemented yet.")
+    end
+end
+
+function compress!(
+    tt::TensorTrain{V, N},
+    method::Symbol=:LU;
+    tolerance::Float64=1e-12,
+    maxbonddim=typemax(Int)
+) where {V, N}
+    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
+        )
+        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
+
+    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
+        )
+        tt.sitetensors[ell] = reshape(right, newbonddim, shaper[2:end]...)
+        shapel = size(tt.sitetensors[ell-1])
+        nexttensor = reshape(tt.sitetensors[ell-1], prod(shapel[1:end-1]), shapel[end]) * left
+        tt.sitetensors[ell-1] = reshape(nexttensor, shapel[1:end-1]..., newbonddim)
+    end
+
+    nothing
+end
+
 
 """
 Fitting data with a TensorTrain object.
@@ -83,7 +141,7 @@ function to_tensors(obj::TensorTrainFit{ValueType}, x::Vector{ValueType}) where
         reshape(
             x[obj.offsets[n]+1:obj.offsets[n+1]],
             size(obj.tt[n])
-            )
+        )
         for n in 1:length(obj.tt)
     ]
 end
@@ -94,5 +152,5 @@ end
 
 function (obj::TensorTrainFit{ValueType})(x::Vector{ValueType}) where {ValueType}
     tensors = to_tensors(obj, x)
-    return sum((abs2(_evaluate(tensors, indexset)  - obj.values[i]) for (i, indexset) in enumerate(obj.indexsets)))
+    return sum((abs2(_evaluate(tensors, indexset) - obj.values[i]) for (i, indexset) in enumerate(obj.indexsets)))
 end

src/util.jl

@@ -107,3 +107,11 @@ function optfirstpivot(
 
     return pivot
 end
+
+function replacenothing(value::Union{T, Nothing}, default::T)::T where {T}
+    if isnothing(value)
+        return default
+    else
+        return value
+    end
+end

test/test_tensortrain.jl

@@ -7,19 +7,33 @@ using Optim
     g(v) = 1 / (sum(v .^ 2) + 1)
     localdims = (6, 6, 6, 6)
     tolerance = 1e-8
+    allindices = CartesianIndices(localdims)
+
     tci, ranks, errors = TCI.crossinterpolate1(Float64, g, localdims; tolerance=tolerance)
     tt = TCI.TensorTrain(tci)
-    @test rank(tci) == rank(tt)
+    @test TCI.rank(tci) == TCI.rank(tt)
     @test TCI.linkdims(tci) == TCI.linkdims(tt)
     gsum = 0.0
-    for i in CartesianIndices(localdims)
+    for i in allindices
         @test TCI.evaluate(tci, i) ≈ TCI.evaluate(tt, i)
         @test tt(i) == TCI.evaluate(tt, i)
         functionvalue = g(Tuple(i))
         @test abs(TCI.evaluate(tt, i) - functionvalue) < tolerance
         gsum += functionvalue
     end
     @test gsum ≈ TCI.sum(tt)
+
+    for method in [:LU, :CI, :SVD]
+        ttcompressed = deepcopy(tt)
+        TCI.compress!(ttcompressed, method; maxbonddim=5)
+        @test TCI.rank(ttcompressed) <= 5
+    end
+
+    for method in [:LU, :CI, :SVD]
+        ttcompressed = deepcopy(tt)
+        TCI.compress!(ttcompressed, method; tolerance=1.0)
+        @test TCI.rank(ttcompressed) <= TCI.rank(tt)
+    end
 end
 
 @testset "batchevaluate" begin