Merge pull request #38 from tensor4all/37-reconstructing-a-full-tensor-using-matrix-multiplications

Vectorize fulltensor()

Author: shinaoka

src/tensortrain.jl

@@ -253,11 +253,17 @@ function (obj::TensorTrainFit{ValueType})(x::Vector{ValueType}) where {ValueType
 end
 
 
-
 function fulltensor(obj::TensorTrain{T,N})::Array{T} where {T,N}
     sitedims_ = sitedims(obj)
     localdims = collect(prod.(sitedims_))
-    r = [obj(collect(Tuple(i))) for i in CartesianIndices(Tuple(localdims))]
+    result::Matrix{T} = reshape(obj.sitetensors[1], localdims[1], :)
+    leftdim = localdims[1]
+    for l in 2:length(obj)
+        nextmatrix = reshape(
+            obj.sitetensors[l], size(obj.sitetensors[l], 1), localdims[l] * size(obj.sitetensors[l])[end])
+        leftdim *= localdims[l]
+        result = reshape(result * nextmatrix, leftdim, size(obj.sitetensors[l])[end])
+    end
     returnsize = collect(Iterators.flatten(sitedims_))
-    return reshape(r, returnsize...)
+    return reshape(result, returnsize...)
 end

test/test_tensortrain.jl

@@ -43,6 +43,25 @@ using Optim
 end
 
 
+function _fulltensor(obj::TCI.TensorTrain{T,N})::Array{T} where {T,N}
+    sitedims_ = TCI.sitedims(obj)
+    localdims = collect(prod.(sitedims_))
+    r = [obj(collect(Tuple(i))) for i in CartesianIndices(Tuple(localdims))]
+    returnsize = collect(Iterators.flatten(sitedims_))
+    return reshape(r, returnsize...)
+end
+
+@testset "TT fulltensor" for T in [Float64, ComplexF64]
+    linkdims = [1, 2, 3, 1]
+    L = length(linkdims) - 1
+    localdims = fill(4, L)
+    tts = TCI.TensorTrain{T,3}([randn(T, linkdims[n], localdims[n], linkdims[n+1]) for n in 1:L])
+    tto = TCI.TensorTrain{4}(tts, fill([2, 2], L))
+
+    @test _fulltensor(tts) ≈ TCI.fulltensor(tts)
+end
+
+
 @testset "TT shape conversion" for T in [Float64, ComplexF64]
     linkdims = [1, 2, 3, 1]
     L = length(linkdims) - 1