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

# # Steadystate of the Bloch-Redfield Master Equation

# In[1]:


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


# In[2]:


import numpy as np
from qutip import *
from IPython.display import display


# ## Compare steadystate of brmesolve and mesolve

# In[15]:


def solve(H, psi0, c_ops, a_ops, e_ops):
    
    result_me = mesolve(H, psi0, times, c_ops, e_ops)
    result_brme = brmesolve(H, psi0, times, a_ops, e_ops)
    
    fig, ax = plot_expectation_values([result_me, result_brme])
    display(fig)
    plt.close(fig)
    
    R, ekets = bloch_redfield_tensor(H, a_ops)
    print("="* 20 + " Bloch-Redfield tensor: ")
    display(R)
    
    
    L = liouvillian(H, c_ops)
    print("="* 20 + " Lindblad liouvilllian: ")
    display(L)
    
    print("="* 20 + " Bloch-Redfield steadystate dm")
    R_rhoss_eb = steadystate(R)
    R_rhoss = R_rhoss_eb.transform(ekets, True)
    display(R_rhoss)
    
    print("="* 20 + " Lindblad steadystate dm")
    L_rhoss = steadystate(L)
    display(L_rhoss)

    print("="* 20 + " Steadystate expectation values")
    print("R_ob: ", [expect(e, R_rhoss) for e in e_ops])
    print("R_eb: ", [expect(e.transform(ekets), R_rhoss_eb) for e in e_ops])
    print("L   : ", [expect(e, L_rhoss) for e in e_ops])

    print("="* 20 + " Dynamics final states")

    print("R: ", [e[-1].real for e in result_brme.expect])
    print("L: ", [e[-1] for e in result_me.expect])


# ## Two-level system

# In[16]:


delta = 0.0 * 2 * np.pi
epsilon = 0.5 * 2 * np.pi
gamma = 0.25
times = np.linspace(0, 50, 100)


# In[17]:


H = delta/2 * sigmay() + epsilon/2 * sigmaz()
psi0 = (2 * basis(2, 0) + basis(2, 1)).unit()
c_ops = [np.sqrt(gamma) * sigmam()]
a_ops = [[sigmax(),lambda w : gamma * (w >= 0)]]
e_ops = [sigmax(), sigmay(), sigmaz()]


# In[18]:


solve(H, psi0, c_ops, a_ops, e_ops)


# ## Harmonic oscillator

# In[25]:


N = 10

w0 = 1.0 * 2 * np.pi
g = 0.05 * w0
kappa = 0.15

times = np.linspace(0, 50, 1000)


# In[26]:


a = destroy(N)
H = w0 * a.dag() * a + g * (a + a.dag())
psi0 = ket2dm((basis(N, 4) + basis(N, 2) + basis(N,0)).unit())
a_ops = [[a + a.dag(),lambda w : kappa * (w >= 0)]]
e_ops = [a.dag() * a, a + a.dag()]


# ### Zero temperature

# In[27]:


c_ops = [np.sqrt(kappa) * a]


# In[28]:


solve(H, psi0, c_ops, a_ops, e_ops)


# ### Finite temperature

# In[29]:


n_th = 1.5
c_ops = [np.sqrt(kappa * (n_th + 1)) * a, np.sqrt(kappa * n_th) * a.dag()]
w_th = w0/np.log(1 + 1/n_th)

def S_w_func(w):
    if w >= 0:
        return (n_th + 1) * kappa
    else:
        return (n_th + 1) * kappa * np.exp(w / w_th)
    
a_ops = [[a + a.dag(), S_w_func]]


# In[30]:


solve(H, psi0, c_ops, a_ops, e_ops)


# ## Atom-Cavity

# In[34]:


N = 10
a = tensor(destroy(N), identity(2))
sm = tensor(identity(N), destroy(2))
psi0 = ket2dm(tensor(basis(N, 1), basis(2, 0)))
a_ops = [[a + a.dag(),lambda w : kappa*(w > 0)]]
e_ops = [a.dag() * a, sm.dag() * sm]


# ### Weak coupling

# In[35]:


w0 = 1.0 * 2 * np.pi
g = 0.05 * 2 * np.pi
kappa = 0.05
times = np.linspace(0, 150 * 2 * np.pi / g, 1000)
c_ops = [np.sqrt(kappa) * a]
H = w0 * a.dag() * a + w0 * sm.dag() * sm + g * (a + a.dag()) * (sm + sm.dag())


# In[36]:


solve(H, psi0, c_ops, a_ops, e_ops)


# ### Strong coupling

# In[37]:


w0 = 1.0 * 2 * np.pi
g = 0.75 * 2 * np.pi
kappa = 0.05
times = np.linspace(0, 150 * 2 * np.pi / g, 1000)


# In[38]:


c_ops = [np.sqrt(kappa) * a]
H = w0 * a.dag() * a + w0 * sm.dag() * sm + g * (a + a.dag()) * (sm + sm.dag())


# In[39]:


solve(H, psi0, c_ops, a_ops, e_ops)


# ## Versions

# In[40]:


from qutip.ipynbtools import version_table

version_table()


# In[ ]: