Implement the classical shadow coding error?

Question

I'm trying to reproduce the basic method of classical shadow, which is based on the tutorial of pennylane. However, I've met some realization problems here when I finish reading the tutorial of pennylane, and trying to finish the method myself, just to check if my understanding is correct, because I'm not sure if my understanding about the inverse of the map $M$ mentioned in the paper is correct.

I will describe the method shortly first, and then showing my code with qiskit, pennylane, and matlab. But all of them failed to have the same effect as the tutorial of pennylane does(increasing the number of measurements, the distance between the state I want to reconstruct and the original state should be more and more close).

The idea of classical shadow(or the process of the algorithm) is kind of simple while the math behind it might be complicated. The process states that for any density matrix $\rho$, we act some unitary matrix $U$ which are chosen randomly from a specific set of the unitary matrix $\mathcal{U}$ on it, i.e., $U\rho U^\dagger$. Then we do one-shot measurement based on the computational basis on $U\rho U^\dagger$. Then the state will collapse into some state $|\hat{b}\rangle$. And then we do the rest of the work in classical data analysis style. First we undo the unitary matrix, i.e., $U^\dagger |\hat{b}\rangle\langle \hat{b}| U$ . Then we do the inverse of the map $\hat{\rho}\equiv M^{-1}(U^\dagger |\hat{b}\rangle\langle \hat{b}| U)$ which can be defined as $M(\rho)\equiv E(U^\dagger |\hat{b}\rangle\langle \hat{b}| U)$, where the $E$ stands for expectation over both unitary matrix and measurement result $|\hat{b}\rangle$. And for a specific choice of unitary set(seems Clifford group, not very clear here), we have a form of the inverse of $M$ states as $$ \hat{\rho} = \bigotimes_{j=1}^n(3U^{\dagger}_j|\hat{b}_j\rangle\langle\hat{b}_j|U_j-\mathbb{I})\tag{1} $$ Then I will introduce my code first. We specify the unitary group into Hadamard gate, phase gate, and identity. Then we do the computational basis measurement and rebuild the classical shadow $\hat{\rho}$ with the help of eq.(1), and then calculate the expectation value of $\hat{\rho}$ by directly divide measurement times.

Following the codes(pennylane, qiskit, Matlab), aiming at construct the classical shadow of bell state:

from networkx.algorithms.centrality import harmonic
from networkx.exception import HasACycle
from networkx.readwrite.sparse6 import write_sparse6
from numpy import dtype
from numpy.random.mtrand import rand
import pennylane as qml
from pennylane import wires
import pennylane.numpy as np
import matplotlib.pyplot as plt
import time

def distance(rho):
    return np.sqrt(np.trace(rho.conjugate().transpose() @ rho))

def my_quantum_function(x, y):
    unitary = [qml.Hadamard, qml.S, qml.Identity]
    qml.Hadamard(wires=0)
    qml.CNOT(wires=[0,1])
        
    unitary[y[0]](wires=0)
    unitary[y[1]](wires=1)

    # all measure in computational basis, i.e., mean value of pauliz, one-shot case
    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]

# one-shot case shots = 1 to simulate the measure in computational basis requirement
dev = qml.device('default.qubit', wires=2, shots=1)
circuit = qml.QNode(my_quantum_function, dev)

# generate random number seed for easy replicate the experiment
np.random.seed(666)

# init
phase_z = np.array([[1, 0], [0, 1j]], dtype=complex)
hadamard = qml.Hadamard(0).matrix
identity = qml.Identity(0).matrix

unitary = [hadamard, phase_z, identity]

snapshot = 1000
state0 = np.array([[1,0],[0,0]])
state1 = np.array([[0,0],[0,1]])
record_rho = np.zeros([4,4])

for i in range(snapshot):
    randnum = np.random.randint(0,3,size=2)
    [res0, res1] = circuit(0,randnum)
    # print(circuit.draw())
    if res0 == 1:
        rho1 = 3*(unitary[randnum[0]].conj().T @ state0 @ unitary[randnum[0]]) - identity
    else:
        rho1 = 3*(unitary[randnum[0]].conj().T @ state1 @ unitary[randnum[0]]) - identity

    if res0 == 1:
        rho2 = 3*(unitary[randnum[1]].conj().T @ state0 @ unitary[randnum[1]]) - identity
    else:
        rho2 = 3*(unitary[randnum[1]].conj().T @ state1 @ unitary[randnum[1]]) - identity

    record_rho = record_rho + np.kron(rho1,rho2)

record_rho = record_rho/snapshot
bell_state = np.array([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]])
print(record_rho)

print(distance(record_rho - bell_state))

from math import exp
from qiskit import *
from qiskit import Aer
import numpy as np
import matplotlib.pyplot as plt
from random import randrange

np.random.seed(222)
def one_shot(operator):
    
    sim = Aer.get_backend('aer_simulator')
    qc = QuantumCircuit(2)
    unitary = [qc.h, qc.sdg, qc.id]
    qc.h(0)
    qc.cx(0,1)
    unitary[operator[0]](0)
    unitary[operator[1]](1)
    qc.measure_all()

    qobj = assemble(qc,shots=1)
    result = sim.run(qobj).result().get_counts()
    return result

def distance(rho):
    '''
    calculate distance of two density matrix
    '''
    return np.sqrt(np.trace(rho.conjugate().transpose().dot(rho)))

hadamard = 1/np.sqrt(2)*np.array([[1,1],[1,-1]])
s_gate = np.array([[1,0],[0,-1j]],dtype=complex)
id = np.identity(2)
unitary = [hadamard,np.dot(hadamard,s_gate),id]

snapshot_num = 1000
state0 = np.array([[1,0],[0,0]])
state1 = np.array([[0,0],[0,1]])
record_rho = np.zeros([4,4])
for i in range(snapshot_num):
    randnum = np.random.randint(0,3,size=2)
    result = one_shot(randnum)
    if result.get('00') == 1:
        rho = np.kron(3*np.dot(unitary[randnum[0]].conj().T,state0).dot(unitary[randnum[0]] - id),3*np.dot(unitary[randnum[1]].conj().T,state0).dot(unitary[randnum[1]]) - id)
    elif result.get('01') == 1:
        rho = np.kron(3*np.dot(unitary[randnum[0]].conj().T,state0).dot(unitary[randnum[0]] - id),3*np.dot(unitary[randnum[1]].conj().T,state1).dot(unitary[randnum[1]]) - id)
    elif result.get('10') == 1:
        rho = np.kron(3*np.dot(unitary[randnum[0]].conj().T,state1).dot(unitary[randnum[0]] - id),3*np.dot(unitary[randnum[1]].conj().T,state0).dot(unitary[randnum[1]]) - id)
    else:
        rho = np.kron(3*np.dot(unitary[randnum[0]].conj().T,state1).dot(unitary[randnum[0]] - id),3*np.dot(unitary[randnum[1]].conj().T,state1).dot(unitary[randnum[1]]) - id)

    record_rho = record_rho + rho
record_rho = record_rho/snapshot_num
bell_state = np.array([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]])
print(record_rho)
print(distance(record_rho - bell_state))

% Pauli matrix as random unitary
pauli = eye(2);
pauli(:,:,2) = [1 0;0 -1j]; pauli(:,:,3) = 1/sqrt(2)*pauli(:,:,2)'*[1 1;1 -1];
% two computational basis |0) and |1)
state = eye(4);
state0 = [1;0];
state1 = [0;1];

psi = 1/sqrt(2)*[1;0;0;1];
rho = psi*psi';
record_rho = zeros(4);
n = 6000;
for i = 1:n
    randnum = randi([1 3],[1 2]);
    rhot = kron(pauli(:,:,randnum(1)),pauli(:,:,randnum(2)))*rho*kron(pauli(:,:,randnum(1)),pauli(:,:,randnum(2)))';
    prob1 = state(:,1)'*rhot*state(:,1);
    prob2 = state(:,2)'*rhot*state(:,2);
    prob3 = state(:,3)'*rhot*state(:,3);
    prob4 = state(:,4)'*rhot*state(:,4);
    
    % Utilizing if to simulate the quantum measurement
    % The inverse process is using the formula of eq(S44) in supplemental
    % material of the original paper
    if rand < prob1
        rhot = 3*pauli(:,:,randnum(1))'*(state0*state0')*pauli(:,:,randnum(1)) - eye(2);
        rhot = kron(rhot,3*pauli(:,:,randnum(2))'*(state0*state0')*pauli(:,:,randnum(2)) - eye(2));
    elseif rand < prob1 + prob2
        rhot = 3*pauli(:,:,randnum(1))'*(state0*state0')*pauli(:,:,randnum(1)) - eye(2);
        rhot = kron(rhot,3*pauli(:,:,randnum(2))'*(state1*state1')*pauli(:,:,randnum(2)) - eye(2));
    elseif rand < prob1 + prob2 + prob3
        rhot = 3*pauli(:,:,randnum(1))'*(state1*state1')*pauli(:,:,randnum(1)) - eye(2);
        rhot = kron(rhot,3*pauli(:,:,randnum(2))'*(state0*state0')*pauli(:,:,randnum(2)) - eye(2));
    else
        rhot = 3*pauli(:,:,randnum(1))'*(state1*state1')*pauli(:,:,randnum(1)) - eye(2);
        rhot = kron(rhot,3*pauli(:,:,randnum(2))'*(state1*state1')*pauli(:,:,randnum(2)) - eye(2));
    end
    record_rho = record_rho + rhot;
end
record_rho = record_rho/n;
sqrt(trace((rho - record_rho)'*(rho - record_rho)))
Fidelity(rho,record_rho)

Answer 1

There's a few bugs in your code as well as a slight misunderstanding about the guarantees of the protocol.

First to clarify some details: The protocol you implement samples $U \in \text{Cl}(2)^{\otimes 2}$ (not the same as $U\in \text{Cl}(2^2)$!) and then applies these operators pre-measurement. Up to global phase this is equivalent to performing local measurements randomly selected to be in either the $x$-, $y$-, or $z$-basis. And so equivalently this can be implemented as randomly applying a unitary $U_j$ drawn from $\{H, HS^\dagger, I\}$, applying to each qubit $j=0, 1$, and then measuring $|b_0 b_1\rangle\langle b_0 b_1|$.

Mathematically the effect of this pre-measurement rotation on an input state averaged over many trials turns out to be the depolarizing channel. This is why you recover a shadow of the input state by applying the inverse of a depolarizing channel \begin{align} \hat{\rho} &= \left(\mathcal{D}_{1/3}^{-1}\right)^{\otimes 2} \circ (U_0^\dagger \otimes U_1^\dagger)|b_0 b_1\rangle\langle b_0 b_1| (U_0 \otimes U_1)\tag{1} \\&=\bigotimes_{j=1}^2(3U^{\dagger}_j|\hat{b}_j\rangle\langle\hat{b}_j|U_j-\mathbb{I}) \end{align}

where the equivalence of these two expressions is derived very nicely in Huang's paper. However, this protocol doesn't actually guarantee tight bounds on $\text{Tr}(\rho \hat{\rho})$. Rather, the guarantee of this protocol is that for $O_i$ taken from the set of $2$-local Paulis, the difference

$$\tag{2} |\text{Tr}(O_i \hat{\rho}) - \text{Tr}(O_i \rho)| \leq \epsilon $$

will be upper bounded according to the shadow norm and accuracy parameters. On the other hand, if you were interested in observables with much lower locality than the size of the system you would still see good convergence re: Equation (2) but the distance between the input state and the classical shadow could be large.

With these factors in mind, here is a working codeblock slightly modified from yours:

from math import exp
from qiskit import *
from qiskit import Aer
import numpy as np
import matplotlib.pyplot as plt
from random import randrange

np.random.seed(222)
def one_shot(operator):
    
    sim = Aer.get_backend('aer_simulator')
    qc = QuantumCircuit(2)
#     unitary = [qc.h, qc.sdg, qc.id] # no! Bug #1
    unitary = [[qc.h], [qc.sdg, qc.h], [qc.id]] 
    qc.h(0)
    qc.cx(0,1)
    # new code:
    for x in unitary[operator[0]]:
        x(0)
    for x in unitary[operator[1]]:
        x(1)
    qc.measure_all()
    qobj = assemble(qc,shots=1)
    result = sim.run(qobj).result().get_counts()
    return result

def distance(rho):
    '''
    calculate distance of two density matrix
    '''
    return np.sqrt(np.trace(rho.conjugate().transpose().dot(rho)))

hadamard = 1/np.sqrt(2)*np.array([[1,1],[1,-1]])
s_gate = np.array([[1,0],[0,-1j]],dtype=complex)
id = np.identity(2) # don't overwrite builtins with variables!
unitary = [hadamard,np.dot(hadamard,s_gate),id]
snapshot_num = 5000 # more shots
state0 = np.array([[1,0],[0,0]])
state1 = np.array([[0,0],[0,1]])
states = [state0, state1]
record_rho = np.zeros([4,4])
for i in range(snapshot_num):
    randnum = np.random.randint(0,3,size=2)
    result = one_shot(randnum)
    
    bit0, bit1 = [int(x) for x in list(result.keys())[0]] # assuming one shot
    U0, U1 = unitary[randnum[0]], unitary[randnum[1]]
    # No! Bug #2: your parentheses were incorrect here
    rhohat = np.kron(3* U0.conj().T @ states[bit0] @ U0 - id , 3* U1.conj().T @ states[bit1] @ U1 - id)
    record_rho = record_rho + rhohat
    
record_rho = record_rho/snapshot_num
bell_state = np.array([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]])

print("State distance")
print(distance(record_rho - bell_state))

# VERIFICATION
print("verify 2-local expecations converge")
I = np.eye(2)
X = np.array([[0, 1], [1, 0]])
Y = np.array([[0, -1j], [1j, 0]])
Z = np.array([[1, 0], [0, -1]])
paulis = [I, X, Y, Z]

for i in range(4):
    for j in range(4):
        twolocal_pauli = np.kron(paulis[i], paulis[j])
        expectation_true = np.trace(bell_state @ twolocal_pauli)
        expectation_rhohat = np.trace(record_rho @ twolocal_pauli)
        print((i,j), abs(expectation_true - expectation_rhohat))

which outputs:

State distance
(0.06070337717129093+0j)
verify 2-local expecations converge
(0, 0) 0.0
(0, 1) 0.045599999999999995
(0, 2) 0.0108
(0, 3) 0.008400000000000074
(1, 0) 0.009000000000000003
(1, 1) 0.03160000000000007
(1, 2) 0.03779999999999999
(1, 3) 0.003599999999999997
(2, 0) 0.006599999999999998
(2, 1) 0.0342
(2, 2) 0.04960000000000009
(2, 3) 0.04319999999999999
(3, 0) 0.013200000000000045
(3, 1) 0.0414
(3, 2) 0.034199999999999994
(3, 3) 0.036799999999999944

So after increasing shots to 5000 can see nice convergence on $\text{Tr}(\hat{\rho} \sigma_i \otimes \sigma_j)$ even if the fidelity $\text{Tr}(\hat{\rho}\rho)$ is increasing at a slower rate$^1$.

The bugs are described below:

1. unitary = [qc.h, qc.sdg, qc.id] in the function one_shot: The $S^\dagger$ is not how you perform a measurement in the Y-basis. You need $H S^\dagger$.
2. rho = np.kron(3*np.dot(unitary[randnum[0]].conj().T,state0).dot(unitary[randnum[0]] - id),3*np.dot(unitary[randnum[1]].conj().T,state0).dot(unitary[randnum[1]]) - id) in the loop over snapshot_num: specifically, .dot(unitary[randnum[0]] - id) has the parentheses in the wrong place. If you had inspected $\rho$ you would have observed it was often equal to the all-zeros matrix which is a bad sign. While we can't rely on $\text{Tr}(\hat{\rho})=1$ in this protocol because $M^{-1}$ is not trace preserving, we can at least expect $\text{Tr}(\hat{\rho})\geq 1$.

---

$^1$ In this special case where the observables are $n$-local rather than $k$-local for $k<n$, I think by triangle inequality the state distance is actually bounded as a sum of the bounds on the observables: \begin{align} \lVert\rho - \hat{\rho}\rVert_1 &= \left\lVert \sum_{i \in \{I, X,Y,Z\}^2 } (\text{Tr}(P_i\rho )- \text{Tr}(P_i\hat{\rho} ) )P_i \right\rVert_1 \\\leq &\sum_{i \in \{I, X,Y,Z\}^2 } \left\lVert(\text{Tr}(P_i\rho )- \text{Tr}(P_i\hat{\rho} ) )P_i \right\rVert_1 \\&\leq 16 \epsilon \end{align} where $P_i = \sigma_{i_1} \otimes \sigma_{i_2}$ and the bound of Equation (2) was applied. This is why the sample complexity of estimating classical shadows for $n$-local Pauli operators (which are a basis for $n$-qubit states) with error parameter $\epsilon$ is much lower than the sample complexity for state tomography with error parameter $\epsilon$.