Source code for toqito.channels.pauli_channel

"""Generates and applies Pauli Channel to a matrix."""

import numpy as np
from scipy import sparse

from toqito.channel_ops.kraus_to_choi import kraus_to_choi
from toqito.helper.update_odometer import update_odometer
from toqito.matrices.pauli import pauli




def pauli_channel(
    prob: int | np.ndarray, return_kraus_ops: bool = False, input_mat: np.ndarray | None = None
) -> np.ndarray | sparse.csc_matrix | tuple:
    r"""Generate and apply a Pauli channel to a matrix.

    Generates the Choi matrix of a Pauli channel with given probabilities and optionally applies it
    to an input matrix. The Pauli channel is defined by the set of Pauli operators weighted by
    the probability vector. For a given probability vector \((p_0, \ldots, p_{4^q -1 })\), the
    channel is defined as shown below. Where, $q$ is the number of qubits.

    \[
        \Phi(\rho) = \sum_{i=0}^{4^q - 1} p_i P_i \rho P_i^*
    \]

    where \(P_i\) are Pauli operators generated by a lexographically increasing sequence of pauli operators of
    length strictly equal to \(q\), and \(p_i\) is the corresponding probability of that operator.
    For example, when \(q = 2\),

    \(P_{0} = I \otimes I\), \(P_{1} = I \otimes X\), \(P_{2} = I \otimes Y\), \(P_{3} = I \otimes Z\),
    \(P_{4} = X \otimes I\), \(P_{5} = X \otimes Y , \ldots P_{15} = Z \otimes Z\)

    If `prob` is a scalar, it generates a random `prob`-qubit Pauli channel.
    The length of the probability vector (if provided) must be \(4^q\) for some
    integer \(q\) (number of qubits).

    Examples:
        Generate a random single-qubit Pauli channel:

        ```python exec="1" source="above"
        from toqito.channels import pauli_channel

        print(pauli_channel(prob=1))
        ```

        Apply a specific two-qubit Pauli channel to an input matrix:

        ```python exec="1" source="above"
        import numpy as np
        from toqito.channels import pauli_channel

        _, output = pauli_channel(
            prob=np.array([0.1, 0.2, 0.3, 0.4]), input_mat=np.eye(2)
        )
        print(output)
        ```

    Raises:
        ValueError: If probabilities are negative or don't sum to 1.
        ValueError: If length of probability vector is not ``4^q`` for some integer ``q``.

    Args:
        prob: Probability vector for Pauli operators. If scalar, generates random probabilities for \(q =\) `prob`
        qubits. The probabilities correspond to Pauli operators in lexographical order of length strictly equal to \(q\)
            ,when `prob` is a vector.
        return_kraus_ops: Flag to return Kraus operators. Default is ``False``.
        input_mat: Optional input matrix to apply the channel to. Default is ``None``.

    Returns:
        The Choi matrix of the channel. If ``input_mat`` is provided, also returns the output matrix.
        If ``return_kraus_ops`` is ``True``, returns Kraus operators as well.

    """
    if not isinstance(prob, np.ndarray):
        if np.isscalar(prob):
            q = prob
            prob = np.random.rand(4**q)
            prob /= np.sum(prob)
        else:
            prob = np.array(prob)

    if np.any(prob < 0) or not np.isclose(np.sum(prob), 1):
        raise ValueError("Probabilities must be non-negative and sum to 1.")

    q = int(np.round(np.log2(len(prob)) / 2))
    if len(prob) != 4**q:
        raise ValueError("The length of the probability vector must be 4^q for some integer q (number of qubits).")

    Phi = sparse.csc_matrix((4**q, 4**q), dtype=complex)

    kraus_operators = []
    ind = np.zeros(q, dtype=int)

    for j in range(len(prob)):
        pauli_op = pauli(ind.tolist())
        kraus_operators.append(np.sqrt(prob[j]) * pauli_op)
        Phi += prob[j] * kraus_to_choi([[pauli_op, pauli_op.conj().T]])
        ind = update_odometer(ind, 4 * np.ones(q, dtype=int))

    output_mat = None
    if input_mat is not None:
        output_mat = sum(k @ input_mat @ k.conj().T for k in kraus_operators)

    if return_kraus_ops:
        return (Phi, output_mat, kraus_operators) if input_mat is not None else (Phi, kraus_operators)

    return (Phi, output_mat) if input_mat is not None else Phi