#!/usr/bin/env python
# coding: utf-8

# # Spin Squeezing in presence of local and collective noise
# 
# Notebook author: Nathan Shammah (nathan.shammah at gmail.com)
# 
# Here we study the effect of collective and local processes on a spin squeezing Hamiltonian. 
# 
# We consider a system of $N$ two-level systems (TLSs) with identical frequency $\omega_{0}$, which can de-excite incoherently or collectively at the rates $\gamma_\text{E}$ and $\gamma_\text{CE}$,
# 
# \begin{eqnarray}
# \dot{\rho} &=&-i\lbrack -i\Lambda\left(J_{+}^2-J_{-}^2\right),\rho \rbrack
# +\frac{\gamma_\text {CE}}{2}\mathcal{L}_{J_{-}}[\rho]
# +\frac{\gamma_\text{E}}{2}\sum_{n=1}^{N}\mathcal{L}_{J_{-,n}}[\rho]
# \end{eqnarray}
# 
# We study the time evolution of the spin squeezing parameter [1-4]
# \begin{eqnarray}
# \xi^2 &=& N\langle\Delta J_y^2\rangle/\left(\langle J_z\rangle^2+\langle J_x\rangle^2\right)
# \end{eqnarray}
# 
# We assess how different dynamical conditions and initial states can be explored to optimize the spin squeezing of a given Dicke state [5-7]. This study can be generalized to other types of local and collective incoherent processes. A table grouping this processes is given below, 
# 
# <table>
#     <tr>
# <td> Keyword</td>
# <td> Rate  $\gamma_j$</td>
# <td> Lindbladian $\mathcal{L}[\rho]$</td>
# </tr>
# 
# <tr>
# <td> $\texttt{emission}$ </td>
# <td> $\gamma_\text{E}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=&\sum_n^N \left(J_{-,n}\rho J_{+,n} - \frac{1}{2}J_{+,n}J_{-,n}\rho - \frac{1}{2}\rho J_{+,n}J_{-,n} \right)\end{eqnarray}</td>
# </tr>
# 
# <tr>
# <td> $\texttt{pumping}$ </td>
# <td> $\gamma_\text{P}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=&\sum_n^N \left(J_{+,n}\rho J_{-,n} - \frac{1}{2}J_{-,n}J_{+,n}\rho - \frac{1}{2}\rho J_{-,n}J_{+,n} \right)\end{eqnarray}</td>
# </tr>
# 
# <tr>
# <td> $\texttt{dephasing}$ </td>
# <td> $\gamma_\text{D}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=&\sum_n^N \left(J_{z,n}\rho J_{z,n} - \frac{1}{2}J_{z,n}J_{z,n}\rho - \frac{1}{2}\rho J_{z,n}J_{z,n} \right)\end{eqnarray}</td>
# </tr>
# 
# <tr>
# <td> $\texttt{collective}\_\texttt{emission}$ </td>
# <td> $\gamma_\text{CE}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=& J_{-}\rho J_{+} - \frac{1}{2}J_{+}J_{-}\rho - \frac{1}{2}\rho J_{+}J_{-} \end{eqnarray}</td>
# </tr>
# 
# <tr>
# <td> $\texttt{collective}\_\texttt{pumping}$ </td>
# <td> $\gamma_\text{CP}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=&  J_{+}\rho J_{-} - \frac{1}{2}J_{-}J_{+}\rho - \frac{1}{2}\rho J_{-}J_{+} \end{eqnarray}</td>
# </tr>
# 
# <tr>
# <td> $\texttt{collective}\_\texttt{dephasing}$ </td>
# <td> $\gamma_\text{CD}$</td>
# <td> \begin{eqnarray}\mathcal{L}[\rho]&=& J_{z}\rho J_{z} - \frac{1}{2}J_{z}^2\rho - \frac{1}{2}\rho J_{z}^2 \end{eqnarray}</td>
# </tr>
# 
# </table>
# 
# 
# Note that in the table above and in $\texttt{qutip.piqs}$ functions, the Lindbladian $\mathcal{L}[\rho]$ is written with a factor 1/2 with respect to $\mathcal{L}_{A}[\rho]$ reported in the LaTeX math equations, in order to have the Lindbladian and full Liouvillian matrix consistently defined by the rates $\gamma_\alpha$. 
# 
# Note also that the *local depolarizing channel* can be written in terms of this Lindbladians as
# \begin{eqnarray}
# \gamma_{Dep}\sum_n^N\left(\mathcal{L}_{J_{x,n}}+\mathcal{L}_{J_{y,n}}+\mathcal{L}_{J_{z,n}}\right)=\gamma_{Dep}\sum_n^N\left(\frac{1}{2}\mathcal{L}_{J_{+,n}}+\frac{1}{2}\mathcal{L}_{J_{-,n}}+ \mathcal{L}_{J_{z,n}}\right).
# \end{eqnarray}
# Similarly, the *collective depolarizing channel* reads
# 
# \begin{eqnarray}
# \gamma_\text{CDep}\left(\mathcal{L}_{J_{x}}+\mathcal{L}_{J_{y}}+\mathcal{L}_{J_{z}}\right)=\gamma_\text{CDep}\left(
# \frac{1}{2}\mathcal{L}_{J_{+}}+\frac{1}{2}\mathcal{L}_{J_{-}}+ \mathcal{L}_{J_{z}}\right).
# \end{eqnarray}

# In[1]:


from time import clock
from scipy.io import mmwrite
import matplotlib.pyplot as plt
from qutip import *
from qutip.piqs import *
from scipy.sparse import load_npz, save_npz


# In[2]:


def isdicke(N, j, m):
    """
    Check if an element in a matrix is a valid element in the Dicke space.
    Dicke row: j value index. Dicke column: m value index. 
    The function returns True if the element exists in the Dicke space and
    False otherwise.

    Parameters
    ----------
    N : int
        The number of two-level systems.    
    j: float
        "j" index of the element in Dicke space which needs to be checked.
    m: float
        "m" index of the element in Dicke space which needs to be checked.
    """
    dicke_row = j
    dicke_col = m
    
    rows = N + 1
    cols = 0

    if (N % 2) == 0:
        cols = int(N/2 + 1)
    else:
        cols = int(N/2 + 1/2)

    if (dicke_row > rows) or (dicke_row < 0):
        return (False)

    if (dicke_col > cols) or (dicke_col < 0):
        return (False)

    if (dicke_row < int(rows/2)) and (dicke_col > dicke_row):
        return False

    if (dicke_row >= int(rows/2)) and (rows - dicke_row <= dicke_col):
        return False

    else:
        return True

def dicke_space(N):
    """
    Generate a matrix to visualize the Dicke space.
    j is on the horizontal axis, increasing right to left.
    m is on the vertical axis, increasing bottom to top.
    It puts 1 in all allowed (j,m) values.
    It puts 0 in all not-allowed (j,m) values.
    Parameters
    ----------
    N : int
        The number of two-level systems.
    Returns
    ----------
    dicke_space : ndarray
        The matrix of all allowed (j,m) pairs.

    """        
    rows = N + 1
    cols = 0

    if (rows % 2) == 0:
        cols = int((rows/2))

    else:
        cols = int((rows + 1)/2)

    dicke_space = np.zeros((rows, cols), dtype = int)

    for (i, j) in np.ndindex(rows, cols):
        dicke_space[i, j] = isdicke(N, i, j)

    return (dicke_space)


# In[3]:


## general parameters
N = 20
ntls = N
nds = num_dicke_states(N)
[jx, jy, jz] = jspin(N)
jp = jspin(N, "+")
jm = jspin(N, "-")
jpjm = jp*jm

Lambda = 1
factor_l = 5

#spin hamiltonian
h = -1j*Lambda * (jp**2-jm**2)
gCE = Lambda/factor_l
gE = Lambda/factor_l

# system with collective emission only
system = Dicke(N=N)
# system2 with local emission only
system2 = Dicke(N=N)
system.collective_emission = gCE
system2.emission = gE
system.hamiltonian = h
system2.hamiltonian = h
liouv = system.liouvillian() 
liouv2 = system2.liouvillian()

print(system)
print(system2)


# # Time evolution of Spin Squuezing Parameter $\xi^2= \frac{N \langle\Delta J_y^2\rangle}{\langle J_z\rangle^2}$

# In[4]:


#set initial state for spins (Dicke basis)
nt = 1001
td0 = 1/(N*Lambda)
tmax =  10 * td0
t = np.linspace(0, tmax, nt)
excited = dicke(N, N/2, N/2)
load_file = False
if load_file == False:
    # cycle over all states in Dicke space
    xi2_1_list = []
    xi2_2_list = []
    xi2_1_min_list = []
    xi2_2_min_list = []

    for j in j_vals(N):
        #for m in m_vals(j):
        m = j
        rho0 = dicke(N, j, m)
        #solve using qutip (Dicke basis)
        # Dissipative dynamics: Only collective emission 
        result = mesolve(liouv, rho0, t, [], 
                         e_ops = [jz, jy, jy**2,jz**2, jx],
                         options = Options(store_states=True))
        rhot = result.states
        jz_t = result.expect[0]
        jy_t = result.expect[1]
        jy2_t = result.expect[2]
        jz2_t = result.expect[3]
        jx_t = result.expect[4]
        Delta_jy = jy2_t - jy_t**2
        xi2_1 = N * Delta_jy / (jz_t**2+jx_t**2)
        # Dissipative dynamics: Only local emission 
        result2 = mesolve(liouv2, rho0, t, [], 
                          e_ops = [jz, jy, jy**2,jz**2, jx],
                          options = Options(store_states=True))
        rhot2 = result2.states
        jz_t2 = result2.expect[0]
        jy_t2 = result2.expect[1]
        jy2_t2 = result2.expect[2]
        jz2_t2 = result2.expect[3]
        jx_t2 = result2.expect[4]
        Delta_jy2 = jy2_t2 - jy_t2**2
        xi2_2 = N * Delta_jy2 / (jz_t2**2+jx_t2**2)

        xi2_1_min = np.min(xi2_1)
        xi2_2_min = np.min(xi2_2)        
        xi2_1_list.append(xi2_1)
        xi2_2_list.append(xi2_2)
        xi2_1_min_list.append(xi2_1_min)
        xi2_2_min_list.append(xi2_2_min)        

        print("|j, m> = ",j,m)


# #### Visualization

# In[5]:


label_size2 = 20
lw = 3
texplot = False
# if texplot == True:
#     plt.rc('text', usetex = True)
#     plt.rc('xtick', labelsize=label_size) 
#     plt.rc('ytick', labelsize=label_size)

fig1 = plt.figure(figsize = (10,6))
for xi2_1 in xi2_1_list:
    plt.plot(t*(N*Lambda), xi2_1, '-', label = r' $\gamma_\Downarrow=0.2$', linewidth = lw)
for xi2_2 in xi2_2_list:
    plt.plot(t*(N*Lambda), xi2_2, '-.', label = r'$\gamma_\downarrow=0.2$')
plt.plot(t*(N*Lambda), 1+0*t, '--k')
plt.xlim([0,3])
plt.ylim([0,8000.5])
plt.ylim([0,2.5])
plt.xlabel(r'$ N \Lambda t$', fontsize = label_size2)
plt.ylabel(r'$\xi^2$', fontsize = label_size2)
#plt.legend(fontsize = label_size2*0.8)
plt.title(r'Spin Squeezing Parameter, $N={}$'.format(N), fontsize = label_size2)
plt.show()
plt.close()


# In[6]:


## Here we find for how long the spin-squeezing parameter, xi2, 
## is less than 1 (non-classical or "quantum" condition), in the two dynamics

dt_quantum_xi1_list = []
dt_quantum_xi2_list = []

dt1_jm =[]
dt2_jm =[]
ds = dicke_space(N)
i = 0
for j in j_vals(N):
    #for m in m_vals(j):
    m = j
    rho0 = dicke(N, j, m)
    quantum_xi1 = xi2_1_list[i][xi2_1_list[i] < 1.0] 
    quantum_xi2 = xi2_2_list[i][xi2_2_list[i] < 1.0]

    # first ensemble
    if len(quantum_xi1)>0:
        dt_quantum_xi1 = len(quantum_xi1)
        dt1_jm.append((dt_quantum_xi1, j, m))

    else:
        dt_quantum_xi1 = 0.0

    # second ensemble
    if len(quantum_xi2)>0:
        dt_quantum_xi2 = len(quantum_xi2)
        dt2_jm.append((dt_quantum_xi2, j, m))
    else:
        dt_quantum_xi2 = 0.0

    dt_quantum_xi1_list.append(dt_quantum_xi1)
    dt_quantum_xi2_list.append(dt_quantum_xi2)

    i = i+1


# In[7]:


print("collective emission: (squeezing time, j, m)")
print(dt1_jm)
print("local emission: (squeezing time, j, m)")
print(dt2_jm)


# #### Visualization

# In[8]:


plt.rc('text', usetex = True)
label_size = 20
label_size2 = 20
label_size3 = 20
plt.rc('xtick', labelsize=label_size) 
plt.rc('ytick', labelsize=label_size)

lw = 3
i0 = -3
i0s=2
fig1 = plt.figure(figsize = (8,5))
# excited state spin squeezing
plt.plot(t*(N*Lambda), xi2_1_list[-1], 'k-', 
         label = r'$|\frac{N}{2},\frac{N}{2}\rangle$, $\gamma_\Downarrow=0.2\Lambda$', 
         linewidth = 0.8)
plt.plot(t*(N*Lambda), xi2_2_list[-1], 'r--',
         label = r'$|\frac{N}{2},\frac{N}{2}\rangle$, $\gamma_\downarrow=0.2\Lambda$',
         linewidth = 0.8)
# state with max time of spin squeezing

plt.plot(t*(N*Lambda), xi2_1_list[i0], 'k-', 
         label = r'$|j,j\rangle$, $\gamma_\Downarrow=0.2\Lambda$', 
         linewidth = 0.8+0.4*i0s*lw)
plt.plot(t*(N*Lambda), xi2_2_list[i0], 'r--',
         label = r'$|j,j\rangle$, $\gamma_\downarrow=0.2\Lambda$',
         linewidth = 0.8+0.4*i0s*lw)
plt.plot(t*(N*Lambda), 1+0*t, '--k')

plt.xlim([0,2.5])
plt.yticks([0,1,2])
plt.ylim([-1,2.])

plt.xlabel(r'$ N \Lambda t$', fontsize = label_size3)
plt.ylabel(r'$\xi^2$', fontsize = label_size3)
plt.legend(fontsize = label_size2*0.8, ncol=2)
fname = 'figures/spin_squeezing_N_{}_states.pdf'.format(N)
plt.title(r'Spin Squeezing Parameter, $N={}$'.format(N), fontsize = label_size2)
plt.show()
plt.close()


# The plot shows the spin squeezing parameter for two different dynamics -- only collective de-excitation, black curves; only local de-excitation, red curves -- and for two different inital states, the maximally excited state (thin curves) and another Dicke state with longer squeezing time (thick curves). This study, performed in Refs. [5,6] for the maximally excited state has been extended to any Dicke state in Ref. [7].

# In[10]:


# plot the dt matrix in the Dicke space
plt.rc('text', usetex = True)
label_size = 20
label_size2 = 20
label_size3 = 20
plt.rc('xtick', labelsize=label_size) 
plt.rc('ytick', labelsize=label_size)

lw = 3
i0 = 7
i0s=2
ratio_squeezing_local = 3
fig1 = plt.figure(figsize = (6,8))
ds = dicke_space(N)
value_excited = 3
ds[0,0]=value_excited
ds[int(N/2-i0),int(N/2-i0)]=value_excited * ratio_squeezing_local
plt.imshow(ds, cmap="inferno_r")
plt.xticks([])
plt.yticks([])
plt.xlabel(r"$j$", fontsize = label_size3)
plt.ylabel(r"$m$", fontsize = label_size3)
plt.title(r"Dicke space $(j,m)$ for $N={}$".format(N), fontsize = label_size3)
plt.show()
plt.close()


# The Plot above shows the two initial states (darker dots) $|\frac{N}{2},\frac{N}{2}\rangle$ (top edge of the Dicke triangle, red dot) and $|j,j\rangle$, with $j=\frac{N}{2}-3=7$ (black dot). A study of the Dicke triangle (dark yellow space) and state engineering is performed in Ref. [8] for different initial state. 

# #### References
# 
# [1] D. J. Wineland, J. J. Bollinger, W. M. Itano, F. L. Moore, and D. J. Heinzen, Spin squeezing and reduced quantum noise in spectroscopy, *Phys. Rev. A* **46**, R6797 (1992)
# 
# [2] M. Kitagawa and M. Ueda, Squeezed spin states, *Phys. Rev. A* **47**, 5138 (1993)
# 
# [3] J. Ma, X. Wang, C.-P. Sun, and F. Nori, Quantum spin squeezing, *Physics Reports* **509**, 89 (2011)
# 
# [4] L. Pezzè, A. Smerzi, M. K. Oberthaler, R. Schmied, and P. Treutlein, Quantum metrology with nonclassical states of atomic ensembles, *Reviews of Modern Physics*, in press (2018)
# 
# [5] B. A. Chase and J. Geremia, Collective processes of an ensemble of spin-1 particles, *Phys. Rev. A* **78**,0521012 (2008)
# 
# [6] B. Q. Baragiola, B. A. Chase, and J. Geremia, Collective uncertainty in partially polarized and partially deco- hered spin-1 systems, *Phys. Rev. A* **81**, 032104 (2010)
# 
# [7] N. Shammah, S. Ahmed, N. Lambert, S. De Liberato, and F. Nori, 
# Open quantum systems with local and collective incoherent processes: Efficient numerical simulation using permutational invariance https://arxiv.org/abs/1805.05129
# 
# [8] N. Shammah, N. Lambert, F. Nori, and S. De Liberato, Superradiance with local phase-breaking effects, *Phys. Rev. A* **96**, 023863 (2017).
# 
# 

# In[ ]:


qutip.about()