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

# # GRAPE calculation of control fields for single-qubit rotation

# 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
from numpy import pi


# In[2]:


from qutip import *
from qutip.control import *


# In[3]:


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


# In[4]:


theta, phi = np.random.rand(2)


# In[5]:


# target unitary transformation (random single qubit rotation)
U = rz(phi) * rx(theta); U


# In[6]:


R = 150
H_ops = [sigmax(), sigmay(), sigmaz()]

H_labels = [r'$u_{x}$',
            r'$u_{y}$',
            r'$u_{z}$',
        ]


# In[7]:


H0 = 0 * pi * sigmaz()


# # GRAPE

# In[8]:


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


# In[9]:


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


# In[10]:


u0 = np.array([np.random.rand(len(times)) * 2 * pi * 0.005 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[11]:


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


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

# In[12]:


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


# In[13]:


# target unitary
U


# In[14]:


# unitary from grape pulse
result.U_f


# In[15]:


# target / result overlap
_overlap(U, result.U_f).real, abs(_overlap(U, result.U_f))**2


# ### Verify correctness of the Hamiltonian pulses by integration

# 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)


# # Bloch sphere dynamics

# In[20]:


psi0 = basis(2, 0)
e_ops = [sigmax(), sigmay(), sigmaz()]


# In[21]:


me_result = mesolve(result.H_t, psi0, times, c_ops, e_ops)


# In[22]:


b = Bloch()

b.add_points(me_result.expect)

b.add_states(psi0)
b.add_states(U * psi0)
b.render()


# # Process tomography

# ## Ideal gate

# In[23]:


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


# In[24]:


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)


# ## Gate calculated using GRAPE

# In[25]:


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

U_ideal = spre(result.U_f) * spost(result.U_f.dag())

chi = qpt(U_ideal, op_basis)

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


# ## Versions

# In[26]:


from qutip.ipynbtools import version_table

version_table()