Magic_State_V.py

"""
Code to understand magic states by visualizing single qubit transformations
around the bloch sphere explicitly with the random circuits of the 
Hadamard and T gate.

Code Written by Andrew Projansky
Project Start Date: 7/21/2022
"""

import numpy as np
import matplotlib.pyplot as plt
import random

"""
Defines all universal gates used
"""

H = np.array([[1 / np.sqrt(2), 1 / np.sqrt(2)], [1 / np.sqrt(2), -1 / np.sqrt(2)]])
S = np.array([[1, 0], [0, 1j]])
T = np.array([[np.exp(-1j * np.pi / 8), 0], [0, np.exp(1j * np.pi / 8)]])
Z = np.array([[1, 0], [0, -1]])
CNot = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]])
I = np.array([[1, 0], [0, 1]])

"""
Defines all useful states and matrix products for simulation
"""

magic_state = np.array([1/np.sqrt(2),1/np.sqrt(2)*np.exp(1j*np.pi/4)])
rho_m = np.tensordot(magic_state, np.conjugate(magic_state),0)
sig_z_p = np.kron(I, Z)
gate_dict = {1: T, 2: H}
proj = {0: np.array([[1,0],[0,0]]), 1: np.array([[0,0],[0,1]])}

def cc(gate):
    """
    Returns complex conjugate of gate

    Parameters
    ----------
    gate: list
        Unitary gate that can be applied on a state
    """

    return np.transpose(np.matrix.conjugate(gate))

def partial_conditional_measure(rho_t):
    """
    Performs partial trace over the ancilla system, then applies
    conditional S gate to the leftover state.

    Parameters
    ----------
    rho_t: list
        Hermetian total density matrix to which operations and
        traces will be applied to
    """

    #one day maybe make projection generators rather than just having a list
    bc = get_bias(rho_t)
    fullp = np.kron(I, proj[bc])
    rhopr = np.matmul(rho_t, fullp)
    rho_p = np.matmul(fullp, rhopr)/(np.trace(rhopr))
    rho_red = partial_trace(rho_p)
    if bc == 1:
        rho_red = np.matmul(S, np.matmul(rho_red, cc(S)))
    return rho_red

def get_bias(rho):
    """
    For a multi_qubit density matrix in which we wish to measure over the
    final qubit, we find the weighted distribution of whether to choose the
    0 or 1 state projector, and biased coin flip to pick one or the other

    Parameters
    ----------
    rho: list
        Hermetian density matrix which we take an expectation value over
        identiy tensor sigma z
    """
    e_v = np.real(np.trace( np.matmul(rho, sig_z_p) ))
    ran = random.random()
    return 0 if ran > (0.5+e_v/2) else 1

def partial_trace(rho):
    """
    partial traces over final qubit
    #should generalize this to any size matrix when I have time
    - involves defining new matrix as being generated by taking
    2x2 traces in tons of section - just have to automate it

    Parameters
    ----------
    rho: density matrix we're going to take the partial trace over
    """

    new_matrix = np.array([[rho[0][0]+rho[1][1],rho[0][2]+rho[1][3]],
                           [rho[2][0]+rho[3][1],rho[2][2]+rho[3][3]]])
    return new_matrix

class Experiment:
    """
    Parameters
    ----------
    num_steps : int, optional
        total number of steps. The default is 1.
    magic: bool, optional
        Boolean value to determine whether to use magic state or not. The
        default is False
    state: list, optional
        Initial state position. The default is the plus Z eigenstate

    Attributes
    ----------
    angles : list
        angles from states Used for plotting
    """

    def __init__(
        self,
        num_steps: int = 1,
        magic: bool = False,
        state: list = np.array([[1,0],[0,0]])

    ):
        self.num_steps = num_steps
        self.magic = magic
        self.state = state


        self.angles = [[],[],[]]

    ################# run functions ##################################
    def run_stepwise(self):
        for i in range(self.num_steps):
            gate_choice = random.randint(1,2)
            gate = gate_dict[gate_choice]
            if (gate_choice == 1) and (self.magic):
                self.magic_step()
            else:
                self.state = np.matmul(gate, np.matmul(self.state, cc(gate)))
            self.__gen_xyz_points()

    def magic_step(self):
        rho_t = np.kron(self.state, rho_m)
        rho_t = np.matmul(CNot, np.matmul(rho_t, cc(CNot)))
        self.state = partial_conditional_measure(rho_t)

    ##############plot functions ########################
    def plot(self):
        fig = plt.figure()
        ax = fig.add_subplot(projection="3d")
        u, v = np.mgrid[0 : 2 * np.pi : 50j, 0 : np.pi : 50j]
        xs = np.cos(u) * np.sin(v)
        ys = np.sin(u) * np.sin(v)
        zs = np.cos(v)
        ax.plot_surface(xs, ys, zs, color="lightgrey", alpha=0.3)
        ax.scatter(self.angles[0], self.angles[1], self.angles[2], marker="o", alpha=0.9)
        plt.show()

    def __gen_xyz_points(self):

        self.angles[0].append(2*np.real(self.state[1][0]))
        self.angles[1].append(2*np.imag(self.state[1][0]))
        self.angles[2].append(2*self.state[0][0]-1)

#%%
exp = Experiment(num_steps=100, magic=True)
exp.run_stepwise()
exp.plot()
del exp
#%%