Qutrit Device #2

doc/releases/changelog-dev.md

@@ -107,8 +107,9 @@
   [(#2709)](https://github.com/PennyLaneAI/pennylane/pull/2709)
 
 * Added `QutritDevice` as an abstract base class for qutrit devices.
-  [(#2781)](https://github.com/PennyLaneAI/pennylane/pull/2781)
-  * Added operation `qml.QutritUnitary` for applying user-specified unitary operations on qutrit devices.
+  ([#2781](https://github.com/PennyLaneAI/pennylane/pull/2781), [#2782](https://github.com/PennyLaneAI/pennylane/pull/2782))
+
+* Added operation `qml.QutritUnitary` for applying user-specified unitary operations on qutrit devices.
   [(#2699)](https://github.com/PennyLaneAI/pennylane/pull/2699)
 
 **Operator Arithmetic:**

pennylane/_qutrit_device.py

@@ -17,7 +17,7 @@
 
 # For now, arguments may be different from the signatures provided in QubitDevice to minimize size of pull request
 # e.g. instead of expval(self, observable, wires, par) have expval(self, observable)
-# pylint: disable=abstract-method, no-value-for-parameter,too-many-instance-attributes,too-many-branches, no-member, bad-option-value, arguments-renamed
+# pylint: disable=arguments-differ, abstract-method, no-value-for-parameter,too-many-instance-attributes,too-many-branches, no-member, bad-option-value, arguments-renamed
 import itertools
 
 import numpy as np
@@ -80,10 +80,6 @@ def capabilities(cls):
         capabilities.update(model="qutrit")
         return capabilities
 
-    def statistics(self, observables, shot_range=None, bin_size=None):
-        # Overloading QubitDevice.statistics() as VnEntropy and MutualInfo not yet supported for QutritDevice
-        raise NotImplementedError
-
     def generate_samples(self):
         r"""Returns the computational basis samples generated for all wires.
 
@@ -152,15 +148,18 @@ def states_to_ternary(samples, num_wires, dtype=np.int64):
     def density_matrix(self, wires):
         """Returns the reduced density matrix prior to measurement.
 
-        .. note::
+        Args:
+            wires (Wires): wires of the reduced system
 
-            Only state vector simulators support this property. Please see the
-            plugin documentation for more details.
+        Raises:
+            QuantumFunctionError: density matrix is currently unsupported on :class:`~.QutritDevice`
         """
         # TODO: Add density matrix support. Currently, qml.math is hard-coded to work only with qubit states,
         # (see `qml.math.reduced_dm()`) so it needs to be updated to be able to handle calculations for qutrits
         # before this method can be implemented.
-        raise NotImplementedError
+        raise qml.QuantumFunctionError(
+            "Unsupported return type specified for observable density matrix"
+        )
 
     def vn_entropy(self, wires, log_base):
         r"""Returns the Von Neumann entropy prior to measurement.
@@ -172,13 +171,15 @@ def vn_entropy(self, wires, log_base):
             wires (Wires): Wires of the considered subsystem.
             log_base (float): Base for the logarithm, default is None the natural logarithm is used in this case.
 
-        Returns:
-            float: returns the Von Neumann entropy
+        Raises:
+            QuantumFunctionError: Von Neumann entropy is currently unsupported on :class:`~.QutritDevice`
         """
         # TODO: Add support for VnEntropy return type. Currently, qml.math is hard coded to calculate this for qubit
         # states (see `qml.math.vn_entropy()`), so it needs to be updated before VnEntropy can be supported for qutrits.
         # For now, if a user tries to request this return type, an error will be raised.
-        raise NotImplementedError
+        raise qml.QuantumFunctionError(
+            "Unsupported return type specified for observable Von Neumann entropy"
+        )
 
     def mutual_info(self, wires0, wires1, log_base):
         r"""Returns the mutual information prior to measurement:
@@ -194,13 +195,15 @@ def mutual_info(self, wires0, wires1, log_base):
             wires1 (Wires): wires of the second subsystem
             log_base (float): base to use in the logarithm
 
-        Returns:
-            float: the mutual information
+        Raises:
+            QuantumFunctionError: Mutual information is currently unsupported on :class:`~.QutritDevice`
         """
         # TODO: Add support for MutualInfo return type. Currently, qml.math is hard coded to calculate this for qubit
         # states (see `qml.math.mutual_info()`), so it needs to be updated before MutualInfo can be supported for qutrits.
         # For now, if a user tries to request this return type, an error will be raised.
-        raise NotImplementedError
+        raise qml.QuantumFunctionError(
+            "Unsupported return type specified for observable mutual information"
+        )
 
     def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         """Return the estimated probability of each computational basis state
@@ -317,10 +320,27 @@ def marginal_prob(self, prob, wires=None):
         perm = basis_states @ powers_of_three
         return self._gather(prob, perm)
 
-    # TODO: Update in next PR to add counts capability for binning
-    def sample(
-        self, observable, shot_range=None, bin_size=None
-    ):  # pylint: disable=arguments-differ
+    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
+        def _samples_to_counts(samples, no_observable_provided):
+            """Group the obtained samples into a dictionary.
+
+            **Example**
+
+                >>> samples
+                tensor([[0, 0, 1],
+                        [0, 0, 1],
+                        [1, 1, 1]], requires_grad=True)
+                >>> self._samples_to_counts(samples)
+                {'111':1, '001':2}
+            """
+            if no_observable_provided:
+                # If we describe a state vector, we need to convert its list representation
+                # into string (it's hashable and good-looking).
+                # Before converting to str, we need to extract elements from arrays
+                # to satisfy the case of jax interface, as jax arrays do not support str.
+                samples = ["".join([str(s.item()) for s in sample]) for sample in samples]
+            states, counts = np.unique(samples, return_counts=True)
+            return dict(zip(states, counts))
 
         # TODO: Add special cases for any observables that require them once list of
         # observables is updated.
@@ -329,10 +349,9 @@ def sample(
         device_wires = self.map_wires(observable.wires)
         name = observable.name  # pylint: disable=unused-variable
         sample_slice = Ellipsis if shot_range is None else slice(*shot_range)
+        no_observable_provided = isinstance(observable, MeasurementProcess)
 
-        if isinstance(
-            observable, MeasurementProcess
-        ):  # if no observable was provided then return the raw samples
+        if no_observable_provided:  # if no observable was provided then return the raw samples
             if (
                 len(observable.wires) != 0
             ):  # if wires are provided, then we only return samples from those wires
@@ -359,9 +378,22 @@ def sample(
                 ) from e
 
         if bin_size is None:
+            if counts:
+                return _samples_to_counts(samples, no_observable_provided)
             return samples
 
-        return samples.reshape((bin_size, -1))
+        num_wires = len(device_wires) if len(device_wires) > 0 else self.num_wires
+        if counts:
+            shape = (-1, bin_size, num_wires) if no_observable_provided else (-1, bin_size)
+            return [
+                _samples_to_counts(bin_sample, no_observable_provided)
+                for bin_sample in samples.reshape(shape)
+            ]
+        return (
+            samples.reshape((num_wires, bin_size, -1))
+            if no_observable_provided
+            else samples.reshape((bin_size, -1))
+        )
 
     # TODO: Implement function. Currently unimplemented due to lack of decompositions available
     # for existing operations and lack of non-parametrized observables.