Merged

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.12"
 
 [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/globalsearch.jl

@@ -80,7 +80,9 @@ function _floatingzone(
             )
             err = vec(abs.(exactdata .- prediction))
             pivot[ipos] = argmax(err)
-            maxerror = maximum(err)
+            # In RHS, we compare the maximum of the error vector with the current maxerror
+            # to make sure that the error does not decrease even if maxerror is close to machine precision.
+            maxerror = max(maximum(err), maxerror)
         end
 
         if maxerror == prev_maxerror || maxerror > earlystoptol # early stop

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
@@ -431,7 +431,8 @@ function makecanonical!(
     abstol::Float64=0.0,
     maxbonddim::Int=typemax(Int)
 ) where {F,ValueType}
-    sweep1site!(tci, f, :forward; reltol, abstol, maxbonddim, updatetensors=false)
+    # The first half-sweep is performed exactly without compression.
+    sweep1site!(tci, f, :forward; reltol=0.0, abstol=0.0, maxbonddim=typemax(Int), updatetensors=false)
     sweep1site!(tci, f, :backward; reltol, abstol, maxbonddim, updatetensors=false)
     sweep1site!(tci, f, :forward; reltol, abstol, maxbonddim, updatetensors=true)
 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}
@@ -262,4 +273,20 @@ 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_))
+    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(result, 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_globalsearch.jl

@@ -32,6 +32,6 @@ import QuanticsGrids as QD
     pivoterrors = TCI.estimatetrueerror(TCI.TensorTrain(tci), f)
 
     errors = [e for (_, e) in pivoterrors] 
-    @test all([abs(f(p) - tci(p)) for (p, _) in pivoterrors] .== errors)
+    @test [abs(f(p) - tci(p)) for (p, _) in pivoterrors] ≈ errors
     @test all(errors[1:end-1] .>= errors[2:end]) # check if errors are sorted in descending order
 end