Performance with long range interactions is really poor

[AndreasChristophilopoulos]

I had some code that I decided to migrate from TenPy, but unfortunately performance is far worse for some reason. The relevant part of the Julia code is this:

disordered_distance(i, j, dx) = abs(j - i + dx[j] - dx[i])

function build_rydberg_hamiltonian(L, Δ, Ω, σ, t, V, seed)
    # Random displacements per site, uniform in [-sigma/2, sigma/2]
    rng = MersenneTwister(seed)
    dx = (rand(rng, L) .- 0.5) .* σ

    Vphys = ℂ^2
    physical_spaces = fill(Vphys, L)

    σx = TensorMap([0 1; 1 0], Vphys ← Vphys)
    σz = TensorMap([1 0; 0 -1], Vphys ← Vphys)
    Id = id(Vphys)
    Nop = (Id + σz) / 2
    Sp = TensorMap([0 1; 0 0], Vphys ← Vphys)
    Sm = TensorMap([0 0; 1 0], Vphys ← Vphys)

    lattice = fill(Vphys, L)

    on_site_terms = [(i,) => -Δ * Nop + Ω * σx for i in 1:L]
    interaction_terms = []
    for i in 1:L, j in 1:L 
        if i < j
            r = disordered_distance(i,j,dx)
            Vij = V/r^6
            tij = t/r^3
            push!(interaction_terms, (i,j) => Vij*Nop⊗Nop + tij * ( Sp ⊗ Sm + Sm⊗Sp)) 
        end
    end
    H = FiniteMPOHamiltonian(lattice, on_site_terms..., interaction_terms...) 

    return H, physical_spaces
end

# Sample values 
Δ = 1.0; Ω = 1.0; V = 1.0; t=1.0;
L = 25
chi_max = 50
psi = FiniteMPS(L, ComplexSpace(2), ComplexSpace(chi_max))

H, physical_spaces = build_rydberg_hamiltonian(L, Δ, Ω, sigma, t, V, seed)

psi, envs, dmrg_info = find_groundstate!(psi, H, DMRG(; maxiter=10))

As you can see this implements a Hamiltonian with long range interactions. For comparison, the TenPy code model definition is:

class RydbergModel(CouplingMPOModel):
    
    def init_sites(self, model_params):
        # symmetry settings
        conserve = model_params.get('conserve', None, str)
        sort_charge = model_params.get('sort_charge', True, bool)

        # define spin-1/2 local Hilbert space
        site = SpinHalfSite(conserve=conserve, sort_charge=sort_charge)
        n = np.diag([1,0])
        # site.add_op('N', 0.5 * site.Sigmaz + 0.5 * site.Id)
        site.add_op('N',n)
        return site
    
    def init_terms(self, model_params):
        counter = 0
        # read parameters
    
        Delta = model_params.get('Delta', 0.0)              
        Omega = model_params.get('Omega', 0.0) 
        sigma = model_params.get('sigma', 0.0) 
        t = model_params.get('t', 0.0)
        V = model_params.get('V', 1.0)
        seed = model_params.get('seed', 0.0)
            
        N = self.lat.N_sites

        L = self.lat.Ls[0] 
        #for dx in range(1, L):
        # randomness on possition dxi:
        # so true position of each atom is no longer i, but i + dxi
        rng = np.random.default_rng(seed)
        dx_randomness = rng.uniform(-0.5, 0.5, size=L) * sigma

        # This is for non-periodic boundary conditions
        for i in range(L):
            for j in range(i+1,L): # This ensures i < j    
                dx = abs((j+dx_randomness[j]) - (i + dx_randomness[i]))
                Vij = V / dx**6
                tij = t / ( dx**3) # run again with this value
                self.add_coupling_term(tij, i,j, 'Sp', 'Sm', plus_hc=True)
                self.add_coupling_term(Vij, i, j,'N','N')


        self.add_onsite(-Delta, 0, 'N')

        self.add_onsite(Omega, 0, 'Sigmax') 

and then for the actual ground state computation I run something like

model = MyTFIChain(model_params)
                psi = MPS.from_lat_product_state(model.lat, [['up']])

                dmrg_params = {
                'mixer': True,  # setting this to True helps to escape local minima
                'max_E_err': 1.e-6,
                'trunc_params': {
                    'chi_max': chi_max,
                    'svd_min': 1.e-10,
                    },
                'verbose': True,
                'combine': True
                }

                # get the data:
                eng = dmrg.TwoSiteDMRGEngine(psi, model, dmrg_params)
                E, psi = eng.run()

The TenPy version runs a lot faster. In my actual code, I perform sweeps over regions of Δ and Ω. The TenPy version took about 4.5 minutes to run a sweep over a 5x5 region, while MPSKit takes almost as much time for even the first DMRG calculation (I don't count the time for compilation).

Am I using the wrong algorithm for something, is it a problem with my code, or is TenPy just faster currently?

[lkdvos]

Hi Andreas, thanks for reporting this!

My biggest hunch here is that our default MPO constructor is not actually trying to do anything smart (in fact, it is doing the most naive thing possible), so I'm guessing that for this model with many long-range terms, TenPy is just generating a hamiltonian that is a lot more efficient.

To confirm this and possibly try improve MPSKit's version, something you could try is to see if doing a (quasi-exact) compression of the MPO yields any improvements:

H_compact = changebonds(H, SVDCut(; trscheme = trunctol(; atol=1e-12)))

A second issue (that I'm currently trying to resolve, see #455) is that TenPy is a little bit smarter at the moment with its use of adaptive tolerances in the 2site scheme, which likely also makes their sweeps a little bit faster.
You could try and play around a bit with the default tolerances in MPSKit right now, but I am hoping to get this fixed soon.

I would definitely be happy to hear if this is improving the runtimes at all, and investigate what is going on otherwise!

[AndreasChristophilopoulos]

Thanks a lot for responding. Alright, for starters I will my full program, the way it was originally. For reference, here is the whole thing, in case someone wants to copy-paste and try it out.

I have your version of compressing the MPO commented out in the "build_rydberg_hamiltonian" function. I ran it with the following command:

julia --project=. main.jl --D -4,0,3 --Omega -4,0,3 --L 25 --sigma 0.0 --seed 1234 --V 1.0 --chi 50 --t 0.0

After it's done I will record the time it took and run it your way next. Here is the full code for replication purposes:

using TensorKit, MPSKit, MPSKitModels
using LinearAlgebra, Random
using Plots
using ArgParse
using CSV, DataFrames
using Dates
using Logging

# -----------------------------
# Helpers
# -----------------------------

"""
    parse_float_list(s)

Parse a comma-separated string of numbers, e.g. "-2,2,9" -> [-2.0, 2.0, 9.0].
"""
parse_float_list(s::AbstractString) = parse.(Float64, split(s, ','))

"""
    logspace_values(spec::AbstractString)

`spec` is "start,stop,count"; returns `10 .^ range(start, stop, length=count)`.
"""
function logspace_values(spec::AbstractString)
    start, stop, count = parse_float_list(spec)
    return 10.0 .^ range(start, stop=stop, length=Int(count))
end

"""
    disordered_distance(i, j, dx)

Effective distance between sites `i` and `j` once the per-site random
displacement `dx` is taken into account.
"""
disordered_distance(i, j, dx) = abs(j - i + dx[j] - dx[i])

# -----------------------------
# Build the Rydberg Hamiltonian as an MPO
# -----------------------------

"""
    build_rydberg_hamiltonian(L, Delta, Omega, sigma, t, V, seed)

Return `(H, physical_spaces)` where `H` is the Rydberg-chain MPO Hamiltonian
    H = -Delta * sum_i n_i + Omega * sum_i sigma^x_i
        + sum_{i<j} [ V/r_ij^6 * n_i n_j + t/r_ij^3 * (S+_i S-_j + S-_i S+_j) ]
with r_ij the (disordered) distance between sites, and `physical_spaces` the
vector of local Hilbert spaces to be used e.g. when building the initial MPS.
"""
function build_rydberg_hamiltonian(L, Δ, Ω, σ, t, V, seed)
    # Random displacements per site, uniform in [-sigma/2, sigma/2]
    rng = MersenneTwister(seed)
    dx = (rand(rng, L) .- 0.5) .* σ

    Vphys = ℂ^2
    physical_spaces = fill(Vphys, L)

    # Pauli / ladder operators. `Id` (not `I`) to avoid shadowing LinearAlgebra.I
    σx = TensorMap([0 1; 1 0], Vphys ← Vphys)
    σz = TensorMap([1 0; 0 -1], Vphys ← Vphys)
    Id = id(Vphys)
    Nop = (Id + σz) / 2
    Sp = TensorMap([0 1; 0 0], Vphys ← Vphys)
    Sm = TensorMap([0 0; 1 0], Vphys ← Vphys)

    lattice = fill(Vphys, L)

    on_site_terms = [(i,) => -Δ * Nop + Ω * σx for i in 1:L]
    interaction_terms = []
    for i in 1:L, j in 1:L 
        if i ≠ j
            r = disordered_distance(i,j,dx)
            Vij = V/r^6
            tij = t/r^3
            push!(interaction_terms, (i,j) => Vij*Nop⊗Nop + tij * ( Sp ⊗ Sm + Sm⊗Sp)) 
        end
    end
    H = FiniteMPOHamiltonian(lattice, on_site_terms..., interaction_terms...) 
    # H = changebonds(H, SVDCut(; trscheme = trunctol(; atol=1e-12)))

    return H, physical_spaces
end

# -----------------------------
# Central charge from entanglement entropy profile
# -----------------------------

"""
    central_charge_from_entropy(S, L)

Fit the bipartite entanglement entropy profile `S` (one value per internal
bond, length `L - 1`) to the CFT prediction
    S(x) = (c / 6) * log((2L/π) sin(πx/L)) + const
and return the extracted central charge `c`.
"""
function central_charge_from_entropy(S::Vector{Float64}, L::Int)
    n = L - 1
    @assert length(S) == n "expected $(n) bond entropies, got $(length(S))"
    x = [log((2L / π) * sin(π * i / L)) for i in 1:n]
    X = hcat(x, ones(n))
    slope, _intercept = X \ S
    return 6 * slope
end

# -----------------------------
# Main
# -----------------------------
function main()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--boundary_conditions"
            help = "finite or infinite (ignored, finite is always used)"
            default = "finite"
        "--D"
            help = "Delta range as \"start,stop,count\" (logspace)"
            arg_type = String
            required = true
        "--Omega"
            help = "Omega range as \"start,stop,count\" (logspace)"
            arg_type = String
            required = true
        "--L"
            help = "chain length"
            arg_type = Int
            required = true
        "--sigma"
            help = "disorder strength"
            arg_type = Float64
            default = 0.0
        "--seed"
            help = "random seed"
            arg_type = Int
            default = 1234
        "--t"
            help = "hopping amplitude"
            arg_type = Float64
            default = 0.0
        "--V"
            help = "interaction strength"
            arg_type = Float64
            default = 1.0
        "--chi"
            help = "maximum bond dimension"
            arg_type = Int
            required = true
    end
    args = parse_args(s)

    Delta_values = logspace_values(args["D"])
    Omega_values = logspace_values(args["Omega"])

    L = args["L"]
    V = args["V"]
    t = args["t"]
    sigma = args["sigma"]
    seed = args["seed"]
    chi_max = args["chi"]

    # ----- logging / output directory -----
    dir_name = "RydbergModel_s=$(sigma)_D=$(args["D"])_Omega=$(args["Omega"])_V=$(V)_L=$(L)"
    if isdir(dir_name)
        print("Directory already exists. Press 'y' if you want to overwrite.")
        input = readline()
        agreed =  strip(input) == "y"
        if agreed
            print("Permission to remove granted.")
            rm(dir_name, recursive=true)
        else
            error("Directory $dir_name already exists. If .")
        end
    end
    mkpath(dir_name)

    info_text = """
    Started: $(now())
    Randomness on site positions (dx), uniform distribution. sigma = $sigma, seed = $seed.
    Delta = logspace($(args["D"])), Omega = logspace($(args["Omega"])), L = $L, V = $V, t = $t
    """
    write(joinpath(dir_name, "information.txt"), info_text)

    max_entropy = Matrix{Float64}(undef, length(Delta_values), length(Omega_values))
    central_charge = similar(max_entropy)

    entropy_csv = joinpath(dir_name, "entropy_results_for_each_run.csv")

    start_time = time()

    open(entropy_csv, "w") do io
        for (i, Δ) in enumerate(Delta_values)
            row_S = Vector{Vector{Float64}}(undef, length(Omega_values))
            for (j, Ω) in enumerate(Omega_values)
                psi = FiniteMPS(L, ComplexSpace(2), ComplexSpace(chi_max))
                @info "Running DMRG" Delta = Δ Omega = Ω

                H, physical_spaces = build_rydberg_hamiltonian(L, Δ, Ω, sigma, t, V, seed)

                psi, envs, dmrg_info = find_groundstate!(psi, H, DMRG(; maxiter=10))

                S = [entropy(psi,i) for i in 2:L]  # one value per internal bond, length L - 1
                max_entropy[i, j] = maximum(S)
                central_charge[i, j] = central_charge_from_entropy(S, L)
                row_S[j] = S

            end
            
            for j in 1:length(Omega_values)
                println(io, join(row_S[j], ","))
            end
            println(io)  # blank line separates each Delta row
        end
    end

    elapsed = time() - start_time
    @info "Sweep complete" elapsed_seconds = elapsed

    # ----- save matrices -----
    param_tag = "D=$(args["D"])_Omega=$(args["Omega"])_V=$(V)_L=$(L)"
    CSV.write(joinpath(dir_name, "MaxEntropyData_$(param_tag).csv"), DataFrame(max_entropy, :auto), writeheader = false)
    CSV.write(joinpath(dir_name, "CentralChargeData_$(param_tag).csv"), DataFrame(central_charge, :auto), writeheader = false)

    # ----- plots -----
    p1 = heatmap(Omega_values, Delta_values, max_entropy,
                 xscale = :log10, yscale = :log10,
                 xlabel = "Omega / V", ylabel = "Delta / V",
                 title = "Max Entropy", colorbar = true)
    savefig(p1, joinpath(dir_name, "MaxEntropyPlot_$(param_tag).png"))

    p2 = heatmap(Omega_values, Delta_values, central_charge,
                 xscale = :log10, yscale = :log10,
                 xlabel = "Omega / V", ylabel = "Delta / V",
                 title = "Central Charge", colorbar = true)
    savefig(p2, joinpath(dir_name, "CentralChargePlot_$(param_tag).png"))

    return nothing
end

main()

[AndreasChristophilopoulos]

Ok, so I ended up running with --D -4,0,2 and --Omega -4,0,2 instead because it was taking too much time. I had also made a couple of typos in the code that I posted, I edited the post to fix it. The first run took about 1009 seconds. Now I am waiting for the second one (that uses your recommendation) to be completed.

[AndreasChristophilopoulos]

Uh oh, SVDCut wasn't found... Does it come from some other package I haven't imported?

[lkdvos]

Oh I'm sorry, it should have been SvdCut (note the capitalization)

[AndreasChristophilopoulos]

That was far better, it took about 214 seconds overall. In fact I think there's some compilation related stuff leaking into it, I'll benchmark it so that it doesn't count that, and then compare to TenPy.

[AndreasChristophilopoulos]

Ok, now that I'm using btime on my function, I see that it actually took about 37 seconds to do the actual computation. Meanwhile, the TenPy version which I have parallelized a bit so that it's a little bit faster than it was took about 15-16 seconds. So, TenPy still seems to be faster, but not by as much.

[AndreasChristophilopoulos]

Upon further testing, I found that after your suggestion, it is approximately twice as slow as the original, un-parallelized TenPy code, which is not as bad as it was at first.

However, another thing I noticed is that increasing the values I sweep over also increases the compilation time, which is confusing to me...

[lkdvos]

That does sound like the remainder might be due to the difference in the eigenvalue settings, especially if just the maximal number of iterations is fixed. I think MPSKit is (overly) cautious by default, and something like this might also produce reasonable results:

alg = DMRG2(;
    maxiter = 10,
    alg_eigsolve = (; krylovdim = 16, maxiter = 1, dynamic_tols = true, tol_min = 1e-8, tol_max = 1e-1, tol_factor = 1e-1)
)

This is hopefully what #455 will handle better though.

(Note that there is also always the BLAS settings, which are somewhat harder to get to be the same, and don't really reflect on the speed of the libraries. I assume you are not setting the number of BLAS threads manually? If you are on a machine that supports this, also using MKL tends to be a bit faster.)

---

Regarding the compilation time, I'm not entirely sure why and how that is happening, my first guess would be to check for type stabilities etc, but I can't spot any immediate issues in your script. If you plan on doing quite a bit of these simulations, it might be worth it to invest in a small startup package to precompile a bunch of this, see e.g. https://julialang.github.io/PrecompileTools.jl/stable/#Tutorial:-local-%22Startup%22-packages for some pointers on that

[AndreasChristophilopoulos]

Oh I'm not manually setting anything BLAS related. I will try out your recommendations tomorrow.