%pylab inline

from qutip import *

N = 15
w0 = 0.5 * 2 * pi
times = linspace(0, 15, 150)
gamma = 0.15

a = destroy(N)

H = w0 * a.dag() * a

rho0 = fock(N, 5)

c_ops = [sqrt(gamma) * a]

e_ops = [a.dag() * a, a + a.dag()]

result_ref = mesolve(H, rho0, times, c_ops, e_ops)

plot_expectation_values(result_ref);

# The argument A in the d1 and d2 callback functions is a list with the following
# precomputed superoperators, where c is the stochastic collapse operator given
# to the solver (called once for each operator, if more than one is given)
#
#     A[0] = spre(c)
#     A[1] = spost(c)
#     A[2] = spre(c.dag())
#     A[3] = spost(c.dag())
#     A[4] = spre(n)
#     A[5] = spost(n)
#     A[6] = (spre(c) * spost(c.dag())
#     A[7] = lindblad_dissipator(c)

def d1_rho_func(A, rho_vec):
    n_sum = A[4] + A[5]
    return - n_sum * rho_vec + expect_rho_vec(n_sum, rho_vec) * rho_vec

def d2_rho_func(A, rho_vec):
    e1 = expect_rho_vec(A[6], rho_vec) + 1e-16  # add a small number to avoid division by zero
    return [(A[6] * rho_vec) / e1 - rho_vec]

result = smesolve(H, rho0, times, c_ops=[], sc_ops=c_ops, e_ops=e_ops, 
                  ntraj=5, nsubsteps=100, d1=d1_rho_func, d2=d2_rho_func,
                  distribution='poisson', store_measurement=True)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.step(times, m)

result = smesolve(H, rho0, times, c_ops=[], sc_ops=c_ops, e_ops=e_ops, 
                  ntraj=5, nsubsteps=100, d1=d1_rho_func, d2=d2_rho_func,
                  distribution='poisson', store_measurement=True, noise=result.noise)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.step(times, m)

result = smesolve(H, rho0, times, [], c_ops, e_ops, ntraj=5, nsubsteps=100,
                  method='photocurrent', store_measurement=True)

plot_expectation_values([result, result_ref]);

result = smesolve(H, rho0, times, c_ops=[], sc_ops=c_ops, e_ops=e_ops, ntraj=5, nsubsteps=100,
                  method='photocurrent', store_measurement=True, noise=result.noise)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.step(times, m)

def d1_rho_func(A, rho_vec):
    return A[7] * rho_vec

def d2_rho_func(A, rho_vec):
    n_sum = A[0] + A[3]
    e1 = expect_rho_vec(n_sum, rho_vec)
    return [n_sum * rho_vec - e1 * rho_vec]

result = smesolve(H, rho0, times, [], c_ops, e_ops, 
                  ntraj=5, nsubsteps=100, d1=d1_rho_func, d2=d2_rho_func,
                  distribution='normal', store_measurement=True)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

result = smesolve(H, rho0, times, [], c_ops, e_ops, 
                  ntraj=5, nsubsteps=100, d1=d1_rho_func, d2=d2_rho_func,
                  distribution='normal', store_measurement=True, noise=result.noise)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

result = smesolve(H, rho0, times, [], c_ops, e_ops, ntraj=5, nsubsteps=100,
                  method='homodyne', store_measurement=True)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

result = smesolve(H, rho0, times, [], c_ops, e_ops, ntraj=5, nsubsteps=100,
                  method='homodyne', store_measurement=True, noise=result.noise)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

def d1_rho_func(A, rho_vec):
    return A[7] * rho_vec

def d2_rho_func(A, rho_vec):

    n_sum = A[0] + A[3]
    e1 = expect_rho_vec(n_sum, rho_vec)
    drho1 = n_sum * rho_vec - e1 * rho_vec

    n_sum = A[0] - A[3]
    e1 = expect_rho_vec(n_sum, rho_vec)
    drho2 = n_sum * rho_vec - e1 * rho_vec

    return [1.0/sqrt(2) * drho1, -1.0j/sqrt(2) * drho2]

result = smesolve(H, rho0, times, [], c_ops, e_ops, 
                  ntraj=5, nsubsteps=100, d1=d1_rho_func, d2=d2_rho_func, d2_len=2,
                  distribution='normal', store_measurement=True)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

result = smesolve(H, rho0, times, [], c_ops, e_ops, ntraj=5, nsubsteps=100,
                  method='heterodyne', store_measurement=True)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

result = smesolve(H, rho0, times, [], c_ops, e_ops, ntraj=5, nsubsteps=100,
                  method='heterodyne', store_measurement=True, noise=result.noise)

plot_expectation_values([result, result_ref]);

fig, ax = subplots()

for m in result.measurement:
    ax.plot(times, m)
    
ax.plot(times, array(result.measurement).mean(axis=0), 'k', lw=5);

from qutip.ipynbtools import version_table

version_table()