performance.py

import numpy as np
from scipy.linalg import sqrtm

from utils import Fock_cutoff_calc, Pauli_eigenstates
from states import coherent_state, cat_lossy
from channel import syndromes, channel_tensor_f1, channel_tensor_f1f2


def trace_distance(rho, sigma):
    """
    calculates the trace distance between two density matrices
    """
    return 0.5 * np.trace(sqrtm((rho - sigma).T.conj() @ (rho - sigma)))


def fidelity(rho, sigma):
    """
    calculates the fidelity of two density matrices
    """
    return np.trace(sqrtm(sqrtm(rho) @ sigma @ sqrtm(rho))) ** 2


def syndrome_Pauli_maps(syndromes):
    """
    takes the set of transformations for (n,m) and spits out the associated Pauli correction for each
    benefits from rounding input

    Args:
        syndromes (array): an array of 2x2 transforms associated with each (n,m)
    Returns:
        X_map (array): 1 at each (n,m) needing a bit-flip, 0 otherwise
        Z_map (array): 1 at each (n,m) needing a phase-flip, 0 otherwise
    """

    # check whether either off-diagonal element is non-zero
    X_map = np.array(
        (np.abs(syndromes[:, :, 0, 1]) != 0) + (np.abs(syndromes[:, :, 1, 0]) != 0), dtype=int
    )
    # check whether the product of the elements is less than zero, indicating overall phase change
    # sum is done assuming transformations are either diagonal or off-diag
    Z_map = np.array(np.prod(np.sum(syndromes, axis=2), axis=2) < 0, dtype=int)

    return X_map, Z_map


def composite_Pauli_maps(L, Gamma, x, max_photon_count, tolerance=10):
    """
    a composite maps of X and Z errors
    constructed by scanning over amplitudes

    """

    max_alpha = int(max_photon_count**0.5)

    alpha_list = [alpha for alpha in range(1, max_alpha, 1)]
    n_cutoffs = np.array([Fock_cutoff_calc(alpha, tolerance) for alpha in alpha_list])
    # limit cutoffs to max photon count
    n_cutoffs[n_cutoffs > max_photon_count] = max_photon_count

    X_map = np.zeros((max_photon_count + 3, max_photon_count + 3), dtype=int)
    Z_map = np.zeros((max_photon_count + 3, max_photon_count + 3), dtype=int)

    for a, alpha in enumerate(alpha_list):
        n = n_cutoffs[a]

        # need to
        S = syndromes(L, alpha, Gamma, x, n).round(10)

        X, Z = syndrome_Pauli_maps(S)

        X_map[: n + 1, : n + 1] += X
        Z_map[: n + 1, : n + 1] += Z

    # resize, re-type
    X_map = X_map[: max_photon_count + 1, : max_photon_count + 1].astype(bool).astype(int)
    Z_map = Z_map[: max_photon_count + 1, : max_photon_count + 1].astype(bool).astype(int)

    return X_map, Z_map


def channel_tensor_Pauli_correction(channel_tensor, X_map, Z_map):
    """
    corrects any Pauli errors in channel_tensor_f1f2

    Args:
        channel_tensor (array): map in (n,m) of channel transformations, derived from channel_tensor_f1f2()
        X_map (array): map in (n,m) of where a bit-flip has occurred
        Z_map (array): map in (n,m) of where a phase-flip has occurred
    Returns:
        nc_pc (n,m): map in (n,m) of transformed states with relevant Pauli errors corrected
    """

    # add dummy indices to match those of dm_nm
    X_map_xtend = np.expand_dims(X_map, axis=(2, 3, 4, 5))
    Z_map_xtend = np.expand_dims(Z_map, axis=(2, 3, 4, 5))

    # logical not of the original maps
    X_map_xtend_not = np.abs(X_map_xtend - 1)
    Z_map_xtend_not = np.abs(Z_map_xtend - 1)

    # produce array where only nc elements on which to apply X remain
    channel_tensor_Xc = X_map_xtend * channel_tensor

    # Pauli X in terms of array operations
    #   flip horizontally, then vertically
    channel_tensor_Xc = np.flip(channel_tensor_Xc, axis=3)
    channel_tensor_Xc = np.flip(channel_tensor_Xc, axis=5)

    # add back in the values to which X was not applied
    channel_tensor_Xc += X_map_xtend_not * channel_tensor

    # produce an array that when multiplied in will produce a Pauli Z
    #   first the Z operation where Z_map is nonzero
    Z_array = Z_map_xtend * np.expand_dims([[1, -1], [-1, 1]], axis=(0, 1, 2, 4)).astype(
        np.complex128
    )
    #   then the identity elsewhere (adding in 1)
    Z_array += Z_map_xtend_not
    # apply to X-corrected map of density matrices
    channel_tensor_pc = Z_array * channel_tensor_Xc

    return channel_tensor_pc


def channel_fidelity_uncorrected(L, alpha_list, Gamma_list, tolerance=10):
    """
    channel fidelity for lossy cat states without correction
    does this over specified lists of amplitude and loss values for efficiency
    
    Args:
        L (int): code order
        alpha_list (list): amplitude values
        Gamma_list (list): loss values
        tolerance (int): tolerance used to specify Fock cutoffs
    Returns:
        fidelities (array): channel fidelity for each alpha, Gamma
    """
    n_cutoffs = [Fock_cutoff_calc(alpha, tolerance) for alpha in alpha_list]

    rho_i_00 = [[1, 0], [0, 0]]
    rho_i_11 = [[0, 0], [0, 1]]
    rho_i_01 = [[0, 1], [0, 0]]
    rho_i_10 = [[0, 0], [1, 0]]

    rho_i = np.array([rho_i_00, rho_i_11, rho_i_01, rho_i_10])

    fidelities = np.empty((len(alpha_list), len(Gamma_list)))

    for a, alpha in enumerate(alpha_list):

        count_to = n_cutoffs[a]

        coherent_states = np.array(
            [coherent_state(m, L, alpha, count_to) for m in range(2 * (L + 1))]
        )
        signs = np.expand_dims([1, -1] * (L + 1), axis=1)

        cat0 = np.sum(coherent_states, axis=0)
        cat0 /= np.linalg.norm(cat0)
        cat1 = np.sum(signs * coherent_states, axis=0)
        cat1 /= np.linalg.norm(cat1)

        state_initial_pure = (
            np.tensordot(cat0, cat0, axes=0) + np.tensordot(cat1, cat1, axes=0)
        ) / 2**0.5
        state_initial_pure = state_initial_pure.flatten()

        rho_i_Fock = []
        for rho in rho_i:
            rho_i_Fock += [cat_lossy(rho, L, alpha, 0, count_to, tolerance)]

        for g, Gamma in enumerate(Gamma_list):

            rho_f_Fock = []
            # rho_i_Fock = []
            for rho in rho_i:
                # transform the state and sum over all (n,m) (gives the average state)
                rho_f_Fock += [cat_lossy(rho, L, alpha, Gamma, count_to, tolerance)]

            rho_entangled = 0
            for i in range(len(rho_i)):
                # construct entangled state with loss channel applied to first mode
                rho_entangled += 0.5 * np.tensordot(rho_f_Fock[i], rho_i_Fock[i], axes=0)

            del rho_f_Fock

            # using np.reshape() does not construct the tensor as it should
            dm = np.block(
                [[rho_entangled[i, j] for j in range(count_to + 1)] for i in range(count_to + 1)]
            )

            fidelities[a, g] = state_initial_pure.T.conj() @ dm @ state_initial_pure

    return fidelities


def channel_fidelity(channel_average, N=10):
    """
    channel fidelity as defined in "Performance and structure of single-mode bosonic codes"

    Args:
        transform_average (array): average of channel dervied from channel_tensor_f1f2()
        N (int): number of successive iterations of the channel
    Returns:
        (array): channel fidelity
    """

    rho_i_00 = [[1, 0], [0, 0]]
    rho_i_11 = [[0, 0], [0, 1]]
    rho_i_01 = [[0, 1], [0, 0]]
    rho_i_10 = [[0, 0], [1, 0]]

    dm_i_list = [rho_i_00, rho_i_11, rho_i_01, rho_i_10]

    state_initial_pure = (
        np.tensordot([1, 0], [1, 0], axes=0) + np.tensordot([0, 1], [0, 1], axes=0)
    ) / 2**0.5
    state_initial_pure = np.reshape(state_initial_pure, (4))

    dm_f_list = np.array(dm_i_list, dtype=np.complex128)

    for Ni in range(N):
        # the final average states after transmission
        dm_f = []
        for dm in dm_f_list:
            dm_f += [np.einsum("cdef,ce->df", channel_average, dm)]

        dm_f_list = np.array(dm_f)

    state_final = np.zeros((2, 2, 2, 2), dtype=np.complex128)
    for d in range(4):
        # construct entangled state with loss channel applied to first mode
        state_final += 0.5 * np.tensordot(dm_f_list[d], dm_i_list[d], axes=0)

    dm_pc = np.block([[state_final[i, j] for j in range(2)] for i in range(2)])

    channel_fidelity_pc = state_initial_pure.T.conj() @ dm_pc @ state_initial_pure

    return channel_fidelity_pc.real


def mean_error_probability(channel_average, N=10):
    """
    mean error probability between pairs of Pauli eigenstates
    for a given average channel and number of iterations

    Args:
        transform_average (array): average of channel channel_tensor_f1f2() with Pauli correction
        N (int): number of successive iterations of the channel
    Returns:
        (array): mean error probability
    """

    dm_list = Pauli_eigenstates()
    dm0_ = dm_list[::2]
    dm1_ = dm_list[1::2]

    trace_distances = np.zeros((3), dtype=float)

    for d in range(3):
        dm0_i = dm0_[d]
        dm1_i = dm1_[d]
        for Ni in range(N):
            dm0_i = np.einsum("cdef,ce->df", channel_average, dm0_i)
            dm1_i = np.einsum("cdef,ce->df", channel_average, dm1_i)

        trace_distances[d] = trace_distance(dm0_i, dm1_i).real

    mean_error_prob = 0.5 * (1 - np.sum(trace_distances) / 3)

    return mean_error_prob


def performance_scan(performance_function, L, alpha_list, Gamma_list, N_list, tolerance=10):
    """
    calculates the performance function over amplitude, loss, and iteration values

    Args:
        performance_function (func): function taking the average channel and number of iterations as arguments
        L (int): code order
        alpha_list (list): alpha values
        Gamma_list (list): total loss values
        N_list (list): numbers of iterations
    Returns:
        (array): output of performance_metric() at specified parameters
    """

    n_cutoffs = [Fock_cutoff_calc(alpha, tolerance) for alpha in alpha_list]

    master_channel_f1 = channel_tensor_f1(L, n_cutoffs[-1])

    master_X_map, master_Z_map = composite_Pauli_maps(L, 0.2, 1, n_cutoffs[-1], tolerance=10)

    performance_metric = np.empty((len(alpha_list), len(Gamma_list), len(N_list)), dtype=float)

    for a, alpha in enumerate(alpha_list):

        Fock_size = n_cutoffs[a] + 1

        channel_f1 = master_channel_f1[:Fock_size, :Fock_size]

        X_map = master_X_map[:Fock_size, :Fock_size]
        Z_map = master_Z_map[:Fock_size, :Fock_size]

        for g, Gamma in enumerate(Gamma_list):
            for n, N in enumerate(N_list):

                gamma = 1 - (1 - Gamma) ** (1 / N)

                channel_f1f2 = channel_tensor_f1f2(channel_f1, alpha, gamma)
                channel_f1f2_Pauli_corrected_average = np.sum(
                    channel_tensor_Pauli_correction(channel_f1f2, X_map, Z_map), axis=(0, 1)
                )

                performance_metric[a, g, n] = performance_function(
                    channel_f1f2_Pauli_corrected_average, N
                )

    return performance_metric


def mean_error_probability_scan_Hastrup(L, alpha_list, Gamma_list, N_list, tolerance=10):
    """
    calculates the mean error probability over amplitude, loss, and iteration values
    uses the same method as in "All-optical cat-code quantum error correction"
    using an extra segment of loss on the state before calculating the error prob

    Args:
        L (int): code order
        alpha_list (list): alpha values
        Gamma_list (list): total loss values
        N_list (list): numbers of iterations
    Returns:
        (array): mean error probability
    """

    dm_list = Pauli_eigenstates()
    dm0_ = dm_list[::2]
    dm1_ = dm_list[1::2]

    n_cutoffs = [Fock_cutoff_calc(alpha, tolerance) for alpha in alpha_list]

    master_channel_f1 = channel_tensor_f1(L, n_cutoffs[-1])

    master_X_map, master_Z_map = composite_Pauli_maps(L, 0.2, 1, n_cutoffs[-1], tolerance=10)

    mean_error_prob = np.empty((len(alpha_list), len(Gamma_list), len(N_list)), dtype=float)

    for a, alpha in enumerate(alpha_list):

        Fock_size = n_cutoffs[a] + 1

        channel_f1 = master_channel_f1[:Fock_size, :Fock_size]

        X_map = master_X_map[:Fock_size, :Fock_size]
        Z_map = master_Z_map[:Fock_size, :Fock_size]

        for g, Gamma in enumerate(Gamma_list):
            for n, N in enumerate(N_list):

                gamma = 1 - (1 - Gamma) ** (1 / (N + 1))

                channel_f1f2 = channel_tensor_f1f2(channel_f1, alpha, gamma)
                channel_f1f2_Pauli_corrected_average = np.sum(
                    channel_tensor_Pauli_correction(channel_f1f2, X_map, Z_map), axis=(0, 1)
                )

                trace_distances = np.zeros((3), dtype=float)

                for d in range(3):
                    dm0_i = dm0_[d]
                    dm1_i = dm1_[d]
                    for Ni in range(N):
                        dm0_i = np.einsum(
                            "cdef,ce->df", channel_f1f2_Pauli_corrected_average, dm0_i
                        )
                        dm1_i = np.einsum(
                            "cdef,ce->df", channel_f1f2_Pauli_corrected_average, dm1_i
                        )

                    # could be faster by saving coherent states in advance...
                    dm0_i_Fock = cat_lossy(dm0_i, L, alpha, gamma, n_cutoffs[a])
                    dm1_i_Fock = cat_lossy(dm1_i, L, alpha, gamma, n_cutoffs[a])

                    trace_distances[d] = trace_distance(dm0_i_Fock, dm1_i_Fock).real

                mean_error_prob[a, g, n] = 0.5 * (1 - np.sum(trace_distances) / 3)

    return mean_error_prob


def target_channel_fidelity_scan(
    fidelity_target, channel_f1, Pauli_maps, alpha_search, Gamma, N_bounds
):
    """
    scans the given amplitude and iteration values until channel fidelity >= fidelity_target

    Args:
        target_fidelity (float): minimum fidelity required
        channel_tensor_f1 (array): channel tensor for function f1 associated with desired order
        Pauli_maps (tuple): (X_map, Z_map) the maps of X and Z errors for the given order
        alpha_search (tuple): (start, stop, step) for alpha scan
        Gamma (float): amount of loss
        N_bounds (tuple): (start, stop) for iteration scan
    Returns:
        (tuple): fidelity, N, alpha
    """

    N_range = np.arange(*N_bounds, 1)
    alpha_range = np.arange(*alpha_search)

    fidelity_array = np.zeros((len(N_range), len(alpha_range)), dtype=float)
    params = []

    for Ni, N in enumerate(N_range):
        for a, alpha in enumerate(alpha_range):

            gamma = 1 - (1 - Gamma) ** (1 / N)

            channel_f1f2 = channel_tensor_f1f2(channel_f1, alpha, gamma)
            channel_f1f2_Pauli_corrected_average = np.sum(
                channel_tensor_Pauli_correction(channel_f1f2, *Pauli_maps), axis=(0, 1)
            )
            fidelity_i = channel_fidelity(channel_f1f2_Pauli_corrected_average, N).real

            fidelity_array[Ni, a] = fidelity_i

            params += [(N, alpha)]

            if fidelity_i >= fidelity_target:
                break
        if fidelity_i >= fidelity_target:
            break

    if fidelity_i < fidelity_target:
        print("unable to find sufficient fidelity")
        print("max fidelity =", np.max(fidelity_array.round(4)))
        print("at (N,alpha) =", params[np.argmax(fidelity_array)])
        fidelity_i = None

    if Ni == 0:
        print("decrease lower iteration bound")
        fidelity_i = None

    return fidelity_i, N, alpha