Dictionaries.jl optimized for PauliRotations

[MSRudolph]

Dictionaries.jl is an interesting library aiming to make dictionaries that behave more like arrays. Such dictionaries have several nice benefits, including:

1. Fast iteration over keys/values/pairs
2. Comfortable interface for common array operations like broadcasting (abs.(dict)), sorting (sort(dict), sortperm(dict), in-place manipulation dict *= 0.9 and tokenization. The latter allows manipulating elements by performing only a single hash. See, for example, this Julia issue: julep/rfc: proposal for unified Dict get/insert/has interface JuliaLang/julia#12157
3. Easy high-level parallelization using Base.Threads and the @threads macro.

Using Dictionaries.jl dicts has been brought up in #58, and would have several obvious benefits for gate such as Clifford gates or Pauli noise gates. At the same time, our most common gate type are PauliRotation gates, and there the advantages are not as clear-cut.

Challenges:
My attempts at writing apply functions for the currently most common gate - PauliRotation - have had limited success (see answer below). At the scale of dozens of Pauli strings, I see Dictionaries.jl dicts starting to be slower than Base.Dict, mostly because insertion of elements can be substantially slower. This likely mean that highly-branching gates are even worse with Dictionaries.jl dicts, but our focus should be on the most common gates. The potential benefits here are great.

Concrete goals:

• Prototype a gate application function such as the one below for PauliRotation gates.
• Showcase that it is not slower at scale than our current Base.Dict code at a variety of circuits. Given my experiences, the trend of not slowing down beyond the current code should hold up to runtimes of tens of minutes.
• Finding trade-offs between Base.Dict and Dictionaries.jl dicts is good. If some cases are faster and others are slower, we can still consider supporting both.

[MSRudolph]

Here is an attempt for using Dictionaries.jl, though I think are probably some problems here:

function Base.merge!(d::AbstractDictionary, pstrs, coeffs)
    ## somehow this made everything worse 
    # s = maximum(length(d.values)+1, length(d) + length(coeffs))
    # sizehint!(d, s)
    
    for (pstr, coeff) in zip(pstrs, coeffs)
        (hasindex, token) = gettoken!(d, pstr)
        if hasindex
            @inbounds settokenvalue!(d, token, gettokenvalue(d, token) + coeff)
        else
            @inbounds settokenvalue!(d, token, coeff)
        end
    end
    return d
end


anticommuting(gate, d) = [pstr for pstr in d.indices if !commutes(gate, pstr)]


function myapply(gate::PauliRotation, psum, theta; min_abs_coeff = 1e-10)

    gate = PauliPropagation._tomaskedpaulirotation(gate, keytype(psum))

    # pre-compute the sine and cosine values because the are used for every Pauli string that does not commute with the gate
    cos_val = cos(theta)
    sin_val = sin(theta)
    
    # find the bitstrings that anticommute
    anticomm_pstrings = anticommuting(gate, psum)

    # create the new value list
    new_coeffs = zeros(length(anticomm_pstrings))
    
    for ii in eachindex(anticomm_pstrings)
        # I get confused what is an `Index` and what is the integer Pauli string
        pstr = anticomm_pstrings[ii]

        # old coefficient
        # coeff = psum[pstr]
        (_, token) = gettoken(psum, pstr)
        coeff = gettokenvalue(psum, token)

        # set the cos coefficient
        # psum[pstr] = coeff * cos_val
        settokenvalue!(psum, token, coeff * cos_val)

        # get the new stuff
        new_pstr, sign = PauliPropagation.getnewpaulistring(gate, pstr)
        new_coeff = coeff * sin_val * sign
        
        new_coeffs[ii] = new_coeff
        anticomm_pstrings[ii] = new_pstr

    end

    # anticomm_pstrings are non-overlapping, so we can put them in arrays
    # merge them into the psum
    merge!(psum, anticomm_pstrings, new_coeffs)

    # truncation function - ignoring general truncations for now.    
    Dictionaries.filter!(coeff -> abs(coeff) >= min_abs_coeff, psum)
    
    return psum
end

[MSRudolph]

We aim to give the reward to the fastest version utilizing Dictionaries.jl that beats the existing Base.Dict performance non-trivially.