# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + slideshow={"slide_type": "-"}
'''This file gives all of the functions used to define states
and implement quantum operations.'''


import cmath, random, numpy
import functools
import  matplotlib.pyplot as plt
import sys
import os
import math
from qutip import*

from sympy import*
#from sympsi import*
from scipy import optimize

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import time
import math
from qutip import *
from qutip.ipynbtools import plot_animation
import numpy as np
import matplotlib.pyplot as plt
import qutip
# %matplotlib inline
import matplotlib.pylab as plt
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from IPython.display import display, Math, Latex
import cmath
from mpl_toolkits.axes_grid1 import AxesGrid
from scipy.special import factorial as fac

xvec = np.arange(-80.,80.)*5./80
yvec = np.arange(-50.,50)*5/40
X,Y = np.meshgrid(xvec, xvec)  ##Some plotting params
X1,Y1 = np.meshgrid(yvec,yvec)
N_dim = 35 ##Dimenstion of the Hilbert spac

'''Annihilation operators.  Each "a" is the same, but we define different ones so
as to make defining different operations on different modes less confusing.'''
a1 = destroy(N_dim) ##This is for single-photon field
a2 = destroy(N_dim) ##for coherent field
a3 = destroy(N_dim) ##for vacuum field

'''Define a function for n choose k'''
def n_choose_k(n,k):
    return fac(n)/(fac(n-k)*fac(k))

'''Applies the displacement operator to a density matrix'''
def D(state,alpha):  
    Rho_new=displace(N_dim,alpha)*state*displace(N_dim,alpha).dag()
    return Rho_new   

'''Define a rotation in phase space, or phase shifter operation'''
def Phase(theta):  
    b=-1j*theta*a1.dag()*a1;
    return b.expm()

'''Squeezing operation, inputs a density matrix and outputs the 
squeezed density matrix for squeezing parameter r'''
def Sq(state,r):
    Rho_new=squeeze(N_dim,r)*state*squeeze(N_dim,r).dag();
    return Rho_new

'''The function below creates a beamsplitter operation that acts on 
two modes. The value for k determines what number Fock state could be
filtered out of the first state based on a single photon input for the
second BS port, followed by single photon detection.'''
def BS_operator_filtering(a1, a2, k):
    theta_k = np.arctan(1/np.sqrt(k))
    T = np.sin(theta_k)*np.sin(theta_k)
    R = np.cos(theta_k)*np.cos(theta_k)
    print('I am filtering', k, 'and:', theta_k*180/math.pi)
    print('BS T is : ', T, 'and : ', R)
    b = theta_k*(tensor(a1,a2.dag()) - tensor(a1.dag(),a2))
    return b.expm()

'''Define Schrodinger cat states'''
def cat_plus(alpha):
    cat = (1/(np.sqrt(2)*np.sqrt(1+np.e**(-alpha*alpha.conj()))))*(coherent(N_dim,-alpha)+(coherent(N_dim,alpha)))
    return cat
def cat_minus(alpha):
    cat = (1/(np.sqrt(2)*np.sqrt(1-np.e**(-alpha*alpha.conj()))))*(-coherent(N_dim,-alpha)+(coherent(N_dim,alpha)))
    return cat

'''Define superpositions of squeezed vacuum (SSV), which are 
superpositions where the vacuum has been squeezed and then 
displaced along the antisqueezed axis of the squeezing, in a
way where the number of elements in the sum gives the degree 
of symmetry of the state. Degree 2 SSV states are equivalent 
to cat states, but with a different squeezing parameter'''
def SSV_plus(r,alpha):
    state = ket2dm((displace(N_dim,alpha)+displace(N_dim,-alpha))
                   *squeeze(N_dim,r)*fock(N_dim,0))
    norm_state = state/state.tr()
    return norm_state
def SSV_minus(r,alpha):
    state = ket2dm((displace(N_dim,alpha)-displace(N_dim,-alpha))
                   *squeeze(N_dim,r)*fock(N_dim,0))
    norm_state = state/state.tr()
    return norm_state

'''General SSV state, where r is the squeezing parameter, amp is the
displacement amplitude, and M is the degree of symmetry.'''
def SSV_gen(r,amp,M):
    state=0;
    for k in range(M):
        state+=displace(N_dim,amp*np.e**(2j*k*np.pi/M))*squeeze(N_dim,r*np.e**(4j*k*np.pi/M))*fock(N_dim,0)
    norm_state=ket2dm(state.unit())
    return norm_state

'''Degree 3 symmetry SSV state'''
def sq_tri_cat(r,amp):
    state=ket2dm((displace(N_dim,amp)*squeeze(N_dim,r)*fock(N_dim,0)
                  +displace(N_dim,amp*np.e**(2j*np.pi/3))
                  *squeeze(N_dim,r*np.e**(4j*np.pi/3))*fock(N_dim,0)
                  +displace(N_dim,amp*np.e**(4j*np.pi/3))
                  *squeeze(N_dim,r*np.e**(2j*np.pi/3))*fock(N_dim,0)))
    norm_state=state/state.tr()
    return norm_state

'''Degree 4 symmetry SSV state'''
def sq_quad_cat(r,amp):
    state=ket2dm((displace(N_dim,amp)*squeeze(N_dim,r)*fock(N_dim,0)
                  +displace(N_dim,amp*np.e**(2j*np.pi/4))
                  *squeeze(N_dim,r*np.e**(4j*np.pi/4))*fock(N_dim,0)
                  +displace(N_dim,amp*np.e**(4j*np.pi/4))
                  *squeeze(N_dim,r*np.e**(8j*np.pi/4))*fock(N_dim,0)
                  +displace(N_dim,amp*np.e**(6j*np.pi/4))
                  *squeeze(N_dim,r*np.e**(12j*np.pi/4))*fock(N_dim,0)))
    norm_state=state/state.tr()
    return norm_state

'''Utilize positive-operator value measurments (POVM) of the detector 
to define a PNR detector with efficiency eta'''
def pnr_resolution_detector(eta, click, n_truc):
    pi_n = 0;
    l = np.arange(click,n_truc)
    for i in l:
        pi_n +=  n_choose_k(i,click)*math.pow((1-eta),(i-click))*math.pow(eta,click)*fock(N_dim,i)*fock(N_dim,i).dag()
        #print("The final Povm element is: ", pi_0)
    return Qobj(pi_n)

'''Performs photon catalysis with Fock state input.  Both inputs 
are density matrices, and the returned output mode is a normalized
density matrix after the PNR detection'''
def Fock_Filter_povm(in_state,in_fock,refl,num_det,eta,n_truc):
    Projector = tensor(pnr_resolution_detector(eta, num_det, n_truc),qeye(N_dim));
    Initial_state=tensor(in_state,ket2dm(fock(N_dim,in_fock)));
    theta_k=np.arccos(np.sqrt(refl));
    
    BS1= ((theta_k)*(tensor(a1,a2.dag()) - tensor(a1.dag(),a2))).expm()
        
    Rho=BS1*Initial_state*BS1.dag();
        
    Rho_filtered = ((Rho*Projector).ptrace(1))/((Rho*Projector).tr())
    '''The operation .ptrace(m) takes the partial trace over every mode 
    EXCEPT m, where the numbering startes at 0.  So .ptrace(1) means
    you keep mode 1, which is actually the 2nd mode'''
    print('BS has reflectivity',refl,' and I am detecting the |',num,
          '> state,where my detector has efficiency', eta)
    return Rho_filtered

'''Performs photon catalysis with Fock state input and calculates 
the probability of success.'''
def Fock_Filter_prob(in_state,in_fock,refl,num_det,eta,n_truc):
    Projector = tensor(pnr_resolution_detector(eta, num_det, n_truc),qeye(N_dim));
    Initial_state=tensor(in_state,ket2dm(fock(N_dim,in_fock)));
    theta_k=np.arccos(np.sqrt(refl));
    
    BS1= ((theta_k)*(tensor(a1,a2.dag()) - tensor(a1.dag(),a2))).expm()
        
    Rho=BS1*Initial_state*BS1.dag();
    P=(Rho*Projector).tr()
    print('The probability of a sucessful detection is:',P)    
    Rho_filtered = ((Rho*Projector).ptrace(1))/((Rho*Projector).tr())
    #Rho_filtered=Rho*Projector
    '''The operation .ptrace(m) takes the partial trace over every mode 
    EXCEPT m, where the numbering startes at 0.  So .ptrace(1) means you
    keep mode 1, which is actually the 2nd mode'''
    print('BS has reflectivity',refl,' and I am detecting the |',num,
          '> state, where my detector has efficiency', eta)
    return Rho_filtered

'''Generic photon catalysis where the two input states are allowed to be 
arbitrary. Takes in two density matrices and returns a density matrix 
along with success probability.'''
def catalysis(in1,in2,refl,num_det,eta,n_truc):
    Projector = tensor(pnr_resolution_detector(eta, num_det, n_truc),qeye(N_dim));
    Initial_state=tensor(in1,in2);
    theta_k=np.arccos(np.sqrt(refl));
    
    BS1= ((theta_k)*(tensor(a1,a2.dag()) - tensor(a1.dag(),a2))).expm()
        
    Rho=BS1*Initial_state*BS1.dag();
    
    P=(Rho*Projector).tr()
    print('The probability of a sucessful detection is:',P)   
    
    Rho_filtered = ((Rho*Projector).ptrace(1))/((Rho*Projector).tr())
    '''The operation .ptrace(m) takes the partial trace over every mode 
    EXCEPT m, where the numbering startes at 0.  So .ptrace(1) means you 
    keep mode 1, which is actually the 2nd mode'''
    print('BS has reflectivity',refl,' and I am detecting the |',num,
          '> state, where my detector has efficiency', eta)
    return Rho_filtered

'''Defines the fidelity between two arbitrary quantum states'''
def fid(state1,state2):
    F=np.absolute((state1.sqrtm()*state2*state1.sqrtm()).sqrtm().tr())
    return F



# +
def my_hex_GKP(mu, d, delta, cutoff, nmax=400):
    r"""Hexagonal GKP code state.
    The Hex GKP state is defined by
    .. math::
        |mu> = \sum_{n_1,n_2=-\infty}^\infty e^{-i(q+\sqrt{3}p)/2}
            \sqrt{4\pi/\sqrt{3}d}(dn_1+\mu) e^{iq\sqrt{4\pi/\sqrt{3}d}n_2}|0>
    where d is the dimension of a code space, \mu=0,1,...,d-1, |0> is the
    vacuum state, and the states are modulated by a Gaussian envelope in the
    case of finite energy:
    ..math:: e^{-\Delta ^2 n}|\mu>
    Args:
        d (int): the dimension of the code space.
        mu (int): mu=0,1,...,d-1.
        delta (float): width of the modulating Gaussian envelope.
        cutoff (int): the Fock basis truncation of the returned density matrix.
        nmax (int): the Hex GKP state |mu> is calculated by performing the
            sum using n1,n1=-nmax,...,nmax.
    Returns:
        Density Matrix: a size [cutoff] complex Qobj.
    """
    n1 = np.arange(-nmax, nmax+1)[:, None]
    n2 = np.arange(-nmax, nmax+1)[None, :]

    n1sq = n1**2
    n2sq = n2**2

    sqrt3 = np.sqrt(3)

    #arg1 = -1j*np.pi*n2*(d*n1+mu)/d
    arg1 = -1j*np.pi*n2*(d*n1+mu)/d
    arg2 = -np.pi*(d**2*n1sq+n2sq-d*n1*(n2-2*mu)-n2*mu+mu**2)/(sqrt3*d)
    arg2 *= 1-np.exp(-2*delta**2)

    #amplitude = (np.exp(arg1)*np.exp(arg2)).flatten()[:, None]
    amplitude = (np.exp(arg1)).flatten()[:, None]


    alpha = np.sqrt(np.pi/(2*sqrt3*d)) * (sqrt3*(d*n1+mu) - 1j*(d*n1-2*n2+mu))
    #alpha *= np.exp(-delta**2)

    alpha = alpha.flatten()[:, None]
    n = np.arange(cutoff)[None, :]
    coherent = np.exp(-0.5*np.abs(alpha)**2)*alpha**n/np.sqrt(fac(n))
    
    #hex_state = np.sum(amplitude*coherent*np.exp(-n*delta**2), axis=0)
    hex_state = ket2dm(Qobj(np.sum(amplitude*coherent*np.exp(-n*delta**2), axis=0)))
    final=hex_state/hex_state.tr()
    return final

def sqrGKP(mu, d, delta, cutoff, nmax=400):  
    """Square grid GKP state
    The GKP state is defined by
    .. math::
        |mu> = \sum_{n_1,n_2=-\infty}^\infty e^{-i(p)
            \sqrt{2\pi/d}(dn_1+\mu)} e^{iq\sqrt{2\pi/d}n_2}|0>
    where d is the dimension of a code space, \mu=0,1,...,d-1, |0> is the
    vacuum state, and the states are modulated by a Gaussian envelope in the
    case of finite energy:
    ..math:: e^{-\Delta ^2 n}|\mu>
    Args:
        d (int): the dimension of the code space.
        mu (int): mu=0,1,...,d-1.
        delta (float): width of the modulating Gaussian envelope.
        cutoff (int): the Fock basis truncation of the returned density matrix.
        nmax (int): the GKP state |mu> is calculated by performing the
            sum using n1,n1=-nmax,...,nmax.
    Returns:
        Density Matrix: a size [cutoff] complex Qobj.
    """
    n1 = np.arange(-nmax, nmax+1)[:, None]
    n2 = np.arange(-nmax, nmax+1)[None, :]

    n1sq = n1**2
    n2sq = n2**2

    sqrt3 = np.sqrt(3)

    #arg1 = -1j*np.pi*n2*(d*n1+mu)/d
    arg1 = 1j*np.pi*n2*(d*n1+mu)/d

    amplitude = (np.exp(arg1)).flatten()[:, None]


    alpha = np.sqrt(np.pi/(d)) * ((d*n1+mu - 1j*n2))

    alpha = alpha.flatten()[:, None]
    n = np.arange(cutoff)[None, :]
    coherent = np.exp(-0.5*np.abs(alpha)**2)*alpha**n/np.sqrt(fac(n))
    
    #hex_state = np.sum(amplitude*coherent*np.exp(-n*delta**2), axis=0)
    sqr_state = ket2dm(Qobj(np.sum(amplitude*coherent*np.exp(-n*delta**2), axis=0)))
    final=sqr_state/sqr_state.tr()
    return final
# -



StarcoderdataJupyterScriptsDedupFiltered