Source code for toqito.states.mutually_unbiased_basis

"""Mutually unbiased basis states.

If a system prepared in an eigenstate of one of the bases gives an equal probability of (1/d) when measured with respect
to the other bases, mutually unbiased basis states are orthonormal bases in the Hilbert space Cᵈ.
"""

import numpy as np
from sympy import isprime, primerange

from toqito.matrices import gen_pauli




def mutually_unbiased_basis(dim: int) -> list[np.ndarray]:
    r"""Generate list of MUBs for a given dimension [@WikiMUB].

    Note that this function only works if the dimension provided is prime or a power of a prime. Otherwise, we don't
    know how to generate general MUBs.

    Examples:
        For the case of dimension 2, the three mutually unbiased bases are provided by:

        \[
            M_0 = \left\{|0\rangle, |1\rangle \right\}, \\
            M_1 = \left\{\frac{|0\rangle + |1\rangle}{\sqrt{2}}, \frac{|0\rangle - |1\rangle}{\sqrt{2}}\right\}
            M_2 = \left\{\frac{|0\rangle + i|1\rangle}{\sqrt{2}}, \frac{|0\rangle - i|1\rangle}{\sqrt{2}}\right\}
        \]

        The six vectors above are obtained accordingly:

        ```python exec="1" source="above" session="mubs"
        from toqito.states import mutually_unbiased_basis
        mubs = mutually_unbiased_basis(2)
        print(len(mubs))
        ```

        ```python exec="1" source="above" session="mubs"
        lst =[vec.shape for vec in mubs]
        print(lst)
        ```
    Args:
        dim: The dimension of the mutually unbiased bases to produce.

    Returns:
        The set of mutually unbiased bases of dimension `dim` (if known).

    """
    # The first basis will always be the standard basis:
    mats = [np.eye(dim)]

    pauli_x = gen_pauli(1, 0, dim)
    if isprime(dim):
        pauli_z = gen_pauli(0, 1, dim)

        for j in range(dim, 0, -1):
            _, eigen_vec = np.linalg.eig(pauli_x @ pauli_z ** (j))
            mats.append(eigen_vec)
    elif _is_prime_power(dim) and not isprime(dim):
        raise ValueError(f"Dimension {dim} is a prime power but not prime (more complicated no support at the moment).")
    else:
        raise ValueError(f"No general construction of MUBs is known for dimension: {dim}.")

    mubs: list[np.ndarray] = []
    for mat in mats:
        nrows, ncols = mat.shape[0], mat.shape[1]
        for row in range(nrows):
            mub: list[np.ndarray] = []
            for col in range(ncols):
                mub.append(mat[col][row])
            mubs.append(np.array(mub))
    return mubs



def _is_prime_power(n: int) -> bool:
    """Determine if a number is a prime power.

    A number is a prime power if it can be written as p^k, where p is a prime number and k is an integer greater than 0.

    Args:
        n: An integer to check for being a prime power.

    Returns:
        True if n is a prime power, False otherwise.

    """
    # 1 is not considered a prime power
    if n == 1:
        return False

    # Iterate over primes using a generator
    for p in primerange(2, int(n**0.5) + 1):
        # If p is a divisor of n
        if n % p == 0:
            # Keep dividing n by p as long as possible
            while n % p == 0:
                n //= p
            # If n becomes 1, then it's a prime power
            return n == 1
    # If n is itself a prime number
    return isprime(n)