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

# In[1]:


# Set up the cluster for parallel computation
from IPython.parallel import Client
client = Client(profile='lsf')


# In[2]:


len(client.ids)


# In[3]:


dview = client[:]
dview.use_dill()
lview = client.load_balanced_view();


# In[1]:


#%%px --local

get_ipython().run_line_magic('matplotlib', 'inline')
import qutip
import numpy as np
import scipy
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from IPython import display
import sys
import time

from itertools import product
from types import FunctionType, BuiltinFunctionType
from functools import partial

def expand_lindblad_operator_to_qutip_std_form(*summands):
    """Expands the superoperator L(A+B+...) to sum of independent dissipator operators.
    
    Warning: only real coefficients are supported.
    
    The input `[matA, coefA], [matB, coefB], ...` represents the Lindblad dissipator
    `L(matA*coefA + matB*coefB + ...)`. `mat*` is an operator, `coef*` is a time
    dependent coefficient (in any of the qutip supported formats).
    
    If a summand is constant just input it on its own,
    e.g. `[matA, coefA], matConst, [matB, coefB]`.
    
    The output is a list of independent dissipator superoperators in form that can be
    used directly with `mesolve`'s `c_ops`."""
    matrix_coeff_pairs = [_ if isinstance(_, list) else [_, None] for _ in summands]

    if all(isinstance(_, np.ndarray) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'array'
    elif all(isinstance(_, str) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'str'
    elif all(isinstance(_, (FunctionType, BuiltinFunctionType, partial)) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'function'
    else:
        raise TypeError("Incorrect format for the coefficients")

    c_ops = sum((properly_conjugated_lindblad(*_, coef_type=coef_type) for _ in product(matrix_coeff_pairs, repeat=2)),[])
    return sorted(c_ops, key=lambda _:isinstance(_, list), reverse=True)

def properly_conjugated_lindblad(left, right, coef_type):
    '''Return the two components of the Lindblad superoperator with properly conjugated coeffs.
    
    For input `(A, ca), (B, cb)` return 
    `[[pre(A)*post(Bdag), ca*conj(cb)],[-0.5*(pre(Adag*B)+post(Adag*B)), conj(ca)*cb]]`.
    
    It supports constant operators (just set `ca` or `cb` to `None`).
    '''
    if coef_type == 'array':
        conj = lambda _: np.conj(_)
        prod = lambda a,b: a*b
    elif coef_type == 'str':
        conj = lambda _: 'conj(%s)'%_
        prod = lambda a,b: '(%s)*(%s)'%(a,b)
    elif coef_type == 'function':
        raise NotImplementedError
        
    m_left, c_left = left
    m_right, c_right = right
    
    A = qutip.spre(m_left) * qutip.spost(m_right.dag())
    ld_r = m_left.dag()*m_right
    B = - 0.5 * qutip.spre(ld_r) - 0.5 * qutip.spost(ld_r)

    if c_left is None and c_right is None:
        return [A+B]
    elif c_left is None:
        return [[A, conj(c_right)],
                [B,      c_right]]
    elif c_right is None:
        return [[A,      c_left],
                [B, conj(c_left)]]
    else:
        return [[A, prod(     c_left,  conj(c_right))],
                [B, prod(conj(c_left),      c_right)]]


# In[19]:


get_ipython().run_cell_magic('px', '--local', 'def two_circles_both_move_radius_one_half(alpha0=2, T=200, samples_factor=1):\n    assert np.abs(alpha0) < 5.1, \'alpha > 6 at any point in time goes out of the cutoff for the hilbert space\'\n    N = 80\n    radius = 0.5\n    samples = int(15000*samples_factor)\n\n    unit_interval = np.linspace(0,1,samples)\n    centered_interval = np.linspace(-1,1,samples)\n    one = np.ones(samples)\n    zero = np.zeros(samples)\n    gaussian = np.exp(-centered_interval**2)\n\n    a_op = qutip.destroy(N)\n    adag_op = qutip.create(N)\n    id_op = qutip.identity(N)\n    zero_op = qutip.zero_oper(N)\n    num_op = qutip.num(N)\n\n    beta0 = -alpha0\n    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()\n    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)\n    betas  = beta0  - radius*(1-np.cos(-unit_interval*2*np.pi)) + radius*1j*np.sin(-unit_interval*2*np.pi)\n\n    c_ops = expand_lindblad_operator_to_qutip_std_form(\n                a_op**2,\n                [-a_op, alphas+betas],\n                [id_op, alphas*betas]\n            )\n    \n    \n    import signal\n\n    def signal_handler(signum, frame):\n        raise Exception("Timed out!")\n\n    signal.signal(signal.SIGALRM, signal_handler)\n    signal.alarm(2*60*60)\n    try:\n        ret = qutip.mesolve(zero_op,init_state,unit_interval*T,\n                            c_ops=c_ops,\n                            e_ops=[],\n                            progress_bar=qutip.ui.progressbar.TextProgressBar())\n        return ret.states[-1], len(ret.states)\n    except Exception:\n        signal.alarm(0)\n        return init_state, 1\n')


# In[2]:


alpha0s = np.linspace(2,5,num=16)
Ts = np.concatenate([np.linspace(20,200,num=10)[:-1],np.linspace(200,2000,num=10)])
args = list(product(alpha0s, Ts))


# In[21]:


get_ipython().run_cell_magic('px', '--local', 'def two_circles_one_moves_radius_one_half_sqrt2(alpha0=2, T=200, samples_factor=1):\n    assert np.abs(alpha0) < 5.1, \'alpha > 6 at any point in time goes out of the cutoff for the hilbert space\'\n    N = 80\n    radius = 0.5*(2**0.5)\n    samples = int(15000*samples_factor)\n\n    unit_interval = np.linspace(0,1,samples)\n    centered_interval = np.linspace(-1,1,samples)\n    one = np.ones(samples)\n    zero = np.zeros(samples)\n    gaussian = np.exp(-centered_interval**2)\n\n    a_op = qutip.destroy(N)\n    adag_op = qutip.create(N)\n    id_op = qutip.identity(N)\n    zero_op = qutip.zero_oper(N)\n    num_op = qutip.num(N)\n\n    beta0 = -alpha0\n    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()\n    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)\n    betas  = beta0 * np.ones_like(alphas)\n\n    c_ops = expand_lindblad_operator_to_qutip_std_form(\n                a_op**2,\n                [-a_op, alphas+betas],\n                [id_op, alphas*betas]\n            )\n    \n    import signal\n\n    def signal_handler(signum, frame):\n        raise Exception("Timed out!")\n\n    signal.signal(signal.SIGALRM, signal_handler)\n    signal.alarm(2*60*60)\n    try:\n        ret = qutip.mesolve(zero_op,init_state,unit_interval*T,\n                            c_ops=c_ops,\n                            e_ops=[],\n                            progress_bar=qutip.ui.progressbar.TextProgressBar())\n        return ret.states[-1], len(ret.states)\n    except Exception:\n        signal.alarm(0)\n        return init_state, 1\n    \ndef two_circles_one_moves_radius_one_half(alpha0=2, T=200, samples_factor=1):\n    assert np.abs(alpha0) < 5.1, \'alpha > 6 at any point in time goes out of the cutoff for the hilbert space\'\n    N = 80\n    radius = 0.5\n    samples = int(15000*samples_factor)\n\n    unit_interval = np.linspace(0,1,samples)\n    centered_interval = np.linspace(-1,1,samples)\n    one = np.ones(samples)\n    zero = np.zeros(samples)\n    gaussian = np.exp(-centered_interval**2)\n\n    a_op = qutip.destroy(N)\n    adag_op = qutip.create(N)\n    id_op = qutip.identity(N)\n    zero_op = qutip.zero_oper(N)\n    num_op = qutip.num(N)\n\n    beta0 = -alpha0\n    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()\n    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)\n    betas  = beta0 * np.ones_like(alphas)\n\n    c_ops = expand_lindblad_operator_to_qutip_std_form(\n                a_op**2,\n                [-a_op, alphas+betas],\n                [id_op, alphas*betas]\n            )\n\n    import signal\n\n    def signal_handler(signum, frame):\n        raise Exception("Timed out!")\n\n    signal.signal(signal.SIGALRM, signal_handler)\n    signal.alarm(2*60*60)\n    try:\n        ret = qutip.mesolve(zero_op,init_state,unit_interval*T,\n                            c_ops=c_ops,\n                            e_ops=[],\n                            progress_bar=qutip.ui.progressbar.TextProgressBar())\n        return ret.states[-1], len(ret.states)\n    except Exception:\n        signal.alarm(0)\n        return init_state, 1\n')


# In[ ]:


results_two_circles_both_move_standard = lview.map(lambda a: two_circles_both_move_radius_one_half(*a), args)
results_two_circles_one_moves_sqrt2_standard = lview.map(lambda a: two_circles_one_moves_radius_one_half_sqrt2(*a), args)
results_two_circles_one_moves_standard = lview.map(lambda a: two_circles_one_moves_radius_one_half(*a), args)


# In[23]:


results_two_circles_both_move_standard.wait_interactive()


# In[24]:


import pickle
with open("results_two_circles_both_move_standard", "wb") as f:
    pickle.dump(results_two_circles_both_move_standard.result, f)


# In[25]:


results_two_circles_one_moves_sqrt2_standard.wait_interactive()


# In[26]:


with open("results_two_circles_one_moves_sqrt2_standard", "wb") as f:
    pickle.dump(results_two_circles_one_moves_sqrt2_standard.result, f)


# In[27]:


results_two_circles_one_moves_standard.wait_interactive()


# In[28]:


with open("results_two_circles_one_moves_standard", "wb") as f:
    pickle.dump(results_two_circles_one_moves_standard.result, f)


# In[3]:


import pickle
with open("results_two_circles_both_move_standard", "rb") as f:
    results_two_circles_both_move_standard = pickle.load(f)
with open("results_two_circles_one_moves_sqrt2_standard", "rb") as f:
    results_two_circles_one_moves_sqrt2_standard = pickle.load(f)
with open("results_two_circles_one_moves_standard", "rb") as f:
    results_two_circles_one_moves_standard = pickle.load(f)


# In[4]:


len([_ for _ in results_two_circles_both_move_standard if _[1]])


# In[55]:


def plot_stuff(results, title=''):
    steps_purity = np.array([(_[0]**2).tr() if _[1]==15000 else 0 for _ in results])
    steps_exchange = np.array([(qutip.coherent(80,-a).dag()*_[0]*qutip.coherent(80,a)).tr() if _[1]==15000 else 0
                               for _, (a, t) in zip(results,args)])
    mask = np.array([_[1] for _ in results])
    steps_purity.shape = len(alpha0s), len(Ts)
    steps_exchange.shape = len(alpha0s), len(Ts)
    mask.shape = len(alpha0s), len(Ts)
    #plt.imshow(mask, interpolation='nearest', origin='lower', cmap=plt.get_cmap('RdBu'))
    #mask = np.triu(mask[:,::-1],k=5)[:,::-1]
    #plt.imshow(mask, interpolation='nearest', origin='lower', cmap=plt.get_cmap('RdBu'))
    steps_purity[mask!=15000] = np.nan
    steps_purity = np.ma.MaskedArray(steps_purity, mask!=15000)
    steps_exchange[mask!=15000] = np.nan
    steps_exchange = np.ma.MaskedArray(steps_exchange, mask!=15000)
    f = plt.figure(figsize=(16,16))
    pvst = f.add_subplot(2,2,1)
    pvst.plot(Ts, steps_purity.T[:,:], '.-')
    pvst.set_xlabel(r'$T$')
    pvst.set_title(r'Purity-vs-$T$ for various $\alpha_0$')
    pvst.legend(list(map(str, alpha0s)), loc=4, title=r'$\alpha_0$')
    pvst.set_ylim(0.9,1)
    pvsa = f.add_subplot(2,2,2)
    pvsa.plot(alpha0s, steps_purity[:,:], '.-')
    pvsa.set_xlabel(r'$\alpha_0$')
    pvsa.set_title(r'Purity-vs-$\alpha_0$ for various $T$')
    pvsa.legend(list(map(str, Ts[:10])), loc=4, title='$T$')
    pvsa.set_ylim(0.9,1)
    pvsboth = f.add_subplot(2,2,3, projection='3d')
    am,tm = np.meshgrid(alpha0s,Ts)
    pvsboth.plot_wireframe(am,tm,steps_purity.T)
    pvsboth.set_ylim(20,200)
    pvsboth.set_xlabel(r'$\alpha_0$')
    pvsboth.set_ylabel(r'$T$')
    pvsboth.set_title(r'Purity vs $\alpha_0$ and $T$')
    pvsboth.set_zlim(0.9,1)
    phasevsboth = f.add_subplot(2,2,4, projection='3d')
    phasevsboth.plot_wireframe(am,tm,np.angle(steps_exchange).T)
    phasevsboth.set_ylim(20,200)
    phasevsboth.set_xlabel(r'$\alpha_0$')
    phasevsboth.set_ylabel(r'$T$')
    phasevsboth.set_title(r'$arg(\langle -\alpha_0 | \rho_{final} | \alpha_0 \rangle)$ vs $\alpha_0$ and $T$')
    f.suptitle(title)
    return f
f = plot_stuff(results_two_circles_both_move_standard, 'Two Blobs: Both Move on Opposite Circles')


# In[56]:


f = plot_stuff(results_two_circles_one_moves_standard, 'Two Blobs: Only One Moves')


# In[57]:


f = plot_stuff(results_two_circles_one_moves_sqrt2_standard, 'Two Blobs: Only One Moves, but on Bigger Circle')


# In[28]:


steps_purity_both = np.array([(_[0]**2).tr() if _[1]==15000 else 0 for _ in results_two_circles_both_move_standard])
mask_both = np.array([_[1] for _ in results_two_circles_both_move_standard])
steps_purity_both.shape = len(alpha0s), len(Ts)
mask_both.shape = len(alpha0s), len(Ts)
steps_purity_both[mask_both!=15000] = np.nan
steps_purity_both = np.ma.MaskedArray(steps_purity_both, mask_both!=15000)

steps_purity_single = np.array([(_[0]**2).tr() if _[1]==15000 else 0 for _ in results_two_circles_one_moves_sqrt2_standard])
mask_single = np.array([_[1] for _ in results_two_circles_one_moves_sqrt2_standard])
steps_purity_single.shape = len(alpha0s), len(Ts)
mask_single.shape = len(alpha0s), len(Ts)
steps_purity_single[mask_single!=15000] = np.nan
steps_purity_single = np.ma.MaskedArray(steps_purity_single, mask_single!=15000)


f = plt.figure()
pvsboth = f.add_subplot(1,1,1)
am,tm = np.meshgrid(alpha0s,Ts)
m = pvsboth.contourf(am,tm,((steps_purity_both-steps_purity_single)/(steps_purity_both+steps_purity_single)).T,20)
pvsboth.set_ylim(20,200)
pvsboth.set_xlabel(r'$\alpha_0$')
pvsboth.set_ylabel(r'$T$')
pvsboth.set_title(r'$\frac{P_{both}-P_{single}}{P_{both}+P_{single}}$'+'\n', fontsize=20)
f.colorbar(m);


# In[6]:


steps_purity_both = np.array([(_[0]**2).tr() if _[1]==15000 else 0 for _ in results_two_circles_both_move_standard])
mask_both = np.array([_[1] for _ in results_two_circles_both_move_standard])
steps_purity_both.shape = len(alpha0s), len(Ts)
mask_both.shape = len(alpha0s), len(Ts)
steps_purity_both[mask_both!=15000] = np.nan
steps_purity_both = np.ma.MaskedArray(steps_purity_both, mask_both!=15000)

steps_purity_single = np.array([(_[0]**2).tr() if _[1]==15000 else 0 for _ in results_two_circles_one_moves_sqrt2_standard])
mask_single = np.array([_[1] for _ in results_two_circles_one_moves_sqrt2_standard])
steps_purity_single.shape = len(alpha0s), len(Ts)
mask_single.shape = len(alpha0s), len(Ts)
steps_purity_single[mask_single!=15000] = np.nan
steps_purity_single = np.ma.MaskedArray(steps_purity_single, mask_single!=15000)


f = plt.figure()
pvsboth = f.add_subplot(1,1,1)#,projection='3d')
am,tm = np.meshgrid(alpha0s,Ts)
m = pvsboth.contour(np.log(am),np.log(tm),np.log(1-steps_purity_both).T,20)
pvsboth.set_ylim(np.log(20),np.log(200))
pvsboth.set_xlabel(r'$log(\alpha_0)$')
pvsboth.set_ylabel(r'$log(T)$')
pvsboth.set_title(r'$log(1-P)$'+'\n', fontsize=20);
f.colorbar(m);


# In[14]:


am,tm = np.meshgrid(alpha0s,Ts[:10])
f = plt.figure()
pvsboth = f.add_subplot(1,1,1,projection='3d')
pvsboth.plot_wireframe(np.log(am),np.log(tm),np.log(1-steps_purity_both[:,:10].T))
pvsboth.set_xlabel(r'$\alpha_0$')
pvsboth.set_ylabel(r'$T$')
pvsboth.set_title(r'loglog Purity vs $\alpha_0$ and $T$');


# ### For both moving

# In[7]:


#masking does not work, so we just cut it
A = np.array(list(product(np.log(alpha0s),np.log(Ts[:10]),[1])))
B = np.log(1-steps_purity_both[:,:10]).flatten()
x = np.linalg.lstsq(A, B)[0]
#plt.plot(A.dot(x)-B)
x


# In[8]:


res = A.dot(x)
res.shape = len(alpha0s), 10
am,tm = np.meshgrid(alpha0s,Ts[:10])
f = plt.figure(figsize=(12,8))
pvsboth = f.add_subplot(1,1,1,projection='3d')
pvsboth.set_xlabel(r'$\log{(\alpha)}$', fontsize=22)
pvsboth.set_ylabel(r'$\log{(T)}$', fontsize=22)
pvsboth.set_zlabel(r'$\log{(1-P)}$', fontsize=22)
pvsboth.set_title(r'log-log $1-P$ vs $\alpha$ and $T$'+'\n'+r'fit: $1-P\propto\alpha^{%.3f}T^{%.3f}$'%(x[0],x[1]),fontsize=24)
pvsboth.plot_wireframe(np.log(am),np.log(tm),np.log(1-steps_purity_both[:,:10].T),label='input',linewidth=1)
pvsboth.plot_wireframe(np.log(am),np.log(tm),res.T,color='r',label='fit',alpha=0.4,linewidth=7)
pvsboth.legend(fontsize=22);


# In[9]:


res = A.dot(x)
res.shape = len(alpha0s), 10
am,tm = np.meshgrid(alpha0s,Ts[:10])
f = plt.figure(figsize=(12,8))
pvsboth = f.add_subplot(1,1,1,projection='3d')
pvsboth.set_xlabel(r'$\log{(\alpha)}$', fontsize=22)
pvsboth.set_ylabel(r'$\log{(T)}$', fontsize=22)
pvsboth.set_zlabel(r'$\log{(\Delta)}$', fontsize=22)
pvsboth.set_title(r'log-log $\Delta$ vs $\alpha$ and $T$'+'\n'+r'fit: $\Delta\propto\alpha^{%.3f}T^{%.3f}$'%(x[0],x[1]),fontsize=24)
pvsboth.plot_wireframe(np.log(am),np.log(tm),np.log(1-steps_purity_both[:,:10].T),label='input',linewidth=1)
pvsboth.plot_wireframe(np.log(am),np.log(tm),res.T,color='r',label='fit',alpha=0.4,linewidth=7)
pvsboth.legend(fontsize=22);


# ### For only one moving

# In[30]:


#masking does not work, so we just cut it
A = np.array(list(product(np.log(alpha0s[:14]),np.log(Ts[:10]),[1])))
B = np.log(1-steps_purity_single[:14,:10]).flatten()
x = np.linalg.lstsq(A, B)[0]
#plt.plot(A.dot(x)-B)
res = A.dot(x)
res.shape = 14, 10
am,tm = np.meshgrid(alpha0s[:14],Ts[:10])
f = plt.figure(figsize=(12,8))
pvsboth = f.add_subplot(1,1,1,projection='3d')
pvsboth.set_xlabel(r'$\log{(\alpha)}$', fontsize=22)
pvsboth.set_ylabel(r'$\log{(T)}$', fontsize=22)
pvsboth.set_zlabel(r'$\log{(1-P)}$', fontsize=22)
pvsboth.set_title(r'log-log $1-P$ vs $\alpha$ and $T$'+'\n'+r'fit: $1-P\propto\alpha^{%.3f}T^{%.3f}$'%(x[0],x[1]),fontsize=24)
pvsboth.plot_wireframe(np.log(am),np.log(tm),np.log(1-steps_purity_single[:14,:10].T),label='input',linewidth=1)
pvsboth.plot_wireframe(np.log(am),np.log(tm),res.T,color='r',label='fit',alpha=0.4,linewidth=7)
pvsboth.legend(fontsize=22);


# In[31]:


f = plt.figure(figsize=(12,8))
pvsboth = f.add_subplot(1,1,1,projection='3d')
pvsboth.set_xlabel(r'$\log{(\alpha)}$', fontsize=22)
pvsboth.set_ylabel(r'$\log{(T)}$', fontsize=22)
pvsboth.set_zlabel(r'$\log{(\Delta)}$', fontsize=22)
pvsboth.set_title(r'log-log $\Delta$ vs $\alpha$ and $T$'+'\n'+r'fit: $\Delta\propto\alpha^{%.3f}T^{%.3f}$'%(x[0],x[1]),fontsize=24)
pvsboth.plot_wireframe(np.log(am),np.log(tm),np.log(1-steps_purity_single[:14,:10].T),label='input',linewidth=1)
pvsboth.plot_wireframe(np.log(am),np.log(tm),res.T,color='r',label='fit',alpha=0.4,linewidth=7)
pvsboth.legend(fontsize=22);


# # Some more tests on smaller Hilbert space

# In[ ]:


get_ipython().run_cell_magic('px', '--local', 'def two_circles_both_move_radius_one_half_small(alpha0=2, T=200, samples_factor=1):\n    assert np.abs(alpha0) < 2.1, \'alpha > 3 at any point in time goes out of the cutoff for the hilbert space\'\n    N = 40\n    radius = 0.5\n    samples = int(20000*samples_factor)\n\n    unit_interval = np.linspace(0,1,samples)\n    centered_interval = np.linspace(-1,1,samples)\n    one = np.ones(samples)\n    zero = np.zeros(samples)\n    gaussian = np.exp(-centered_interval**2)\n\n    a_op = qutip.destroy(N)\n    adag_op = qutip.create(N)\n    id_op = qutip.identity(N)\n    zero_op = qutip.zero_oper(N)\n    num_op = qutip.num(N)\n\n    beta0 = -alpha0\n    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()\n    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)\n    betas  = beta0  - radius*(1-np.cos(-unit_interval*2*np.pi)) + radius*1j*np.sin(-unit_interval*2*np.pi)\n\n    c_ops = expand_lindblad_operator_to_qutip_std_form(\n                a_op**2,\n                [-a_op, alphas+betas],\n                [id_op, alphas*betas]\n            )\n    \n    \n    import signal\n\n    def signal_handler(signum, frame):\n        raise Exception("Timed out!")\n\n    signal.signal(signal.SIGALRM, signal_handler)\n    signal.alarm(3*60*60)\n    try:\n        ret = qutip.mesolve(zero_op,init_state,unit_interval*T,\n                            c_ops=c_ops,\n                            e_ops=[],\n                            progress_bar=qutip.ui.progressbar.TextProgressBar())\n        return ret.states[-1], len(ret.states)\n    except Exception:\n        signal.alarm(0)\n        return init_state, 1\n    \ndef two_circles_one_moves_radius_one_half_sqrt2_small(alpha0=2, T=200, samples_factor=1):\n    assert np.abs(alpha0) < 2.1, \'alpha > 3 at any point in time goes out of the cutoff for the hilbert space\'\n    N = 40\n    radius = 0.5*(2**0.5)\n    samples = int(20000*samples_factor)\n\n    unit_interval = np.linspace(0,1,samples)\n    centered_interval = np.linspace(-1,1,samples)\n    one = np.ones(samples)\n    zero = np.zeros(samples)\n    gaussian = np.exp(-centered_interval**2)\n\n    a_op = qutip.destroy(N)\n    adag_op = qutip.create(N)\n    id_op = qutip.identity(N)\n    zero_op = qutip.zero_oper(N)\n    num_op = qutip.num(N)\n\n    beta0 = -alpha0\n    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()\n    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)\n    betas  = beta0 * np.ones_like(alphas)\n\n    c_ops = expand_lindblad_operator_to_qutip_std_form(\n                a_op**2,\n                [-a_op, alphas+betas],\n                [id_op, alphas*betas]\n            )\n    \n    import signal\n\n    def signal_handler(signum, frame):\n        raise Exception("Timed out!")\n\n    signal.signal(signal.SIGALRM, signal_handler)\n    signal.alarm(3*60*60)\n    try:\n        ret = qutip.mesolve(zero_op,init_state,unit_interval*T,\n                            c_ops=c_ops,\n                            e_ops=[],\n                            progress_bar=qutip.ui.progressbar.TextProgressBar())\n        return ret.states[-1], len(ret.states)\n    except Exception:\n        signal.alarm(0)\n        return init_state, 1\n')


# In[10]:


alpha0s = [2]
Ts = np.concatenate([np.linspace(20,500,num=25)[:-1],np.linspace(500,2000,num=24)])
args = list(product(alpha0s, Ts))


# In[10]:


results_two_circles_both_move_small = lview.map(lambda a: two_circles_both_move_radius_one_half_small(*a), args)
results_two_circles_one_moves_sqrt2_small = lview.map(lambda a: two_circles_one_moves_radius_one_half_sqrt2_small(*a), args)


# In[11]:


results_two_circles_both_move_small.wait_interactive()


# In[12]:


results_two_circles_one_moves_sqrt2_small.wait_interactive()


# In[13]:


import pickle
with open("results_two_circles_both_move_small", "wb") as f:
    pickle.dump(results_two_circles_both_move_small.result, f)
with open("results_two_circles_one_moves_sqrt2_small", "wb") as f:
    pickle.dump(results_two_circles_one_moves_sqrt2_small.result, f)


# In[7]:


import pickle
with open("results_two_circles_both_move_small", "rb") as f:
    results_two_circles_both_move_small = pickle.load(f)
with open("results_two_circles_one_moves_sqrt2_small", "rb") as f:
    results_two_circles_one_moves_sqrt2_small = pickle.load(f)


# In[8]:


steps_purity_both = np.array([(_[0]**2).tr() for _ in results_two_circles_both_move_small])
steps_exchange_both = np.array([(qutip.coherent(40,-2).dag()*_[0]*qutip.coherent(40,2)).tr()
                                for _ in results_two_circles_both_move_small])
steps_purity_single = np.array([(_[0]**2).tr() for _ in results_two_circles_one_moves_sqrt2_small])
steps_exchange_single = np.array([(qutip.coherent(40,-2).dag()*_[0]*qutip.coherent(40,2)).tr()
                                for _ in results_two_circles_one_moves_sqrt2_small])


# In[22]:


plt.plot(Ts, np.abs(np.angle(steps_exchange_both)), 'o-', label="both move")
plt.plot(Ts, np.abs(np.angle(steps_exchange_single)), 'v-', label="only one moves")
plt.legend(loc=4)
plt.title(r'$arg(\langle -\alpha_0 | \rho_{final} | \alpha_0 \rangle)$ vs $T$ for $\alpha_0=2$');


# In[11]:


plt.plot(Ts, steps_purity_both, 'o-', label="both move")
plt.plot(Ts, steps_purity_single, 'v-', label="only one moves")
plt.legend(loc=4)
plt.title(r'Purity');


# In[ ]: