Merge branch '46-conversion-between-different-tensortrain-objects' into 'main'

Conversion TensorTrain -> TensorCI2

Closes #46

See merge request tensors4fields/TensorCrossInterpolation.jl!89

Author: rittermarc

src/conversion.jl

@@ -69,3 +69,108 @@ function TensorCI2{ValueType}(tci1::TensorCI1{ValueType}) where {ValueType}
     tci2.maxsamplevalue = tci1.maxsamplevalue
     return tci2
 end
+
+function sweep1sitegetindices!(
+    tt::TensorTrain{ValueType,3}, forwardsweep::Bool,
+    spectatorindices::Vector{Vector{MultiIndex}}=Vector{MultiIndex}[];
+    maxbonddim=typemax(Int), tolerance=0.0
+) where {ValueType}
+    indexset = Vector{MultiIndex}[MultiIndex[[]]]
+    pivoterrorsarray = zeros(rank(tt) + 1)
+
+    function groupindices(T::AbstractArray, next::Bool)
+        shape = size(T)
+        if forwardsweep != next
+            reshape(T, prod(shape[1:end-1]), shape[end])
+        else
+            reshape(T, shape[1], prod(shape[2:end]))
+        end
+    end
+
+    function splitindices(T::AbstractArray, shape, newbonddim, next::Bool)
+        if forwardsweep != next
+            newshape = (shape[1:end-1]..., newbonddim)
+        else
+            newshape = (newbonddim, shape[2:end]...)
+        end
+        reshape(T, newshape)
+    end
+
+    L = length(tt)
+    for i in 1:L-1
+        ell = forwardsweep ? i : L - i + 1
+        ellnext = forwardsweep ? i + 1 : L - i
+        shape = size(tt.sitetensors[ell])
+        shapenext = size(tt.sitetensors[ellnext])
+
+        luci = MatrixLUCI(
+            groupindices(tt.sitetensors[ell], false), leftorthogonal=forwardsweep,
+            abstol=tolerance, maxrank=maxbonddim
+        )
+
+        if forwardsweep
+            push!(indexset, kronecker(last(indexset), shape[2])[rowindices(luci)])
+            if !isempty(spectatorindices)
+                spectatorindices[ell] = spectatorindices[ell][colindices(luci)]
+            end
+        else
+            push!(indexset, kronecker(shape[2], last(indexset))[colindices(luci)])
+            if !isempty(spectatorindices)
+                spectatorindices[ell] = spectatorindices[ell][rowindices(luci)]
+            end
+        end
+
+
+        tt.sitetensors[ell] = splitindices(
+            forwardsweep ? left(luci) : right(luci),
+            shape, npivots(luci), false
+        )
+
+        nexttensor = (
+            forwardsweep
+            ? right(luci) * groupindices(tt.sitetensors[ellnext], true)
+            : groupindices(tt.sitetensors[ellnext], true) * left(luci)
+        )
+
+        tt.sitetensors[ellnext] = splitindices(nexttensor, shapenext, npivots(luci), true)
+        pivoterrorsarray[1:npivots(luci) + 1] = max.(pivoterrorsarray[1:npivots(luci) + 1], pivoterrors(luci))
+    end
+
+    if forwardsweep
+        return indexset, pivoterrorsarray
+    else
+        return reverse(indexset), pivoterrorsarray
+    end
+end
+
+function TensorCI2{ValueType}(
+    tt::TensorTrain{ValueType,3}; tolerance=1e-12, maxbonddim=typemax(Int), maxiter=3
+) where {ValueType}
+    local pivoterrors::Vector{Float64}
+
+    Iset, = sweep1sitegetindices!(tt, true; maxbonddim, tolerance)
+    Jset, pivoterrors = sweep1sitegetindices!(tt, false; maxbonddim, tolerance)
+
+    for iter in 3:maxiter
+        if isodd(iter)
+            Isetnew, pivoterrors = sweep1sitegetindices!(tt, true, Jset)
+            if Isetnew == Iset
+                break
+            end
+        else
+            Jsetnew, pivoterrors = sweep1sitegetindices!(tt, false, Iset)
+            if Jsetnew == Jset
+                break
+            end
+        end
+    end
+
+    tci2 = TensorCI2{ValueType}(first.(sitedims(tt)))
+    tci2.Iset = Iset
+    tci2.Jset = Jset
+    tci2.sitetensors = sitetensors(tt)
+    tci2.pivoterrors = pivoterrors
+    tci2.maxsamplevalue = maximum(maximum.(abs, tci2.sitetensors))
+
+    return tci2
+end

src/tensortrain.jl

@@ -195,6 +195,12 @@ function Base.:/(tt::TensorTrain{V,N}, a) where {V,N}
     return divide(tt, a)
 end
 
+function Base.reverse(tt::AbstractTensorTrain{V}) where {V}
+    return tensortrain(reverse([
+        permutedims(T, (ndims(T), (2:ndims(T)-1)..., 1)) for T in sitetensors(tt)
+    ]))
+end
+
 """
 Fitting data with a TensorTrain object.
 This may be useful when the interpolated function is noisy.

test/test_conversion.jl

@@ -71,3 +71,28 @@ end
         @test tci2(v) ≈ f(v)
     end
 end
+
+@testset "Conversion between TT and TCI2" begin
+    f(v) = (1.0 + 2.0im) ./ (sum(v .^ 2) + 1)
+    tci, = crossinterpolate2(ComplexF64, f, fill(4, 4); tolerance=1e-14, maxbonddim=5)
+    tt = tensortrain(tci)
+    tcib = TensorCI2{ComplexF64}(tt; tolerance=1e-14)
+
+    @test rank(tt) == 5
+    @test linkdims(tt) == linkdims(tci)
+    @test sitedims(tt) == fill([4], 4)
+
+    @test rank(tcib) == 5
+    @test linkdims(tcib) == linkdims(tt)
+    @test sitedims(tcib) == fill([4], 4)
+
+    for v in Iterators.product([1:4 for _ in 1:4]...)
+        @test abs(tt(v) - tci(v)) < 1e-13
+        @test abs(tcib(v) - tci(v)) < 1e-13
+    end
+
+    optimize!(tcib, f; tolerance=1e-14)
+    for v in Iterators.product([1:4 for _ in 1:4]...)
+        @test abs(tcib(v) - f(v)) < 1e-13
+    end
+end

test/test_tensortrain.jl

@@ -10,17 +10,24 @@ using Optim
 
     tci, ranks, errors = TCI.crossinterpolate1(ComplexF64, g, localdims; tolerance=tolerance)
     tt = TCI.TensorTrain(tci)
+    ttr = TCI.reverse(tt)
     @test TCI.rank(tci) == TCI.rank(tt)
+    @test TCI.rank(tci) == TCI.rank(ttr)
     @test TCI.linkdims(tci) == TCI.linkdims(tt)
+    @test TCI.linkdims(tci) == reverse(TCI.linkdims(ttr))
     gsum = ComplexF64(0.0)
     for i in allindices
         @test TCI.evaluate(tci, i) ≈ TCI.evaluate(tt, i)
+        @test TCI.evaluate(tci, i) ≈ TCI.evaluate(ttr, reverse(Tuple(i)))
         @test tt(i) == TCI.evaluate(tt, i)
+        @test tt(i) ≈ TCI.evaluate(tt, reverse(Tuple(i)))
         functionvalue = g(Tuple(i))
         @test abs(TCI.evaluate(tt, i) - functionvalue) < tolerance
+        @test abs(TCI.evaluate(ttr, reverse(Tuple(i))) - functionvalue) < tolerance
         gsum += functionvalue
     end
     @test gsum ≈ TCI.sum(tt)
+    @test gsum ≈ TCI.sum(ttr)
 
     for method in [:LU, :CI, :SVD]
         ttcompressed = deepcopy(tt)
@@ -46,7 +53,7 @@ end
 
     for n in 1:L
         @test all(tts[n] .== tts_reconst[n])
-    end 
+    end
 
     @test_throws ErrorException TCI.TensorTrain{4}(tts, fill([2,3], L)) # Wrong shape
     @test_throws ErrorException TCI.TensorTrain{4}(tts, fill([1,2,3], L)) # Wrong shape