%pylab inline

from qutip import *
import qutip.ipynbtools

import time

qutip.ipynbtools.version_table()

delay_times = numpy.random.rand(10)

delay_times

def task(delay):
    import os, time

    t0 = time.time()
    pid = os.getpid()
    time.sleep(delay)
    t1 = time.time()
    
    return (t1-t0)  # return the actual delay

result = map(task, delay_times)

result

result = parfor(task, delay_times)

result

result = qutip.ipynbtools.parfor(task, delay_times, show_scheduling=True, show_progressbar=True)

result

def visualize_results(g_vec, n_vec):
    
    fig, ax = subplots()
    ax.plot(g_vec, n_vec, lw=2)
    ax.set_xlabel('Coupling strength (g)')
    ax.set_ylabel('Photon Number')
    ax.set_title('# of photons in the steady state')

def compute_task(g, args):
    wc, wa, N, kappa, n_th = args['wc'], args['wa'], args['N'], args['kappa'], args['n_th']
    
    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    nc = a.dag() * a
    na = sm.dag() * sm

    c_ops = [sqrt(kappa * (1 + n_th)) * a, sqrt(kappa * n_th) * a.dag()]

    H0 = wc * nc + wa * na
    H1 = (a.dag() + a) * (sm + sm.dag())
    H = H0 + g * H1
    
    rho_ss = steadystate(H, c_ops)
    return expect(nc, rho_ss)

# problem parameters
args = {'wc': 1.0 * 2 * pi,  # cavity frequency
        'wa': 1.0 * 2 * pi,  # atom frequency
        'N':  25,            # number of cavity fock states
        'kappa': 0.05,
        'n_th':  0.5,
        }

g_vec = linspace(0, 2.5, 50) * 2 * pi

# serial calculation
t0 = time.time()
n_vec = array([compute_task(g, args) for g in g_vec])
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)

# parallel calculation using qutip parfor
t0 = time.time()
n_vec = parfor(compute_task, g_vec, args=args)
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)

# parallel calculation using qutip IPython.parallel parfor
t0 = time.time()
n_vec = qutip.ipynbtools.parfor(compute_task, g_vec, args=args, show_scheduling=True, show_progressbar=True)
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)

def compute_task(g, args):
    H0, H1, c_ops, nc = args['H0'], args['H1'], args['c_ops'], args['nc']
    H = H0 + g * H1
    rho_ss = steadystate(H, c_ops)
    return expect(nc, rho_ss)

# problem parameters
wc = 1.0 * 2 * pi
wa = 1.0 * 2 * pi
N = 75
kappa = 0.05
n_th = 0.5

a = tensor(destroy(N), qeye(2))
sm = tensor(qeye(N), destroy(2))
nc = a.dag() * a
na = sm.dag() * sm

c_ops = [sqrt(kappa * (1 + n_th)) * a, sqrt(kappa * n_th) * a.dag()]

H0 = wc * nc + wa * na
H1 = (a.dag() + a) * (sm + sm.dag())

args = {'H0': H0,
        'H1': H1,
        'c_ops': c_ops,
        'nc': nc
        }

g_vec = linspace(0, 2.5, 50) * 2 * pi

# serial calculation
t0 = time.time()
n_vec = array([compute_task(g, args) for g in g_vec])
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)

# parallel calculation using qutip parfor
t0 = time.time()
n_vec = parfor(compute_task, g_vec, args=args)
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)

# parallel calculation
t0 = time.time()
n_vec = qutip.ipynbtools.parfor(compute_task, g_vec, args=args, show_scheduling=True, show_progressbar=True)
t1 = time.time()

print "elapsed =", (t1-t0)

visualize_results(g_vec / (2 * pi), n_vec)