Repository: PennyLaneAI/pennylane
Pull request #3551 (2022-12-16 15:21:13)

Add `purity` to measurements 
**Context:**

`purity` was added to the `qinfo` module in a previous PR, this PR completes the feature with the purity measurement process

**Description of the Change:**

`purity` added to the measurement module, implementing `purity` and `PurityMP`

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

https://github.com/PennyLaneAI/pennylane/issues/3139


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3551](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2c66f99) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/594fef59cd1f0fe198550da6bdc21afa4b5bdfeb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (594fef5) will **decrease** coverage by `0.01%`.
> The diff coverage is `86.84%`.

```diff
@@            Coverage Diff             @@
##           master    #3551      +/-   ##
==========================================
- Coverage   99.87%   99.85%   -0.02%     
==========================================
  Files         311      312       +1     
  Lines       27903    27941      +38     
==========================================
+ Hits        27867    27900      +33     
- Misses         36       41       +5     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `96.70% <ø> (ø)` | |
| [pennylane/measurements/purity.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9wdXJpdHkucHk=) | `85.29% <85.29%> (ø)` | |
| [pennylane/measurements/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3551/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

astralcai: @AlbertMitjans I just made the update

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5c4a03bf7..8b51fd70e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,9 @@
 
 <h3>New features since last release</h3>
 
+* `qml.purity` is added as a measurement process for purity 
+  [(#3551)](https://github.com/PennyLaneAI/pennylane/pull/3551)
+
 * Added a new template that implements a canonical 2-complete linear (2-CCL) swap network
   described in [arXiv:1905.05118](https://arxiv.org/abs/1905.05118).
   [(#3447)](https://github.com/PennyLaneAI/pennylane/pull/3447)
@@ -130,6 +133,7 @@ This release contains contributions from (in alphabetical order):
 
 Ikko Ashimine
 Utkarsh Azad
+Astral Cai
 Lillian M. A. Frederiksen
 Soran Jahangiri
 Albert Mitjans Coma
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 92eb82c8a..1a1222ec1 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -59,6 +59,7 @@ from pennylane.measurements import (
     state,
     var,
     vn_entropy,
+    purity,
     mutual_info,
     classical_shadow,
     shadow_expval,
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 677d41fdb..297fee07a 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -214,6 +214,7 @@ from .measurements import (
 )
 from .mid_measure import MeasurementValue, MeasurementValueError, MidMeasureMP, measure
 from .mutual_info import MutualInfoMP, mutual_info
+from .purity import PurityMP, purity
 from .probs import ProbabilityMP, probs
 from .sample import SampleMP, sample
 from .state import StateMP, density_matrix, state
diff --git a/pennylane/measurements/purity.py b/pennylane/measurements/purity.py
new file mode 100644
index 000000000..704e9c2d7
--- /dev/null
+++ b/pennylane/measurements/purity.py
@@ -0,0 +1,103 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module contains the qml.purity measurement.
+"""
+
+from typing import Sequence
+import pennylane as qml
+from pennylane.wires import Wires
+
+from .measurements import StateMeasurement
+
+
+def purity(wires):
+    r"""The purity of the system prior to measurement.
+
+    .. math::
+        \gamma = \text{Tr}(\rho^2)
+
+    where :math:`\rho` is the density matrix. The purity of a normalized quantum state satisfies
+    :math:`\frac{1}{d} \leq \gamma \leq 1`, where :math:`d` is the dimension of the Hilbert space.
+    A pure state has a purity of 1.
+
+    Args:
+        wires (Sequence[int] or int): The wires of the subsystem
+
+    Returns:
+        PurityMP: Measurement process instance
+
+    **Example**
+
+    .. code-block:: python3
+
+        dev = qml.device("default.mixed", wires=2)
+
+        @qml.qnode(dev)
+        def circuit_purity(p):
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.BitFlip(p, wires=0)
+            qml.BitFlip(p, wires=1)
+            return qml.purity(wires=[0,1])
+
+    >>> circuit_purity(0.1)
+    array(0.7048)
+
+    .. seealso:: :func:`pennylane.qinfo.transforms.purity` and :func:`pennylane.math.purity`
+    """
+    wires = Wires(wires)
+    return PurityMP(wires=wires)
+
+
+class PurityMP(StateMeasurement):
+    """Measurement process that computes the purity of the system prior to measurement.
+
+    Please refer to :func:`purity` for detailed documentation.
+
+    Args:
+        obs (.Observable): The observable that is to be measured as part of the
+            measurement process. Not all measurement processes require observables (for
+            example ``Probability``); this argument is optional.
+        wires (.Wires): The wires the measurement process applies to.
+            This can only be specified if an observable was not provided.
+        eigvals (array): A flat array representing the eigenvalues of the measurement.
+            This can only be specified if an observable was not provided.
+        id (str): custom label given to a measurement instance, can be useful for some
+            applications where the instance has to be identified
+    """
+
+    @property
+    def numeric_type(self):
+        return float
+
+    def shape(self, device=None):
+        if qml.active_return():
+            return self._shape_new(device)
+        if device is None or device.shot_vector is None:
+            return (1,)
+        num_shot_elements = sum(s.copies for s in device.shot_vector)
+        return (num_shot_elements,)
+
+    def _shape_new(self, device=None):
+        if device is None or device.shot_vector is None:
+            return ()
+        num_shot_elements = sum(s.copies for s in device.shot_vector)
+        return tuple(() for _ in range(num_shot_elements))
+
+    def process_state(self, state: Sequence[complex], wire_order: Wires):
+        wire_map = dict(zip(wire_order, list(range(len(wire_order)))))
+        indices = [wire_map[w] for w in self.wires]
+        return qml.math.purity(state, indices=indices, c_dtype=state.dtype)


Test patch:
diff --git a/tests/measurements/test_purity_measurement.py b/tests/measurements/test_purity_measurement.py
new file mode 100644
index 000000000..bd9b99f6a
--- /dev/null
+++ b/tests/measurements/test_purity_measurement.py
@@ -0,0 +1,340 @@
+import pytest
+
+import numpy as np
+import pennylane as qml
+
+
+def expected_purity_ising_xx(param):
+    """Returns the analytical purity for subsystems of the IsingXX"""
+
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    return eig_1**2 + eig_2**2
+
+
+def expected_purity_grad_ising_xx(param):
+    """The analytic gradient purity for the IsingXX"""
+
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    grad_expected_purity = (
+        2
+        * eig_1
+        * (np.sin(param / 2) ** 3 * np.cos(param / 2) - np.sin(param / 2) * np.cos(param / 2) ** 3)
+        / np.sqrt(1 - 4 * np.sin(param / 2) ** 2 * np.cos(param / 2) ** 2)
+    ) + (
+        2
+        * eig_2
+        * (
+            np.sin(param / 2)
+            * np.cos(param / 2)
+            * (np.cos(param / 2) ** 2 - np.sin(param / 2) ** 2)
+        )
+        / np.sqrt(1 - 4 * np.sin(param / 2) ** 2 * np.cos(param / 2) ** 2)
+    )
+    return grad_expected_purity
+
+
+class TestPurity:
+    """Tests for purity measurements"""
+
+    devices = ["default.qubit", "lightning.qubit", "default.mixed"]
+    grad_supported_devices = ["default.qubit", "default.mixed"]
+    mix_supported_devices = ["default.mixed"]
+
+    diff_methods = ["backprop", "finite-diff"]
+
+    parameters = np.linspace(0, 2 * np.pi, 3)
+    probs = np.array([0.001, 0.01, 0.1, 0.2])
+
+    wires_list = [([0], True), ([1], True), ([0, 1], False)]
+
+    @pytest.mark.parametrize("shots, shape", [(None, (1,)), (10, (1,)), ((1, 10), (2,))])
+    def test_shape(self, shots, shape):
+        """Test the ``shape`` method."""
+        meas = qml.purity(wires=0)
+        dev = qml.device("default.qubit", wires=1, shots=shots)
+        assert meas.shape(dev) == shape
+
+    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
+    def test_shape_new(self, shots, shape):
+        """Test the ``shape_new`` method."""
+        qml.enable_return()
+        meas = qml.purity(wires=0)
+        dev = qml.device("default.qubit", wires=1, shots=shots)
+        assert meas.shape(dev) == shape
+        qml.disable_return()
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity(self, device, param, wires, is_partial):
+        """Tests purity for a qnode"""
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(param)
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.parametrize("device", mix_supported_devices)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("param", probs)
+    def test_bit_flip_qnode_purity(self, device, wires, param, is_partial):
+        """Tests purity for a qnode on a noisy device with bit flips"""
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev)
+        def circuit(p):
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.BitFlip(p, wires=0)
+            qml.BitFlip(p, wires=1)
+            return qml.purity(wires=wires)
+
+        purity = circuit(param)
+        expected_purity = (
+            0.5
+            if is_partial
+            else 4 * (0.5 - (1 - param) * param) ** 2 + 4 * (1 - param) ** 2 * param**2
+        )
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.parametrize("device", grad_supported_devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad(self, device, param, wires, is_partial, diff_method):
+        """Tests purity for a qnode"""
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = qml.grad(circuit)(param)
+        expected_grad = expected_purity_grad_ising_xx(param) if is_partial else 0
+        assert qml.math.allclose(grad_purity, expected_grad, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.parametrize("device", mix_supported_devices)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("param", probs)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_bit_flip_qnode_purity_grad(self, device, wires, param, is_partial, diff_method):
+        """Tests gradient of purity for a qnode on a noisy device with bit flips"""
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev)
+        def circuit(p):
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.BitFlip(p, wires=0)
+            qml.BitFlip(p, wires=1)
+            return qml.purity(wires=wires)
+
+        purity_grad = qml.grad(circuit)(param)
+        expected_purity_grad = 0 if is_partial else 32 * (param - 0.5) ** 3
+        assert qml.math.allclose(purity_grad, expected_purity_grad, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity_jax(self, device, param, wires, is_partial):
+        """Test purity for a QNode with jax interface."""
+
+        import jax.numpy as jnp
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(jnp.array(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("device", grad_supported_devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad_jax(self, device, param, wires, is_partial, diff_method):
+        """Test purity for a QNode gradient with Jax."""
+
+        import jax
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="jax", diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = jax.grad(circuit)(jax.numpy.array(param))
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity_jax_jit(self, device, param, wires, is_partial):
+        """Test purity for a QNode with jax interface."""
+
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="jax-jit")
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = jax.jit(circuit)(jnp.array(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("device", grad_supported_devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad_jax_jit(self, device, param, wires, is_partial, diff_method):
+        """Test purity for a QNode gradient with Jax."""
+
+        import jax
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="jax-jit", diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = jax.jit(jax.grad(circuit))(jax.numpy.array(param))
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity_torch(self, device, param, wires, is_partial):
+        """Tests purity for a qnode"""
+
+        import torch
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(torch.tensor(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("device", grad_supported_devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad_torch(self, device, param, wires, is_partial, diff_method):
+        """Test purity for a QNode gradient with torch."""
+
+        import torch
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="torch", diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        expected_grad = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        param = torch.tensor(param, dtype=torch.float64, requires_grad=True)
+        purity = circuit(param)
+        purity.backward()
+        grad_purity = param.grad
+
+        assert qml.math.allclose(grad_purity, expected_grad, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity_tf(self, device, param, wires, is_partial):
+        """Tests purity for a qnode"""
+
+        import tensorflow as tf
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(tf.Variable(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("device", grad_supported_devices)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad_tf(self, device, param, wires, is_partial, diff_method):
+        """Test purity for a QNode gradient with tf."""
+
+        import tensorflow as tf
+
+        dev = qml.device(device, wires=2)
+
+        @qml.qnode(dev, interface="tf", diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        param = tf.Variable(param)
+        with tf.GradientTape() as tape:
+            purity = circuit(param)
+
+        grad_purity = tape.gradient(purity, param)
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("param", parameters)
+    def test_qnode_entropy_custom_wires(self, device, param):
+        """Test that purity can be returned with custom wires."""
+
+        dev = qml.device(device, wires=["a", 1])
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.IsingXX(x, wires=["a", 1])
+            return qml.purity(wires=["a"])
+
+        purity = circuit(param)
+        expected_purity = expected_purity_ising_xx(param)
+        assert qml.math.allclose(purity, expected_purity)