Allow classical shadow measurements in new device (#4162)

* Support classical shadow measurements

* Add shot vector tests

* Address comments

* Add comment for diag_list

* Apply suggestions from code review

Co-authored-by: Christina Lee <christina@xanadu.ai>

* Fix

* Add unit tests for process_state_with_shots

* Update preprocessing

* Apply suggestions from code review

Co-authored-by: Frederik Wilde <42576579+frederikwilde@users.noreply.github.com>

* Address PR comments

* set rng for test

---------

Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: Frederik Wilde <42576579+frederikwilde@users.noreply.github.com>

doc/releases/changelog-dev.md

@@ -148,6 +148,7 @@
 * Added a function `measure_with_samples` that returns a sample-based measurement result given a state
   [(#4083)](https://github.com/PennyLaneAI/pennylane/pull/4083)
   [(#4093)](https://github.com/PennyLaneAI/pennylane/pull/4093)
+  [(#4162)](https://github.com/PennyLaneAI/pennylane/pull/4162)
   [(#4254)](https://github.com/PennyLaneAI/pennylane/pull/4254)
 
 * Wrap all objects being queued in an `AnnotatedQueue` so that `AnnotatedQueue` is not dependent on

pennylane/devices/qubit/preprocess.py

@@ -23,7 +23,14 @@
 import pennylane as qml
 
 from pennylane.operation import Tensor
-from pennylane.measurements import MidMeasureMP, StateMeasurement, SampleMeasurement, ExpectationMP
+from pennylane.measurements import (
+    MidMeasureMP,
+    StateMeasurement,
+    SampleMeasurement,
+    ExpectationMP,
+    ClassicalShadowMP,
+    ShadowExpvalMP,
+)
 from pennylane.typing import ResultBatch, Result
 from pennylane import DeviceError
 
@@ -187,9 +194,9 @@ def validate_measurements(
             )
 
         for m in circuit.measurements:
-            if not isinstance(m, SampleMeasurement):
+            if not isinstance(m, (SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP)):
                 raise DeviceError(
-                    f"Circuits with finite shots must only contain SampleMeasurements; got {m}"
+                    f"Circuits with finite shots must only contain SampleMeasurements, ClassicalShadowMP, or ShadowExpvalMP; got {m}"
                 )
 
 

pennylane/devices/qubit/sampling.py

@@ -12,25 +12,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Functions to sample a state."""
+from typing import Union
 
 import numpy as np
 import pennylane as qml
 from pennylane.ops import Sum, Hamiltonian
-from pennylane.measurements import SampleMeasurement, Shots, ExpectationMP
+from pennylane.measurements import (
+    SampleMeasurement,
+    Shots,
+    ExpectationMP,
+    ClassicalShadowMP,
+    ShadowExpvalMP,
+)
 from pennylane.typing import TensorLike
 from .apply_operation import apply_operation
 
 
 def measure_with_samples(
-    mp: SampleMeasurement, state: np.ndarray, shots: Shots, is_state_batched: bool = False, rng=None
+    mp: Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP],
+    state: np.ndarray,
+    shots: Shots,
+    is_state_batched: bool = False,
+    rng=None,
 ) -> TensorLike:
     """
     Returns the samples of the measurement process performed on the given state.
     This function assumes that the user-defined wire labels in the measurement process
     have already been mapped to integer wires used in the device.
 
     Args:
-        mp (~.measurements.SampleMeasurement): The sample measurement to perform
+        mp (Union[~.measurements.SampleMeasurement, ~.measurements.ClassicalShadowMP, ~.measurements.ShadowExpvalMP]):
+            The sample measurement to perform
         state (np.ndarray[complex]): The state vector to sample from
         shots (~.measurements.Shots): The number of samples to take
         is_state_batched (bool): whether the state is batched or not
@@ -71,6 +83,9 @@ def _sum_for_single_shot(s):
             unsqueezed_results = tuple(_sum_for_single_shot(Shots(s)) for s in shots)
             return unsqueezed_results if shots.has_partitioned_shots else unsqueezed_results[0]
 
+    if isinstance(mp, (ClassicalShadowMP, ShadowExpvalMP)):
+        return _measure_classical_shadow(mp, state, shots, rng=rng)
+
     # measure with the usual method (rotate into the measurement basis)
     return _measure_with_samples_diagonalizing_gates(
         mp, state, shots, is_state_batched=is_state_batched, rng=rng
@@ -139,6 +154,35 @@ def _measure_with_samples_diagonalizing_gates(
     return processed
 
 
+def _measure_classical_shadow(
+    mp: Union[ClassicalShadowMP, ShadowExpvalMP], state: np.ndarray, shots: Shots, rng=None
+):
+    """
+    Returns the result of a classical shadow measurement on the given state.
+
+    A classical shadow measurement doesn't fit neatly into the current measurement API
+    since different diagonalizing gates are used for each shot. Here it's treated as a
+    state measurement with shots instead of a sample measurement.
+
+    Args:
+        mp (~.measurements.SampleMeasurement): The sample measurement to perform
+        state (np.ndarray[complex]): The state vector to sample from
+        shots (~.measurements.Shots): The number of samples to take
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+
+    Returns:
+        TensorLike[Any]: Sample measurement results
+    """
+    wires = qml.wires.Wires(range(len(state.shape)))
+
+    if shots.has_partitioned_shots:
+        return tuple(mp.process_state_with_shots(state, wires, s, rng=rng) for s in shots)
+
+    return mp.process_state_with_shots(state, wires, shots.total_shots, rng=rng)
+
+
 def sample_state(
     state, shots: int, is_state_batched: bool = False, wires=None, rng=None
 ) -> np.ndarray:

pennylane/measurements/classical_shadow.py

@@ -17,6 +17,7 @@
 import copy
 from collections.abc import Iterable
 from typing import Optional, Union, Sequence
+from string import ascii_letters as ABC
 
 import numpy as np
 
@@ -309,6 +310,120 @@ def process(self, tape, device):
 
         return qml.math.cast(qml.math.stack([outcomes, recipes]), dtype=np.int8)
 
+    def process_state_with_shots(
+        self, state: Sequence[complex], wire_order: Wires, shots: int, rng=None
+    ):
+        """Process the given quantum state with the given number of shots
+
+        Args:
+            state (Sequence[complex]): quantum state vector given as a rank-N tensor, where
+                each dim has size 2 and N is the number of wires.
+            wire_order (Wires): wires determining the subspace that ``state`` acts on; a matrix of
+                dimension :math:`2^n` acts on a subspace of :math:`n` wires
+            shots (int): The number of shots
+            rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+                seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+                If no value is provided, a default RNG will be used. The random measurement outcomes
+                in the form of bits will be generated from this argument, while the random recipes will be
+                created from the ``seed`` argument provided to ``.ClassicalShadowsMP``.
+
+        Returns:
+            tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
+            the measured bits and the second represents the recipes used.
+        """
+        wire_map = {w: i for i, w in enumerate(wire_order)}
+        mapped_wires = [wire_map[w] for w in self.wires]
+        n_qubits = len(mapped_wires)
+        num_dev_qubits = len(state.shape)
+
+        # seed the random measurement generation so that recipes
+        # are the same for different executions with the same seed
+        recipe_rng = np.random.RandomState(self.seed)
+        recipes = recipe_rng.randint(0, 3, size=(shots, n_qubits))
+
+        bit_rng = np.random.default_rng(rng)
+
+        obs_list = np.stack(
+            [
+                qml.PauliX.compute_matrix(),
+                qml.PauliY.compute_matrix(),
+                qml.PauliZ.compute_matrix(),
+            ]
+        )
+
+        # the diagonalizing matrices corresponding to the Pauli observables above
+        diag_list = np.stack(
+            [
+                qml.Hadamard.compute_matrix(),
+                qml.Hadamard.compute_matrix() @ qml.RZ.compute_matrix(-np.pi / 2),
+                qml.Identity.compute_matrix(),
+            ]
+        )
+        obs = obs_list[recipes]
+        diagonalizers = diag_list[recipes]
+
+        # There's a significant speedup if we use the following iterative
+        # process to perform the randomized Pauli measurements:
+        #   1. Randomly generate Pauli observables for all snapshots for
+        #      a single qubit (e.g. the first qubit).
+        #   2. Compute the expectation of each Pauli observable on the first
+        #      qubit by tracing out all other qubits.
+        #   3. Sample the first qubit based on each Pauli expectation.
+        #   4. For all snapshots, determine the collapsed state of the remaining
+        #      qubits based on the sample result.
+        #   4. Repeat iteratively until no qubits are remaining.
+        #
+        # Observe that after the first iteration, the second qubit will become the
+        # "first" qubit in the process. The advantage to this approach as opposed to
+        # simulataneously computing the Pauli expectations for each qubit is that
+        # the partial traces are computed over iteratively smaller subsystems, leading
+        # to a significant speed-up.
+
+        # transpose the state so that the measured wires appear first
+        unmeasured_wires = [i for i in range(num_dev_qubits) if i not in mapped_wires]
+        transposed_state = np.transpose(state, axes=mapped_wires + unmeasured_wires)
+
+        outcomes = np.zeros((shots, n_qubits))
+        stacked_state = np.repeat(transposed_state[np.newaxis, ...], shots, axis=0)
+
+        for active_qubit in range(n_qubits):
+            # stacked_state loses a dimension each loop
+
+            # trace out every qubit except the first
+            num_remaining_qubits = num_dev_qubits - active_qubit
+            conj_state_first_qubit = ABC[num_remaining_qubits]
+            stacked_dim = ABC[num_remaining_qubits + 1]
+
+            state_str = f"{stacked_dim}{ABC[:num_remaining_qubits]}"
+            conj_state_str = f"{stacked_dim}{conj_state_first_qubit}{ABC[1:num_remaining_qubits]}"
+            target_str = f"{stacked_dim}a{conj_state_first_qubit}"
+
+            first_qubit_state = np.einsum(
+                f"{state_str},{conj_state_str}->{target_str}",
+                stacked_state,
+                np.conj(stacked_state),
+            )
+
+            # sample the observables on the first qubit
+            probs = (np.einsum("abc,acb->a", first_qubit_state, obs[:, active_qubit]) + 1) / 2
+            samples = bit_rng.random(size=probs.shape) > probs
+            outcomes[:, active_qubit] = samples
+
+            # collapse the state of the remaining qubits; the next qubit in line
+            # becomes the first qubit for the next iteration
+            rotated_state = np.einsum(
+                "ab...,acb->ac...", stacked_state, diagonalizers[:, active_qubit]
+            )
+            stacked_state = rotated_state[np.arange(shots), samples.astype(np.int8)]
+
+            # re-normalize the collapsed state
+            sum_indices = tuple(range(1, num_remaining_qubits))
+            state_squared = np.abs(stacked_state) ** 2
+            norms = np.sqrt(np.sum(state_squared, sum_indices, keepdims=True))
+            stacked_state /= norms
+
+        return np.stack([outcomes, recipes]).astype(np.int8)
+
     @property
     def samples_computational_basis(self):
         return False
@@ -373,6 +488,29 @@ def process(self, tape, device):
         shadow = qml.shadows.ClassicalShadow(bits, recipes, wire_map=self.wires.tolist())
         return shadow.expval(self.H, self.k)
 
+    def process_state_with_shots(
+        self, state: Sequence[complex], wire_order: Wires, shots: int, rng=None
+    ):
+        """Process the given quantum state with the given number of shots
+
+        Args:
+            state (Sequence[complex]): quantum state
+            wire_order (Wires): wires determining the subspace that ``state`` acts on; a matrix of
+                dimension :math:`2^n` acts on a subspace of :math:`n` wires
+            shots (int): The number of shots
+            rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+                seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+                If no value is provided, a default RNG will be used.
+
+        Returns:
+            float: The estimate of the expectation value.
+        """
+        bits, recipes = qml.classical_shadow(
+            wires=self.wires, seed=self.seed
+        ).process_state_with_shots(state, wire_order, shots, rng=rng)
+        shadow = qml.shadows.ClassicalShadow(bits, recipes, wire_map=self.wires.tolist())
+        return shadow.expval(self.H, self.k)
+
     @property
     def samples_computational_basis(self):
         return False

tests/devices/experimental/test_default_qubit_2.py

@@ -1230,6 +1230,108 @@ def test_complex_hamiltonian(self):
         assert np.allclose(res, expected, atol=0.001)
 
 
+class TestClassicalShadows:
+    """Test that classical shadow measurements works with the new device"""
+
+    @pytest.mark.parametrize("n_qubits", [1, 2, 3])
+    def test_shape_and_dtype(self, n_qubits):
+        """Test that the shape and dtype of the measurement is correct"""
+        dev = DefaultQubit2()
+
+        ops = [qml.Hadamard(i) for i in range(n_qubits)]
+        qs = qml.tape.QuantumScript(ops, [qml.classical_shadow(range(n_qubits))], shots=100)
+        res = dev.execute(qs)
+
+        assert res.shape == (2, 100, n_qubits)
+        assert res.dtype == np.int8
+
+        # test that the bits are either 0 and 1
+        assert np.all(np.logical_or(res[0] == 0, res[0] == 1))
+
+        # test that the recipes are either 0, 1, or 2 (X, Y, or Z)
+        assert np.all(np.logical_or(np.logical_or(res[1] == 0, res[1] == 1), res[1] == 2))
+
+    def test_expval(self):
+        """Test that shadow expval measurements work as expected"""
+        dev = DefaultQubit2(seed=100)
+
+        ops = [qml.Hadamard(0), qml.Hadamard(1)]
+        meas = [qml.shadow_expval(qml.PauliX(0) @ qml.PauliX(1), seed=200)]
+        qs = qml.tape.QuantumScript(ops, meas, shots=1000)
+        res = dev.execute(qs)
+
+        assert res.shape == ()
+        assert np.allclose(res, 1.0, atol=0.05)
+
+    def test_reconstruct_bell_state(self):
+        """Test that a bell state can be faithfully reconstructed"""
+        dev = DefaultQubit2(seed=100)
+
+        ops = [qml.Hadamard(0), qml.CNOT([0, 1])]
+        meas = [qml.classical_shadow(wires=[0, 1], seed=200)]
+        qs = qml.tape.QuantumScript(ops, meas, shots=1000)
+
+        # should prepare the bell state
+        bits, recipes = dev.execute(qs)
+        shadow = qml.shadows.ClassicalShadow(bits, recipes)
+        global_snapshots = shadow.global_snapshots()
+
+        state = np.sum(global_snapshots, axis=0) / shadow.snapshots
+        bell_state = np.array([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]])
+        assert qml.math.allclose(state, bell_state, atol=0.1)
+
+        # reduced state should yield maximally mixed state
+        local_snapshots = shadow.local_snapshots(wires=[0])
+        assert qml.math.allclose(np.mean(local_snapshots, axis=0)[0], 0.5 * np.eye(2), atol=0.05)
+
+        # alternative computation
+        ops = [qml.Hadamard(0), qml.CNOT([0, 1])]
+        meas = [qml.classical_shadow(wires=[0], seed=200)]
+        qs = qml.tape.QuantumScript(ops, meas, shots=1000)
+        bits, recipes = dev.execute(qs)
+
+        shadow = qml.shadows.ClassicalShadow(bits, recipes)
+        global_snapshots = shadow.global_snapshots()
+        local_snapshots = shadow.local_snapshots(wires=[0])
+
+        state = np.sum(global_snapshots, axis=0) / shadow.snapshots
+        assert qml.math.allclose(state, 0.5 * np.eye(2), atol=0.1)
+        assert np.all(local_snapshots[:, 0] == global_snapshots)
+
+    @pytest.mark.parametrize("n_qubits", [1, 2, 3])
+    @pytest.mark.parametrize(
+        "shots",
+        [
+            [1000, 1000],
+            [(1000, 2)],
+            [1000, 2000],
+            [(1000, 2), 2000],
+            [(1000, 3), 2000, (3000, 2)],
+        ],
+    )
+    def test_shot_vectors(self, n_qubits, shots):
+        """Test that classical shadows works when given a shot vector"""
+        dev = DefaultQubit2()
+        shots = qml.measurements.Shots(shots)
+
+        ops = [qml.Hadamard(i) for i in range(n_qubits)]
+        qs = qml.tape.QuantumScript(ops, [qml.classical_shadow(range(n_qubits))], shots=shots)
+        res = dev.execute(qs)
+
+        assert isinstance(res, tuple)
+        assert len(res) == len(list(shots))
+
+        for r, s in zip(res, shots):
+            assert r.shape == (2, s, n_qubits)
+            assert r.dtype == np.int8
+
+            # test that the bits are either 0 and 1
+            assert np.all(np.logical_or(r[0] == 0, r[0] == 1))
+
+            # test that the recipes are either 0, 1, or 2 (X, Y, or Z)
+            assert np.all(np.logical_or(np.logical_or(r[1] == 0, r[1] == 1), r[1] == 2))
+
+
 def test_broadcasted_parameter():
     """Test that DefaultQubit2 handles broadcasted parameters as expected."""
     dev = DefaultQubit2()