start setting up benchmark suite

Author: lkdvos

benchmark/MPSKitBenchmarks/MPSKitBenchmarks.jl

@@ -14,9 +14,11 @@ BenchmarkTools.DEFAULT_PARAMETERS.memory_tolerance = 0.01
 const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json")
 const SUITE = BenchmarkGroup()
 const MODULES = Dict{String, Symbol}(
-    "derivatives" => :Derivatives
+    "derivatives" => :DerivativesBenchmarks
 )
 
+include("derivatives/DerivativesBenchmarks.jl")
+
 load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...)
 
 function load!(group::BenchmarkGroup, id::AbstractString; tune::Bool = false)

benchmark/MPSKitBenchmarks/derivatives/AC2_benchmarks.jl

@@ -0,0 +1,85 @@
+struct AC2Spec{S <: ElementarySpace} # <: BenchmarkSpec
+    physicalspaces::NTuple{2, S}
+    mps_virtualspaces::NTuple{3, S}
+    mpo_virtualspaces::NTuple{3, SumSpace{S}}
+    nonzero_keys::NTuple{2, Vector{Tuple{Int, Int}}}
+end
+
+function AC2Spec(mps, mpo; site = length(mps) ÷ 2)
+    physicalspaces = (physicalspace(mps, site), physicalspace(mps, site + 1))
+    mps_virtualspaces = (left_virtualspace(mps, site), right_virtualspace(mps, site), right_virtualspace(mps, site + 1))
+    mpo_virtualspaces = (left_virtualspace(mpo, site), right_virtualspace(mpo, site), right_virtualspace(mpo, site + 1))
+    ks = (
+        map(x -> (x.I[1], x.I[4]), nonzero_keys(mpo[site])),
+        map(x -> (x.I[1], x.I[4]), nonzero_keys(mpo[site])),
+    )
+    return AC2Spec(physicalspaces, mps_virtualspaces, mpo_virtualspaces, ks)
+end
+
+# Benchmarks
+# ----------
+function MPSKit.MPO_AC2_Hamiltonian(spec::AC2Spec{S}; T::Type = Float64) where {S}
+    GL = randn(T, spec.mps_virtualspaces[1] ⊗ spec.mpo_virtualspaces[1]' ← spec.mps_virtualspaces[1])
+    GR = randn(T, spec.mps_virtualspaces[3] ⊗ spec.mpo_virtualspaces[3] ← spec.mps_virtualspaces[3])
+    W1 = JordanMPOTensor{T, S}(undef, spec.mpo_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.physicalspaces[1] ⊗ spec.mpo_virtualspaces[2])
+    for (r, c) in spec.nonzero_keys[1]
+        r == c == 1 && continue
+        r == size(W1, 1) && c == size(W1, 4) && continue
+        W1[r, 1, 1, c] = randn!(W1[r, 1, 1, c])
+    end
+    W2 = JordanMPOTensor{T, S}(undef, spec.mpo_virtualspaces[2] ⊗ spec.physicalspaces[2] ← spec.physicalspaces[2] ⊗ spec.mpo_virtualspaces[3])
+    for (r, c) in spec.nonzero_keys[2]
+        r == c == 1 && continue
+        r == size(W2, 1) && c == size(W2, 4) && continue
+        W2[r, 1, 1, c] = randn!(W2[r, 1, 1, c])
+    end
+
+    return MPSKit.MPO_AC2_Hamiltonian(GL, W1, W2, GR)
+end
+
+function contraction_benchmark(spec::AC2Spec; T::Type = Float64)
+    AA = randn(T, spec.mps_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.mps_virtualspaces[3] ⊗ spec.physicalspaces[2]')
+    H_eff = MPSKit.MPO_AC2_Hamiltonian(spec; T)
+    H_prep, x_prep = MPSKit.prepare_operator!!(H_eff, AA)
+    init() = randn!(similar(x_prep))
+
+    return @benchmarkable $H_prep * x setup = (x = $init())
+end
+
+function preparation_benchmark(spec::AC2Spec; T::Type = Float64)
+    init() = randn(T, spec.mps_virtualspaces[1] ⊗ spec.physicalspaces[1] ← spec.mps_virtualspaces[3] ⊗ spec.physicalspaces[2]')
+    H_eff = MPSKit.MPO_AC2_Hamiltonian(spec; T)
+
+    return @benchmarkable begin
+        O′, x′ = MPSKit.prepare_operator!!($H_eff, x)
+        y = MPSKit.unprepare_operator!!(x′, O′, x)
+    end setup = (x = $init())
+end
+
+# Converters
+# ----------
+function tomlify(spec::AC2Spec)
+    return Dict(
+        "physicalspaces" => collect(tomlify.(spec.physicalspaces)),
+        "mps_virtualspaces" => collect(tomlify.(spec.mps_virtualspaces)),
+        "mpo_virtualspaces" => collect(tomlify.(spec.mpo_virtualspaces)),
+        "nonzero_keys" => collect(map(Base.Fix1(map, collect), spec.nonzero_keys))
+    )
+end
+
+function untomlify(::Type{AC2Spec}, x)
+    physicalspaces = Tuple(map(untomlify, x["physicalspaces"]))
+    mps_virtualspaces = Tuple(map(untomlify, x["mps_virtualspaces"]))
+    mpo_virtualspaces = Tuple(map(untomlify, x["mpo_virtualspaces"]))
+    nonzero_keys = Tuple(map(Base.Fix1(map, Base.Fix1(map, Tuple)), x["nonzero_keys"]))
+    return AC2Spec(physicalspaces, mps_virtualspaces, mpo_virtualspaces, nonzero_keys)
+end
+
+function Base.convert(::Type{AC2Spec{S₁}}, spec::AC2Spec{S₂}) where {S₁, S₂}
+    return S₁ === S₂ ? spec : AC2Spec(
+            S₁.(spec.physicalspaces),
+            S₁.(spec.mps_virtualspaces),
+            SumSpace{S₁}.(spec.mpo_virtualspaces),
+            spec.nonzero_keys
+        )
+end

benchmark/MPSKitBenchmarks/derivatives/DerivativesBenchmarks.jl

@@ -0,0 +1,17 @@
+module DerivativesBenchmarks
+
+export AC2Spec
+
+using BenchmarkTools
+using TensorKit
+using BlockTensorKit
+using MPSKit
+using ..BenchUtils
+import ..BenchUtils: tomlify, untomlify
+
+const SUITE = BenchmarkGroup()
+
+include("AC2_benchmarks.jl")
+
+
+end

benchmark/MPSKitBenchmarks/utils/BenchUtils.jl

@@ -1,3 +1,27 @@
 module BenchUtils
 
+export tomlify, untomlify
+
+using TensorKit
+using BlockTensorKit
+using TOML
+
+tomlify(x::VectorSpace) = string(x)
+untomlify(::Type{<:VectorSpace}, x) = eval(Meta.parse(x))
+
+
+# Type piracy but oh well
+TensorKit.ComplexSpace(V::ElementarySpace) = ComplexSpace(dim(V), isdual(V))
+
+function TensorKit.U1Space(V::SU2Space)
+    dims = TensorKit.SectorDict{U1Irrep, Int}()
+    for c in sectors(V), m in (-c.j):(c.j)
+        u1 = U1Irrep(m)
+        dims[u1] = get(dims, u1, 0) + 1
+    end
+    return U1Space(dims; dual = isdual(V))
+end
+
+BlockTensorKit.SumSpace{S}(V::SumSpace) where {S} = SumSpace(map(S, V.spaces))
+
 end

src/MPSKit.jl

@@ -65,6 +65,7 @@ using MatrixAlgebraKit: TruncationStrategy, PolarViaSVD, LAPACK_SVDAlgorithm
 using BlockTensorKit
 using BlockTensorKit: TensorMapSumSpace
 using TensorOperations
+using TensorOperations: AbstractBackend, DefaultBackend, DefaultAllocator
 using KrylovKit
 using KrylovKit: KrylovAlgorithm
 using OptimKit

src/algorithms/derivatives/derivatives.jl

@@ -213,3 +213,31 @@ const DerivativeOrMultiplied{D <: DerivativeOperator} = Union{MultipliedOperator
 (x::LazySum{<:DerivativeOrMultiplied})(y, t::Number) = sum(O -> O(y, t), x)
 (x::LazySum{<:DerivativeOrMultiplied})(y) = sum(O -> O(y), x)
 Base.:*(h::LazySum{<:Union{DerivativeOrMultiplied}}, v) = h(v)
+
+# Operator preparation
+# --------------------
+"""
+    prepare_operator!!(O, x, [backend], [allocator]) -> O′, x′
+
+Given an operator and vector, try to construct a more efficient representation of that operator for repeated application.
+This should always be used in conjunction with [`unprepare_operator!!`](@ref).
+"""
+function prepare_operator!!(
+        O, x,
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
+    )
+    return O, x
+end
+
+"""
+    unprepare_operator!!(y, O, x, [backend], [allocator]) -> y′
+
+Given the result `y` of applying a prepared operator `O` on vectors like `x`, undo the preparation work on the vector, and clean up the operator.
+This should always be used in conjunction with [`prepare_operator!!`](@ref).
+"""
+function unprepare_operator!!(
+        y, O, x,
+        backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()
+    )
+    return y
+end

src/algorithms/derivatives/hamiltonian_derivatives.jl

@@ -75,13 +75,13 @@ const _HAM_MPS_TYPES = Union{
     InfiniteMPS{<:MPSTensor},
 }
 
-function AC_hamiltonian(
-        site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:JordanMPOTensor},
-        above::_HAM_MPS_TYPES, envs
+function prepare_operator!!(
+        O::MPO_AC_Hamiltonian{<:Any, <:JordanMPOTensor, <:Any}, x::MPSTensor,
+        backend::AbstractBackend, allocator
     )
-    GL = leftenv(envs, site, below)
-    GR = rightenv(envs, site, below)
-    W = operator[site]
+    GL = O.leftenv
+    W = only(O.operators)
+    GR = O.rightenv
 
     # starting
     if nonzero_length(W.C) > 0
@@ -117,7 +117,7 @@ function AC_hamiltonian(
     end
 
     # not_started
-    if isfinite(operator) && site == length(operator)
+    if size(W, 4) == 1
         not_started = missing
     elseif !ismissing(starting)
         I = id(storagetype(GR[1]), physicalspace(W))
@@ -128,7 +128,7 @@ function AC_hamiltonian(
     end
 
     # finished
-    if isfinite(operator) && site == 1
+    if size(W, 1) == 1
         finished = missing
     elseif !ismissing(ending)
         I = id(storagetype(GL[end]), physicalspace(W))
@@ -142,19 +142,20 @@ function AC_hamiltonian(
     A = W.A
     continuing = (GL[2:(end - 1)], A, GR[2:(end - 1)])
 
-    return JordanMPO_AC_Hamiltonian(
+    O′ = JordanMPO_AC_Hamiltonian(
         onsite, not_started, finished, starting, ending, continuing
     )
+
+    return O′, x
 end
 
-function AC2_hamiltonian(
-        site::Int, below::_HAM_MPS_TYPES, operator::MPOHamiltonian{<:JordanMPOTensor},
-        above::_HAM_MPS_TYPES, envs
+function prepare_operator!!(
+        O::MPO_AC2_Hamiltonian{<:Any, <:JordanMPOTensor, <:JordanMPOTensor, <:Any}, x::MPOTensor,
+        backend::AbstractBackend, allocator
     )
-    GL = leftenv(envs, site, below)
-    GR = rightenv(envs, site + 1, below)
-    W1 = operator[site]
-    W2 = operator[site + 1]
+    GL = O.leftenv
+    W1, W2 = O.operators
+    GR = O.rightenv
 
     # starting left - continuing right
     if nonzero_length(W1.C) > 0 && nonzero_length(W2.A) > 0
@@ -256,7 +257,7 @@ function AC2_hamiltonian(
     end
 
     # finished
-    if isfinite(operator) && site + 1 == length(operator)
+    if size(W2, 4) == 1
         II = missing
     elseif !ismissing(IC)
         I = id(storagetype(GR[1]), physicalspace(W2))
@@ -271,7 +272,7 @@ function AC2_hamiltonian(
     end
 
     # unstarted
-    if isfinite(operator) && site == 1
+    if size(W1, 1) == 1
         EE = missing
     elseif !ismissing(BE)
         I = id(storagetype(GL[end]), physicalspace(W1))
@@ -289,7 +290,8 @@ function AC2_hamiltonian(
     # TODO: MPODerivativeOperator code reuse + optimization
     AA = (GL[2:(end - 1)], W1.A, W2.A, GR[2:(end - 1)])
 
-    return JordanMPO_AC2_Hamiltonian(II, IC, ID, CB, CA, AB, AA, BE, DE, EE)
+    O′ = JordanMPO_AC2_Hamiltonian(II, IC, ID, CB, CA, AB, AA, BE, DE, EE)
+    return O′, x
 end
 
 # Actions

src/algorithms/fixedpoint.jl

@@ -8,26 +8,30 @@ Compute the fixedpoint of a linear operator `A` using the specified eigensolver
 fixedpoint is assumed to be unique.
 """
 function fixedpoint(A, x₀, which::Symbol, alg::Lanczos)
-    vals, vecs, info = eigsolve(A, x₀, 1, which, alg)
+    A′, x₀′ = prepare_operator!!(A, x₀)
+    vals, vecs, info = eigsolve(A′, x₀′, 1, which, alg)
 
-    if info.converged == 0
+    info.converged == 0 &&
         @warnv 1 "fixedpoint not converged after $(info.numiter) iterations: normres = $(info.normres[1])"
-    end
 
-    return vals[1], vecs[1]
+    λ = vals[1]
+    v = unprepare_operator!!(vecs[1], A′, x₀)
+
+    return λ, v
 end
 
 function fixedpoint(A, x₀, which::Symbol, alg::Arnoldi)
-    TT, vecs, vals, info = schursolve(A, x₀, 1, which, alg)
+    A′, x₀′ = prepare_operator!!(A, x₀)
+    TT, vecs, vals, info = schursolve(A′, x₀′, 1, which, alg)
 
-    if info.converged == 0
+    info.converged == 0 &&
         @warnv 1 "fixedpoint not converged after $(info.numiter) iterations: normres = $(info.normres[1])"
-    end
-    if size(TT, 2) > 1 && TT[2, 1] != 0
-        @warnv 1 "non-unique fixedpoint detected"
-    end
+    size(TT, 2) > 1 && TT[2, 1] != 0 && @warnv 1 "non-unique fixedpoint detected"
+
+    λ = vals[1]
+    v = unprepare_operator!!(vecs[1], A′, x₀)
 
-    return vals[1], vecs[1]
+    return λ, v
 end
 
 function fixedpoint(A, x₀, which::Symbol; kwargs...)