Lightning qubit2 add measurements (#627)

* lightning interface for new device api * Auto update version * Merge Master * update version * update version * update lightning_qubit_2 * add LightningQubit2 to init * add LightningStateVector class * add LightningMeasurements class * add new QuantumScriptSerializer class * allow lightning.qubit2 to be tested within our suite * add tests and CI workflow for lightning_qubit_2 * update CI * update CI * add wire mapping, black * add tests for custom wires * add tests for custom wires * add review suggestions * format * remove python class to reverse order of PRs * update simulate to get a LightningStateVector * add reset state * update simulate * update docs * add Result import * Update pennylane_lightning/lightning_qubit/_measurements.py Co-authored-by: Christina Lee <[email protected]> * Update pennylane_lightning/lightning_qubit/_measurements.py Co-authored-by: Christina Lee <[email protected]> * fix reset state * Update pennylane_lightning/lightning_qubit/_state_vector.py Co-authored-by: Christina Lee <[email protected]> * Update pennylane_lightning/lightning_qubit/_state_vector.py Co-authored-by: Christina Lee <[email protected]> * Update pennylane_lightning/lightning_qubit/_state_vector.py Co-authored-by: Christina Lee <[email protected]> * remove LightningQubit2 references * remove unnecessary modules * merging Serializer classes * update serialize tests * update measurements with new serialize class * remove outdated test * remove obsolete tests * remove unused dtype input from simulate * update measurements * update state_vector * update lightning_qubit2 * format * pylint * Update pennylane_lightning/lightning_qubit/_state_vector.py Co-authored-by: Christina Lee <[email protected]> * remove old comment * some review suggestions * Auto update version * remove print * update state vector class * add state vector class tests * adding measurement tests * update state vector and tests * move and rename test files, and format * Auto update version * skip measurements class for other devices and in the absence of binaries * format * update measurements class * expand measurement class testing * garbage collection * typo * update coverage and StateVector class * expand measurements class coverage * Auto update version * add coverage for n-controlled operations * add map to standard wires to get_final_state for safety * update jax config import * Auto update version * trigger CI * update state vector class and tests for improved coverage * update measurement class tests * update dev version * remove device definition * update dev version * clean test_measurements_class.py * isort+black * review suggestion * fix docs * Add qml.var support. * Add probs support. * increase tolerance * Auto update version * isort * Add double-obs tests. * Pin pytest version (#624) * update dev version * update changelog * pin pytest version in requirement files * add a requirements file for tests against Pennylane master * update wheels' workflows * Version Bump (#626) * post release version bump * trigger CI --------- Co-authored-by: AmintorDusko <[email protected]> Co-authored-by: AmintorDusko <[email protected]> * increase tolerance * Introduce isort. (#623) * Introduce isort. * Auto update version * Update changelog * Auto update version * Update changelog. * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> * Auto update version * isort * Add qml.var support. * Add probs support. * Add measurement tests with wires. * review suggestions * remove unused imports * remove diagonalization gate application from state vector * pytest.skip tests * Auto update version * Fix format * Update tests/lightning_qubit/test_measurements_class.py Co-authored-by: Amintor Dusko <[email protected]> * Update pennylane_lightning/lightning_qubit/_measurements.py Co-authored-by: Christina Lee <[email protected]> * Remove pylint: disable=protected-access * trigger ci * Add var tests * update changelog --------- Co-authored-by: albi3ro <[email protected]> Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: AmintorDusko <[email protected]> Co-authored-by: Amintor Dusko <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: AmintorDusko <[email protected]>
PennyLaneAI · Mar 7, 2024 · cdb098a · cdb098a
1 parent 22d04fe
commit cdb098a
Show file tree

Hide file tree

Showing 5 changed files with 269 additions and 22 deletions.
diff --git a/.github/CHANGELOG.md b/.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

diff --git a/pennylane_lightning/core/_version.py b/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"
diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/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}")

diff --git a/pennylane_lightning/lightning_qubit/_measurements.py b/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])