%pylab inline

from qutip import *

import time

def solve_dynamics(H, args, psi0, tlist, T, c_ops, e_ops, fc_ops, J_cbs):
    """
    Solve the dynamics for a quantum system with a bunch of different solvers
    and return the results. The idea is that the results should be quite similar,
    even though some differences can be expected due to the different situation
    (approximations) the different solvers are applicable for. Also, some solvers
    here solve for unitary evolution and some the dissipative evolution, so
    depending on the values c_ops and fc_ops the dynamics can be very different.
    """
    
    sol_list = []
    
    #
    # Solve with unitary floquet solver by going through all steps manually
    #
    sol = Odedata()
    sol.times = tlist
    sol.solver = "Floquet WF"
    sol.expect = zeros((len(e_ops), len(tlist)))

    f_modes_0, f_energies = floquet_modes(H, T, args)
    f_coeff = floquet_state_decomposition(f_modes_0, f_energies, psi0)
    f_modes_table_t = floquet_modes_table(f_modes_0, f_energies, tlist, H, T, args)

    for n, t in enumerate(tlist):
        f_modes_t = floquet_modes_t_lookup(f_modes_table_t, t, T)
        psi_t = floquet_wavefunction(f_modes_t, f_energies, f_coeff, t)
        for e_idx, e in enumerate(e_ops):
            sol.expect[e_idx, n] = expect(e, psi_t)
    sol_list.append(sol)

    #
    # Use fsesolve to do the same as above with a single function call
    #
    sol = fsesolve(H, psi0, tlist, e_ops, T, args)
    sol_list.append(sol)
    
    #
    # Solve the floquet-markov master equation by requesting density matrices,
    # and then loop through the states and calculate the expectation values.
    # Do this with and without asking the solver to transform states back to the
    # computational basis.
    #
    sol = fmmesolve(H, psi0, tlist, fc_ops, [], J_cbs, T, args, floquet_basis=True)
    sol.solver = "Floquet ME states #1"
    sol.expect = zeros((len(e_ops), len(tlist)))
    for n, t in enumerate(sol.times):
        f_modes_t = floquet_modes_t_lookup(f_modes_table_t, t, T)
        for e_idx, e in enumerate(e_ops):
            sol.expect[e_idx, n] = expect(e, sol.states[n].transform(f_modes_t, False))
    sol_list.append(sol)

    sol = fmmesolve(H, psi0, tlist, fc_ops, [], J_cbs, T, args, floquet_basis=False)
    sol.solver = "Floquet ME states #2"
    sol.expect = zeros((len(e_ops), len(tlist)))
    for n, t in enumerate(sol.times):
        for e_idx, e in enumerate(e_ops):
            sol.expect[e_idx, n] = expect(e, sol.states[n])
    sol_list.append(sol)
    
    #
    # Solve the floquet-markov master equation by requesting expectation values
    #
    sol = fmmesolve(H, psi0, tlist, fc_ops, e_ops, J_cbs, T, args)
    sol_list.append(sol)

    # 
    # Solve with mesolve
    #
    sol = mesolve(H, psi0, tlist, c_ops, e_ops, args)
    sol_list.append(sol)

    return sol_list

def visualize_solutions(sol_list, split=False):

    if split:
        fig, axes = subplots(len(sol_list), 1, figsize=(12, 3 * len(sol_list)), squeeze=False)
    else:
        fig, axes = subplots(1, 1, figsize=(12, 6), squeeze=False)
        
    for idx, sol in enumerate(sol_list):
        if not split:
            idx = 0

        for e_idx, e in enumerate(sol.expect):
            axes[idx,0].plot(sol.times, real(e), label="%s e_%d" % (sol.solver, e_idx))

        axes[idx,0].set_xlabel('Time')
        axes[idx,0].set_ylabel('Occupation probability')
        axes[idx,0].legend(loc=0)

delta = 0.45 * 2*pi
eps0  = 1.0 * 2*pi
A     = 0.3 * 2*pi
omega = sqrt(delta**2 + eps0**2)
T      = (2*pi)/omega
tlist  = linspace(0.0, 10 * T, 501)
psi0   = rand_ket(2)

H0 = - eps0/2.0 * sigmaz() - delta/2.0 * sigmax()
H1 = A/2.0 * sigmax()
args = {'w': omega}
H = [H0, [H1, lambda t,args: sin(args['w'] * t)]]

gamma1 = 0.005
def noise_spectrum(omega):
    return 0.25 * gamma1 * omega /(2*pi)

c_ops = [sqrt(gamma1) * sigmam()]
e_ops = [num(2)]
fc_ops = [sigmax()]
J_cbs = [noise_spectrum]

t = time.time()
sol_list = solve_dynamics(H, args, psi0, tlist, T, c_ops, e_ops, fc_ops, J_cbs)
print "elapsed =", time.time() - t

visualize_solutions(sol_list)

H = [H0, [H1, "sin(w * t)"]]

gamma1 = 0.075
def noise_spectrum(omega):
    return 0.25 * gamma1 * omega /(2*pi)

J_cbs = [noise_spectrum]
c_ops = [sqrt(gamma1) * sigmam()]

t = time.time()
sol_list = solve_dynamics(H, args, psi0, tlist, T, c_ops, e_ops, fc_ops, J_cbs)
print "elapsed =", time.time() - t

visualize_solutions(sol_list)

delta = 0.0 * 2*pi
eps0 = 1.0 * 2*pi
A = 0.4 * 2*pi
w0 = 2.0 * 2 * pi
g = 0.15 * 2 * pi

omega = w0 - sqrt(delta**2 + eps0**2)
T      = (2*pi)/omega
tlist  = linspace(0.0, 10 * T, 1001)
N = 2
a  = tensor(destroy(N), qeye(2))
sx = tensor(qeye(N), sigmax())
sz = tensor(qeye(N), sigmaz())
sn = tensor(qeye(N), num(2))
sm = tensor(qeye(N), destroy(2))
n = a.dag() * a

psi0 = tensor(basis(N, 1), basis(2,0))

H0 = w0 * a.dag() * a - eps0/2.0 * sz - delta/2.0 * sx + g * (a + a.dag()) * sx
H1 = A/2.0 * sz
args = {'w': omega}
H = [H0, [H1, lambda t,args: sin(args['w'] * t)]]

gamma1 = 0.001
def noise_spectrum(omega):
    return 0.25 * gamma1 * omega /(2*pi)

c_ops = [sqrt(gamma1) * sm]
e_ops = [sn, n]
fc_ops = [sx]
J_cbs = [noise_spectrum]

t = time.time()
sol_list = solve_dynamics(H, args, psi0, tlist, T, c_ops, e_ops, fc_ops, J_cbs)
print "elapsed =", time.time() - t

visualize_solutions(sol_list)

from qutip.ipynbtools import version_table
version_table()