From 8d03ebe242efd3699b79e0595d4785258ab24852 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 1 May 2026 12:40:57 -0400 Subject: [PATCH 1/8] add benchmark code for make_time_mpo --- .../MPSKitBenchmarks/MPSKitBenchmarks.jl | 4 ++- .../timestep/TimestepBenchmarks.jl | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl diff --git a/benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl b/benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl index c70a01f70..9546df2ba 100644 --- a/benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl +++ b/benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl @@ -14,10 +14,12 @@ BenchmarkTools.DEFAULT_PARAMETERS.memory_tolerance = 0.01 const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json") const SUITE = BenchmarkGroup() const MODULES = Dict{String, Symbol}( - "derivatives" => :DerivativesBenchmarks + "derivatives" => :DerivativesBenchmarks, + "timestep" => :TimestepBenchmarks, ) include("derivatives/DerivativesBenchmarks.jl") +include("timestep/TimestepBenchmarks.jl") load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...) diff --git a/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl b/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl new file mode 100644 index 000000000..ce2bb4395 --- /dev/null +++ b/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl @@ -0,0 +1,28 @@ +module TimestepBenchmarks + +using BenchmarkTools +using MPSKit +using MPSKitModels + +const SUITE = BenchmarkGroup() +const dt = 0.05 + +const HAMILTONIANS = ( + "tfim" => transverse_field_ising(), + "heis_xxx" => heisenberg_XXX(), +) + +# Steady-state benchmark: time `make_time_mpo` after compilation has settled. +# Keep the matrix small — we want a stable baseline, not an exhaustive sweep. +let g_top = addgroup!(SUITE, "make_time_mpo") + for (Hname, H) in HAMILTONIANS + g = addgroup!(g_top, Hname) + for N in (1, 2, 3) + alg = TaylorCluster(; N) + g["N=$N"] = @benchmarkable make_time_mpo($H, $dt, $alg) + end + g["WI"] = @benchmarkable make_time_mpo($H, $dt, $(MPSKit.WI)) + end +end + +end # module From 40c47de88f2ec65117d4b0bb2c945a536e31414c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 1 May 2026 12:41:22 -0400 Subject: [PATCH 2/8] refactor make_time_mpo: use function barrier for `N` --- src/algorithms/timestep/taylorcluster.jl | 220 +++++++++++++++-------- 1 file changed, 150 insertions(+), 70 deletions(-) diff --git a/src/algorithms/timestep/taylorcluster.jl b/src/algorithms/timestep/taylorcluster.jl index a24142431..a0d299518 100644 --- a/src/algorithms/timestep/taylorcluster.jl +++ b/src/algorithms/timestep/taylorcluster.jl @@ -31,54 +31,128 @@ function make_time_mpo( H::MPOHamiltonian, dt::Number, alg::TaylorCluster; tol = eps(real(scalartype(H))), imaginary_evolution::Bool = false ) - N = alg.N + return _make_time_mpo( + H, dt, Val(alg.N), alg.extension, alg.compression; + tol, imaginary_evolution + ) +end + +# Function-barrier body: `N` is now a type parameter, so every `ntuple(_, N)`, +# `LinearIndices`, and `CartesianIndex` downstream has a statically-known rank. +function _make_time_mpo( + H::MPOHamiltonian, dt::Number, ::Val{N}, + extension::Bool, compression::Bool; + tol, imaginary_evolution::Bool + ) where {N} τ = imaginary_evolution ? -dt : -1im * dt - # start with H^N - H_n = H^N + H_n, virtual_sz, linds = _taylor_setup(H, Val(N)) + + if extension + _taylor_extension!(H_n, H, virtual_sz, linds, Val(N), τ) + end + + mpo = MPO(map(SparseBlockTensorMap, parent(H_n))) + _taylor_loopback!(mpo, virtual_sz, linds, Val(N), τ) + _taylor_remove_equivalents!(mpo, virtual_sz, linds) + + if compression + _taylor_compression!(mpo, virtual_sz, linds, Val(N), τ) + end + + return remove_orphans!(mpo; tol) +end + +# `H^N` unrolled with `Val(N)` so the result type of `H_n` is statically known. +# Mirrors `Base.power_by_squaring`'s tree (and copy-on-N=1 behaviour) so the +# bond dimensions match `H^N` exactly. +@inline _pow(H, ::Val{1}) = copy(H) +@inline function _pow(H, ::Val{N}) where {N} + half = _pow(H, Val(N ÷ 2)) + return iseven(N) ? half * half : half * half * H +end + +# Insert `v` at position `i` of tuple `t`, preserving rank in the result type. +@inline function _insert(t::NTuple{M, T}, i::Integer, v::T) where {M, T} + return ntuple(Val(M + 1)) do j + j < i ? t[j] : (j == i ? v : t[j - 1]) + end +end + +# Stable partition: elements equal to `sentinel` first, others after, +# preserving relative order within each group. Tuple-in / tuple-out so the +# resulting `CartesianIndex` keeps a known rank — this replaces +# `CartesianIndex(sort(collect(t); by = (!=(sentinel)))...)`. +@inline function _partition_first(t::NTuple{M, Int}, sentinel::Int) where {M} + n_match = count(==(sentinel), t) + return ntuple(Val(M)) do j + if j <= n_match + _kth_match(t, sentinel, j, true) + else + _kth_match(t, sentinel, j - n_match, false) + end + end +end + +@inline function _kth_match(t::NTuple{M, Int}, sentinel::Int, k::Int, want_match::Bool) where {M} + cnt = 0 + for x in t + if (x == sentinel) == want_match + cnt += 1 + cnt == k && return x + end + end + return 0 +end + +function _taylor_setup(H::MPOHamiltonian, ::Val{N}) where {N} + H_n = _pow(H, Val(N)) virtual_sz = map(0:length(H)) do i return i == 0 ? size(H[1], 1) : size(H[i], 4) end linds = map(virtual_sz) do sz - return LinearIndices(ntuple(Returns(sz), N)) + return LinearIndices(ntuple(Returns(sz), Val(N))) end + return H_n, virtual_sz, linds +end - # extension step: Algorithm 3 - # incorporate higher order terms +# Algorithm 3 of Van Damme et al. (2024): incorporate higher-order corrections +# from `H_next = H_n * H` into `H_n` in place. +function _taylor_extension!(H_n, H, virtual_sz, linds, ::Val{N}, τ::Number) where {N} # TODO: don't need to fully construct H_next... - if alg.extension - H_next = H_n * H - for (i, slice) in enumerate(parent(H_n)) - linds_left = linds[i] - linds_right = linds[i + 1] - V_left = virtual_sz[i] - V_right = virtual_sz[i + 1] - linds_next_left = LinearIndices(ntuple(Returns(V_left), N + 1)) - linds_next_right = LinearIndices(ntuple(Returns(V_right), N + 1)) - - for a in CartesianIndices(linds_left), b in CartesianIndices(linds_right) - all(>(1), b.I) || continue - all(in((1, V_left)), a.I) && any(==(V_left), a.I) && continue - - n1 = count(==(1), a.I) + 1 - n3 = count(==(V_right), b.I) + 1 - factor = τ * factorial(N) / (factorial(N + 1) * n1 * n3) - - for c in 1:(N + 1), d in 1:(N + 1) - aₑ = insert!([a.I...], c, 1) - bₑ = insert!([b.I...], d, V_right) - - # TODO: use VectorInterface for memory efficiency - slice[linds_left[a], 1, 1, linds_right[b]] += factor * - H_next[i][linds_next_left[aₑ...], 1, 1, linds_next_right[bₑ...]] - end + H_next = H_n * H + for (i, slice) in enumerate(parent(H_n)) + linds_left = linds[i] + linds_right = linds[i + 1] + V_left = virtual_sz[i] + V_right = virtual_sz[i + 1] + linds_next_left = LinearIndices(ntuple(Returns(V_left), Val(N + 1))) + linds_next_right = LinearIndices(ntuple(Returns(V_right), Val(N + 1))) + + for a in CartesianIndices(linds_left), b in CartesianIndices(linds_right) + all(>(1), b.I) || continue + all(in((1, V_left)), a.I) && any(==(V_left), a.I) && continue + + n1 = count(==(1), a.I) + 1 + n3 = count(==(V_right), b.I) + 1 + factor = τ * factorial(N) / (factorial(N + 1) * n1 * n3) + + for c in 1:(N + 1), d in 1:(N + 1) + aₑ = _insert(a.I, c, 1) + bₑ = _insert(b.I, d, V_right) + + # TODO: use VectorInterface for memory efficiency + slice[linds_left[a], 1, 1, linds_right[b]] += factor * + H_next[i][linds_next_left[aₑ...], 1, 1, linds_next_right[bₑ...]] end end end + return H_n +end - # loopback step: Algorithm 1 - # constructing the Nth order time evolution MPO - mpo = MPO(map(SparseBlockTensorMap, parent(H_n))) +# Algorithm 1: project the auxiliary virtual-bond directions onto the physical +# block, completing the Nth-order time-evolution MPO. +function _taylor_loopback!(mpo, virtual_sz, linds, ::Val{N}, τ::Number) where {N} for (i, slice) in enumerate(parent(mpo)) V_right = virtual_sz[i + 1] linds_right = linds[i + 1] @@ -95,14 +169,18 @@ function make_time_mpo( end end end + return mpo +end - # Remove equivalent rows and columns: Algorithm 2 +# 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) for (i, slice) in enumerate(parent(mpo)) V_left = virtual_sz[i] linds_left = linds[i] for c in CartesianIndices(linds_left) c_lin = linds_left[c] - s_c = CartesianIndex(sort(collect(c.I); by = (!=(1)))...) + s_c = CartesianIndex(_partition_first(c.I, 1)) n1 = count(==(1), c.I) n3 = count(==(V_left), c.I) @@ -122,7 +200,7 @@ function make_time_mpo( linds_right = linds[i + 1] for c in CartesianIndices(linds_right) c_lin = linds_right[c] - s_r = CartesianIndex(sort(collect(c.I); by = (!=(V_right)))...) + s_r = CartesianIndex(_partition_first(c.I, V_right)) n1 = count(==(1), c.I) n3 = count(==(V_right), c.I) @@ -138,44 +216,46 @@ function make_time_mpo( end end end + return mpo +end - # Approximate compression step: Algorithm 4 - if alg.compression - for (i, slice) in enumerate(parent(mpo)) - V_right = virtual_sz[i + 1] - linds_right = linds[i + 1] - for a in CartesianIndices(linds_right) - all(>(1), a.I) || continue - a_lin = linds_right[a] - n1 = count(==(V_right), a.I) - n1 == 0 && continue - b = CartesianIndex(replace(a.I, V_right => 1)) - b_lin = linds_right[b] - factor = τ^n1 * factorial(N - n1) / factorial(N) - for k in 1:size(slice, 1) - I = CartesianIndex(k, 1, 1, a_lin) - if I in nonzero_keys(slice) - slice[k, 1, 1, b_lin] += factor * slice[I] - delete!(slice, I) - end - end - end - V_left = virtual_sz[i] - linds_left = linds[i] - for a in CartesianIndices(linds_left) - all(>(1), a.I) || continue - a_lin = linds_left[a] - n1 = count(==(V_left), a.I) - n1 == 0 && continue - for k in 1:size(slice, 4) - I = CartesianIndex(a_lin, 1, 1, k) +# Algorithm 4: approximate compression — fold the right-going `V_right` +# columns back into the `1`-column with the matching Taylor weight, then drop +# the residual `V_left` rows. +function _taylor_compression!(mpo, virtual_sz, linds, ::Val{N}, τ::Number) where {N} + for (i, slice) in enumerate(parent(mpo)) + V_right = virtual_sz[i + 1] + linds_right = linds[i + 1] + for a in CartesianIndices(linds_right) + all(>(1), a.I) || continue + a_lin = linds_right[a] + n1 = count(==(V_right), a.I) + n1 == 0 && continue + b = CartesianIndex(map(x -> x == V_right ? 1 : x, a.I)) + b_lin = linds_right[b] + factor = τ^n1 * factorial(N - n1) / factorial(N) + for k in 1:size(slice, 1) + I = CartesianIndex(k, 1, 1, a_lin) + if I in nonzero_keys(slice) + slice[k, 1, 1, b_lin] += factor * slice[I] delete!(slice, I) end end end + V_left = virtual_sz[i] + linds_left = linds[i] + for a in CartesianIndices(linds_left) + all(>(1), a.I) || continue + a_lin = linds_left[a] + n1 = count(==(V_left), a.I) + n1 == 0 && continue + for k in 1:size(slice, 4) + I = CartesianIndex(a_lin, 1, 1, k) + delete!(slice, I) + end + end end - - return remove_orphans!(mpo; tol) + return mpo end # Hack to treat FiniteMPOhamiltonians as Infinite From 848bb2ec46ad93ce35b27b370794610d19c2f4a8 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 1 May 2026 13:36:27 -0400 Subject: [PATCH 3/8] Union-split small orders --- src/algorithms/timestep/taylorcluster.jl | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/algorithms/timestep/taylorcluster.jl b/src/algorithms/timestep/taylorcluster.jl index a0d299518..d69f907fc 100644 --- a/src/algorithms/timestep/taylorcluster.jl +++ b/src/algorithms/timestep/taylorcluster.jl @@ -31,9 +31,15 @@ function make_time_mpo( H::MPOHamiltonian, dt::Number, alg::TaylorCluster; tol = eps(real(scalartype(H))), imaginary_evolution::Bool = false ) - return _make_time_mpo( - H, dt, Val(alg.N), alg.extension, alg.compression; - tol, imaginary_evolution + # Manual union-split for typical small N: emits a direct call into the specialised + # `_make_time_mpo(H, dt, ::Val{N}, ...)` instead of paying a dynamic dispatch on `Val(alg.N)`. + # Falls back to the generic path for N > 4. + n = alg.N + return Base.Cartesian.@nif( + 4, + d -> n == d, + d -> _make_time_mpo(H, dt, Val(d), alg.extension, alg.compression; tol, imaginary_evolution), + d -> _make_time_mpo(H, dt, Val(n), alg.extension, alg.compression; tol, imaginary_evolution), ) end From 5b9827eb913964e69d0d285b60416a9be55661a9 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Fri, 1 May 2026 13:40:36 -0400 Subject: [PATCH 4/8] simplifications --- src/algorithms/timestep/taylorcluster.jl | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/algorithms/timestep/taylorcluster.jl b/src/algorithms/timestep/taylorcluster.jl index d69f907fc..81fb46520 100644 --- a/src/algorithms/timestep/taylorcluster.jl +++ b/src/algorithms/timestep/taylorcluster.jl @@ -54,17 +54,13 @@ function _make_time_mpo( H_n, virtual_sz, linds = _taylor_setup(H, Val(N)) - if extension - _taylor_extension!(H_n, H, virtual_sz, linds, Val(N), τ) - end + extension && _taylor_extension!(H_n, H, virtual_sz, linds, Val(N), τ) mpo = MPO(map(SparseBlockTensorMap, parent(H_n))) _taylor_loopback!(mpo, virtual_sz, linds, Val(N), τ) _taylor_remove_equivalents!(mpo, virtual_sz, linds) - if compression - _taylor_compression!(mpo, virtual_sz, linds, Val(N), τ) - end + compression && _taylor_compression!(mpo, virtual_sz, linds, Val(N), τ) return remove_orphans!(mpo; tol) end @@ -78,13 +74,6 @@ end return iseven(N) ? half * half : half * half * H end -# Insert `v` at position `i` of tuple `t`, preserving rank in the result type. -@inline function _insert(t::NTuple{M, T}, i::Integer, v::T) where {M, T} - return ntuple(Val(M + 1)) do j - j < i ? t[j] : (j == i ? v : t[j - 1]) - end -end - # Stable partition: elements equal to `sentinel` first, others after, # preserving relative order within each group. Tuple-in / tuple-out so the # resulting `CartesianIndex` keeps a known rank — this replaces @@ -144,8 +133,8 @@ function _taylor_extension!(H_n, H, virtual_sz, linds, ::Val{N}, τ::Number) whe factor = τ * factorial(N) / (factorial(N + 1) * n1 * n3) for c in 1:(N + 1), d in 1:(N + 1) - aₑ = _insert(a.I, c, 1) - bₑ = _insert(b.I, d, V_right) + aₑ = TT.insertafter(a.I, c - 1, (1,)) + bₑ = TT.insertafter(b.I, d - 1, (V_right,)) # TODO: use VectorInterface for memory efficiency slice[linds_left[a], 1, 1, linds_right[b]] += factor * From dc70e9091af96c3cddbda38658cd2566d3df446c Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 2 May 2026 19:41:05 -0400 Subject: [PATCH 5/8] refactor WII to improve TTFX --- src/algorithms/timestep/wii.jl | 129 ++++++++++++++++++--------------- 1 file changed, 72 insertions(+), 57 deletions(-) diff --git a/src/algorithms/timestep/wii.jl b/src/algorithms/timestep/wii.jl index 742992e7a..5c39500fd 100644 --- a/src/algorithms/timestep/wii.jl +++ b/src/algorithms/timestep/wii.jl @@ -23,67 +23,82 @@ function make_time_mpo( H::InfiniteMPOHamiltonian, dt::Number, alg::WII; imaginary_evolution::Bool = false ) - WA = H.A - WB = H.B - WC = H.C - WD = H.D - - # needs to be complex if negative because of square root - δ = imaginary_evolution ? -complex(dt) : (-1im * dt) + T = complex(promote_type(scalartype(H), typeof(dt))) + δ = imaginary_evolution ? T(-dt) : T(-im * dt) exp_alg = Arnoldi(; alg.tol, alg.maxiter) - Wnew = map(1:length(H)) do i - for j in 2:(size(H[i], 1) - 1), k in 2:(size(H[i], 4) - 1) - init = ( - isometry(storagetype(WD[i]), codomain(WD[i]), domain(WD[i])), - zero(H[i][1, 1, 1, k]), zero(H[i][j, 1, 1, end]), zero(H[i][j, 1, 1, k]), - ) - - y, convhist = exponentiate(1.0, init, exp_alg) do (x1, x2, x3, x4) - @plansor y1[-1 -2; -3 -4] := δ * x1[-1 1; -3 -4] * - H[i][1, 1, 1, end][2 3; 1 4] * τ[-2 4; 2 3] - - @plansor y2[-1 -2; -3 -4] := δ * x2[-1 1; -3 -4] * - H[i][1, 1, 1, end][2 3; 1 4] * τ[-2 4; 2 3] + - sqrt(δ) * x1[1 2; -3 4] * H[i][1, 1, 1, k][-1 -2; 3 -4] * τ[3 4; 1 2] - - @plansor y3[-1 -2; -3 -4] := δ * x3[-1 1; -3 -4] * - H[i][1, 1, 1, end][2 3; 1 4] * τ[-2 4; 2 3] + - sqrt(δ) * x1[1 2; -3 4] * H[i][j, 1, 1, end][-1 -2; 3 -4] * τ[3 4; 1 2] - - @plansor y4[-1 -2; -3 -4] := ( - δ * x4[-1 1; -3 -4] * H[i][1, 1, 1, end][2 3; 1 4] * τ[-2 4; 2 3] + - x1[1 2; -3 4] * H[i][j, 1, 1, k][-1 -2; 3 -4] * τ[3 4; 1 2] - ) + ( - sqrt(δ) * x2[1 2; -3 -4] * H[i][j, 1, 1, end][-1 -2; 3 4] * τ[3 4; 1 2] - ) + (sqrt(δ) * x3[-1 4; -3 3] * H[i][1, 1, 1, k][2 -2; 1 -4] * τ[3 4; 1 2]) - - return y1, y2, y3, y4 - end - - convhist.converged == 0 && - @warn "failed to exponentiate $(convhist.normres)" - - WA[i][j - 1, 1, 1, k - 1] = y[4] - WB[i][j - 1, 1, 1, 1] = y[3] - WC[i][1, 1, 1, k - 1] = y[2] - WD[i] = y[1] - end - - Vₗ = left_virtualspace(H, i)[1:(end - 1)] - Vᵣ = right_virtualspace(H, i)[1:(end - 1)] - P = physicalspace(H, i) - - h′ = similar(H[i], Vₗ ⊗ P ← P ⊗ Vᵣ) - h′[2:end, 1, 1, 2:end] = WA[i] - h′[2:end, 1, 1, 1] = WB[i] - h′[1, 1, 1, 2:end] = WC[i] - h′[1, 1, 1, 1] = WD[i] - - return h′ + O = tmap(parent(parent(H)); scheduler = Defaults.scheduler[]) do W + return _make_time_mpo(W, δ, exp_alg) + end + return InfiniteMPO(PeriodicArray(O)) +end + +function _make_time_mpo(W::JordanMPOTensor, δ, exp_alg) + # Allocate output + Vₗ = left_virtualspace(W)[1:(end - 1)] + Vᵣ = right_virtualspace(W)[1:(end - 1)] + P = physicalspace(W) + O = similar(W, storagetype(W), Vₗ ⊗ P ← P ⊗ Vᵣ) + + # fill onsite first: treated exactly + expD = exp!(scale(only(W.D), δ)) + O[1, 1, 1, 1] = insertrightunit(insertleftunit(expD, 1), 3) + + # remaining entries: + D = W[1, 1, 1, end] + for j in 2:(size(W, 1) - 1), k in 2:(size(W, 4) - 1) + B = W[j, 1, 1, end] + C = W[1, 1, 1, k] + A = W[j, 1, 1, k] + O[j, 1, 1, k], O[j, 1, 1, 1], O[1, 1, 1, k], _ = + wii_solve_block!(W, A, B, C, D, δ, exp_alg) end - return InfiniteMPO(PeriodicArray(Wnew)) + return O +end + +function wii_solve_block!( + W, A, B, C, D, δ::T, exp_alg + ) where {T <: Number} + init = (zero(A), zero(B), zero(C), isometry(storagetype(D), space(D))) + op = WIIStep(A, B, C, D, δ) + (yᴬ, yᴮ, yᶜ, yᴰ), convhist = exponentiate(op, one(real(scalartype(W))), init, exp_alg) + convhist.converged == 0 && @warn "failed to exponentiate $(convhist.normres)" + return yᴬ, yᴮ, yᶜ, yᴰ +end + +# Replaces the `do`-block closure passed to `exponentiate`. +struct WIIStep{HA, HB, HC, HD, T <: Number} + A::HA # interior / propagation block + B::HB # right boundary / operator end + C::HC # left boundary / operator start + D::HD # onsite / diagonal + δ::T +end + +function (op::WIIStep)((xᴬ, xᴮ, xᶜ, xᴰ)) + δ = op.δ + sqrtδ = sqrt(δ) + + @plansor yᴰ[-1 -2; -3 -4] := δ * xᴰ[-1 1; -3 -4] * + op.D[2 3; 1 4] * τ[-2 4; 2 3] + + @plansor yᶜ[-1 -2; -3 -4] := δ * xᶜ[-1 1; -3 -4] * + op.D[2 3; 1 4] * τ[-2 4; 2 3] + + sqrtδ * xᴰ[1 2; -3 4] * op.C[-1 -2; 3 -4] * τ[3 4; 1 2] + + @plansor yᴮ[-1 -2; -3 -4] := δ * xᴮ[-1 1; -3 -4] * + op.D[2 3; 1 4] * τ[-2 4; 2 3] + + sqrtδ * xᴰ[1 2; -3 4] * op.B[-1 -2; 3 -4] * τ[3 4; 1 2] + + @plansor yᴬ[-1 -2; -3 -4] := ( + δ * xᴬ[-1 1; -3 -4] * op.D[2 3; 1 4] * τ[-2 4; 2 3] + + xᴰ[1 2; -3 4] * op.A[-1 -2; 3 -4] * τ[3 4; 1 2] + ) + ( + sqrtδ * xᶜ[1 2; -3 -4] * op.B[-1 -2; 3 4] * τ[3 4; 1 2] + ) + (sqrtδ * xᴮ[-1 4; -3 3] * op.C[2 -2; 1 -4] * τ[3 4; 1 2]) + + return yᴬ, yᴮ, yᶜ, yᴰ end # Hack to treat FiniteMPOhamiltonians as Infinite From 3c89639419e50033633795301e971366dcb55836 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 2 May 2026 19:43:37 -0400 Subject: [PATCH 6/8] update benchmarks --- benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl b/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl index ce2bb4395..fdc71afc5 100644 --- a/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl +++ b/benchmark/MPSKitBenchmarks/timestep/TimestepBenchmarks.jl @@ -12,8 +12,6 @@ const HAMILTONIANS = ( "heis_xxx" => heisenberg_XXX(), ) -# Steady-state benchmark: time `make_time_mpo` after compilation has settled. -# Keep the matrix small — we want a stable baseline, not an exhaustive sweep. let g_top = addgroup!(SUITE, "make_time_mpo") for (Hname, H) in HAMILTONIANS g = addgroup!(g_top, Hname) @@ -21,7 +19,7 @@ let g_top = addgroup!(SUITE, "make_time_mpo") alg = TaylorCluster(; N) g["N=$N"] = @benchmarkable make_time_mpo($H, $dt, $alg) end - g["WI"] = @benchmarkable make_time_mpo($H, $dt, $(MPSKit.WI)) + g["WII"] = @benchmarkable make_time_mpo($H, $dt, $(MPSKit.WII())) end end From ccb159a6ba52ca267d23779fc2d8a29758cab7b3 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 2 May 2026 19:50:15 -0400 Subject: [PATCH 7/8] cleaning --- src/algorithms/timestep/taylorcluster.jl | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/src/algorithms/timestep/taylorcluster.jl b/src/algorithms/timestep/taylorcluster.jl index 81fb46520..c8945aabd 100644 --- a/src/algorithms/timestep/taylorcluster.jl +++ b/src/algorithms/timestep/taylorcluster.jl @@ -27,13 +27,14 @@ First order Taylor expansion for a time-evolution MPO. """ const WI = TaylorCluster(; N = 1, extension = false, compression = false) +# For type stability reasons, we add a function barrier here and dispatch on Val(N). +# `LinearIndices` and `CartesianIndex{N}` are otherwise abstract. +# However, to not incur a dynamic dispatch for typical cases (N ≤ 4), we manually union-split + function make_time_mpo( H::MPOHamiltonian, dt::Number, alg::TaylorCluster; tol = eps(real(scalartype(H))), imaginary_evolution::Bool = false ) - # Manual union-split for typical small N: emits a direct call into the specialised - # `_make_time_mpo(H, dt, ::Val{N}, ...)` instead of paying a dynamic dispatch on `Val(alg.N)`. - # Falls back to the generic path for N > 4. n = alg.N return Base.Cartesian.@nif( 4, @@ -43,8 +44,6 @@ function make_time_mpo( ) end -# Function-barrier body: `N` is now a type parameter, so every `ntuple(_, N)`, -# `LinearIndices`, and `CartesianIndex` downstream has a statically-known rank. function _make_time_mpo( H::MPOHamiltonian, dt::Number, ::Val{N}, extension::Bool, compression::Bool; @@ -65,19 +64,15 @@ function _make_time_mpo( return remove_orphans!(mpo; tol) end -# `H^N` unrolled with `Val(N)` so the result type of `H_n` is statically known. -# Mirrors `Base.power_by_squaring`'s tree (and copy-on-N=1 behaviour) so the -# bond dimensions match `H^N` exactly. +# `H^N` unrolled with `Val(N)`. +# Mirrors `Base.power_by_squaring`'s tree (and copy-on-N=1 behaviour) so the bond dimensions match `H^N` exactly. @inline _pow(H, ::Val{1}) = copy(H) @inline function _pow(H, ::Val{N}) where {N} half = _pow(H, Val(N ÷ 2)) return iseven(N) ? half * half : half * half * H end -# Stable partition: elements equal to `sentinel` first, others after, -# preserving relative order within each group. Tuple-in / tuple-out so the -# resulting `CartesianIndex` keeps a known rank — this replaces -# `CartesianIndex(sort(collect(t); by = (!=(sentinel)))...)`. +# Stable partition: elements equal to `sentinel` first, others after, preserving relative order within each group. @inline function _partition_first(t::NTuple{M, Int}, sentinel::Int) where {M} n_match = count(==(sentinel), t) return ntuple(Val(M)) do j From ecbb0b4e8e21e12bda367eaf21f0056bff8c9d01 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 2 May 2026 20:34:09 -0400 Subject: [PATCH 8/8] some more cleaning --- src/algorithms/timestep/taylorcluster.jl | 41 ++++++------------------ 1 file changed, 10 insertions(+), 31 deletions(-) diff --git a/src/algorithms/timestep/taylorcluster.jl b/src/algorithms/timestep/taylorcluster.jl index c8945aabd..9e5943288 100644 --- a/src/algorithms/timestep/taylorcluster.jl +++ b/src/algorithms/timestep/taylorcluster.jl @@ -29,29 +29,27 @@ const WI = TaylorCluster(; N = 1, extension = false, compression = false) # For type stability reasons, we add a function barrier here and dispatch on Val(N). # `LinearIndices` and `CartesianIndex{N}` are otherwise abstract. -# However, to not incur a dynamic dispatch for typical cases (N ≤ 4), we manually union-split function make_time_mpo( H::MPOHamiltonian, dt::Number, alg::TaylorCluster; tol = eps(real(scalartype(H))), imaginary_evolution::Bool = false ) - n = alg.N - return Base.Cartesian.@nif( - 4, - d -> n == d, - d -> _make_time_mpo(H, dt, Val(d), alg.extension, alg.compression; tol, imaginary_evolution), - d -> _make_time_mpo(H, dt, Val(n), alg.extension, alg.compression; tol, imaginary_evolution), - ) + τ = imaginary_evolution ? -dt : -1im * dt + return _make_time_mpo(H, τ, Val(alg.N), alg.extension, alg.compression; tol, imaginary_evolution) end function _make_time_mpo( - H::MPOHamiltonian, dt::Number, ::Val{N}, + H::MPOHamiltonian, τ::Number, ::Val{N}, extension::Bool, compression::Bool; tol, imaginary_evolution::Bool ) where {N} - τ = imaginary_evolution ? -dt : -1im * dt - - H_n, virtual_sz, linds = _taylor_setup(H, Val(N)) + H_n = H^N + virtual_sz = map(0:length(H)) do i + return i == 0 ? size(H[1], 1) : size(H[i], 4) + end + linds = map(virtual_sz) do sz + return LinearIndices(ntuple(Returns(sz), Val(N))) + end extension && _taylor_extension!(H_n, H, virtual_sz, linds, Val(N), τ) @@ -64,14 +62,6 @@ function _make_time_mpo( return remove_orphans!(mpo; tol) end -# `H^N` unrolled with `Val(N)`. -# Mirrors `Base.power_by_squaring`'s tree (and copy-on-N=1 behaviour) so the bond dimensions match `H^N` exactly. -@inline _pow(H, ::Val{1}) = copy(H) -@inline function _pow(H, ::Val{N}) where {N} - half = _pow(H, Val(N ÷ 2)) - return iseven(N) ? half * half : half * half * H -end - # Stable partition: elements equal to `sentinel` first, others after, preserving relative order within each group. @inline function _partition_first(t::NTuple{M, Int}, sentinel::Int) where {M} n_match = count(==(sentinel), t) @@ -95,17 +85,6 @@ end return 0 end -function _taylor_setup(H::MPOHamiltonian, ::Val{N}) where {N} - H_n = _pow(H, Val(N)) - virtual_sz = map(0:length(H)) do i - return i == 0 ? size(H[1], 1) : size(H[i], 4) - end - linds = map(virtual_sz) do sz - return LinearIndices(ntuple(Returns(sz), Val(N))) - end - return H_n, virtual_sz, linds -end - # Algorithm 3 of Van Damme et al. (2024): incorporate higher-order corrections # from `H_next = H_n * H` into `H_n` in place. function _taylor_extension!(H_n, H, virtual_sz, linds, ::Val{N}, τ::Number) where {N}