[BUG] Correlation function is not fully compatible with Blockade subspace

[hxwang-amazon]

Describe the bug
When using blockade_subspace, the rydberg_corr function has bugs for Op.X and Op.Y
The rydberg_corr function returns a correlation matrix. Theoretically, the correlation matrix should have diagonal entries always equal to ones (i.e., diagonal = [1,1,1...1]), since X*X=I, Y*Y=I and Z*Z=I. When using the full Hilbert space registers, this is indeed the case. However, when using the blockade_subspace registers, the diagonal entries of the correlation matrix (for X and Y) are not always equal to ones (e.g., sometimes the diagonal entries are approximately [0,1,0,1,...]).

To Reproduce

using Bloqade
using KrylovKit
using SparseArrays
using ProgressBars
using BitBasis
using Statistics
using Plots

densities = []
x_corrs = [] # correlation matrix for X
y_corrs = []
z_corrs = []

nsites = 9
scale = 5.72
atoms = generate_sites(ChainLattice(), nsites, scale=scale)
space = blockade_subspace(atoms, scale + 0.01);
total_time = 3.0;
time_step = 0.1;
Ω_max = 2π * 4;
Ω = piecewise_linear(clocks=[0.0, 0.1, 2.1, 2.2, total_time], values=[0.0, Ω_max, Ω_max, 0, 0]);
U1 = -2π * 10;
U2 = 2π * 10;
Δ = piecewise_linear(clocks=[0.0, 0.6, 2.1, total_time], values=[U1, U1, U2, U2]);

h = rydberg_h(atoms; Δ, Ω);

## ---- Full Hilbert Space ----
reg = zero_state(nsites)
## --- Blockade Subspace
# reg = zero_state(space) # just uncomment this to use Blockage subspace registers

prob = SchrodingerProblem(reg, total_time, h);
integrator = init(prob, Vern8());

times = 0.0:time_step:total_time
for _ in TimeChoiceIterator(integrator, tqdm(times))
    push!(densities, rydberg_density(reg))
    push!(x_corrs, rydberg_corr(Op.X, reg))
    push!(y_corrs, rydberg_corr(Op.Y, reg))
    push!(z_corrs, rydberg_corr(Op.Z, reg))
end


D = hcat(densities...);

# convert lists to 3D arrays
# e.g., Xc is arrays of all Pauli-X correlation matrices obtained during this adiabatic evolution
Xc = permutedims(cat(x_corrs...; dims=3), (3, 1, 2))
Yc = permutedims(cat(y_corrs...; dims=3), (3, 1, 2))
Zc = permutedims(cat(z_corrs...; dims=3), (3, 1, 2))

# Visualize a Pauli-X Correlation Matrix
time_step = 18 # overall we have 31 time steps here (0:0.1:3.0), and we take the step 18 for visualization
# You can change Xc to Yc or Zc for visualization of the Pauli-Y or Pauli-Z correlation matrix
heatmap(real.(Xc[time_step, :, :]), yflip=true,
    clim=(0.0, 1.0), c=:thermal)

Expected behavior
When using blockade_subspace, the correlation matrix (for Op.X or Op.Y) should have diagonal entries = [1,1,...,1].

Screenshots
If applicable, add screenshots to help explain your problem.

A Pauli-X Correlation matrix obtained by full-Hilbert-space registers (Correct one)

A Pauli-X Correlation matrix obtained by Blockade subspace registers (Wrong one)

Version Info

• Julia version: 1.7.3 (Linux/Ubuntu) and 1.8.0 (OS X w/ ARM CPU)
• Package Version: Bloqade 0.1.7

Additional context
I think there are two ways to fix this

1. Create a function that converts Blockade subspace registers to full-Hilbert-space registers, and then just add a step of converting registers in the rydberg_corr function. Notice: It seems this new feature mentioned in [Feature Request] conversion from subspace register to fullspace #364 (by @Roger-luo ) might be able to fix the bug in this suggested way.
2. In the rydberg_corr function, for each pair of qubits of interests (to compute the two-point correlation function), get the reduced density matrix of the two-qubit subsystem and use it to compute the expectation value of the two-point correlation function (for X, Y, Z or Op.n).

[Roger-luo]

I think this is not necessarily be fixed by #373 I think the symbolic expression should at least guarantee the diag component is correct, but I could be wrong. I need to take a deeper look into this.

[weinbe58]

I think the issue is here. The problem is at intermediate points of the calculation for off diagonal observables you might activate a register that doesn't obey the Rydberg constraint, as such when you search for that Fock state in the subspace you get -1.

This is an issue when dealing with subspaces in ED methods. individual operators might create intermediate states that do not live in the subspace, but after applying the entire chain of operators will you get a valid state again.

you might fix the diagonal part of the correlation function symbolically, however the off diagonal correlations will not be correct because of this issue.

Would it be possible to prevent the projection back to the subspace before applying the second operator?

EDIT: A short term but not necessarily the best solution would be to implement suggestion 2. its not ideal because it requires an SVD but its an OK patch for the short term. if I have time I can try to implement this for YaoSubspaceArrayReg

[weinbe58]

The issue is that when using a local operator, sometimes the intermediate state falls outside of the Blockade subspace. Effecitvely the current implementation measures pBpAp where p and A/B are the operators is the projection operator onto the blockade subspace, but this is incorrect we need to measure pBAp so one needs to build build the AB blocks directly using the Yao API.

I propose we include the following operators:

XX_0101
XX_0r0r
XX_1r1r
YY_0101
YY_0r0r
YY_1r1r
ZZ_0000
ZZ_1111
ZZ_rrrr

where:
Z: Pauli-Z
X: Pauli-X
Y: Pauli-Y

and the subscripts correspond to which levels are coupled inside the qubit/qudit