Source code for toqito.state_props.entanglement_of_formation

"""Computes the entanglement of formation of a bipartite quantum state."""

import numpy as np
import scipy

from toqito.matrix_ops import partial_trace
from toqito.state_props import concurrence, von_neumann_entropy




def entanglement_of_formation(rho: np.ndarray, dim: list[int] | int | None = None) -> float:
    r"""Compute entanglement-of-formation of a bipartite quantum state [@Quantiki_EOF].

    Entanglement-of-formation is the entropy of formation of the bipartite
    quantum state `rho`. Note that this function currently only supports
    `rho` being a pure state or a 2-qubit state: it is not known how to
    compute the entanglement-of-formation of higher-dimensional mixed states.

    This function was adapted from QETLAB.

    Examples:
        Compute the entanglement-of-formation of a Bell state.

        Let \(u = \frac{1}{\sqrt{2}} \left(|00\rangle + |11\rangle \right)\)
        and let

        \[
            \rho = uu^* = \frac{1}{2}\begin{pmatrix}
                                        1 & 0 & 0 & 1 \\
                                        0 & 0 & 0 & 0 \\
                                        0 & 0 & 0 & 0 \\
                                        1 & 0 & 0 & 1
                                     \end{pmatrix}.
        \]

        The entanglement-of-formation of \(\rho\) is equal to 1.

        ```python exec="1" source="above"
        import numpy as np
        from toqito.state_props import entanglement_of_formation
        from toqito.states import bell
        u_vec = bell(0)
        rho = u_vec @ u_vec.conj().T
        print(entanglement_of_formation(rho))
        ```

    Raises:
        ValueError: If matrices have improper dimension.

    Args:
        rho: A matrix or vector.
        dim: The default has both subsystems of equal dimension.

    Returns:
        A value between 0 and 1 that corresponds to the entanglement-of-formation of `rho`.

    """
    dim_x, dim_y = rho.shape
    round_dim = int(np.round(np.sqrt(max(dim_x, dim_y))))

    if dim is None:
        dim_int = round_dim
    elif isinstance(dim, int):
        dim_int = dim
    else:
        dim_int = None

    # User can specify dimension as integer.
    if dim_int is not None:
        dim_arr = np.array([dim_int, max(dim_x, dim_y) / dim_int], dtype=int)
        dim_arr[1] = np.round(dim_arr[1])
    else:
        dim_arr = np.array(dim)

    if np.prod(dim_arr) != max(dim_x, dim_y):
        raise ValueError("Invalid dimension: Please provide local dimensions that match the size of `rho`.")
    # If `rho` is a rank-1 density matrix, turn it into a vector instead
    # so we can compute the entanglement-of-formation easily.
    tmp_rho = scipy.linalg.orth(rho)
    if dim_x == dim_y and tmp_rho.shape[1] == 1:
        rho = tmp_rho
        dim_y = 1

    # Start computing entanglement-of-formation.
    if min(dim_x, dim_y) == 1:
        rho = rho[:]
        dim_list = [int(x) for x in dim_arr]
        return von_neumann_entropy(partial_trace(rho @ rho.conj().T, [1], dim_list))

    # Case: `rho` is a density matrix.
    if dim_x == dim_y:
        # In the two-qubit case, we know how to compute the
        # entanglement-of-formation exactly.
        if dim_x == 4:
            rho_c = concurrence(rho)

            rho_c1 = (1 + np.sqrt(1 - rho_c**2)) / 2
            rho_c2 = (1 - np.sqrt(1 - rho_c**2)) / 2

            rho_c1_log2 = 0 if rho_c1 == 0 else np.log2(rho_c1)
            rho_c2_log2 = 0 if rho_c2 == 0 else np.log2(rho_c2)

            return -rho_c1 * rho_c1_log2 - rho_c2 * rho_c2_log2
        raise ValueError(
            "Invalid input: It is presently only known how to compute "
            "the entanglement-of-formation for two-qubit states and pure "
            "states."
        )
    raise ValueError("Invalid dimension: `rho` must be either a vector or square matrix.")