Support wire labels in qinfo transforms

doc/releases/changelog-dev.md

@@ -9,6 +9,9 @@
 * `PauliWord` sparse matrices are much faster, which directly improves `PauliSentence`.
   [#4272](https://github.com/PennyLaneAI/pennylane/pull/4272)
 
+* QNode transforms in `qml.qinfo` now support custom wire labels.
+  [#4331](https://github.com/PennyLaneAI/pennylane/pull/4331)
+
 <h3>Breaking changes 💔</h3>
 
 * The `do_queue` keyword argument in `qml.operation.Operator` has been removed. Instead of
@@ -46,10 +49,14 @@
 * Raise a warning if control indicators are hidden when calling `qml.draw_mpl`
   [(#4295)](https://github.com/PennyLaneAI/pennylane/pull/4295)
 
+* `qml.qinfo.purity` now produces correct results with custom wire labels.
+  [#4331](https://github.com/PennyLaneAI/pennylane/pull/4331)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
 
+Edward Jiang,
 Christina Lee,
 Borja Requena,
 Matthew Silverman

pennylane/qinfo/transforms.py

@@ -51,6 +51,8 @@ def circuit(x):
 
     .. seealso:: :func:`pennylane.density_matrix` and :func:`pennylane.math.reduce_dm`
     """
+    wire_map = {w: i for i, w in enumerate(qnode.device.wires)}
+    indices = [wire_map[w] for w in wires]
 
     def wrapper(*args, **kwargs):
         qnode.construct(args, kwargs)
@@ -66,7 +68,7 @@ def wrapper(*args, **kwargs):
 
         # TODO: optimize given the wires by creating a tape with relevant operations
         state_built = qnode(*args, **kwargs)
-        density_matrix = dm_func(state_built, indices=wires, c_dtype=qnode.device.C_DTYPE)
+        density_matrix = dm_func(state_built, indices=indices, c_dtype=qnode.device.C_DTYPE)
         return density_matrix
 
     return wrapper
@@ -121,6 +123,8 @@ def circuit(x):
 
     .. seealso:: :func:`pennylane.math.purity`
     """
+    wire_map = {w: i for i, w in enumerate(qnode.device.wires)}
+    indices = [wire_map[w] for w in wires]
 
     def wrapper(*args, **kwargs):
         # Construct tape
@@ -141,7 +145,7 @@ def wrapper(*args, **kwargs):
         if not dm_measurement:
             state_built = qml.math.dm_from_state_vector(state_built)
 
-        return qml.math.purity(state_built, wires, c_dtype=qnode.device.C_DTYPE)
+        return qml.math.purity(state_built, indices, c_dtype=qnode.device.C_DTYPE)
 
     return wrapper
 
@@ -184,6 +188,8 @@ def circuit(x):
 
     .. seealso:: :func:`pennylane.math.vn_entropy` and :func:`pennylane.vn_entropy`
     """
+    wire_map = {w: i for i, w in enumerate(qnode.device.wires)}
+    indices = [wire_map[w] for w in wires]
 
     density_matrix_qnode = qml.qinfo.reduced_dm(qnode, qnode.device.wires)
 
@@ -205,11 +211,13 @@ def wrapper(*args, **kwargs):
                 return 0.0
 
             density_matrix = qnode(*args, **kwargs)
-            entropy = qml.math.vn_entropy(density_matrix, wires, base, c_dtype=qnode.device.C_DTYPE)
+            entropy = qml.math.vn_entropy(
+                density_matrix, indices, base, c_dtype=qnode.device.C_DTYPE
+            )
             return entropy
 
         density_matrix = density_matrix_qnode(*args, **kwargs)
-        entropy = qml.math.vn_entropy(density_matrix, wires, base, c_dtype=qnode.device.C_DTYPE)
+        entropy = qml.math.vn_entropy(density_matrix, indices, base, c_dtype=qnode.device.C_DTYPE)
         return entropy
 
     return wrapper
@@ -263,12 +271,15 @@ def circuit(x):
 
     .. seealso:: :func:`~.qinfo.vn_entropy`, :func:`pennylane.math.mutual_info` and :func:`pennylane.mutual_info`
     """
+    wire_map = {w: i for i, w in enumerate(qnode.device.wires)}
+    indices0 = [wire_map[w] for w in wires0]
+    indices1 = [wire_map[w] for w in wires1]
 
     density_matrix_qnode = qml.qinfo.reduced_dm(qnode, qnode.device.wires)
 
     def wrapper(*args, **kwargs):
         density_matrix = density_matrix_qnode(*args, **kwargs)
-        entropy = qml.math.mutual_info(density_matrix, wires0, wires1, base=base)
+        entropy = qml.math.mutual_info(density_matrix, indices0, indices1, base=base)
         return entropy
 
     return wrapper

tests/measurements/test_mutual_info.py

@@ -179,6 +179,28 @@ def circuit_state(params):
 
         assert np.allclose(actual, expected)
 
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    def test_mutual_info_wire_labels(self, device):
+        """Test that mutual_info is correct with custom wire labels"""
+        param = np.array([0.678, 1.234])
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(param):
+            qml.RY(param, wires=wires[0])
+            qml.CNOT(wires=wires)
+            return qml.state()
+
+        actual = qml.qinfo.mutual_info(circuit, wires0=[wires[0]], wires1=[wires[1]])(param)
+
+        # compare transform results with analytic values
+        expected = -2 * np.cos(param / 2) ** 2 * np.log(np.cos(param / 2) ** 2) - 2 * np.sin(
+            param / 2
+        ) ** 2 * np.log(np.sin(param / 2) ** 2)
+
+        assert np.allclose(actual, expected)
+
     @pytest.mark.jax
     @pytest.mark.parametrize("params", np.linspace(0, 2 * np.pi, 8))
     def test_qnode_state_jax_jit(self, params):

tests/qinfo/test_entropies.py

@@ -363,6 +363,30 @@ def circuit_state(x):
 
         assert qml.math.allclose(entropy, expected_entropy)
 
+    @pytest.mark.parametrize("device", devices)
+    def test_entropy_wire_labels(self, device, tol):
+        """Test that vn_entropy is correct with custom wire labels"""
+        param = np.array(1.234)
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.PauliX(wires=wires[0])
+            qml.IsingXX(x, wires=wires)
+            return qml.state()
+
+        entropy0 = qml.qinfo.vn_entropy(circuit, wires=[wires[0]])(param)
+        eigs0 = [np.sin(param / 2) ** 2, np.cos(param / 2) ** 2]
+        exp0 = -np.sum(eigs0 * np.log(eigs0))
+
+        entropy1 = qml.qinfo.vn_entropy(circuit, wires=[wires[1]])(param)
+        eigs1 = [np.cos(param / 2) ** 2, np.sin(param / 2) ** 2]
+        exp1 = -np.sum(eigs1 * np.log(eigs1))
+
+        assert np.allclose(exp0, entropy0, atol=tol)
+        assert np.allclose(exp1, entropy1, atol=tol)
+
 
 class TestRelativeEntropy:
     """Tests for the mutual information functions"""
@@ -764,6 +788,33 @@ def circuit2(param=0):
 
         assert np.allclose(actual, expected)
 
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    def test_entropy_wire_labels(self, device, tol):
+        """Test that relative_entropy is correct with custom wire labels"""
+        param = np.array([0.678, 1.234])
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(param):
+            qml.RY(param, wires=wires[0])
+            qml.CNOT(wires=wires)
+            return qml.state()
+
+        rel_ent_circuit = qml.qinfo.relative_entropy(circuit, circuit, [wires[0]], [wires[1]])
+        actual = rel_ent_circuit((param[0],), (param[1],))
+
+        # compare transform results with analytic results
+        first_term = np.cos(param[0] / 2) ** 2 * (
+            np.log(np.cos(param[0] / 2) ** 2) - np.log(np.cos(param[1] / 2) ** 2)
+        )
+        second_term = np.sin(param[0] / 2) ** 2 * (
+            np.log(np.sin(param[0] / 2) ** 2) - np.log(np.sin(param[1] / 2) ** 2)
+        )
+        expected = first_term + second_term
+
+        assert np.allclose(actual, expected, atol=tol)
+
 
 @pytest.mark.parametrize("device", ["default.qubit", "default.mixed"])
 class TestBroadcasting:

tests/qinfo/test_fidelity.py

@@ -163,7 +163,6 @@ def circuit1(x, y):
     def test_fidelity_qnodes_rx_pauliz(self, device, param, wire):
         """Test the fidelity between Rx and PauliZ circuits."""
         dev = qml.device(device, wires=[wire])
-        print(dev.wires)
 
         @qml.qnode(dev)
         def circuit0(x):
@@ -175,10 +174,7 @@ def circuit1():
             qml.PauliZ(wires=wire)
             return qml.state()
 
-        # todo: Once #4318 is closed, will need to change the wires0, wires1 arguments
-        # to the commented version
-        fid = qml.qinfo.fidelity(circuit0, circuit1, wires0=[0], wires1=[0])((param))
-        # fid = qml.qinfo.fidelity(circuit0, circuit1, wires0=[wire], wires1=[wire])((param))
+        fid = qml.qinfo.fidelity(circuit0, circuit1, wires0=[wire], wires1=[wire])((param))
         expected_fid = expected_fidelity_rx_pauliz(param)
         assert qml.math.allclose(fid, expected_fid)
 
@@ -205,6 +201,28 @@ def circuit1():
         expected_fid = expected_grad_fidelity_rx_pauliz(param)
         assert qml.math.allclose(fid_grad, expected_fid)
 
+    @pytest.mark.parametrize("device", devices)
+    def test_fidelity_wire_labels(self, device, tol):
+        """Test that fidelity is correct with custom wire labels"""
+        param = np.array([0.678, 1.234])
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.PauliX(wires=wires[0])
+            qml.IsingXX(x, wires=wires)
+            return qml.state()
+
+        fid_circuit = qml.qinfo.fidelity(circuit, circuit, [wires[0]], [wires[1]])
+        actual = fid_circuit((param[0],), (param[1],))
+
+        expected = (
+            np.sin(param[0] / 2) * np.cos(param[1] / 2)
+            + np.sin(param[1] / 2) * np.cos(param[0] / 2)
+        ) ** 2
+        assert np.allclose(actual, expected, atol=tol)
+
     interfaces = ["auto", "autograd"]
 
     @pytest.mark.autograd

tests/qinfo/test_purity.py

@@ -21,10 +21,7 @@
 
 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
+    return np.cos(param / 2) ** 4 + np.sin(param / 2) ** 4
 
 
 def expected_purity_grad_ising_xx(param):
@@ -378,6 +375,26 @@ def circuit_state(x):
 
         assert qml.math.allclose(grad_purity, grad_expected_purity)
 
+    @pytest.mark.parametrize("device", devices)
+    def test_purity_wire_labels(self, device, tol):
+        """Test that purity is correct with custom wire labels"""
+        param = np.array(1.234)
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit_state(x):
+            qml.PauliX(wires=wires[0])
+            qml.IsingXX(x, wires=wires)
+            return qml.state()
+
+        purity0 = qml.qinfo.purity(circuit_state, wires=[wires[0]])(param)
+        purity1 = qml.qinfo.purity(circuit_state, wires=[wires[1]])(param)
+        expected = expected_purity_ising_xx(param)
+
+        assert qml.math.allclose(purity0, expected, atol=tol)
+        assert qml.math.allclose(purity1, expected, atol=tol)
+
 
 @pytest.mark.parametrize("device", ["default.qubit", "default.mixed"])
 def test_broadcasting(device):

tests/qinfo/test_reduced_dm.py

@@ -81,6 +81,28 @@ def expected_density_matrix(x, wires):
 
         assert np.allclose(expected_density_matrix(angle, wires), density_matrix, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("angle", angle_values)
+    def test_density_matrix_wire_labels(self, device, angle, tol):
+        """Test that density matrix is correct with custom wire labels"""
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.PauliX(wires=wires[0])
+            qml.IsingXX(x, wires=wires)
+            return qml.state()
+
+        dm0 = qml.qinfo.reduced_dm(circuit, wires=[wires[0]])(angle)
+        dm1 = qml.qinfo.reduced_dm(circuit, wires=[wires[1]])(angle)
+
+        exp0 = np.array([[np.sin(angle / 2) ** 2, 0], [0, np.cos(angle / 2) ** 2]])
+        exp1 = np.array([[np.cos(angle / 2) ** 2, 0], [0, np.sin(angle / 2) ** 2]])
+
+        assert np.allclose(exp0, dm0, atol=tol)
+        assert np.allclose(exp1, dm1, atol=tol)
+
     def test_qnode_not_returning_state(self):
         """Test that the QNode of reduced_dm function must return state."""
         dev = qml.device("default.qubit", wires=1)

tests/qinfo/test_trace_distance.py

@@ -110,6 +110,28 @@ def circuit1(x, y):
         assert qml.math.allclose(td_arg_kwarg, 0.0)
         assert qml.math.allclose(td_kwargs, 0.0)
 
+    @pytest.mark.parametrize("device", devices)
+    def test_trace_distance_wire_labels(self, device, tol):
+        """Test that trace_distance is correct with custom wire labels"""
+        param = np.array([0.678, 1.234])
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.PauliX(wires=wires[0])
+            qml.IsingXX(x, wires=wires)
+            return qml.state()
+
+        td_circuit = qml.qinfo.trace_distance(circuit, circuit, [wires[0]], [wires[1]])
+        actual = td_circuit((param[0],), (param[1],))
+
+        expected = 0.5 * (
+            np.abs(np.cos(param[0] / 2) ** 2 - np.sin(param[1] / 2) ** 2)
+            + np.abs(np.cos(param[1] / 2) ** 2 - np.sin(param[0] / 2) ** 2)
+        )
+        assert np.allclose(actual, expected, atol=tol)
+
     # 0 and 2 * pi are removed to avoid nan values in the gradient
     parameters = np.linspace(0, 2 * np.pi, 20)[1:-1]
     wires = [1, 2]