Repository: PennyLaneAI/pennylane
Pull request #5830 (2024-06-10 18:41:33)

Process density matrix
**Context:**
Current Pennylane already defined the measurement for pure states. However, the same interface for density matrices is yet implemented. Here in this Pull Request, we aim to provide this feature.

**Description of the Change:**
- Implemented `process_density_matrix` in both `ProbabilityMP` and `DensityMatrixMP` to correctly handle density matrices, accommodating specific subsystems or rearranged wire orders.
- Raise error if `StateMP.process_density_matrix` is called, as suggested at [PennyLane Issue #5628](https://github.com/PennyLaneAI/pennylane/issues/5628#issuecomment-2161383960).
- Kept the intrinsic behaviour of the two new methods as close to other implementations that already existed as possible. Specifically, the `ProbabilityMP.process_density_matrix` is implemented by reading out the full prob dist from the dm and then constructing a new 'pseudo state' to be processed by `process_state`.
- Added multiple test cases.

**Benefits:**
- **Flexibility and Robustness:** Provides versatile management of quantum state configurations, crucial for scenarios involving partial traces and wire reordering.
- **Enhanced Performance:** Kept extra computation introduced in the implementations as little as possible.

**Possible Drawbacks:**
- **Potential Coupled Issue:** `ProbabilityMP.process_density_matrix` calls `process_state`, which means these two methods are coupled and the latter's performance might affect the former.

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5628


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5830?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`2100476`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/21004763ba1dddb59898275b9d8e8df847ae9ef6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`01d6e5b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/01d6e5b9bdfed4fd830967e303d81feffbd76af6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 4 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5830      +/-   ##
==========================================
+ Coverage   99.51%   99.66%   +0.14%     
==========================================
  Files         415      415              
  Lines       39985    39512     -473     
==========================================
- Hits        39791    39379     -412     
+ Misses        194      133      -61     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5830?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: Thanks @JerryChen97! Let us know when this is ready for review or if you need a hand with anything.
JerryChen97: > Thanks @JerryChen97! Let us know when this is ready for review or if you need a hand with anything.

Sure! Still working on some testcases. Will mention you onfe fin
JerryChen97: > Thanks @JerryChen97! Let us know when this is ready for review or if you need a hand with anything.

Hi Thomas, on my side for now I don't have more to add or modify. Please review if available and hmu wherever potential issues arise
JerryChen97: @trbromley Hi Thomas would you like to review this PR? The auto-checks always take too long but they passed anyway
JerryChen97: Note that the version of `master` changed so I filled this pr into the new blank dev log.
I've also marked the discussion initiated by @trbromley regarding type hints as resolved, following @albi3ro 's indication that there will be a future review focused on improving our type hinting practices across the codebase.
Unless there are any further suggestions or objections, I recommend we proceed with merging this into the master branch @Alex-Preciado 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4f23ec990..22f268d7b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -3,7 +3,9 @@
 # Release 0.38.0-dev (development release)
 
 <h3>New features since last release</h3>
-
+* A new method `process_density_matrix` has been added to the `ProbabilityMP` and `DensityMatrixMP` classes, allowing for more efficient handling of quantum density matrices, particularly with batch processing support. This method simplifies the calculation of probabilities from quantum states represented as density matrices.
+  [(#5830)](https://github.com/PennyLaneAI/pennylane/pull/5830)
+  
 <h3>Improvements 🛠</h3>
 
 <h3>Breaking changes 💔</h3>
@@ -17,3 +19,5 @@
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
+
+Yushao Chen.
\ No newline at end of file
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index c24c6aee2..ccf87e59c 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -676,6 +676,22 @@ class StateMeasurement(MeasurementProcess):
                 dimension :math:`2^n` acts on a subspace of :math:`n` wires
         """
 
+    def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
+        """
+        Process the given density matrix.
+
+        Args:
+            density_matrix (TensorLike): The density matrix representing the (mixed) quantum state,
+                which may be single or batched. For a single matrix, the shape should be ``(2^n, 2^n)``
+                where `n` is the number of wires the matrix acts upon. For batched matrices, the shape
+                should be ``(batch_size, 2^n, 2^n)``.
+            wire_order (Wires): The wires determining the subspace that the ``density_matrix`` acts on.
+                A matrix of dimension :math:`2^n` acts on a subspace of :math:`n` wires. This parameter specifies
+                the mapping of matrix dimensions to physical qubits, allowing the function to correctly
+                trace out the subsystems not involved in the measurement or operation.
+        """
+        raise NotImplementedError
+
 
 class MeasurementTransform(MeasurementProcess):
     """Measurement process that applies a transform into the given quantum tape. This transform
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index c1c0dc45c..3359b0696 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -19,6 +19,7 @@ from typing import Sequence, Tuple
 import numpy as np
 
 import pennylane as qml
+from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
 from .measurements import Probability, SampleMeasurement, StateMeasurement
@@ -260,6 +261,18 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
 
         return prob_vector
 
+    def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
+        if len(np.shape(density_matrix)) == 2:
+            prob = qml.math.diagonal(density_matrix)
+        else:
+            prob = qml.math.array(
+                [qml.math.diagonal(density_matrix[i]) for i in range(np.shape(density_matrix)[0])]
+            )
+
+        # Since we only care about the probabilities, we can simplify the task here by creating a 'pseudo-state' to carry the diagonal elements and reuse the process_state method
+        p_state = np.sqrt(prob)
+        return self.process_state(p_state, wire_order)
+
     @staticmethod
     def _count_samples(indices, batch_size, dim):
         """Count the occurrences of sampled indices and convert them to relative
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index d97fcf8a9..7962206fd 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -17,6 +17,7 @@ This module contains the qml.state measurement.
 from typing import Optional, Sequence
 
 import pennylane as qml
+from pennylane.typing import TensorLike
 from pennylane.wires import WireError, Wires
 
 from .measurements import State, StateMeasurement
@@ -191,6 +192,10 @@ class StateMP(StateMeasurement):
         state = qml.math.reshape(state, flat_shape)
         return qml.math.cast(state, "complex128") if is_tf_interface else state + 0.0j
 
+    def process_density_matrix(self, density_matrix: Sequence[complex], wire_order: Wires):
+        # pylint:disable=redefined-outer-name
+        raise ValueError("Processing from density matrix to state is not supported.")
+
 
 class DensityMatrixMP(StateMP):
     """Measurement process that returns the quantum state in the computational basis.
@@ -238,3 +243,14 @@ class DensityMatrixMP(StateMP):
         if not qml.math.is_abstract(state) and qml.math.any(qml.math.iscomplex(state)):
             kwargs["c_dtype"] = state.dtype
         return qml.math.reduce_statevector(state, **kwargs)
+
+    def process_density_matrix(self, density_matrix: TensorLike, wire_order: Wires):
+        # pylint:disable=redefined-outer-name
+        wire_map = dict(zip(wire_order, range(len(wire_order))))
+        mapped_wires = [wire_map[w] for w in self.wires]
+        kwargs = {"indices": mapped_wires, "c_dtype": "complex128"}
+        if not qml.math.is_abstract(density_matrix) and qml.math.any(
+            qml.math.iscomplex(density_matrix)
+        ):
+            kwargs["c_dtype"] = density_matrix.dtype
+        return qml.math.reduce_dm(density_matrix, **kwargs)


Test patch:
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 7398a9794..20c789fe2 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -648,6 +648,19 @@ class TestStateMeasurement:
         ):
             circuit()
 
+    def test_state_measurement_process_density_matrix_not_implemented(self):
+        """Test that the process_density_matrix method of StateMeasurement raises
+        NotImplementedError."""
+
+        class MyMeasurement(StateMeasurement):
+            def process_state(self, state, wire_order):
+                return qml.math.sum(state)
+
+        with pytest.raises(NotImplementedError):
+            MyMeasurement().process_density_matrix(
+                density_matrix=qml.math.array([[1, 0], [0, 0]]), wire_order=Wires([0, 1])
+            )
+
 
 class TestMeasurementTransform:
     """Tests for the MeasurementTransform class."""
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 014465652..ea6b6784a 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -211,6 +211,123 @@ class TestProbs:
         assert subset_probs.shape == qml.math.shape(expected)
         assert qml.math.allclose(subset_probs, expected)
 
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_process_density_matrix_basic(self, interface):
+        """Test that process_density_matrix returns correct probabilities from a density matrix."""
+        dm = qml.math.array([[0.5, 0], [0, 0.5]], like=interface)
+        wires = qml.wires.Wires(range(1))
+        expected = qml.math.array([0.5, 0.5], like=interface)
+        calculated_probs = qml.probs().process_density_matrix(dm, wires)
+        assert qml.math.allclose(calculated_probs, expected)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize(
+        "subset_wires, expected",
+        [
+            ([0], [0.5, 0.5]),
+            ([1], [0.25, 0.75]),
+            ([1, 0], [0.15, 0.1, 0.35, 0.4]),
+            ([0, 1], [0.15, 0.35, 0.1, 0.4]),
+        ],
+    )
+    def test_process_density_matrix_subsets(self, interface, subset_wires, expected):
+        """Test processing of density matrix with subsets of wires."""
+        dm = qml.math.array(
+            [[0.15, 0, 0.1, 0], [0, 0.35, 0, 0.4], [0.1, 0, 0.1, 0], [0, 0.4, 0, 0.4]],
+            like=interface,
+        )
+        wires = qml.wires.Wires(range(2))
+        subset_probs = qml.probs(wires=subset_wires).process_density_matrix(dm, wires)
+        assert subset_probs.shape == qml.math.shape(expected)
+        assert qml.math.allclose(subset_probs, expected)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize(
+        "subset_wires, expected",
+        [
+            ([0], [[1, 0], [0.5, 0.5], [0.5, 0.5]]),
+            ([1], [[1, 0], [0.5, 0.5], [0.5, 0.5]]),
+            ([1, 0], [[1, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.5, 0, 0, 0.5]]),
+            ([0, 1], [[1, 0, 0, 0], [0.25, 0.25, 0.25, 0.25], [0.5, 0, 0, 0.5]]),
+        ],
+    )
+    def test_process_density_matrix_batched(self, interface, subset_wires, expected):
+        """Test processing of a batch of density matrices."""
+        # Define a batch of density matrices
+        dm_batch = qml.math.array(
+            [
+                # Pure state |00⟩
+                np.outer([1, 0, 0, 0], [1, 0, 0, 0]),
+                # Maximally mixed state
+                np.identity(4) / 4,
+                # Bell state |Φ+⟩ = 1/√2(|00⟩ + |11⟩)
+                0.5 * np.outer([1, 0, 0, 1], [1, 0, 0, 1]),
+            ],
+            like=interface,
+        )
+
+        wires = qml.wires.Wires(range(2))
+        # Process the entire batch of density matrices
+        subset_probs = qml.probs(wires=subset_wires).process_density_matrix(dm_batch, wires)
+
+        expected = qml.math.array(expected, like=interface)
+        # Check if the calculated probabilities match the expected values
+        assert (
+            subset_probs.shape == expected.shape
+        ), f"Shape mismatch: expected {expected.shape}, got {subset_probs.shape}"
+        assert qml.math.allclose(
+            subset_probs, expected
+        ), f"Value mismatch: expected {expected.tolist()}, got {subset_probs.tolist()}"
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize(
+        "subset_wires",
+        [([3, 1, 0])],
+    )
+    def test_process_density_matrix_medium(self, interface, subset_wires):
+        """Test processing of a random generated, medium-sized density matrices."""
+        # Define a density matrix for a 4-qubit system
+        size = 16
+        B = np.random.rand(size, size)
+        dm_np = B + B.conjugate().T
+        dm_np = dm_np / np.trace(dm_np)
+        dm = qml.math.array(
+            dm_np,
+            like=interface,
+        )
+
+        wires = qml.wires.Wires(range(4))
+        # Process the entire batch of density matrices
+        subset_probs = qml.probs(wires=subset_wires).process_density_matrix(dm, wires)
+
+        # Trace out the second qubit (qubit indexed 2) and calculate probabilities for qubits 3, 1, 0
+        # We need to sum over the indices corresponding to the second qubit, which are in positions 4, 5, 6, 7, ...
+        # and their strides in the flattened array are 4 (since 2^2)
+        reshaped_dm = dm_np.reshape(
+            (2, 2, 2, 2, 2, 2, 2, 2)
+        )  # Reshape to (2^8) with each qubit getting a dimension
+        reduced_dm = np.einsum(
+            "ijklmnkp->ljipnm", reshaped_dm
+        )  # Sum over the second qubit (k, m -> trace out)
+
+        # Reshape back to a 8x8 matrix to get the density matrix for qubits 3, 1, 0
+        reduced_dm = reduced_dm.reshape((8, 8))
+
+        # Extract probabilities (diagonal elements of the reduced density matrix)
+        expected = np.diag(reduced_dm)
+        expected = qml.math.array(expected, like=interface)
+        # Check if the calculated probabilities match the expected values
+        assert (
+            subset_probs.shape == expected.shape
+        ), f"Shape mismatch: expected {expected.shape}, got {subset_probs.shape}"
+        assert qml.math.allclose(
+            subset_probs, expected
+        ), f"Value mismatch: expected {expected.tolist()}, got {subset_probs.tolist()}"
+
     def test_integration(self, tol):
         """Test the probability is correct for a known state preparation."""
         dev = qml.device("default.qubit", wires=2)
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index b131dc229..5c580c1eb 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -129,6 +129,24 @@ class TestStateMP:
         with pytest.raises(WireError, match=r"Unexpected unique wires <Wires = \[0, 1, 2\]> found"):
             StateMP(wires=[0, 1]).process_state([1, 0], wire_order=Wires(2))
 
+    @pytest.mark.parametrize(
+        "dm",
+        [(np.array([[1, 0, 1, 23], [0, 0, 0, 0], [1, 0, 1, 23], [23, 0, 23, 529]]) / 531 + 0.0j)],
+    )
+    def test_process_density_matrix(self, dm):
+        """Test the processing of a state vector."""
+
+        mp = StateMP(wires=None)
+        assert mp.return_type == State
+        assert mp.numeric_type is complex
+
+        # Expecting a NotImplementedError to be raised when process_density_matrix is called
+
+        with pytest.raises(
+            ValueError, match="Processing from density matrix to state is not supported."
+        ):
+            mp.process_density_matrix(dm, [0, 1])
+
 
 class TestDensityMatrixMP:
     """Tests for the DensityMatrix measurement process"""
@@ -158,9 +176,6 @@ class TestDensityMatrixMP:
             exp = reduce_statevector(vec, wires)
         assert qml.math.allclose(processed, exp)
 
-    @pytest.mark.xfail(
-        reason="DensityMatrixMP.process_state no longer supports density matrix parameters"
-    )
     @pytest.mark.parametrize(
         "mat, wires",
         [
@@ -181,7 +196,7 @@ class TestDensityMatrixMP:
 
         num_wires = int(np.log2(len(mat)))
         order = list(range(num_wires))
-        processed = mp.process_state(mat, order)
+        processed = mp.process_density_matrix(mat, order)
         assert qml.math.shape(processed) == (2 ** len(wires), 2 ** len(wires))
         if len(wires) == num_wires:
             exp = _permute_dense_matrix(mat, wires, order, None)