fix adaptive_patching logic

Author: SamuelBadr

src/patching.jl

@@ -81,29 +81,18 @@ If the bond dimension of a SubDomainMPS exceeds `maxdim`, perform patching.
 function adaptive_patching(
     subdmps::SubDomainMPS, patchorder; cutoff=0.0, maxdim=typemax(Int)
 )::Vector{SubDomainMPS}
-    if maxbonddim(subdmps) <= maxdim
+    prjidx = _next_projindex(subdmps.projector, patchorder)
+    if maxbonddim(subdmps) <= maxdim || prjidx === nothing
         return [subdmps]
     end
 
-    # If the bond dimension exceeds maxdim, perform patching
-    refined_subdmpss = SubDomainMPS[]
-    nextprjidx = _next_projindex(subdmps.projector, patchorder)
-    if nextprjidx === nothing
-        return [subdmps]
-    end
-
-    for prjval in 1:ITensors.dim(nextprjidx)
-        prj_ = subdmps.projector & Projector(nextprjidx => prjval)
-        subdmps_ = truncate(project(subdmps, prj_); cutoff, maxdim)
-        if maxbonddim(subdmps_) <= maxdim
-            push!(refined_subdmpss, subdmps_)
-        else
-            append!(
-                refined_subdmpss, adaptive_patching(subdmps_, patchorder; cutoff, maxdim)
-            )
-        end
+    # The bond dimension exceeds maxdim and there's an index to be projected, so perform patching
+    prjvals = 1:ITensors.dim(prjidx)
+    return mapreduce(vcat, prjvals) do prjval
+        prj_ = subdmps.projector & Projector(prjidx => prjval)
+        subdmps_ = truncate(project(subdmps, prj_); cutoff)
+        adaptive_patching(subdmps_, patchorder; cutoff, maxdim)
     end
-    return refined_subdmpss
 end
 
 """
@@ -113,14 +102,14 @@ Do patching recursively to reduce the bond dimension.
 If the bond dimension of a SubDomainMPS exceeds `maxdim`, perform patching.
 """
 function adaptive_patching(
-    prjmpss::PartitionedMPS, patchorder; cutoff=0.0, maxdim=typemax(Int)
+    prjmps::PartitionedMPS, patchorder; cutoff=0.0, maxdim=typemax(Int)
 )::PartitionedMPS
     return PartitionedMPS(
         collect(
-            Iterators.flatten((
-                apdaptive_patching(prjmps; cutoff, maxdim, patchorder) for
-                prjmps in values(prjmpss)
-            )),
+            Iterators.flatten(
+                adaptive_patching(subdmps, patchorder; cutoff, maxdim) for
+                subdmps in values(prjmps)
+            ),
         ),
     )
 end

test/patching_tests.jl

@@ -1,9 +1,55 @@
 using Test
 import PartitionedMPSs:
-    PartitionedMPSs, Projector, project, SubDomainMPS, adaptive_patching, PartitionedMPS
+    PartitionedMPSs,
+    Projector,
+    project,
+    SubDomainMPS,
+    adaptive_patching,
+    PartitionedMPS,
+    maxbonddim
 import FastMPOContractions as FMPOC
 using ITensors
+using ITensorMPS
 using Random
+import TensorCrossInterpolation as TCI
+import QuanticsGrids as QG
+
+include("_util.jl")
+
+# first example from the patching paper
+function examplefunc(r)
+    r1 = (0.2, 0.2)
+    r2 = (0.8, 0.8)
+    A = 10^4
+    σ = 1e-1
+    k = 10^3
+    return A * exp(-norm(r .- r1)^2 / (2σ^2)) * sin(k * norm(r .- r1)) +
+           A / 2 * exp(-norm(r .- r2)^2 / (σ^2 / 2)) * sin(k * norm(r .- r2) / 2)
+end
+
+function tci_examplefunc(R)
+    grid = QG.DiscretizedGrid{2}(R; unfoldingscheme=:interleaved)
+    qf(x) = examplefunc(QG.quantics_to_origcoord(grid, x))
+    localdims = fill(2, 2 * R)
+
+    Npivots = 5
+    Random.seed!(1)
+    firstpivots = [[rand(1:dim) for dim in localdims] for _ in 1:Npivots]
+    firstpivots = map(p -> TCI.optfirstpivot(qf, localdims, p), firstpivots)
+    tci_func, _ = TCI.crossinterpolate2(Float64, qf, localdims, firstpivots; tolerance=1e-7)
+
+    sitesx = [Index(2, "Qubit,x=$n") for n in 1:R]
+    sitesy = [Index(2, "Qubit,y=$n") for n in 1:R]
+
+    sites_ = collect.(zip(sitesx, sitesy))
+    sites = collect(Iterators.flatten(sites_))
+
+    mps = MPS(TCI.TensorTrain(tci_func); sites)
+    sdmps = SubDomainMPS(mps)
+    pmps = PartitionedMPS([sdmps])
+
+    return sites, pmps
+end
 
 @testset "patching.jl" begin
     @testset "adaptive_patching" begin
@@ -26,4 +72,14 @@ using Random
 
         @test MPS(partmps) ≈ MPS(subdmps) rtol = 1e-12
     end
+
+    @testset "adaptive_patching PartitionedMPS" begin
+        R = 15
+        sites, pmps = tci_examplefunc(R)
+        pmps_partitioned = adaptive_patching(pmps, sites; maxdim=100, cutoff=1e-10)
+
+        @test maximum(maxbonddim, values(pmps_partitioned)) <= 100
+        @test length(values(pmps_partitioned)) > 1
+        @test MPS(pmps_partitioned) ≈ MPS(pmps) rtol = 6e-5
+    end
 end