CMP #5: Add Sample MP class

pennylane/circuit_graph.py

@@ -18,11 +18,10 @@
 # pylint: disable=too-many-branches,too-many-arguments,too-many-instance-attributes
 from collections import namedtuple
 
-import retworkx as rx
 import numpy as np
+import retworkx as rx
 
 import pennylane as qml
-
 from pennylane.wires import Wires
 
 
@@ -125,7 +124,7 @@ def __init__(self, ops, obs, wires, par_info=None, trainable_params=None):
 
                 elif op.return_type is qml.measurements.Sample and op.wires == Wires([]):
                     # Sampling without specifying wires is treated as sampling all wires
-                    op = qml.measurements.MeasurementProcess(qml.measurements.Sample, wires=wires)
+                    op = qml.sample(wires=wires)
 
                 op.queue_idx = k
 

pennylane/measurements/classical_shadow.py

@@ -16,6 +16,7 @@
 This module contains the qml.classical_shadow measurement.
 """
 from collections.abc import Iterable
+from typing import Sequence, Tuple
 
 import numpy as np
 
@@ -293,3 +294,8 @@ def __copy__(self):
         obj.H = self.H
         obj.k = self.k
         return obj
+
+    def process(
+        self, samples: Sequence[complex], shot_range: Tuple[int] = None, bin_size: int = None
+    ):
+        """TODO: Implement this method."""

pennylane/measurements/measurements.py

@@ -21,12 +21,14 @@
 import contextlib
 import copy
 import functools
+from abc import abstractmethod
 from enum import Enum
+from typing import Sequence, Tuple, Union
 
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import Operator
+from pennylane.operation import Observable
 from pennylane.wires import Wires
 
 # =============================================================================
@@ -106,7 +108,7 @@ class MeasurementShapeError(ValueError):
     quantum tape."""
 
 
-class MeasurementProcess:
+class MeasurementProcess:  # TODO: Inherit from ABC
     """Represents a measurement process occurring at the end of a
     quantum variational circuit.
 
@@ -131,7 +133,7 @@ class MeasurementProcess:
     def __init__(
         self,
         return_type: ObservableReturnTypes,
-        obs: Operator = None,
+        obs: Union[Observable, None] = None,
         wires=None,
         eigvals=None,
         id=None,
@@ -692,3 +694,18 @@ def map_wires(self, wire_map: dict):
         if self.obs is not None:
             new_measurement.obs = self.obs.map_wires(wire_map=wire_map)
         return new_measurement
+
+    @abstractmethod
+    def process(
+        self, samples: Sequence[complex], shot_range: Tuple[int] = None, bin_size: int = None
+    ):
+        """Process the given samples.
+
+        Args:
+            samples (Sequence[complex]): computational basis samples generated for all wires
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+        """

pennylane/measurements/sample.py

@@ -16,13 +16,18 @@
 This module contains the qml.sample measurement.
 """
 import warnings
+from typing import Tuple, Union
 
+import numpy as np
+
+import pennylane as qml
+from pennylane.operation import Observable
 from pennylane.wires import Wires
 
 from .measurements import MeasurementProcess, Sample
 
 
-def sample(op=None, wires=None):
+def sample(op: Union[Observable, None] = None, wires=None):
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device,
     returning raw samples. If no observable is provided then basis state samples are returned
@@ -106,4 +111,47 @@ def circuit(x):
             )
         wires = Wires(wires)
 
-    return MeasurementProcess(Sample, obs=op, wires=wires)
+    return _Sample(Sample, obs=op, wires=wires)
+
+
+# TODO: Make public when removing the ObservableReturnTypes enum
+class _Sample(MeasurementProcess):
+    """Measurement process that returns the samples of a given observable."""
+
+    def process(self, samples: np.ndarray, shot_range: Tuple[int] = None, bin_size: int = None):
+        name = self.obs.name if self.obs is not None else None
+        # Select the samples from samples that correspond to ``shot_range`` if provided
+        if shot_range is not None:
+            # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last
+            # colon (:) is required because shots is the second-to-last axis and the
+            # Ellipsis (...) otherwise would take up broadcasting and shots axes.
+            samples = samples[..., slice(*shot_range), :]
+
+        if len(self.wires) != 0:
+            # if wires are provided, then we only return samples from those wires
+            samples = samples[..., np.array(self.wires)]
+
+        num_wires = samples.shape[-1]  # wires is the last dimension
+
+        if self.obs is None:
+            # if no observable was provided then return the raw samples
+            return samples if bin_size is None else samples.reshape(num_wires, bin_size, -1)
+
+        if name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
+            # Process samples for observables with eigenvalues {1, -1}
+            samples = 1 - 2 * qml.math.squeeze(samples)
+        else:
+            # Replace the basis state in the computational basis with the correct eigenvalue.
+            # Extract only the columns of the basis samples required based on ``wires``.
+            powers_of_two = 2 ** np.arange(num_wires)[::-1]
+            indices = samples @ powers_of_two
+            indices = np.array(indices)  # Add np.array here for Jax support.
+            try:
+                samples = self.obs.eigvals()[indices]
+            except qml.operation.EigvalsUndefinedError as e:
+                # if observable has no info on eigenvalues, we cannot return this measurement
+                raise qml.operation.EigvalsUndefinedError(
+                    f"Cannot compute samples of {self.obs.name}."
+                ) from e
+
+        return samples if bin_size is None else samples.reshape((bin_size, -1))

pennylane/ops/functions/equal.py

@@ -16,6 +16,7 @@
 """
 # pylint: disable=too-many-arguments,too-many-return-statements
 from typing import Union
+
 import pennylane as qml
 from pennylane.measurements import MeasurementProcess, ShadowMeasurementProcess
 from pennylane.operation import Operator
@@ -85,12 +86,12 @@ def equal(
     if op1.__class__ is not op2.__class__:
         return False
 
-    if op1.__class__ is MeasurementProcess:
-        return equal_measurements(op1, op2)
-
-    if op1.__class__ is ShadowMeasurementProcess:
+    if isinstance(op1, ShadowMeasurementProcess):
         return equal_shadow_measurements(op1, op2)
 
+    if isinstance(op1, MeasurementProcess):
+        return equal_measurements(op1, op2)
+
     return equal_operator(op1, op2, check_interface, check_trainability, rtol, atol)