Computing a linear combination of all reduced density matrices of a given pure state

Hi all! I have been struggling for a while on the following calculation:
(For the simulations, I’m using Yao)

Imagine you are given a pure state \psi of N qubits. And you want to compute
where A runs over all subregions of {1, …, N}, and \sigma_A refers to the reduced density matrix of the given state \psi, associated with region A. There is, however, a caveat that the reduced density matrix has to be embedded into the original Hilbert space of N qubits.

Do you have any suggestions on how to do this in general? I have a code in Julia which computes the reduced density and then embedded it back into the full Hilbert space but it is quite inefficient.

Any help/advice with this is highly appreciated. :slight_smile:

1 Like

To speed up the computation of the reduced density matrix you might try Tullio, reshape your density matrix into a tensor and then sum over the indices you don’t need.

When you embed it back into the full Hilbert space do you mutate an existing matrix or do you create a new matrix?

By embedding it back into the full Hilbert space I mean something like the following
image. At the moment, I’m creating a new matrix each time

I think at least you could preallocate the matrix for storing your results? Otherwise it could be slow. I would suggest running a profile to actually see what is causing the slow down

So, I just found that my Julia code for computing the embedded back reduced density matrix has some issues. I have a function written in Mathematica that basically does what I want correctly but I want to work in Julia so I basically “translated” it into Julia code. I think the issue is associated with Julia being column-major while Mathematica is row-major. This is the function I have written at the moment but it is not even returning a Hermitian matrix

using TensorCore, Yao, QuantumInformation, LinearAlgebra

function eye(n::Int)
    m = zeros(Float64, n, n)
    for i in 1:n
        m[i, i] = 1.0
    return m

function embed_rdmat(rho_init, N::Int, rA::Vector{Int})
    sizeA = length(rA);
    rB = setdiff(collect(1:N), rA);
    sizeB = length(rB)
    rho_tensor_A = reshape(rho_init, Tuple(map(x->2, collect(1:2*sizeA))))
    id_tensor_B = reshape(eye(2^sizeB), Tuple(map(x->2, collect(1:2*sizeB))))
    rho = reshape(tensor(rho_tensor_A ,id_tensor_B), Tuple(map(x->2, 1:2*N)))

    indices_to_permute = zeros(Int, 2, N)
    indices_to_permute[1, :] = union(rA, rB)
    indices_to_permute[2, :] = union(rA, rB) .+ N

    rho = reshape(permutedims(rho, reduce(hcat, indices_to_permute)[1,:]), (2^N, 2^N))

    return rho

N = 6;
ψ = rand_state(N);
rA = [1, 3]
rB = [2, 4, 5, 6]

ψvec = ψ |> state;

rhoA =  QuantumInformation.ptrace(ψvec*ψvec', map(x->2, collect(1:N)), rB);

rhoA_fullH = embed_rdmat(rhoA, N, rA)


Any suggestion on how to tackle this issue is more than welcome!

Can you write out mathematically what you want to accomplish? I think the embed function from QuantumOptics might solve your problem

Wow, thanks for pointing this out. That’s exactly what I want, although I don’t understand well the source code of the embed function yet. I want to write the reduced density matrix as an operator in the full Hilbert space ( by adding identity in the sites in B), i.e., \rho_A \to rhoFull. The idea is that if I trace out the region B once again I get 'rho_A. My issue is with how to order the indices. I did it in Mathematica but since Julia is column major the code does not work. And I honestly don’t know how to do it in that way.

Hi once again. Besides the function from QuantumOptics I manage with the help of a coworker to implement the following solution to the above issue. It uses the "parallel bit deposit and extract” instructions, which are part of BMI2. See this. This is the working code: it receives the reduced density matrix (2^nA x 2^nA) as well as another matrix ( 2^N x 2^N ) which is then modified:

function get_mask(rA::Vector{Int}, N::Int)
    mask = UInt64(0)
    for a in rA
        mask = mask | (1 << (N-a))
    return mask

pdep(x::UInt64, y::UInt64) = ccall("llvm.x86.bmi.pdep.64", llvmcall, UInt64, (UInt64, UInt64), x, y)

function embed_rdm(rho:: Matrix{ComplexF64}, N::Int, rA::Vector{Int}, rhoToModify::Matrix{ComplexF64})
    nA = length(rA)
    nB = N - nA;

    maskA = get_mask(rA, N);
    maskB = get_mask(setdiff(collect(1:N), rA), N)

    for iB in 0:(2^nB - 1)
        sB = pdep(UInt64(iB), maskB)
        for iA1 in 0:(2^nA - 1)
            sA1 = pdep(UInt64(iA1), maskA)
            for iA2 in 0:(2^nA - 1)
                sA2 = pdep(UInt64(iA2), maskA)
                rhoToModify[ (sB | sA1) + 1, (sB | sA2) + 1] += rho[iA1 + 1, iA2 + 1]/(2^nB)

I hope is useful to someone. Thanks for all the help :slight_smile:

One simple improvement to keep in mind in general is to use I(n) (from LinearAlgebra) instead of your eye function. It’s pretty much strictly better.


Just saw you question, it seems like you want the following function

julia> using Yao

julia> reg = rand_state(18);

julia> dm = density_matrix(reg, (3,))
DensityMatrix{2, ComplexF64, Matrix{ComplexF64}}(ComplexF64[0.4999692201666899 + 0.0im -0.0005383780430281125 + 0.0009932672235014744im; -0.0005383780430281125 - 0.0009932672235014744im 0.5000307798333101 + 0.0im])

julia> mat(put(18, 2=>matblock(dm.state))) |> size
(262144, 262144)

julia> @time mat(put(18, 3=>matblock(dm.state)));
  0.003275 seconds (15 allocations: 14.001 MiB)
1 Like

This is quite nice! Thanks a lot for sharing. Maybe a silly question on my side but is density_matrix in the Yao doc. I didn’t manage to find it