added density_matrix_qnode, negativity_cost_fn, and partial_transpose…

… functions with test cases

updated functionality of partial transpose so it works with any input size

src/qnetvo/cost/__init__.py

@@ -6,3 +6,4 @@
 from .chsh_inequality import *
 from .linear_inequalities import *
 from .mutual_info import *
+from .negativity import *

src/qnetvo/cost/negativity.py

@@ -0,0 +1,60 @@
+import pennylane as qml
+from pennylane import math
+import numpy as np
+from ..qnodes import density_matrix_qnode
+from ..utilities import partial_transpose
+
+
+def negativity_cost_fn(network_ansatz, m, n, wires, qnode_kwargs={}):
+    """Constructs an ansatz-specific negativity cost function.
+
+    Negativity can be used to identify if two subsystems :math:`A` and :math:`B` are
+    entangled, through the PPT criterion. Negativity is an upper bound for distillable entanglement.
+
+    This entanglement measure is expressed as
+
+    .. math::
+
+            \\mathcal{N}(\\rho) = |\\sum_{\\lambda_i < 0}\\lambda_i|,
+
+    where :math:`\\rho^{\\Gamma_B}` is the partial transpose of the joint state with respect to
+    the :math:`B` party, and :math:`\\lambda_i` are all of the eigenvalues of :math:`\\rho^{\\Gamma_B}`.
+
+    For more information on negativity and its applications in quantum information theory,
+    see [Vidal and Werner (2001)](https://arxiv.org/pdf/quant-ph/0102117).
+
+    :param ansatz: The ansatz circuit on which the negativity is evaluated.
+    :type ansatz: NetworkAnsatz
+
+    :param m: The size of the :math:`A` subsystem.
+    :type m: int
+
+    :param n: The size of the :math:`B` subsystem.
+    :type n: int
+
+    :param wires: The wires which define the joint state.
+    :type wires: list[int]
+
+    :param qnode_kwargs: Keyword arguments passed to the execute qnodes.
+    :type qnode_kwargs: dictionary
+
+    :returns: A cost function ``negativity_cost(*network_settings)`` parameterized by
+              the ansatz-specific scenario settings.
+    :rtype: Function
+
+    :raises ValueError: If the sum of the sizes of the two subsystems (``m + n``) does not match the length of ``wires``.
+    """
+
+    if len(wires) != m + n:
+        raise ValueError(f"Sum of sizes of two subsystems should be {len(wires)}; got {m+n}.")
+
+    density_qnode = density_matrix_qnode(network_ansatz, wires, **qnode_kwargs)
+
+    def negativity_cost(*network_settings):
+        dm = density_qnode(network_settings)
+        dm_pt = partial_transpose(dm, 2**m, 2**n)
+        eigenvalues = math.eigvalsh(dm_pt)
+        negativity = np.sum(np.abs(eigenvalues[eigenvalues < 0]))
+        return -negativity
+
+    return negativity_cost

src/qnetvo/qnodes.py

@@ -95,3 +95,35 @@ def circuit(settings):
         return qml.probs(wires=network_ansatz.layers_wires[-1])
 
     return circuit
+
+
+def density_matrix_qnode(network_ansatz, wires=None, **qnode_kwargs):
+    """
+    Constructs a qnode that computes the density matrix in the computational basis
+    across specified wires, or across all wires if no specific wires are provided.
+
+    :param network_ansatz: A ``NetworkAnsatz`` class specifying the quantum network simulation.
+    :type network_ansatz: NetworkAnsatz
+
+    :param wires: The wires on which the node operates. If None, the density matrix will be
+                  computed across all wires in the network ansatz.
+    :type wires: list[int] or None
+
+    :returns: A qnode called as ``qnode(settings)`` for evaluating the (reduced) density matrix
+              of the network ansatz.
+    :rtype: ``pennylane.QNode``
+
+    :raises ValueError: If the specified wires are not a subset of the wires in the network ansatz.
+    """
+
+    wires = network_ansatz.layers_wires[-1] if wires is None else wires
+
+    if not set(wires).issubset(network_ansatz.layers_wires[-1]):
+        raise ValueError("Specified wires must be a subset of the wires in the network ansatz.")
+
+    @qml.qnode(qml.device(**network_ansatz.dev_kwargs), **qnode_kwargs)
+    def circuit(settings):
+        network_ansatz.fn(settings)
+        return qml.density_matrix(wires)
+
+    return circuit

src/qnetvo/utilities.py

@@ -187,3 +187,39 @@ def ragged_reshape(input_list, list_dims):
         start_id = end_id
 
     return output_list
+
+
+def partial_transpose(dm, d1, d2):
+    """
+    Computes the partial transpose of a density matrix with respect to the second subsystem.
+
+    :param dm: The density matrix to be partially transposed.
+    :type dm: np.array
+
+    :param d1: The dimension of the first subsystem (e.g., 2^m where m is the number of qubits in the first subsystem).
+    :type d1: int
+
+    :param d2: The dimension of the second subsystem (e.g., 2^n where n is the number of qubits in the second subsystem).
+    :type d2: int
+
+    :returns: The partially transposed density matrix.
+    :rtype: np.array
+
+    :raises ValueError: If the product of `d1` and `d2` does not match the size of the density matrix.
+    """
+
+    if d1 * d2 != dm.shape[0]:
+        raise ValueError(
+            "The dimensions of the subsystems do not match the size of the density matrix."
+        )
+
+    bfm = np.empty((d2, d2), dtype=dm.dtype)
+    trm = np.empty((d2, d2), dtype=dm.dtype)
+
+    for i in range(d1):
+        for j in range(d1):
+            bfm = dm[i * d2 : (i + 1) * d2, j * d2 : (j + 1) * d2]
+            np.copyto(trm, bfm.T)
+            dm[i * d2 : (i + 1) * d2, j * d2 : (j + 1) * d2] = trm
+
+    return dm

test/cost/negativity_test.py

@@ -0,0 +1,45 @@
+import pytest
+from pennylane import numpy as np
+
+import qnetvo as qnet
+
+
+class TestNegativityCost:
+    def test_negativity_cost_fn(self):
+
+        prep_nodes = [
+            qnet.PrepareNode(1, [0, 1], qnet.bell_state_copies, 2),
+        ]
+
+        ansatz = qnet.NetworkAnsatz(prep_nodes)
+
+        negativity_cost = qnet.negativity_cost_fn(ansatz, m=1, n=1, wires=[0, 1])
+
+        zero_settings = ansatz.zero_network_settings()
+
+        negativity_value = negativity_cost(*zero_settings)
+
+        expected_negativity = -0.5
+        assert np.isclose(
+            negativity_value, expected_negativity
+        ), f"Expected {expected_negativity}, but got {negativity_value}"
+
+        separable_prep_nodes = [
+            qnet.PrepareNode(4, [0, 1], qnet.local_RY, 2),
+        ]
+
+        separable_ansatz = qnet.NetworkAnsatz(separable_prep_nodes)
+
+        separable_negativity_cost = qnet.negativity_cost_fn(
+            separable_ansatz, m=1, n=1, wires=[0, 1]
+        )
+
+        separable_negativity_value = separable_negativity_cost(*zero_settings)
+
+        expected_separable_negativity = 0
+        assert np.isclose(
+            separable_negativity_value, expected_separable_negativity
+        ), f"Expected {expected_separable_negativity}, but got {separable_negativity_value}"
+
+        with pytest.raises(ValueError, match="Sum of sizes of two subsystems should be"):
+            qnet.negativity_cost_fn(ansatz, m=2, n=1, wires=[0, 1])

test/cost/qnodes_test.py

@@ -101,3 +101,46 @@ def test_joint_probs_qnode(self):
         assert np.allclose(qnode([np.pi, 0, 0]), [0, 0, 0, 0, 1, 0, 0, 0])
         assert np.allclose(qnode([0, np.pi, 0]), [0, 0, 1, 0, 0, 0, 0, 0])
         assert np.allclose(qnode([0, 0, np.pi]), [0, 1, 0, 0, 0, 0, 0, 0])
+
+    def test_density_matrix_qnode(self):
+        prep_nodes = [
+            qnet.PrepareNode(1, [0, 1, 2], qnet.W_state, 3),
+        ]
+
+        ansatz = qnet.NetworkAnsatz(prep_nodes)
+        qnode = qnet.density_matrix_qnode(ansatz)
+
+        zero_settings = ansatz.zero_network_settings()
+        density_matrix = qnode(zero_settings)
+
+        expected_density_matrix = np.array(
+            [
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 1 / 3, 1 / 3, 0, 1 / 3, 0, 0, 0],
+                [0, 1 / 3, 1 / 3, 0, 1 / 3, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 1 / 3, 1 / 3, 0, 1 / 3, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+            ],
+            dtype=np.complex128,
+        )
+
+        assert np.allclose(
+            density_matrix, expected_density_matrix
+        ), "Density matrix did not match expected Bell state density matrix."
+
+        qnode_subset_wires = qnet.density_matrix_qnode(ansatz, wires=[0])
+
+        density_matrix_subset = qnode_subset_wires(zero_settings)
+        expected_density_matrix_subset = np.array([[2 / 3, 0], [0, 1 / 3]])
+
+        assert np.allclose(
+            density_matrix_subset, expected_density_matrix_subset
+        ), "Reduced density matrix did not match expected result for wire 0."
+
+        with pytest.raises(
+            ValueError, match="Specified wires must be a subset of the wires in the network ansatz."
+        ):
+            qnet.density_matrix_qnode(ansatz, wires=[3])

test/utilities_test.py

@@ -309,3 +309,59 @@ def test_ragged_reshape_error(self, input, list_dims):
             match=r"`len\(input_list\)` must match the sum of `list_dims`\.",
         ):
             qnetvo.ragged_reshape(input, list_dims)
+
+    def test_partial_transpose(self):
+        dm = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+
+        expected_result = np.array([[1, 0, 0, 1], [0, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 1]])
+
+        result = qnetvo.partial_transpose(dm, d1=2, d2=2)
+        assert np.allclose(
+            result, expected_result
+        ), "Partial transpose did not return the expected result."
+
+        dm = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+
+        expected_result = np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+
+        result = qnetvo.partial_transpose(dm, d1=1, d2=4)
+        assert np.allclose(
+            result, expected_result
+        ), "Partial transpose did not return the expected result."
+
+        dm = np.array(
+            [
+                [0, 1, 2, 3, 4, 5, 6, 7],
+                [8, 9, 10, 11, 12, 13, 14, 15],
+                [16, 17, 18, 19, 20, 21, 22, 23],
+                [24, 25, 26, 27, 28, 29, 30, 31],
+                [32, 33, 34, 35, 36, 37, 38, 39],
+                [40, 41, 42, 43, 44, 45, 46, 47],
+                [48, 49, 50, 51, 52, 53, 54, 55],
+                [56, 57, 58, 59, 60, 61, 62, 63],
+            ]
+        )
+
+        expected_result = np.array(
+            [
+                [0, 8, 16, 24, 4, 12, 20, 28],
+                [1, 9, 17, 25, 5, 13, 21, 29],
+                [2, 10, 18, 26, 6, 14, 22, 30],
+                [3, 11, 19, 27, 7, 15, 23, 31],
+                [32, 40, 48, 56, 36, 44, 52, 60],
+                [33, 41, 49, 57, 37, 45, 53, 61],
+                [34, 42, 50, 58, 38, 46, 54, 62],
+                [35, 43, 51, 59, 39, 47, 55, 63],
+            ]
+        )
+
+        result = qnetvo.partial_transpose(dm, d1=2, d2=4)
+        assert np.allclose(
+            result, expected_result
+        ), "Partial transpose did not return the expected result."
+
+        with pytest.raises(
+            ValueError,
+            match="The dimensions of the subsystems do not match the size of the density matrix.",
+        ):
+            qnetvo.partial_transpose(dm, d1=3, d2=3)