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

# # Speed test of stochastic solvers 
# ## Based on development-smesolve-milstein notebook
# 
# Denis V. Vasilyev
# 
# 30 september 2013
# 
# 
# Modified by Eric Giguere, March 2018  
# Include new solvers from the pull request #815 for qutip

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')


# In[2]:


from qutip import *
from numpy import log2, cos, sin
from scipy.integrate import odeint
from qutip.cy.spmatfuncs import cy_expect_rho_vec, spmv


# # Single stochastic operator

# In[3]:


th = 0.1 # Interaction parameter
alpha = cos(th)
beta = sin(th)
gamma = 1

# Exact steady state solution for Vc
Vc = (alpha*beta - gamma + sqrt((gamma-alpha*beta)**2 + 4*gamma*alpha**2))/(4*alpha**2)

#********* Model ************
NN = 200
tlist = linspace(0,50,NN)
Nsub = 200
N = 20
Id = qeye(N)
a = destroy(N)
s = 0.5*((alpha+beta)*a + (alpha-beta)*a.dag())
x = (a + a.dag())/sqrt(2)
H = Id
c_op = [sqrt(gamma)*a]
sc_op = [s]
e_op = [x, x*x]
rho0 = fock_dm(N,0) #initial vacuum state


# In[4]:


# Solution of the differential equation for the variance Vc
y0 = 0.5
def func(y, t):
    return -(gamma - alpha*beta)*y - 2*alpha*alpha*y*y + 0.5*gamma
y = odeint(func, y0, tlist)


# In[5]:


# list solver
help(stochastic_solvers)


# In[6]:


# Euler-Maruyama
sol_euler = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='euler-maruyama',
                   options=Odeoptions(store_states=True, average_states=False))


# In[7]:


# Milstein
sol_mil = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='milstein',
                   options=Odeoptions(store_states=True, average_states=False))


# In[8]:


# predictor-corrector
sol_pc_euler = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='pc-euler',
                   options=Odeoptions(store_states=True, average_states=False))


# In[9]:


# predictor-corrector-2
sol_pc_euler2 = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='pred-corr-2',
                   options=Odeoptions(store_states=True, average_states=False))


# In[10]:


# Default: taylor1.5
sol_taylor15 = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='taylor1.5',
                   options=Odeoptions(store_states=True, average_states=False))


# In[11]:


# Taylor 2.0
sol_taylor20 = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='taylor2.0',
                   options=Odeoptions(store_states=True, average_states=False))


# ## Variance $V_{\mathrm{c}}$ as a function of time

# In[12]:


plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2, label='Euler')
plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2, label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2, label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2, label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2, label='taylor1.5')
plot(tlist,sol_taylor20.expect[1] - abs(sol_taylor20.expect[0])**2, label='taylor2.0')
plot(tlist,Vc*ones(NN),"k:", label='exact steady state solution')
plot(tlist,y,"k", label="exact solution")
legend();
show()

#plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2, label='Euler')
plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2, label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2, label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2, label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2, label='taylor1.5')
plot(tlist,sol_taylor20.expect[1] - abs(sol_taylor20.expect[0])**2, label='taylor2.0')
plot(tlist,Vc*ones(NN),"k:", label='exact steady state solution')
plot(tlist,y,"k", label="exact solution")
legend();
xlim([0,25])
ylim([0.3238,0.325])
show()

#plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2, label='Euler')
#plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2, label='Milstein')
#plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2, label='pc-euler')
#plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2, label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2, label='taylor1.5')
plot(tlist,sol_taylor20.expect[1] - abs(sol_taylor20.expect[0])**2, label='taylor2.0')
plot(tlist,Vc*ones(NN),"k:", label='exact steady state solution')
plot(tlist,y,"k", label="exact solution")
legend();
xlim([0,25])
ylim([0.3241,0.32418])
show()


# # Two stochastic operators

# In[13]:


th = 0.1
alpha = cos(th)
beta = sin(th)
gamma = 1
eps = 0.3

VcEps = ((1-2*eps)*alpha*beta - gamma + \
         sqrt((gamma-alpha*beta)**2 + 4*gamma*alpha*((1-eps)*alpha + eps*beta)))/(4*(1-eps)*alpha**2)
UcEps = (-(1-2*eps)*alpha*beta - gamma + \
         sqrt((gamma-alpha*beta)**2 + 4*eps*beta*gamma*(beta-alpha)))/(4*eps*beta**2)

NN = 200
tlist = linspace(0,3,NN)
Nsub = 200
N = 20
Id = qeye(N)
a = destroy(N)
s = 0.5*((alpha+beta)*a + (alpha-beta)*a.dag())
x = (a + a.dag())/sqrt(2)
H = Id
c_op = [sqrt(gamma)*a]
sc_op = [sqrt(1-eps)*s, sqrt(eps)*1j*s]
e_op = [x, x*x]
rho0 = fock_dm(N,0)
y0 = 0.5


# In[14]:


def func(y, t):
    return -(gamma - (1-2*eps)*alpha*beta)*y - 2*(1-eps)*alpha*alpha*y*y + 0.5*(gamma + eps*beta*beta)
y = odeint(func, y0, tlist)

def funcZ(z, t):
    return -(gamma + (1-2*eps)*alpha*beta)*z - 2*eps*beta*beta*z*z + 0.5*(gamma + (1-eps)*alpha*alpha)
z = odeint(funcZ, y0, tlist)


# In[15]:


# Euler-Maruyama
sol_euler = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='euler-maruyama',
                   options=Odeoptions(store_states=True, average_states=False))


# In[16]:


# Milstein
sol_mil = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='milstein',
                   options=Odeoptions(store_states=True, average_states=False))


# In[17]:


# predictor-corrector
sol_pc_euler = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='pc-euler',
                   options=Odeoptions(store_states=True, average_states=False))


# In[18]:


# predictor-corrector-2
sol_pc_euler2 = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='pred-corr-2',
                   options=Odeoptions(store_states=True, average_states=False))


# In[19]:


# Default: taylor1.5
sol_taylor15 = smesolve(H, rho0, tlist, c_op, sc_op, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='taylor1.5',
                   options=Odeoptions(store_states=True, average_states=False))


# ## Variance $V_{\mathrm{c}}$ as a function of time

# In[20]:


plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2, label='Euler')
plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2, label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2, label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2, label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2, label='taylor1.5')
plot(tlist,Vc*ones(NN), label='exact steady state solution')
plot(tlist,y, label="exact solution")
legend();
show()

plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2, label='Euler')
plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2, label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2, label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2, label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2, label='taylor1.5')
plot(tlist,Vc*ones(NN), label='exact steady state solution')
plot(tlist,y, label="exact solution")
legend();
xlim([1.5,3.0])
ylim([0.347,0.356])


# In[21]:


plot(tlist,sol_euler.expect[1] - abs(sol_euler.expect[0])**2-y.transpose()[0], label='Euler')
plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2-y.transpose()[0], label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2-y.transpose()[0], label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2-y.transpose()[0], label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2-y.transpose()[0], label='taylor1.5')
legend()
show()

plot(tlist,sol_mil.expect[1] - abs(sol_mil.expect[0])**2-y.transpose()[0], label='Milstein')
plot(tlist,sol_pc_euler.expect[1] - abs(sol_pc_euler.expect[0])**2-y.transpose()[0], label='pc-euler')
plot(tlist,sol_pc_euler2.expect[1] - abs(sol_pc_euler2.expect[0])**2-y.transpose()[0], label='pc-euler-2')
plot(tlist,sol_taylor15.expect[1] - abs(sol_taylor15.expect[0])**2-y.transpose()[0], label='taylor1.5')
legend()
show()


# #Multiple stochastic operators

# In[22]:


sc_ops_multiple = [sc_op[0]*0.5**0.5, sc_op[0]*0.5**0.5]
#Euler-Maruyama
sol_eul = smesolve(H, rho0, tlist, c_op, sc_ops_multiple, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='euler',
                   options=Odeoptions(store_states=True, average_states=False))


# In[23]:


#Milstein
sol_milstein = smesolve(H, rho0, tlist, c_op, sc_ops_multiple, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='milstein',
                   options=Odeoptions(store_states=True, average_states=False))


# In[24]:


#taylor1.5
sol_taylor = smesolve(H, rho0, tlist, c_op, sc_ops_multiple, e_op,
                   nsubsteps=Nsub, method='homodyne', solver='taylor1.5',
                   options=Odeoptions(store_states=True, average_states=False))


# In[25]:


plot(tlist,sol_eul.expect[1] - abs(sol_eul.expect[0])**2, label='Euler')
plot(tlist,sol_milstein.expect[1] - abs(sol_milstein.expect[0])**2, label='Milstein')
plot(tlist,sol_taylor.expect[1] - abs(sol_taylor.expect[0])**2, label='taylor1.5')
plot(tlist,Vc*ones(NN), label='exact steady state solution')
plot(tlist,y, label="exact solution")
legend()
show()


# ### Software versions

# In[26]:


from qutip.ipynbtools import version_table

version_table()


# In[ ]: