Implement pivots recycling

Author: gianlucagrosso

src/crossinterpolate.jl

@@ -122,6 +122,7 @@ mutable struct TCI2PatchCreator{T} <: AbstractPatchCreator{T,TensorTrainState{T}
     checkbatchevaluatable::Bool
     loginterval::Int
     initialpivots::Vector{MultiIndex} # Make it to Vector{MMultiIndex}?
+    recyclepivots::Bool
 end
 
 function Base.show(io::IO, obj::TCI2PatchCreator{T}) where {T}
@@ -141,6 +142,7 @@ function TCI2PatchCreator{T}(obj::TCI2PatchCreator{T})::TCI2PatchCreator{T} wher
         obj.checkbatchevaluatable,
         obj.loginterval,
         obj.initialpivots,
+        obj.recyclepivots,
     )
 end
 
@@ -157,7 +159,8 @@ function TCI2PatchCreator(
     ninitialpivot=5,
     checkbatchevaluatable=false,
     loginterval=10,
-    initialpivots=Vector{MultiIndex}[],
+    initialpivots=MultiIndex[],
+    recyclepivots=false,
 )::TCI2PatchCreator{T} where {T}
     #t1 = time_ns()
     if projector === nothing
@@ -183,6 +186,7 @@ function TCI2PatchCreator(
         checkbatchevaluatable,
         loginterval,
         initialpivots,
+        recyclepivots,
     )
 end
 
@@ -206,6 +210,7 @@ function _crossinterpolate2!(
     verbosity::Int=0,
     checkbatchevaluatable=false,
     loginterval=10,
+    recyclepivots=false,
 ) where {T}
     ncheckhistory = 3
     ranks, errors = TCI.optimize!(
@@ -231,13 +236,45 @@ function _crossinterpolate2!(
     ncheckhistory_ = min(ncheckhistory, length(errors))
     maxbonddim_hist = maximum(ranks[(end - ncheckhistory_ + 1):end])
 
-    return PatchCreatorResult{T,TensorTrain{T,3}}(
-        TensorTrain(tci), TCI.maxbonderror(tci) < tolerance && maxbonddim_hist < maxbonddim
-    )
+    if recyclepivots
+        return PatchCreatorResult{T,TensorTrain{T,3}}(
+            TensorTrain(tci),
+            TCI.maxbonderror(tci) < tolerance && maxbonddim_hist < maxbonddim,
+            _globalpivots(tci),
+        )
+
+    else
+        return PatchCreatorResult{T,TensorTrain{T,3}}(
+            TensorTrain(tci),
+            TCI.maxbonderror(tci) < tolerance && maxbonddim_hist < maxbonddim,
+        )
+    end
+end
+
+# Generating global pivots from local ones
+function _globalpivots(
+    tci::TCI.TensorCI2{T}; onlydiagonal=true
+)::Vector{MultiIndex} where {T}
+    Isets = tci.Iset
+    Jsets = tci.Jset
+    L = length(Isets)
+    p = Set{MultiIndex}()
+    # Pivot matrices
+    for bondindex in 1:(L - 1)
+        if onlydiagonal
+            for (x, y) in zip(Isets[bondindex + 1], Jsets[bondindex])
+                push!(p, vcat(x, y))
+            end
+        else
+            for x in Isets[bondindex + 1], y in Jsets[bondindex]
+                push!(p, vcat(x, y))
+            end
+        end
+    end
+    return collect(p)
 end
 
 function createpatch(obj::TCI2PatchCreator{T}) where {T}
-    proj = obj.projector
     fsubset = _FuncAdapterTCI2Subset(obj.f)
 
     tci = if isapproxttavailable(obj.f)
@@ -253,21 +290,25 @@ function createpatch(obj::TCI2PatchCreator{T}) where {T}
         end
         tci
     else
-        # Random initial pivots
         initialpivots = MultiIndex[]
-        let
-            mask = [!isprojectedat(proj, n) for n in 1:length(proj)]
-            for idx in obj.initialpivots
-                idx_ = [[i] for i in idx]
-                if idx_ <= proj
-                    push!(initialpivots, idx[mask])
-                end
+        if obj.recyclepivots
+            # First patching iteration: random pivots
+            if length(fsubset.localdims) == length(obj.localdims)
+                initialpivots = union(
+                    obj.initialpivots,
+                    findinitialpivots(fsubset, fsubset.localdims, obj.ninitialpivot),
+                )
+                # Next iterations: recycle previously generated pivots
+            else
+                initialpivots = copy(obj.initialpivots)
             end
+        else
+            initialpivots = union(
+                obj.initialpivots,
+                findinitialpivots(fsubset, fsubset.localdims, obj.ninitialpivot),
+            )
         end
-        append!(
-            initialpivots,
-            findinitialpivots(fsubset, fsubset.localdims, obj.ninitialpivot),
-        )
+
         if all(fsubset.(initialpivots) .== 0)
             return PatchCreatorResult{T,TensorTrainState{T}}(nothing, true)
         end
@@ -282,6 +323,7 @@ function createpatch(obj::TCI2PatchCreator{T}) where {T}
         verbosity=obj.verbosity,
         checkbatchevaluatable=obj.checkbatchevaluatable,
         loginterval=obj.loginterval,
+        recyclepivots=obj.recyclepivots,
     )
 end
 
@@ -296,14 +338,14 @@ function adaptiveinterpolate(
 end
 
 function adaptiveinterpolate(
-    f::ProjectableEvaluator{T},
-    pordering::PatchOrdering=PatchOrdering(collect(1:length(f.sitedims)));
+    f::ProjectableEvaluator{T};
+    pordering::PatchOrdering=PatchOrdering(collect(1:length(f.sitedims))),
     verbosity=0,
     maxbonddim=typemax(Int),
     tolerance=1e-8,
-    initialpivots=Vector{MultiIndex}[], # Make it to Vector{MMultiIndex}?
+    initialpivots=MultiIndex[], # Make it to Vector{MMultiIndex}?
+    recyclepivots=false,
 )::ProjTTContainer{T} where {T}
-    t1 = time_ns()
     creator = TCI2PatchCreator(
         T,
         f,
@@ -313,6 +355,7 @@ function adaptiveinterpolate(
         verbosity,
         ntry=10,
         initialpivots=initialpivots,
+        recyclepivots=recyclepivots,
     )
     tmp = adaptiveinterpolate(creator, pordering; verbosity)
     return reshape(tmp, f.sitedims)

src/patching.jl

@@ -38,6 +38,19 @@ abstract type AbstractPatchCreator{T,M} end
 mutable struct PatchCreatorResult{T,M}
     data::Union{M,Nothing}
     isconverged::Bool
+    resultpivots::Vector{MultiIndex}
+
+    function PatchCreatorResult{T,M}(
+        data::Union{M,Nothing}, isconverged::Bool, resultpivots::Vector{MultiIndex}
+    )::PatchCreatorResult{T,M} where {T,M}
+        return new{T,M}(data, isconverged, resultpivots)
+    end
+
+    function PatchCreatorResult{T,M}(
+        data::Union{M,Nothing}, isconverged::Bool
+    )::PatchCreatorResult{T,M} where {T,M}
+        return new{T,M}(data, isconverged, MultiIndex[])
+    end
 end
 
 function _reconst_prefix(projector::Projector, pordering::PatchOrdering)
@@ -63,10 +76,19 @@ function __taskfunc(creator::AbstractPatchCreator{T,M}, pordering; verbosity=0)
         for ic in 1:creator.localdims[pordering.ordering[length(prefix) + 1]]
             prefix_ = vcat(prefix, ic)
             projector_ = makeproj(pordering, prefix_, creator.localdims)
-            #if verbosity > 0
-            ##println("Creating a task for $(prefix_) ...")
-            #end
-            push!(newtasks, project(creator, projector_))
+
+            # Pivots are shorter, pordering index is in a different position
+            active_dims_ = findall(x -> x == [0], creator.projector.data)
+            pos_ = findfirst(x -> x == pordering.ordering[length(prefix) + 1], active_dims_)
+            pivots_ = [
+                copy(piv) for piv in filter(piv -> piv[pos_] == ic, patch.resultpivots)
+            ]
+
+            if !isempty(pivots_)
+                deleteat!.(pivots_, pos_)
+            end
+
+            push!(newtasks, project(creator, projector_; pivots=pivots_))
         end
         return nothing, newtasks
     end
@@ -77,14 +99,19 @@ function _zerott(T, prefix, po::PatchOrdering, localdims::Vector{Int})
     return TensorTrain([zeros(T, 1, d, 1) for d in localdims_])
 end
 
-function project(obj::AbstractPatchCreator{T,M}, projector::Projector) where {T,M}
+function project(
+    obj::AbstractPatchCreator{T,M},
+    projector::Projector;
+    pivots::Vector{MultiIndex}=MultiIndex[],
+) where {T,M}
     projector <= obj.projector || error(
         "Projector $projector is not a subset of the original projector $(obj.f.projector)",
     )
 
     obj_copy = TCI2PatchCreator{T}(obj) # shallow copy
     obj_copy.projector = deepcopy(projector)
     obj_copy.f = project(obj_copy.f, projector)
+    obj_copy.initialpivots = deepcopy(pivots)
     return obj_copy
 end
 

test/crossinterpolate_tests.jl

@@ -89,7 +89,7 @@ using Random
             ntry=10,
         )
 
-        obj = TCIA.adaptiveinterpolate(creator, pordering; verbosity=2)
+        obj = TCIA.adaptiveinterpolate(creator, pordering; verbosity=0)
 
         points = [(rand() * 10 - 5, rand() * 10 - 5) for i in 1:100]
 
@@ -140,7 +140,7 @@ using Random
 
         ptt = TCIA.project(TCIA.ProjTensorTrain(tt), p)
 
-        obj = TCIA.adaptiveinterpolate(ptt; verbosity=1, maxbonddim=5)
+        obj = TCIA.adaptiveinterpolate(ptt; verbosity=0, maxbonddim=5)
 
         @test vec(TCIA.fulltensor(obj)) ≈ vec(TCIA.fulltensor(ptt))
     end
@@ -175,9 +175,9 @@ using Random
             verbosity=0,
             ntry=10,
         )
-        obj = TCIA.adaptiveinterpolate(creator, pordering; verbosity=2)
+        obj = TCIA.adaptiveinterpolate(creator, pordering; verbosity=0)
 
-        points = [(rand() * 2*pi - pi, rand() * 2*pi - pi) for i in 1:1000]
+        points = [(rand() * 2 * pi - pi, rand() * 2 * pi - pi) for i in 1:1000]
 
         @assert length(obj) > localdims[1]
 
@@ -186,6 +186,85 @@ using Random
             [qf(QG.origcoord_to_quantics(grid, p)) for p in points];
             atol=1e-4,
         )
+    end
+
+    @testset "recyclepivots" begin
+        Random.seed!(1234)
+
+        function linear_gaussians(
+            x::Float64, y::Float64; centers::Vector{Vector{Float64}}=[[0.0, 0.0]]
+        )
+            N_gauss = length(centers)
+            input_vec = [x, y]
+            σ = [2.0^(-j - 1) for j in 1:N_gauss]
+            return sum(exp(-norm(input_vec - centers[j])^2 / σ[j]^2) for j in 1:N_gauss)
+        end
+
+        R = 30
+        grid = QG.DiscretizedGrid{2}(R, (-3.0, -3.0), (3.0, 3.0))
+        localdims = fill(4, R)
+        sitedims = [[2, 2] for _ in 1:R]
+
+        gauss_centers = [[-2.0, 2.0], [2.0, -2.0], [-2.0, -2.0], [2.0, 2.0]]
+        qf =
+            x -> linear_gaussians(
+                QG.quantics_to_origcoord(grid, x)...; centers=gauss_centers
+            )
+
+        pordering = TCIA.PatchOrdering(collect(1:R))
+        tol = 1e-7
+        mb = 20
+
+        # Provide very good pivots on the maximums of the function 
+        initialpivots = [
+            TCI.optfirstpivot(qf, localdims, QG.origcoord_to_quantics(grid, Tuple(c))) for
+            c in gauss_centers
+        ]
+
+        creator_recycle = TCIA.TCI2PatchCreator(
+            Float64,
+            TCIA.makeprojectable(Float64, qf, localdims),
+            localdims,
+            ;
+            maxbonddim=mb,
+            tolerance=tol,
+            recyclepivots=true,
+            initialpivots=initialpivots,
+            ninitialpivot=0,
+        )
 
+        creator_no_recycle = TCIA.TCI2PatchCreator(
+            Float64,
+            TCIA.makeprojectable(Float64, qf, localdims),
+            localdims,
+            ;
+            maxbonddim=mb,
+            tolerance=tol,
+            recyclepivots=false,
+            initialpivots=initialpivots,
+            ninitialpivot=0,
+        )
+
+        # Perform ony one patch iteration
+        _, recycle_patches = TCIA.__taskfunc(creator_recycle, pordering)
+        _, no_recycle_patches = TCIA.__taskfunc(creator_no_recycle, pordering)
+
+        # Check if the "good" pivots are conserved in each patch
+        for patch in recycle_patches
+            prj_val = collect(Iterators.flatten(patch.projector))[1]
+            patch_pivots = vcat.([prj_val], patch.initialpivots)
+
+            @test intersect(patch_pivots, initialpivots) ==
+                filter(piv -> piv[1] == prj_val, initialpivots)
+        end
+
+        # Check that the information is lost
+        for patch in no_recycle_patches
+            prj_val = collect(Iterators.flatten(patch.projector))[1]
+            patch_pivots = vcat.([prj_val], patch.initialpivots)
+
+            @test intersect(patch_pivots, initialpivots) !==
+                filter(piv -> piv[1] == prj_val, initialpivots)
+        end
     end
 end