Lightning qubit2 add measurements

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add analytic-mode `qml.probs` and `qml.var` support in `lightning.qubit2`.
+  [(#627)](https://github.com/PennyLaneAI/pennylane-lightning/pull/627)
+
 ### Breaking changes
 
 ### Improvements

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev2"
+__version__ = "0.36.0-dev3"

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -65,18 +65,20 @@
     from pennylane.wires import Wires
 
     # pylint: disable=import-error, no-name-in-module, ungrouped-imports
-    from pennylane_lightning.core._serialize import (
-        QuantumScriptSerializer,
-        global_phase_diagonal,
-    )
-    from pennylane_lightning.core._version import __version__
     from pennylane_lightning.lightning_kokkos_ops.algorithms import (
         AdjointJacobianC64,
         AdjointJacobianC128,
         create_ops_listC64,
         create_ops_listC128,
     )
 
+    # pylint: disable=import-error, no-name-in-module, ungrouped-imports
+    from pennylane_lightning.core._serialize import (
+        QuantumScriptSerializer,
+        global_phase_diagonal,
+    )
+    from pennylane_lightning.core._version import __version__
+
     def _kokkos_dtype(dtype):
         if dtype not in [np.complex128, np.complex64]:  # pragma: no cover
             raise ValueError(f"Data type is not supported for state-vector computation: {dtype}")

pennylane_lightning/lightning_qubit/_measurements.py

@@ -28,7 +28,13 @@
 
 import numpy as np
 import pennylane as qml
-from pennylane.measurements import ExpectationMP, MeasurementProcess, StateMeasurement
+from pennylane.measurements import (
+    ExpectationMP,
+    MeasurementProcess,
+    ProbabilityMP,
+    StateMeasurement,
+    VarianceMP,
+)
 from pennylane.tape import QuantumScript
 from pennylane.typing import Result, TensorLike
 from pennylane.wires import Wires
@@ -87,7 +93,6 @@ def state_diagonalizing_gates(self, measurementprocess: StateMeasurement) -> Ten
         """
         diagonalizing_gates = measurementprocess.diagonalizing_gates()
         self._qubit_state.apply_operations(measurementprocess.diagonalizing_gates())
-
         state_array = self._qubit_state.state
         wires = Wires(range(self._qubit_state.num_wires))
 
@@ -133,6 +138,60 @@ def expval(self, measurementprocess: MeasurementProcess):
             measurementprocess.obs.name, measurementprocess.obs.wires
         )
 
+    def probs(self, measurementprocess: MeasurementProcess):
+        """Probabilities of the supplied observable or wires contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Probabilities of the supplied observable or wires
+        """
+        diagonalizing_gates = measurementprocess.diagonalizing_gates()
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(diagonalizing_gates)
+        results = self._measurement_lightning.probs(measurementprocess.wires.tolist())
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(
+                [qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)]
+            )
+        return results
+
+    def var(self, measurementprocess: MeasurementProcess):
+        """Variance of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Variance of the observable
+        """
+
+        if measurementprocess.obs.name == "SparseHamiltonian":
+            # ensuring CSR sparse representation.
+            CSR_SparseHamiltonian = measurementprocess.obs.sparse_matrix(
+                wire_order=list(range(self._qubit_state.num_wires))
+            ).tocsr(copy=False)
+            return self._measurement_lightning.var(
+                CSR_SparseHamiltonian.indptr,
+                CSR_SparseHamiltonian.indices,
+                CSR_SparseHamiltonian.data,
+            )
+
+        if (
+            measurementprocess.obs.name in ["Hamiltonian", "Hermitian"]
+            or (measurementprocess.obs.arithmetic_depth > 0)
+            or isinstance(measurementprocess.obs.name, List)
+        ):
+            ob_serialized = QuantumScriptSerializer(
+                self._qubit_state.device_name, self.dtype == np.complex64
+            )._ob(measurementprocess.obs)
+            return self._measurement_lightning.var(ob_serialized)
+
+        return self._measurement_lightning.var(
+            measurementprocess.obs.name, measurementprocess.obs.wires
+        )
+
     def get_measurement_function(
         self, measurementprocess: MeasurementProcess
     ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
@@ -153,6 +212,16 @@ def get_measurement_function(
                     return self.state_diagonalizing_gates
                 return self.expval
 
+            if isinstance(measurementprocess, ProbabilityMP):
+                return self.probs
+
+            if isinstance(measurementprocess, VarianceMP):
+                if measurementprocess.obs.name in [
+                    "Identity",
+                    "Projector",
+                ]:
+                    return self.state_diagonalizing_gates
+                return self.var
             if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
                 return self.state_diagonalizing_gates
 
@@ -182,9 +251,6 @@ def measure_final_state(self, circuit: QuantumScript) -> Result:
             Tuple[TensorLike]: The measurement results
         """
 
-        if circuit.shots:
-            raise NotImplementedError
-        # analytic case
         if len(circuit.measurements) == 1:
             return self.measurement(circuit.measurements[0])