Merge pull request #20 from tensor4all/19-improvement-of-interface-of-tensor-trains

Make interface of contract allow type conversion

Author: shinaoka

src/contraction.jl

@@ -487,14 +487,14 @@ end
 
 """
     function contract(
-        A::TensorTrain{ValueType,4},
-        B::TensorTrain{ValueType,4};
+        A::TensorTrain{V1,4},
+        B::TensorTrain{V2,4};
         algorithm::Symbol=:TCI,
         tolerance::Float64=1e-12,
         maxbonddim::Int=typemax(Int),
         f::Union{Nothing,Function}=nothing,
         kwargs...
-    ) where {ValueType}
+    ) where {V1,V2}
 
 Contract two tensor trains `A` and `B`.
 
@@ -513,45 +513,48 @@ Arguments:
 - `kwargs...` are forwarded to [`crossinterpolate2`](@ref) if `algorithm=:TCI`.
 """
 function contract(
-    A::TensorTrain{ValueType,4},
-    B::TensorTrain{ValueType,4};
+    A::TensorTrain{V1,4},
+    B::TensorTrain{V2,4};
     algorithm::Symbol=:TCI,
     tolerance::Float64=1e-12,
     maxbonddim::Int=typemax(Int),
     f::Union{Nothing,Function}=nothing,
     kwargs...
-) where {ValueType}
+)::TensorTrain{promote_type(V1,V2),4} where {V1,V2}
+    Vres = promote_type(V1, V2)
+    A_ = TensorTrain{Vres,4}(A)
+    B_ = TensorTrain{Vres,4}(B)
     if algorithm === :TCI
-        return contract_TCI(A, B; tolerance=tolerance, maxbonddim=maxbonddim, f=f, kwargs...)
+        return contract_TCI(A_, B_; tolerance=tolerance, maxbonddim=maxbonddim, f=f, kwargs...)
     elseif algorithm === :naive
         if f !== nothing
             error("Naive contraction implementation cannot contract matrix product with a function. Use algorithm=:TCI instead.")
         end
-        return contract_naive(A, B; tolerance=tolerance, maxbonddim=maxbonddim)
+        return contract_naive(A_, B_; tolerance=tolerance, maxbonddim=maxbonddim)
     elseif algorithm === :zipup
         if f !== nothing
             error("Zipup contraction implementation cannot contract matrix product with a function. Use algorithm=:TCI instead.")
         end
-        return contract_zipup(A, B; tolerance, maxbonddim)
+        return contract_zipup(A_, B_; tolerance, maxbonddim)
     else
         throw(ArgumentError("Unknown algorithm $algorithm."))
     end
 end
 
 function contract(
     A::Union{TensorCI1{V},TensorCI2{V},TensorTrain{V,3}},
-    B::TensorTrain{V,4};
+    B::TensorTrain{V2,4};
     kwargs...
-)::TensorTrain{V,3} where {V}
+)::TensorTrain{promote_type(V,V2),3} where {V,V2}
     tt = contract(TensorTrain{4}(A, [(1, s...) for s in sitedims(A)]), B; kwargs...)
     return TensorTrain{3}(tt, prod.(sitedims(tt)))
 end
 
 function contract(
     A::TensorTrain{V,4},
-    B::Union{TensorCI1{V},TensorCI2{V},TensorTrain{V,3}};
+    B::Union{TensorCI1{V2},TensorCI2{V2},TensorTrain{V2,3}};
     kwargs...
-)::TensorTrain{V,3} where {V}
+)::TensorTrain{promote_type(V,V2),3} where {V,V2}
     tt = contract(A, TensorTrain{4}(B, [(s..., 1) for s in sitedims(B)]); kwargs...)
     return TensorTrain{3}(tt, prod.(sitedims(tt)))
 end

src/tensortrain.jl

@@ -30,6 +30,17 @@ mutable struct TensorTrain{ValueType,N} <: AbstractTensorTrain{ValueType}
     end
 end
 
+function Base.show(io::IO, obj::TensorTrain{V,N}) where {V,N}
+    print(
+        io,
+        "$(typeof(obj)) of rank $(maximum(linkdims(obj)))"
+    )
+end
+
+function TensorTrain{V2,N}(tt::TensorTrain{V})::TensorTrain{V2,N} where {V,V2,N}
+    return TensorTrain{V2,N}(Array{V2}.(sitetensors(tt)))
+end
+
 """
     function TensorTrain(sitetensors::Vector{Array{V, 3}}) where {V}
 
@@ -53,7 +64,7 @@ function TensorTrain(tci::AbstractTensorTrain{V})::TensorTrain{V,3} where {V}
 end
 
 """
-    function TensorTrain{N}(tci::AbstractTensorTrain{V}) where {V,N}
+    function TensorTrain{V2,N}(tci::AbstractTensorTrain{V}) where {V,V2,N}
 
 Convert a tensor-train-like object into a tensor train.
 
@@ -62,15 +73,15 @@ Arguments:
 - `localdims`: a vector of local dimensions for each tensor in the tensor train. A each element
   of `localdims` should be an array-like object of `N-2` integers.
 """
-function TensorTrain{V,N}(tt::AbstractTensorTrain{V}, localdims)::TensorTrain{V,N} where {V,N}
+function TensorTrain{V2,N}(tt::AbstractTensorTrain{V}, localdims)::TensorTrain{V2,N} where {V,V2,N}
     for d in localdims
         length(d) == N - 2 || error("Each element of localdims be a list of N-2 integers.")
     end
     for n in 1:length(tt)
         prod(size(tt[n])[2:end-1]) == prod(localdims[n]) || error("The local dimensions at n=$n must match the tensor sizes.")
     end
-    return TensorTrain{V,N}(
-        [reshape(t, size(t, 1), localdims[n]..., size(t)[end]) for (n, t) in enumerate(sitetensors(tt))])
+    return TensorTrain{V2,N}(
+        [reshape(Array{V2}(t), size(t, 1), localdims[n]..., size(t)[end]) for (n, t) in enumerate(sitetensors(tt))])
 end
 
 function TensorTrain{N}(tt::AbstractTensorTrain{V}, localdims)::TensorTrain{V,N} where {V,N}
@@ -239,4 +250,14 @@ 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)))
+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))]
+    returnsize = collect(Iterators.flatten(sitedims_))
+    return reshape(r, returnsize...)
 end

test/runtests.jl

@@ -20,4 +20,4 @@ include("test_tensortrain.jl")
 include("test_conversion.jl")
 include("test_contraction.jl")
 include("test_integration.jl")
-include("test_globalsearch.jl")
+include("test_globalsearch.jl")

test/test_tensortrain.jl

@@ -48,15 +48,15 @@ end
     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))
+    tto = TCI.TensorTrain{4}(tts, fill([2, 2], L))
     tts_reconst = TCI.TensorTrain{3}(tto, localdims)
 
     for n in 1:L
         @test all(tts[n] .== tts_reconst[n])
     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
+    @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
 end
 
 @testset "batchevaluate" begin
@@ -176,3 +176,19 @@ end
     indicesmultileg = @. collect(zip(indices, indices))
     @test ttmultileg2.(indicesmultileg) ≈ 2 .* ttmultileg.(indicesmultileg)
 end
+
+@testset "tensor train cast" begin
+    Random.seed!(10)
+    localdims = [2, 2, 2]
+    linkdims_ = [1, 2, 3, 1]
+    L = length(localdims)
+
+    tt1 = TCI.TensorTrain{Float64,3}([randn(Float64, linkdims_[n], localdims[n], linkdims_[n+1]) for n in 1:L])
+
+    tt2 = TCI.TensorTrain{ComplexF64,3}(tt1, localdims)
+    @test TCI.fulltensor(tt1) ≈ TCI.fulltensor(tt2)
+
+    tt3 = TCI.TensorTrain{Float64,3}(tt2, localdims)
+    @test TCI.fulltensor(tt1) ≈ TCI.fulltensor(tt3)
+end
+