FixTaylorCluster for non-square (finite) Hamiltonians

docs/src/changelog.md

@@ -121,6 +121,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
   from `make_time_mpo`) into an MPS, since it braided the physical leg directly against a
   `SumSpace` virtual leg. It now densifies the tensor first, which the function already did to its
   return value regardless, so this adds no extra cost. ([#509](https://github.com/QuantumKitHub/MPSKit.jl/pull/509))
+- `make_time_mpo` with `TaylorCluster` on a Hamiltonian whose virtual bond dimension varies along
+  the chain are now correctly handled. ([#511](https://github.com/QuantumKitHub/MPSKit.jl/pull/511))
 
 ### Performance
 

src/algorithms/timestep/taylorcluster.jl

@@ -130,21 +130,57 @@ function _taylor_loopback!(mpo, virtual_sz, linds, ::Val{N}, τ::Number) where {
         V_right = virtual_sz[i + 1]
         linds_right = linds[i + 1]
         cinds_right = CartesianIndices(linds_right)
-        for b in cinds_right[2:end]
-            all(in((1, V_right)), b.I) || continue
 
-            b_lin = linds_right[b]
-            a = count(==(V_right), b.I)
+        threelevel_linds_right = _get_threelevel_linds(V_right, linds_right)
+        for c in threelevel_linds_right
+            c_cart = cinds_right[c]
+            a = count(==(V_right), c_cart.I)
             factor = τ^a * factorial(N - a) / factorial(N)
-            slice[:, 1, 1, 1] = slice[:, 1, 1, 1] + factor * slice[:, 1, 1, b_lin]
-            for I in nonzero_keys(slice)
-                (I[1] == b_lin || I[4] == b_lin) && delete!(slice, I)
-            end
+            slice[:, 1, 1, 1] = slice[:, 1, 1, 1] + factor * slice[:, 1, 1, c]
         end
     end
+    _remove_threelevels!(mpo, virtual_sz, linds)
     return mpo
 end
 
+function _get_threelevel_linds(virtual_sz, linds)
+    row_threelevels_linds = Int[]
+    for c in CartesianIndices(linds)
+        c_lin = linds[c]
+        cI = c.I
+        all(in((1, virtual_sz)), cI) && !all(==(1), cI) && push!(row_threelevels_linds, c_lin)
+    end
+    return row_threelevels_linds
+end
+
+# removes the rows and columns that are labled by level-labels made up of 1 and "3" only, but not all 1's.
+function _remove_threelevels!(mpo, virtual_sz, linds)
+    for (i, slice) in enumerate(parent(mpo))
+        # three-level rows
+        row_threelevel_linds = _get_threelevel_linds(virtual_sz[i], linds[i])
+
+        # three-level columns
+        col_threelevel_linds = _get_threelevel_linds(virtual_sz[i + 1], linds[i + 1])
+
+        # build the set of all CartesianIndices in a three-level row or column
+        d1, _, _, d4 = size(slice)
+        threelevel_cinds = Set{CartesianIndex{4}}()
+        for r in row_threelevel_linds
+            union!(threelevel_cinds, CartesianIndices((r:r, 1:1, 1:1, 1:d4)))
+        end
+        for c in col_threelevel_linds
+            union!(threelevel_cinds, CartesianIndices((1:d1, 1:1, 1:1, c:c)))
+        end
+
+        # remove the nonzero entries that fall in a three-level row or column
+        # this avoids a lot of misses wrt looping over all nonzero_keys
+        for I in intersect(nonzero_keys(slice), threelevel_cinds)
+            delete!(slice, I)
+        end
+    end
+    return
+end
+
 # Algorithm 2: collapse rows and columns that are equivalent under the
 # permutation symmetry of the Taylor expansion.
 function _taylor_remove_equivalents!(mpo, virtual_sz, linds)

test/timeevolution/taylorcluster.jl

@@ -94,7 +94,14 @@ using TensorKit: ℙ
     end
 
     L = 4
-    Hs = [transverse_field_ising(; L), heisenberg_XXX(; L)]
+    h1 = S_z_S_z(ComplexF64, Trivial; spin = 1 // 2) * 4
+    h2 = S_x_S_x(ComplexF64, Trivial; spin = 1 // 2) * 4
+    rectangular_lattice = fill(space(h1, 1), L)
+    rectangular_ising = FiniteMPOHamiltonian(
+        rectangular_lattice, ((i, i + 1) => h1 for i in 1:(L - 1))..., (1, 3) => h2
+    )
+
+    Hs = [transverse_field_ising(; L), heisenberg_XXX(; L), rectangular_ising]
 
     Ns = [1, 2, 3]
     dts = [1.0e-2, 1.0e-3]