Merge branch 'main' into 17-improvements-for-compress-function-using-svd

Author: shinaoka

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorCrossInterpolation"
 uuid = "b261b2ec-6378-4871-b32e-9173bb050604"
 authors = ["Ritter.Marc <[email protected]>, Hiroshi Shinaoka <[email protected]> and contributors"]
-version = "0.9.7"
+version = "0.9.10"
 
 [deps]
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"

src/batcheval.jl

@@ -95,6 +95,9 @@ struct ThreadedBatchEvaluator{T} <: BatchEvaluator{T}
     end
 end
 
+function (obj::ThreadedBatchEvaluator{T})(indexset::Vector{Int})::T where {T}
+    return obj.f(indexset)
+end
 
 # Batch evaluation (loop over all index sets)
 function (obj::ThreadedBatchEvaluator{T})(leftindexset::Vector{Vector{Int}}, rightindexset::Vector{Vector{Int}}, ::Val{M})::Array{T,M + 2} where {T,M}

src/cachedfunction.jl

@@ -90,6 +90,9 @@ function _batcheval_imp_default(f::CachedFunction{V,K},
     rightindexset::AbstractVector{MultiIndex},
     ::Val{M}
 )::Array{V,M + 2} where {V,K,M}
+    if length(leftindexset) * length(rightindexset) == 0
+        return Array{V,M + 2}(undef, ntuple(d -> 0, M + 2)...)
+    end
     nl = length(first(leftindexset))
     nr = length(first(rightindexset))
     L = length(f.localdims)
@@ -113,6 +116,9 @@ function _batcheval_imp_for_batchevaluator(f::CachedFunction{V,K},
     rightindexset::AbstractVector{MultiIndex},
     ::Val{M}
 )::Array{V,M + 2} where {V,K,M}
+    if length(leftindexset) * length(rightindexset) == 0
+        return Array{V,M + 2}(undef, ntuple(d -> 0, M + 2)...)
+    end
     nl = length(first(leftindexset))
     nr = length(first(rightindexset))
     L = length(f.localdims)

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/tensorci2.jl

@@ -47,7 +47,7 @@ function TensorCI2{ValueType}(
 ) where {F,ValueType,N}
     tci = TensorCI2{ValueType}(localdims)
     addglobalpivots!(tci, initialpivots)
-    tci.maxsamplevalue = maximum(abs, func.(initialpivots))
+    tci.maxsamplevalue = maximum(abs, (func(x) for x in initialpivots))
     abs(tci.maxsamplevalue) > 0.0 || error("maxsamplevalue is zero!")
     invalidatesitetensors!(tci)
     return tci

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}
@@ -262,4 +273,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_batcheval.jl

@@ -65,4 +65,21 @@ end
 
         @test result ≈ ref
     end
-end
+
+    @testset "ThreadedBatchEvaluator (from Matsuura)" begin
+        function f(x)
+            sleep(1e-3)
+            return sum(x)
+        end
+
+        L = 20
+        localdims = fill(2, L)
+        parf = TCI.ThreadedBatchEvaluator{Float64}(f, localdims)
+
+        tci, ranks, errors = TCI.crossinterpolate2(Float64, parf, localdims)
+
+        tci_ref, ranks_ref, errors_ref = TCI.crossinterpolate2(Float64, f, localdims)
+
+        @test TCI.fulltensor(TCI.TensorTrain(tci)) ≈ TCI.fulltensor(TCI.TensorTrain(tci_ref))
+    end
+end

test/test_tensortrain.jl

@@ -178,7 +178,6 @@ end
     @test ttmultileg2.(indicesmultileg) ≈ 2 .* ttmultileg.(indicesmultileg)
 end
 
-
 @testset "norm" begin
     T = Float64
     sitedims_ = [[2], [2], [2]]
@@ -217,4 +216,20 @@ end
     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
+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
+