Should `JordanMPOTensor` be reworked?

[kshyatt]

This grew out of a Zulip discussion with @lkdvos and I thought I should create an issue in order to not lose track of it and get feedback from others.

Right now, JordanMPOTensor wraps a SparseBlockTensorMap which itself has a Union{BraidingTensor, TensorMap} for its element type. This leads to lots of annoying storage type inference problems and many other sad situations.

It might rather be better to simply store either:

• two separate SparseBlockTensorMaps
• or one for the TensorMaps and just an array of scalars (this was Lukas' suggestion and I am too group-retreat drained to comment much on what it would look like)

[lkdvos]

Some historic context: before the whole BlockTensorKit machinery existed, we actually stored an Array{Union{TensorMap, Number}, 2} and only inserted/instantiated the BraidingTensors where appropriate.
I think the Union{TensorMap, BraidingTensor} was a reasonable attempt, but given the storagetype inference misery this is clearly not that compatible with GPU support.
A secondary reason to go back to storing scalars is that we did actually lose some expressivity: the BraidingTensor does not hold on to a scalar, so for exponential MPOs we have to instantiate an actual tensor anyways.
Ultimately, it could even be reasonable to ask if this optimization is really worth it at all, but the issue is that we still need to be able to detect when to do the renormalization parts, which e.g. for variance is not only the "first/last level"

I think in the long run it might be reasonable to really just have something like this:

struct JordanMPOTensor
    A
    B
    C
    D
    scalars::Matrix{<:Number}
end

and then see if we can make that behave as it should.

[lkdvos]

I did some minor experiments today with the existing benchmarks, where I replaced the BraidingTensor entries with dense ones to see if we are actually gaining anything. Specifically for the cylinder tests, I locally get the following results:

Symmetry	mps_dim	sparse	dense	ratio (dense/sparse)
Trivial	3	0.34 ms	0.31 ms	0.92×
Trivial	16	0.58 ms	0.54 ms	0.94×
Trivial	43	2.24 ms	2.20 ms	0.98×
Trivial	110	12.7 ms	12.6 ms	0.99×
Trivial	269	79.8 ms	79.1 ms	0.99×
Trivial	668	774 ms	822 ms	1.06× ⚠
Irrep[U₁]	3	0.58 ms	0.31 ms	0.53×
Irrep[U₁]	16	0.94 ms	0.62 ms	0.66×
Irrep[U₁]	43	1.30 ms	0.98 ms	0.75×
Irrep[U₁]	110	2.87 ms	2.58 ms	0.90×
Irrep[U₁]	269	13.3 ms	13.1 ms	0.98×
Irrep[U₁]	668	78.3 ms	77.8 ms	0.99×
Irrep[U₁]	1646	614 ms	579 ms	0.94×
Irrep[U₁]	4059	6.22 s	6.16 s	0.99×
Irrep[SU₂]	3	0.46 ms	0.23 ms	0.49×
Irrep[SU₂]	16	0.68 ms	0.41 ms	0.60×
Irrep[SU₂]	43	0.85 ms	0.56 ms	0.66×
Irrep[SU₂]	110	1.08 ms	0.79 ms	0.73×
Irrep[SU₂]	269	1.92 ms	1.61 ms	0.84×
Irrep[SU₂]	668	6.32 ms	5.92 ms	0.94×
Irrep[SU₂]	1646	32.0 ms	31.3 ms	0.98×
Irrep[SU₂]	4059	200 ms	202 ms	1.01×
Irrep[SU₂]	9999	1.74 s	1.75 s	1.01×

Interesting here is that at large bond dimensions, there is no difference, while at small bond dimensions there is even quite a few cases where using the BraidingTensor is actually slowing us down. Considering how the current BlockTensorKit is structured, this is actually not all that surprising, I don't think the new contraction routines ever get to use these specializations. With this in mind, I would actually be in favor of removing the BraidingTensor union altogether, and trying to fix the variance through a different mechanism.

benchmark script:

Details

using MPSKit
using TensorKit
using BlockTensorKit
using BenchmarkTools
using TOML
using LinearAlgebra: norm
using Printf
using TensorKit: BraidingTensor
using BlockTensorKit: SparseBlockTensorMap, undef_blocks

include(joinpath(@__DIR__, "MPSKitBenchmarks", "utils", "BenchUtils.jl"))
using .BenchUtils

const CYL_TOML = joinpath(
    @__DIR__, "MPSKitBenchmarks", "derivatives", "heisenberg_cylinder.toml"
)

# --- spec parsing ------------------------------------------------------------

struct Spec
    Vphys::Vector
    Vmps::Vector
    Vmpo::Vector
    keys::Vector{Tuple{Int, Int}}
end

function parse_spec(d)
    parse_V(s) = eval(Meta.parse(s))
    Vphys = parse_V.(d["physicalspaces"])
    Vmps = parse_V.(d["mps_virtualspaces"])
    Vmpo = parse_V.(d["mpo_virtualspaces"])
    keys1 = [Tuple(k) for k in d["nonzero_keys"][1]]
    return Spec(Vphys, Vmps, Vmpo, keys1)
end

# --- operator builders -------------------------------------------------------

function build_random_op_data(::Type{T}, spec::Spec, identity_keys) where {T}
    sp = spec.Vmpo[1] ⊗ spec.Vphys[1] ← spec.Vphys[1] ⊗ spec.Vmpo[2]
    S = spacetype(spec.Vphys[1])
    TM = TensorMap{T, S, 2, 2, Vector{T}}
    proto = SparseBlockTensorMap{TM}(undef_blocks, sp)
    data = Dict{Tuple{Int, Int}, TM}()
    for (r, c) in spec.keys
        (r, c) in identity_keys && continue
        data[(r, c)] = randn(T, eachspace(proto)[r, 1, 1, c])
    end
    return data
end

function build_W_sparse(::Type{T}, spec::Spec, identity_keys, op_data) where {T}
    sp = spec.Vmpo[1] ⊗ spec.Vphys[1] ← spec.Vphys[1] ⊗ spec.Vmpo[2]
    S = spacetype(spec.Vphys[1])
    TM = TensorMap{T, S, 2, 2, Vector{T}}
    Tτ = BraidingTensor{T, S}
    W = SparseBlockTensorMap{Union{TM, Tτ}}(undef_blocks, sp)
    for (r, c) in spec.keys
        V_block = eachspace(W)[r, 1, 1, c]
        if (r, c) in identity_keys
            W[r, 1, 1, c] = BraidingTensor{T}(V_block)
        else
            W[r, 1, 1, c] = op_data[(r, c)]
        end
    end
    return W
end

function build_W_dense(::Type{T}, spec::Spec, identity_keys, op_data) where {T}
    sp = spec.Vmpo[1] ⊗ spec.Vphys[1] ← spec.Vphys[1] ⊗ spec.Vmpo[2]
    S = spacetype(spec.Vphys[1])
    TM = TensorMap{T, S, 2, 2, Vector{T}}
    W = SparseBlockTensorMap{TM}(undef_blocks, sp)
    for (r, c) in spec.keys
        V_block = eachspace(W)[r, 1, 1, c]
        if (r, c) in identity_keys
            τ = BraidingTensor{T}(V_block)
            W[r, 1, 1, c] = densify(τ)
        else
            W[r, 1, 1, c] = op_data[(r, c)]
        end
    end
    return W
end

# materialize a BraidingTensor as a dense TensorMap with the same planar action
densify(τ::BraidingTensor) = convert(TensorMap, τ)

# --- environment builders ----------------------------------------------------

function build_envs(::Type{T}, spec::Spec) where {T}
    GL = randn(T,spec.Vmps[1] ⊗ spec.Vmpo[1]' ← spec.Vmps[1])
    GR = randn(T,spec.Vmps[2] ⊗ spec.Vmpo[2] ← spec.Vmps[2])
    return GL, GR
end

# --- driver ------------------------------------------------------------------

function bench_one(spec::Spec; T = Float64, secs = 3.0)
    identity_keys = [(r, c) for (r, c) in spec.keys if r == c]
    op_data = build_random_op_data(T, spec, identity_keys)
    W_sparse = build_W_sparse(T, spec, identity_keys, op_data)
    W_dense  = build_W_dense(T, spec, identity_keys, op_data)

    GL, GR = build_envs(T, spec)
    H_sparse = MPSKit.MPO_AC_Hamiltonian(GL, W_sparse, GR)
    H_dense  = MPSKit.MPO_AC_Hamiltonian(GL, W_dense, GR)

    x = randn(T,spec.Vmps[1] ⊗ spec.Vphys[1] ← spec.Vmps[2])

    # correctness check
    y_sparse = H_sparse(x)
    y_dense  = H_dense(x)
    rel_err  = norm(y_sparse - y_dense) / norm(y_sparse)

    b_sparse = @benchmarkable $H_sparse($x)
    b_dense  = @benchmarkable $H_dense($x)
    t_sparse = run(b_sparse; seconds = secs)
    t_dense  = run(b_dense; seconds = secs)

    return (; t_sparse, t_dense, rel_err)
end

function summarize(symmetry, idx, spec, r)
    ts = minimum(r.t_sparse).time / 1e6
    td = minimum(r.t_dense).time / 1e6
    ms = BenchmarkTools.memory(minimum(r.t_sparse)) / 1024
    md = BenchmarkTools.memory(minimum(r.t_dense)) / 1024
    return @sprintf(
        "%-12s spec %d  mps=%4d mpo=%2d  | sparse %.2f ms (mem %.1f KiB)  dense %.2f ms (mem %.1f KiB)  | ratio %.2fx  err %.1e",
        symmetry, idx,
        dim(spec.Vmps[1]), dim(spec.Vmpo[1]),
        ts, ms, td, md, td / ts, r.rel_err
    )
end

function main(; symmetries = nothing, max_mpsdim = 200, T = Float64)
    params = TOML.parsefile(CYL_TOML)
    chosen = symmetries === nothing ? sort(collect(keys(params))) : symmetries
    for symmetry in chosen
        haskey(params, symmetry) || continue
        specs = params[symmetry]
        println("\n=== Symmetry: $symmetry ===")
        for (i, d) in enumerate(specs)
            spec = parse_spec(d)
            if dim(spec.Vmps[1]) > max_mpsdim
                @printf(
                    "%-12s spec %d  mps=%4d mpo=%2d  | SKIPPED (mps_dim > %d)\n",
                    symmetry, i, dim(spec.Vmps[1]), dim(spec.Vmpo[1]), max_mpsdim
                )
                continue
            end
            try
                r = bench_one(spec; T)
                println(summarize(symmetry, i, spec, r))
            catch e
                @printf("%-12s spec %d  mps=%4d  | FAILED: %s\n",
                    symmetry, i, dim(spec.Vmps[1]), sprint(showerror, e))
            end
        end
    end
    return nothing
end

if abspath(PROGRAM_FILE) == @__FILE__
    sizecap = parse(Int, get(ENV, "AC_BENCH_MAXDIM", "200"))
    syms = haskey(ENV, "AC_BENCH_SYMS") ? split(ENV["AC_BENCH_SYMS"], ",") : nothing
    main(; max_mpsdim = sizecap, symmetries = syms)
end

[VictorVanthilt]

Can you remind me in what situations we would need to keep track of scalars? In Hamiltonians the scalars just get absorbed by D and either the C or B block, what would the "scalars" matrix keep track of in your JordanMPOTensor example above?

[lkdvos]

This is for long range operators: $X_1X_5$ is actually $X_1I_2I_3I_4X_5$, and this would be how it is filled out. In principle you don't have to actually contract with the identities, and you can get away with a simple permutation, but since we are not actually contracting the networks as a whole anymore (by first contracting env with hamiltonian), this doesn't actually help that much.
The scalars could be convenient for a term like $X_iX_j\lambda^{|i - j|}$, but otherwise also that doesn't change that much.
The biggest problem to fix is the variance, since this has identities on the diagonal that we have to treat separately to properly regularize:

MPSKit.jl/src/environments/infinite_envs.jl

Lines 193 to 205 in 98aab36

	if isidentitylevel(operator, i) # identity matrices; do the hacky renormalization
	T = regularize(TransferMatrix(above.AL, below.AL), ρ_left, ρ_right)
	GLs[1][i], convhist = linsolve(flip(T), GLs[1][i], prev, alg, 1, -1)
	convhist.converged == 0 &&
	@warn "GL$i failed to converge: normres = $(convhist.normres)"
	(L > 1) && left_cyclethrough!(i, GLs, below, operator, above)
	# go through the unitcell, again subtracting fixpoints
	for site in 1:L
	@plansor GLs[site][i][-1 -2; -3] -= GLs[site][i][1 -2; 2] *
	r_LL(above, site - 1)[2; 1] * l_LL(above, site)[-1; -3]
	end