Add a CTMRG benchmark suite

Project.toml

@@ -4,7 +4,7 @@ version = "0.7.0"
 authors = ["Paul Brehmer", "Lander Burgelman", "Zhengyuan Yue", "Lukas Devos <ldevos98@gmail.com>"]
 
 [workspace]
-projects = ["test", "docs"]
+projects = ["test", "docs", "benchmark"]
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

benchmark/PEPSKitBenchmarks/PEPSKitBenchmarks.jl

@@ -0,0 +1,56 @@
+module PEPSKitBenchmarks
+
+using BenchmarkTools
+using PEPSKit
+using TOML
+
+include("utils/BenchUtils.jl")
+
+BenchmarkTools.DEFAULT_PARAMETERS.seconds = 60.0
+BenchmarkTools.DEFAULT_PARAMETERS.samples = 5
+BenchmarkTools.DEFAULT_PARAMETERS.evals = 1
+
+const PARAMS_PATH = joinpath(@__DIR__, "etc", "params.json")
+const SUITE = BenchmarkGroup()
+const MODULES = Dict{String, Symbol}(
+    "ctmrg" => :CTMRGBenchmarks,
+)
+
+include("ctmrg/CTMRGBenchmarks.jl")
+
+load!(id::AbstractString; kwargs...) = load!(SUITE, id; kwargs...)
+
+function load!(group::BenchmarkGroup, id::AbstractString; tune::Bool = false)
+    modsym = MODULES[id]
+    mod = Core.eval(@__MODULE__, modsym)
+    modsuite = @invokelatest getglobal(mod, :SUITE)
+    group[id] = modsuite
+    if tune
+        results = BenchmarkTools.load(PARAMS_PATH)[1]
+        haskey(results, id) && loadparams!(modsuite, results[id], :evals)
+    end
+    return group
+end
+
+loadall!(; kwargs...) = loadall!(SUITE; kwargs...)
+
+function loadall!(group::BenchmarkGroup; verbose::Bool = true, tune::Bool = false)
+    for id in keys(MODULES)
+        if verbose
+            print("loading group $(repr(id))... ")
+            time = @elapsed load!(group, id; tune = false)
+            println("done (took $time seconds)")
+        else
+            load!(group, id; tune = false)
+        end
+    end
+    if tune
+        results = BenchmarkTools.load(PARAMS_PATH)[1]
+        for (id, suite) in group
+            haskey(results, id) && loadparams!(suite, results[id], :evals)
+        end
+    end
+    return group
+end
+
+end # module

benchmark/PEPSKitBenchmarks/ctmrg/CTMRGBenchmarks.jl

@@ -0,0 +1,42 @@
+module CTMRGBenchmarks
+
+export CTMRGSpec
+
+using BenchmarkTools
+using TOML
+using TensorKit
+using PEPSKit
+using ..BenchUtils
+import ..BenchUtils: tomlify, untomlify
+
+const SUITE = BenchmarkGroup()
+
+include("ctmrg_iteration_benchmarks.jl")
+
+const allparams = Dict(
+    "default" => TOML.parsefile(joinpath(@__DIR__, "default.toml")),
+    "su3_hubbard" => TOML.parsefile(joinpath(@__DIR__, "su3_hubbard.toml")),
+)
+
+const RESERVED_KEYS = ("unitcell", "algorithms")
+
+T = ComplexF64
+
+for (scenario, params) in allparams
+    unitcell = (Int(params["unitcell"][1]), Int(params["unitcell"][2]))
+    algs = [untomlify(PEPSKit.CTMRGAlgorithm, a) for a in params["algorithms"]]
+    g_scen = addgroup!(SUITE, scenario)
+    for alg in algs
+        g_alg = addgroup!(g_scen, algname(alg))
+        for (sym_name, specs) in params
+            sym_name in RESERVED_KEYS && continue
+            g_sym = addgroup!(g_alg, sym_name)
+            for spec_dict in specs
+                spec = untomlify(CTMRGSpec, spec_dict; unitcell)
+                g_sym[benchname(spec)] = ctmrg_iteration_benchmark(spec, alg; T)
+            end
+        end
+    end
+end
+
+end # module

benchmark/PEPSKitBenchmarks/ctmrg/ctmrg_iteration_benchmarks.jl

@@ -0,0 +1,124 @@
+struct CTMRGSpec{S <: ElementarySpace}
+    Pspaces::Matrix{S}
+    Nspaces::Matrix{S}
+    Espaces::Matrix{S}
+    chi_north::Matrix{S}
+    chi_east::Matrix{S}
+    chi_south::Matrix{S}
+    chi_west::Matrix{S}
+end
+
+unitcell(spec::CTMRGSpec) = size(spec.Pspaces)
+
+function benchname(spec::CTMRGSpec)
+    Dmin = min(minimum(dim, spec.Nspaces), minimum(dim, spec.Espaces))
+    chimin = min(
+        minimum(dim, spec.chi_north), minimum(dim, spec.chi_east),
+        minimum(dim, spec.chi_south), minimum(dim, spec.chi_west)
+    )
+    return "D$(Dmin)_chi$(chimin)"
+end
+
+algname(alg::PEPSKit.CTMRGAlgorithm) =
+    "$(typeof(alg).name.name)_$(typeof(alg.projector_alg).name.name)"
+
+function setup_problem(spec::CTMRGSpec; T::Type = ComplexF64)
+    peps = InfinitePEPS(randn, T, spec.Pspaces, spec.Nspaces, spec.Espaces)
+    env = CTMRGEnv(
+        randn, T, peps,
+        spec.chi_north, spec.chi_east, spec.chi_south, spec.chi_west,
+    )
+    network = InfiniteSquareNetwork(peps)
+    return network, env
+end
+
+function ctmrg_iteration_benchmark(spec::CTMRGSpec, alg; T::Type = ComplexF64)
+    network, env = setup_problem(spec; T)
+    return @benchmarkable PEPSKit.ctmrg_iteration($network, $env, $alg)
+end
+
+# Convert a TOML value (scalar string or 2D array of strings) to a Matrix sized to `unitcell`.
+function _to_space_matrix(val::AbstractString, unitcell::Tuple{Int, Int}, key::AbstractString)
+    V = untomlify(VectorSpace, val)
+    return fill(V, unitcell)
+end
+function _to_space_matrix(val::AbstractVector, unitcell::Tuple{Int, Int}, key::AbstractString)
+    rows, cols = unitcell
+    length(val) == rows || throw(
+        ArgumentError(
+            "Field `$key` has $(length(val)) rows but unitcell expects $rows."
+        )
+    )
+    for (r, row) in enumerate(val)
+        length(row) == cols || throw(
+            ArgumentError(
+                "Field `$key` row $r has $(length(row)) cols but unitcell expects $cols."
+            )
+        )
+    end
+    return [untomlify(VectorSpace, val[r][c]) for r in 1:rows, c in 1:cols]
+end
+
+function _get_matrix(d, key, unitcell)
+    haskey(d, key) || throw(ArgumentError("Spec is missing required key `$key`."))
+    return _to_space_matrix(d[key], unitcell, key)
+end
+
+function tomlify(spec::CTMRGSpec)
+    return Dict(
+        "Pspaces" => _matrix_tomlify(spec.Pspaces),
+        "Nspaces" => _matrix_tomlify(spec.Nspaces),
+        "Espaces" => _matrix_tomlify(spec.Espaces),
+        "chi_north" => _matrix_tomlify(spec.chi_north),
+        "chi_east" => _matrix_tomlify(spec.chi_east),
+        "chi_south" => _matrix_tomlify(spec.chi_south),
+        "chi_west" => _matrix_tomlify(spec.chi_west),
+    )
+end
+
+_matrix_tomlify(M::AbstractMatrix{<:ElementarySpace}) =
+    [[tomlify(M[r, c]) for c in axes(M, 2)] for r in axes(M, 1)]
+
+function untomlify(::Type{CTMRGSpec}, d; unitcell::Tuple{Int, Int})
+    Pspaces = _get_matrix(d, "Pspaces", unitcell)
+    Nspaces = _get_matrix(d, "Nspaces", unitcell)
+    Espaces = _get_matrix(d, "Espaces", unitcell)
+
+    chi_north = _get_matrix(d, "chi_north", unitcell)
+    chi_east = _get_matrix(d, "chi_east", unitcell)
+    chi_south = _get_matrix(d, "chi_south", unitcell)
+    chi_west = _get_matrix(d, "chi_west", unitcell)
+
+    return CTMRGSpec(
+        Pspaces, Nspaces, Espaces,
+        chi_north, chi_east, chi_south, chi_west,
+    )
+end
+
+# One tomlify method per concrete CTMRG algorithm settings type. Each produces a
+# uniform two-key Dict that round-trips via `untomlify(PEPSKit.CTMRGAlgorithm, d)`.
+
+function tomlify(alg::SequentialCTMRG)
+    return Dict(
+        "type" => "SequentialCTMRG",
+        "projector_alg" => string(typeof(alg.projector_alg).name.name),
+    )
+end
+
+function tomlify(alg::SimultaneousCTMRG)
+    return Dict(
+        "type" => "SimultaneousCTMRG",
+        "projector_alg" => string(typeof(alg.projector_alg).name.name),
+    )
+end
+
+function untomlify(::Type{<:PEPSKit.CTMRGAlgorithm}, d)
+    t = d["type"]
+    pa = Symbol(d["projector_alg"])
+    if t == "SequentialCTMRG"
+        return SequentialCTMRG(; projector_alg = pa)
+    elseif t == "SimultaneousCTMRG"
+        return SimultaneousCTMRG(; projector_alg = pa)
+    end
+    throw(ArgumentError("Unknown CTMRG algorithm type: $(t)"))
+end

benchmark/PEPSKitBenchmarks/ctmrg/default.toml

@@ -0,0 +1,50 @@
+unitcell = [1, 1]
+
+[[algorithms]]
+type = "SimultaneousCTMRG"
+projector_alg = "HalfInfiniteProjector"
+
+[[Trivial]]
+Pspaces = "ℂ^2"
+Nspaces = "ℂ^2"
+Espaces = "ℂ^2"
+chi_north = "ℂ^8"
+chi_east = "ℂ^8"
+chi_south = "ℂ^8"
+chi_west = "ℂ^8"
+
+[[Trivial]]
+Pspaces = "ℂ^2"
+Nspaces = "ℂ^2"
+Espaces = "ℂ^2"
+chi_north = "ℂ^12"
+chi_east = "ℂ^12"
+chi_south = "ℂ^12"
+chi_west = "ℂ^12"
+
+[[Trivial]]
+Pspaces = "ℂ^2"
+Nspaces = "ℂ^3"
+Espaces = "ℂ^3"
+chi_north = "ℂ^12"
+chi_east = "ℂ^12"
+chi_south = "ℂ^12"
+chi_west = "ℂ^12"
+
+[[Trivial]]
+Pspaces = "ℂ^2"
+Nspaces = "ℂ^3"
+Espaces = "ℂ^3"
+chi_north = "ℂ^18"
+chi_east = "ℂ^18"
+chi_south = "ℂ^18"
+chi_west = "ℂ^18"
+
+[[Trivial]]
+Pspaces = "ℂ^2"
+Nspaces = "ℂ^4"
+Espaces = "ℂ^4"
+chi_north = "ℂ^16"
+chi_east = "ℂ^16"
+chi_south = "ℂ^16"
+chi_west = "ℂ^16"