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

# # GRAPE calculation of control fields for iSWAP implementation

# Robert Johansson (robert@riken.jp)

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import time
import numpy as np


# In[2]:


from qutip import *
from qutip.control import *


# In[3]:


T = 1
times = np.linspace(0, T, 100)


# In[4]:


U = iswap()
R = 50
H_ops = [#tensor(sigmax(), identity(2)),
         #tensor(sigmay(), identity(2)),
         #tensor(sigmaz(), identity(2)),
         #tensor(identity(2), sigmax()),
         #tensor(identity(2), sigmay()),
         #tensor(identity(2), sigmaz()),
         tensor(sigmax(), sigmax()),
         tensor(sigmay(), sigmay()),
         tensor(sigmaz(), sigmaz())]

H_labels = [#r'$u_{1x}$',
            #r'$u_{1y}$',
            #r'$u_{1z}$',
            #r'$u_{2x}$',
            #r'$u_{2y}$',
            #r'$u_{2z}$',
            r'$u_{xx}$',
            r'$u_{yy}$',
            r'$u_{zz}$',
        ]


# In[5]:


H0 = 0 * np.pi * (tensor(sigmaz(), identity(2)) + tensor(identity(2), sigmaz()))


# # GRAPE

# In[6]:


from qutip.control.grape import plot_grape_control_fields, _overlap, grape_unitary_adaptive, cy_grape_unitary


# In[7]:


from scipy.interpolate import interp1d
from qutip.ui.progressbar import TextProgressBar


# In[8]:


u0 = np.array([np.random.rand(len(times)) * (2 * np.pi / T) * 0.01 for _ in range(len(H_ops))])

u0 = [np.convolve(np.ones(10)/10, u0[idx, :], mode='same') for idx in range(len(H_ops))]


# In[9]:


result = cy_grape_unitary(U, H0, H_ops, R, times, u_start=u0, eps=2*np.pi/T,
                          progress_bar=TextProgressBar())


# In[10]:


#result = grape_unitary(U, H0, H_ops, R, times, u_start=u0, eps=2*np.pi/T,
#                       progress_bar=TextProgressBar())


# ## Plot control fields for iSWAP gate in the presense of single-qubit tunnelling

# In[11]:


plot_grape_control_fields(times, result.u / (2 * np.pi), H_labels, uniform_axes=True);


# In[12]:


# compare to the analytical results
np.mean(result.u[-1,0,:]), np.mean(result.u[-1,1,:]), np.pi/(4 * T)


# ## Fidelity

# In[13]:


U


# In[14]:


result.U_f.tidyup(1e-2)


# In[15]:


_overlap(U, result.U_f).real


# ## Test numerical integration of GRAPE pulse

# In[16]:


c_ops = []


# In[17]:


U_f_numerical = propagator(result.H_t, times[-1], c_ops, args={})


# In[18]:


U_f_numerical


# In[19]:


_overlap(U, U_f_numerical).real


# # Process tomography

# ## Ideal iSWAP gate

# In[20]:


op_basis = [[qeye(2), sigmax(), sigmay(), sigmaz()]] * 2
op_label = [["i", "x", "y", "z"]] * 2


# In[21]:


fig = plt.figure(figsize=(8,6))

U_ideal = spre(U) * spost(U.dag())

chi = qpt(U_ideal, op_basis)

fig = qpt_plot_combined(chi, op_label, fig=fig, threshold=0.001)


# ## iSWAP gate calculated using GRAPE

# In[22]:


fig = plt.figure(figsize=(8,6))

U_ideal = to_super(result.U_f)

chi = qpt(U_ideal, op_basis)

fig = qpt_plot_combined(chi, op_label, fig=fig, threshold=0.001)


# ## Versions

# In[23]:


from qutip.ipynbtools import version_table

version_table()