Add partial_transpose to quantum_info.linalg_operations

doc/source/api-reference/qibo.rst

@@ -1963,6 +1963,12 @@ Partial trace
 .. autofunction:: qibo.quantum_info.partial_trace
 
 
+Partial transpose
+"""""""""""""""""
+
+.. autofunction:: qibo.quantum_info.partial_transpose
+
+
 Matrix exponentiation
 """""""""""""""""""""
 

src/qibo/quantum_info/linalg_operations.py

@@ -97,7 +97,9 @@ def anticommutator(operator_1, operator_2):
     return operator_1 @ operator_2 + operator_2 @ operator_1
 
 
-def partial_trace(state, traced_qubits: Union[List[int], Tuple[int]], backend=None):
+def partial_trace(
+    state, traced_qubits: Union[List[int], Tuple[int, ...]], backend=None
+):
     """Returns the density matrix resulting from tracing out ``traced_qubits`` from ``state``.
 
     Total number of qubits is inferred by the shape of ``state``.
@@ -159,6 +161,83 @@ def partial_trace(state, traced_qubits: Union[List[int], Tuple[int]], backend=No
     return backend.np.einsum("abac->bc", state)
 
 
+def partial_transpose(
+    operator, partition: Union[List[int], Tuple[int, ...]], backend=None
+):
+    """Return matrix after the partial transposition of ``partition`` qubits in ``operator``.
+
+    Given a :math:`n`-qubit operator :math:`O \\in \\mathcal{H}_{A} \\otimes \\mathcal{H}_{B}`,
+    the partial transpose with respect to ``partition`` :math:`B` is given by
+
+    .. math::
+        \\begin{align}
+        O^{T_{B}} &= \\sum_{jklm} \\, O_{lm}^{jk} \\, \\ketbra{j}{k} \\otimes
+            \\left(\\ketbra{l}{m}\\right)^{T} \\\\
+        &= \\sum_{jklm} \\, O_{lm}^{jk} \\, \\ketbra{j}{k} \\otimes \\ketbra{m}{l} \\\\
+        &= \\sum_{jklm} \\, O_{ml}^{jk} \\, \\ketbra{j}{k} \\otimes \\ketbra{l}{m} \\, ,
+        \\end{align}
+
+    where the superscript :math:`T` indicates the transposition operation,
+    and :math:`T_{B}` indicates transposition on ``partition`` :math:`B`.
+    The total number of qubits is inferred by the shape of ``operator``.
+
+    Args:
+        operator (ndarray): :math:`1`- or :math:`2`-dimensional operator, or an array of
+            :math:`1`- or :math:`2`-dimensional operators,
+        partition (Union[List[int], Tuple[int, ...]]): indices of qubits to be transposed.
+        backend (:class:`qibo.backends.abstract.Backend`, optional): backend
+            to be used in the execution. If ``None``, it uses
+            :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
+
+    Returns:
+        ndarray: Partially transposed operator(s) :math:`\\O^{T_{B}}`.
+    """
+    backend = _check_backend(backend)
+
+    shape = operator.shape
+    nstates = shape[0]
+    dims = shape[-1]
+    nqubits = math.log2(dims)
+
+    if not nqubits.is_integer():
+        raise_error(
+            ValueError,
+            f"dimensions of ``state`` (or states in a batch) must be a power of 2.",
+        )
+
+    if (len(shape) > 3) or (nstates == 0) or (len(shape) == 2 and nstates != dims):
+        raise_error(
+            TypeError,
+            "``operator`` must have dims either (k,), (k, k), (N, 1, k) or (N, k, k), "
+            + f"but has dims {shape}.",
+        )
+
+    nqubits = int(nqubits)
+
+    if len(shape) == 1:
+        operator = backend.np.outer(operator, backend.np.conj(operator.T))
+    elif len(shape) == 3 and shape[1] == 1:
+        operator = backend.np.einsum(
+            "aij,akl->aijkl", operator, backend.np.conj(operator)
+        ).reshape(nstates, dims, dims)
+
+    new_shape = list(range(2 * nqubits + 1))
+    for ind in partition:
+        ind += 1
+        new_shape[ind] = ind + nqubits
+        new_shape[ind + nqubits] = ind
+    new_shape = tuple(new_shape)
+
+    reshaped = backend.np.reshape(operator, [-1] + [2] * (2 * nqubits))
+    reshaped = backend.np.transpose(reshaped, new_shape)
+
+    final_shape = (dims, dims)
+    if len(operator.shape) == 3:
+        final_shape = (nstates,) + final_shape
+
+    return backend.np.reshape(reshaped, final_shape)
+
+
 def matrix_exponentiation(
     phase: Union[float, complex],
     matrix,

tests/test_quantum_info_operations.py

@@ -7,6 +7,7 @@
     commutator,
     matrix_power,
     partial_trace,
+    partial_transpose,
 )
 from qibo.quantum_info.metrics import purity
 from qibo.quantum_info.random_ensembles import random_density_matrix, random_statevector
@@ -113,6 +114,93 @@ def test_partial_trace(backend, density_matrix):
     backend.assert_allclose(traced, Id)
 
 
+def _werner_state(p, backend):
+    zero, one = np.array([1, 0], dtype=complex), np.array([0, 1], dtype=complex)
+    psi = (np.kron(zero, one) - np.kron(one, zero)) / np.sqrt(2)
+    psi = np.outer(psi, np.conj(psi.T))
+    psi = backend.cast(psi, dtype=psi.dtype)
+
+    state = p * psi + (1 - p) * backend.identity_density_matrix(2, normalize=True)
+
+    # partial transpose of two-qubit werner state is known analytically
+    transposed = (1 / 4) * np.array(
+        [
+            [1 - p, 0, 0, -2 * p],
+            [0, p + 1, 0, 0],
+            [0, 0, p + 1, 0],
+            [-2 * p, 0, 0, 1 - p],
+        ],
+        dtype=complex,
+    )
+    transposed = backend.cast(transposed, dtype=transposed.dtype)
+
+    return state, transposed
+
+
+@pytest.mark.parametrize("batch", [False, True])
+@pytest.mark.parametrize("statevector", [False, True])
+@pytest.mark.parametrize("p", [1 / 5, 1 / 3, 1.0])
+def test_partial_transpose(backend, p, statevector, batch):
+    with pytest.raises(ValueError):
+        state = random_density_matrix(3, backend=backend)
+        test = partial_transpose(state, [0], backend)
+    with pytest.raises(TypeError):
+        state = np.random.rand(2, 2, 2, 2).astype(complex)
+        state += 1j * np.random.rand(2, 2, 2, 2)
+        state = backend.cast(state, dtype=state.dtype)
+        test = partial_transpose(state, [1], backend=backend)
+
+    if statevector:
+        zero, one = np.array([1, 0], dtype=complex), np.array([0, 1], dtype=complex)
+        psi = (np.kron(zero, one) - np.kron(one, zero)) / np.sqrt(2)
+
+        # testing statevector
+        target = np.zeros((4, 4), dtype=complex)
+        target[0, 3] = -1 / 2
+        target[1, 1] = 1 / 2
+        target[2, 2] = 1 / 2
+        target[3, 0] = -1 / 2
+        target = backend.cast(target, dtype=target.dtype)
+
+        psi = backend.cast(psi, dtype=psi.dtype)
+
+        if batch:
+            # the inner cast is required because of torch
+            psi = backend.cast([backend.cast([psi]) for _ in range(2)])
+
+        transposed = partial_transpose(psi, [0], backend=backend)
+
+        if batch:
+            for j in range(2):
+                backend.assert_allclose(transposed[j], target)
+        else:
+            backend.assert_allclose(transposed, target)
+    else:
+        state, target = _werner_state(p, backend)
+        if batch:
+            state = backend.cast([state for _ in range(2)])
+
+        # partial transpose of two-qubit werner state is known analytically
+        target = (1 / 4) * np.array(
+            [
+                [1 - p, 0, 0, -2 * p],
+                [0, p + 1, 0, 0],
+                [0, 0, p + 1, 0],
+                [-2 * p, 0, 0, 1 - p],
+            ],
+            dtype=complex,
+        )
+        target = backend.cast(target, dtype=target.dtype)
+
+        transposed = partial_transpose(state, [1], backend)
+
+        if batch:
+            for j in range(2):
+                backend.assert_allclose(transposed[j], target)
+        else:
+            backend.assert_allclose(transposed, target)
+
+
 @pytest.mark.parametrize("power", [2, 2.0, "2"])
 def test_matrix_power(backend, power):
     nqubits = 2