%pylab inline

from qutip.ipynbtools import version_table
version_table()

import cloud
def add(x, y):
    return x + y

jid = cloud.call(add, 1, 2)

cloud.result(jid)


import numpy as np
jid = cloud.call(add, np.arange(9.0), np.arange(9.0))

cloud.result(jid)

jid = cloud.call(version_table, _env='/mbaden/precise_with_qutip')
cloud.result(jid)

from qutip import *
import time

wc = 1.0 * 2 * pi # cavity frequency
wa = 1.0 * 2 * pi # atom frequency
N = 25 # number of cavity fock states
kappa = 0.05 # cavity decay rate
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
}

# Create a list of 400 coupling strengths If you have a slow computer
# you want to decrease that number to ~50 first.
g_vec = linspace(0, 2.5, 400) * 2 * pi

def compute_task(g_vec, args):
    H0, H1, c_ops = args['H0'], args['H1'], args['c_ops']
    n_vec = zeros(g_vec.shape)
    for k, g in enumerate(g_vec):
	H = H0 + g * H1
    
	rho_ss = steadystate(H, c_ops)
	n_vec[k] = expect(nc, rho_ss)
    return n_vec

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')

t0 = time.time()
n_vec = compute_task(g_vec, args)
t1 = time.time()

print "elapsed =", (t1-t0)

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

t0 = time.time()
no_chunks = 10
g_chunks = array_split(g_vec, no_chunks)
single_variable_task = lambda g: compute_task(g, args=args)
jids = cloud.map(single_variable_task, g_chunks,
		 _env='/mbaden/precise_with_qutip',
		 _type='c2')
n_chunks = cloud.result(jids)
n_vec = concatenate(n_chunks)
t1 = time.time()

print "elapsed =", (t1-t0)


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