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

# ## QobjEvo usage example
# Features of QobjEvo for the users.  
# Made by Eric Giguere

# In[1]:


from qutip import *
import time
import numpy as np


# ## Definition of time-dependant Qobj
# 
# QobjEvo are definied from list of Qobj:  
# [Qobj0, [Qobj1, coeff1], [Qobj2, coeff2]]  
# coeff can be one of: 
# - function
# - string
# - np.array

# In[2]:


# Definition of base Qobj and 
N = 4

def sin_w(t, args):
    return np.cos(args["w"]*t)

def cos_w(t, args):
    return np.cos(args["w"]*t)

tlist = np.linspace(0,10,10000)
tlistlog = np.logspace(-3,1,10000)


# In[3]:


# constant QobjEvo
cte_QobjEvo = QobjEvo(destroy(N))
cte_QobjEvo(1)


# In[4]:


# QobjEvo with function based coeff
func_QobjEvo = QobjEvo([destroy(N),[qeye(N),cos_w]],args={"w":2})
func_QobjEvo(1)


# In[5]:


# QobjEvo with sting based coeff
str_QobjEvo = QobjEvo([destroy(N),[qeye(N),"cos(w*t)"]],args={"w":2})
str_QobjEvo(1)


# In[6]:


# QobjEvo with array based coeff
array_QobjEvo = QobjEvo([destroy(N),[qeye(N),np.cos(2*tlist)]],tlist=tlist)
array_QobjEvo(1)


# In[7]:


# QobjEvo with array based coeff, log timescale
Log_array_QobjEvo = QobjEvo([destroy(N),[qeye(N),np.cos(2*tlistlog)]],tlist=tlistlog)
Log_array_QobjEvo(1)


# In[8]:


# Reference
destroy(N) + qeye(N) * np.cos(2)


# ## Mathematic
# - addition (QobjEvo, Qobj)
# - substraction (QobjEvo, Qobj)
# - product (QobjEvo, Qobj, scalar)
# - division (scalar)
# 
# The examples are done with function type coefficients only, but work for any type of coefficient.  
# Mixing coefficients type is possible, however this support would be removed if QobjEvo * QobjEvo is to be implemented.

# In[9]:


# Build objects
o1 = QobjEvo([qeye(N),[destroy(N),sin_w]],args={"w":2})
o2 = QobjEvo([qeye(N),[create(N),cos_w]],args={"w":2})
t = np.random.random()*10


# In[10]:


# addition and subtraction 
o3 = o1 + o2
print(o3(t) == o1(t) + o2(t))
o3 = o1 - o2
print(o3(t) == o1(t) - o2(t))
o3 = o1 + destroy(N)
print(o3(t) == o1(t) + destroy(N))
o3 = o1 - destroy(N)
print(o3(t) == o1(t) - destroy(N))


# In[11]:


# product
oc = QobjEvo([qeye(N)])
o3 = o1 * destroy(N)
print(o3(t) == o1(t) * destroy(N))
o3 = o1 * (0.5+0.5j)
print(o3(t) == o1(t) * (0.5+0.5j))
o3 = o1 / (0.5+0.5j)
print(o3(t) == o1(t) / (0.5+0.5j))
o3 = o1 * oc
print(o3(t) == o1(t) * oc(t))
o3 = oc * o1
print(o3(t) == oc(t) * o1(t))
o3 = o1 * o2
print(o3(t) == o1(t) * o2(t))


# In[12]:


o1 = QobjEvo([qeye(N),[destroy(N),sin_w]],args={"w":2})
o2 = QobjEvo([qeye(N),[create(N),cos_w]],args={"w":2})
o1 += o2
print(o1(t) == (qeye(N)*2 + destroy(N)*sin_w(t,args={"w":2}) + create(N)*cos_w(t,args={"w":2})))


# In[13]:


o1 = QobjEvo([qeye(N),[destroy(N),sin_w]],args={"w":2})
o2 = QobjEvo([qeye(N),[create(N),cos_w]],args={"w":2})
o1 -= o2
print(o1(t) == (destroy(N)*sin_w(t,args={"w":2}) - create(N)*cos_w(t,args={"w":2})))

o1 = QobjEvo([qeye(N),[destroy(N),sin_w]],args={"w":2})
o2 = QobjEvo([qeye(N),[create(N),cos_w]],args={"w":2})
o1 += -o2
print(o1(t) == (destroy(N)*sin_w(t,args={"w":2}) - create(N)*cos_w(t,args={"w":2})))


# In[14]:


o1 = QobjEvo([qeye(N),[destroy(N),sin_w]],args={"w":2})
o1 *= destroy(N)
print(o1(t) == (destroy(N) + destroy(N)*destroy(N)*sin_w(t,args={"w":2})))


# ## Unitary operations: 
# - conj
# - dag 
# - trans 
# - _cdc: QobjEvo.dag * QobjEvo

# In[15]:


o_real = QobjEvo([qeye(N),[destroy(N), sin_w]], args={"w":2})
o_cplx = QobjEvo([qeye(N),[create(N), cos_w]], args={"w":-1j})

print(o_real(t).trans() == o_real.trans()(t))
print(o_real(t).conj() == o_real.conj()(t))
print(o_real(t).dag() == o_real.dag()(t))

print(o_cplx(t).trans() == o_cplx.trans()(t))
print(o_cplx(t).conj() == o_cplx.conj()(t))
print(o_cplx(t).dag() == o_cplx.dag()(t))


# In[16]:


# the operator norm correspond to c.dag * c.
td_cplx_f0 = qobjevo.QobjEvo([qeye(N)])
td_cplx_f1 = qobjevo.QobjEvo([qeye(N),[destroy(N)*create(N),sin_w]], args={'w':2.+0.001j})
td_cplx_f2 = qobjevo.QobjEvo([qeye(N),[destroy(N),cos_w]], args={'w':2.+0.001j})
td_cplx_f3 = qobjevo.QobjEvo([qeye(N),[create(N),1j*np.sin(tlist)]], tlist=tlist)

print(td_cplx_f0(t).dag()*td_cplx_f0(t) == td_cplx_f0._cdc()(t))
print(td_cplx_f1(t).dag()*td_cplx_f1(t) == td_cplx_f1._cdc()(t))
print(td_cplx_f2(t).dag()*td_cplx_f2(t) == td_cplx_f2._cdc()(t))
print(td_cplx_f3(t).dag()*td_cplx_f3(t) == td_cplx_f3._cdc()(t))


# ## Liouvillian and lindblad dissipator, to use in solver
# 
# Functions in qutip.superoperator can be used for QobjEvo.

# In[17]:


td_L = liouvillian(H=func_QobjEvo)
L = liouvillian(H=func_QobjEvo(t))
td_L(t) == L


# In[18]:


td_cplx_f0 = qobjevo.QobjEvo([qeye(N)])
td_cplx_f1 = qobjevo.QobjEvo([[destroy(N)*create(N),sin_w]], args={'w':2.})

td_L = liouvillian(H=func_QobjEvo,c_ops=[td_cplx_f0,td_cplx_f1])
L = liouvillian(H=func_QobjEvo(t),c_ops=[td_cplx_f0(t),td_cplx_f1(t)])
print(td_L(t) == L)


# In[19]:


td_P = spre(td_cplx_f1)
P = spre(td_cplx_f1(t))
print(td_P(t) == P)


# ## Getting the list back for the object

# In[20]:


print(td_L.to_list())


# ## Arguments modification
# 
# To change the args: qobjevo.arguments(new_args)
# 
# Call with other arguments without changing them: qobjevo.with_args(t, new_args)
# 

# In[21]:


def Args(t, args):
    return args['w']
td_args = qobjevo.QobjEvo([qeye(N), Args],args={'w':1.}) 
print(td_args(t) == qeye(N))
td_args.arguments({'w':2.})
print(td_args(t) == qeye(N)*2)
print(td_args(t,args={'w':3.}) == qeye(N)*3)


# When summing QobjEvo that have an arguments in common, only one is kept.

# In[22]:


td_args_1 = qobjevo.QobjEvo([qeye(N), [destroy(N), Args]],args={'w':1.})
td_args_2 = qobjevo.QobjEvo([qeye(N), [destroy(N), Args]],args={'w':2.})
td_str_sum = td_args_1 + td_args_2

# Only one value for args is kept
print(td_str_sum(t) == td_args_1(t) + td_args_2(t)) 
print(td_str_sum(t) == 2*td_args_2(t)) 

# Updating args affect all part
td_str_sum.arguments({'w':1.})
print(td_str_sum(t) == 2*td_args_1(t))


# Argument with different names are fine.

# In[23]:


def Args2(t, args):
    return args['x']
td_args_1 = qobjevo.QobjEvo([qeye(N), [destroy(N), cos_w]],args={'w':1.})
td_args_2 = qobjevo.QobjEvo([qeye(N), [destroy(N), Args2]],args={'x':2.})
td_str_sum = td_args_1 + td_args_2

# Only one value for args is kept
print(td_str_sum(t) == td_args_1(t) + td_args_2(t))


# ## Other

# In[24]:


# Obtain the sparce matrix at a time t instead of a Qobj
str_QobjEvo(1, data=True)


# In[25]:


# Test is the QobjEvo does depend on time
print(cte_QobjEvo.const)
print(str_QobjEvo.const)


# In[26]:


# Obtain the size, shape, oper flag etc:
# The QobjEvo.cte always exist and contain the constant part of the QobjEvo
# It can be used to get the shape, etc. since the QobjEvo do not directly have them.
td_cplx_f1 = qobjevo.QobjEvo([[destroy(N)*create(N),sin_w]], args={'w':2.})
print(td_cplx_f1.cte.dims)
print(td_cplx_f1.cte.shape)
print(td_cplx_f1.cte.isoper)
print(td_cplx_f1.cte)


# In[27]:


# Creating a copy
str_QobjEvo_2 = str_QobjEvo.copy()
str_QobjEvo_2 += 1
str_QobjEvo_2(1) -  str_QobjEvo(1)


# In[28]:


about()


# In[ ]: