Development notebook: Tests for QuTiP's stochastic Schrödinger equation solver

Copyright (C) 2011 and later, Paul D. Nation & Robert J. Johansson

In this notebook we test the qutip stochastic Schrödinger equation solver (ssesolve) with a few examples. The same examples are solved using the stochastic master equation solver in the notebook development-smesolve-tests.

In [1]:
%matplotlib inline
import matplotlib.pyplot as plt
In [2]:
from qutip import *
import numpy as np

Photo-count detection

Theory

Stochastic Schrödinger equation for photocurrent detection, in Milburn's formulation, is

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt + \left(\frac{a}{\sqrt{\langle a^\dagger a \rangle_\psi(t)}} - 1\right) |\psi(t)\rangle dN(t) - \frac{1}{2}\gamma\left(\langle a^\dagger a \rangle_\psi(t) - a^\dagger a \right) |\psi(t)\rangle dt $$

and $dN(t)$ is a Poisson distributed increment with $E[dN(t)] = \gamma \langle a^\dagger a\rangle (t)$.

Formulation in QuTiP

In QuTiP we write the stochastic Schrödinger equation on the form:

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt + D_{1}[c] |\psi(t)\rangle dt + D_{2}[c] |\psi(t)\rangle dW $$

where $c = \sqrt{\gamma} a$ is the collapse operator including the rate of the process as a coefficient in the operator. We can identify

$$ D_{1}[c] |\psi(t)\rangle = - \frac{1}{2}\left(\langle c^\dagger c \rangle_\psi(t) - c^\dagger c \right) $$$$ D_{2}[c] |\psi(t)\rangle = \frac{c}{\sqrt{\langle c^\dagger c \rangle_\psi(t)}} - 1 $$$$dW = dN(t)$$

Reference solution: deterministic master equation

In [3]:
N = 10
w0 = 0.5 * 2 * np.pi
times = np.linspace(0, 15, 150)
dt = times[1] - times[0]
gamma = 0.25
A = 2.5
ntraj = 50
nsubsteps = 100
In [4]:
a = destroy(N)
x = a + a.dag()
In [5]:
H = w0 * a.dag() * a
In [6]:
psi0 = fock(N, 5)
In [7]:
sc_ops = [np.sqrt(gamma) * a]
In [8]:
e_ops = [a.dag() * a, a + a.dag(), (-1j)*(a - a.dag())]
In [9]:
result_ref = mesolve(H, psi0, times, sc_ops, e_ops)
In [10]:
plot_expectation_values(result_ref);

Solve using stochastic master equation

$\displaystyle D_{1}[a, |\psi\rangle] = \frac{1}{2} ( \langle \psi|a^\dagger a |\psi\rangle - a^\dagger a )|\psi\rangle $

$\displaystyle D_{2}[a, |\psi\rangle] = \frac{a|\psi\rangle}{\sqrt{\langle \psi| a^\dagger a|\psi\rangle}} - |\psi\rangle $

Using QuTiP built-in photo-current detection functions for $D_1$ and $D_2$

In [11]:
result = photocurrent_sesolve(H, psi0, times, sc_ops, e_ops,
                  ntraj=ntraj*5, nsubsteps=nsubsteps, store_measurement=True, normalize=0)
10.0%. Run time:  23.07s. Est. time left: 00:00:03:27
20.0%. Run time:  50.54s. Est. time left: 00:00:03:22
30.0%. Run time:  81.23s. Est. time left: 00:00:03:09
40.0%. Run time: 112.12s. Est. time left: 00:00:02:48
50.0%. Run time: 141.57s. Est. time left: 00:00:02:21
60.0%. Run time: 169.49s. Est. time left: 00:00:01:52
70.0%. Run time: 195.67s. Est. time left: 00:00:01:23
80.0%. Run time: 223.69s. Est. time left: 00:00:00:55
90.0%. Run time: 243.49s. Est. time left: 00:00:00:27
Total run time: 268.70s
In [12]:
plot_expectation_values([result, result_ref]);
In [13]:
for m in result.measurement:
    plt.step(times, dt * m.real)

Homodyne detection

In [14]:
ntraj = 50
nsubsteps = 200
In [15]:
H = w0 * a.dag() * a + A * (a + a.dag())
In [16]:
result_ref = mesolve(H, psi0, times, sc_ops, e_ops)

Theory

Stochastic master equation for homodyne can be written as (see Gardiner and Zoller, Quantum Noise)

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt + \frac{1}{2}\gamma \left( \langle a + a^\dagger \rangle_\psi a - a^\dagger a - \frac{1}{4}\langle a + a^\dagger \rangle_\psi^2 \right)|\psi(t)\rangle dt + \sqrt{\gamma}\left( a - \frac{1}{2} \langle a + a^\dagger \rangle_\psi \right)|\psi(t)\rangle dW $$

where $dW(t)$ is a normal distributed increment with $E[dW(t)] = \sqrt{dt}$.

In QuTiP format we have:

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt + D_{1}[c, |\psi(t)\rangle] dt + D_{2}[c, |\psi(t)\rangle] dW $$

where $c = \sqrt{\gamma} a$, so we can identify

$$ D_{1}[c, |\psi\rangle] = \frac{1}{2}\left( \langle c + c^\dagger \rangle_\psi c - c^\dagger c - \frac{1}{4}\langle c + c^\dagger \rangle_\psi^2 \right) |\psi(t)\rangle $$
In [17]:
op = sc_ops[0]
opp = (op + op.dag()).data
opn = (op.dag()*op).data
op = op.data
Hd = H.data * -1j
def d1_psi_func(t, psi):
    e_x = cy.cy_expect_psi(opp, psi, 0)
    return  cy.spmv(Hd, psi) + 0.5 * (e_x * cy.spmv(op, psi) - cy.spmv(opn, psi) - 0.25 * e_x ** 2 * psi)
$$ D_{2}[c,|\psi\rangle] = \left(c - \frac{1}{2} \langle c + c^\dagger \rangle_\psi \right)|\psi(t)\rangle $$
In [18]:
def d2_psi_func(t, psi):
    out = np.zeros((1,len(psi)), dtype=complex)
    e_x = cy.cy_expect_psi(opp, psi, 0)
    out[0,:] = cy.spmv(op,psi)
    out -= 0.5 * e_x * psi
    return out
In [19]:
result = general_stochastic(psi0, times, d1=d1_psi_func, d2=d2_psi_func, 
                            e_ops=e_ops, ntraj=ntraj, 
                            m_ops=[a + a.dag()], dW_factors=[1/np.sqrt(gamma)],
                            nsubsteps=nsubsteps, store_measurement=True)
10.0%. Run time:  45.09s. Est. time left: 00:00:06:45
20.0%. Run time:  94.20s. Est. time left: 00:00:06:16
30.0%. Run time: 140.39s. Est. time left: 00:00:05:27
40.0%. Run time: 191.46s. Est. time left: 00:00:04:47
50.0%. Run time: 233.52s. Est. time left: 00:00:03:53
60.0%. Run time: 268.13s. Est. time left: 00:00:02:58
70.0%. Run time: 307.96s. Est. time left: 00:00:02:11
80.0%. Run time: 360.42s. Est. time left: 00:00:01:30
90.0%. Run time: 404.45s. Est. time left: 00:00:00:44
Total run time: 446.63s
In [20]:
plot_expectation_values([result, result_ref]);
In [21]:
for m in result.measurement:
    plt.plot(times, m[:, 0].real, 'b', alpha=0.1)

plt.plot(times,  np.array(result.measurement).mean(axis=0)[:,0].real, 'r', lw=2)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.ylim(-15, 15);

Solve problem again, this time with a specified noise (from previous run)

In [22]:
result = general_stochastic(psi0, times, d1=d1_psi_func, d2=d2_psi_func, 
                            e_ops=e_ops, ntraj=ntraj, noise=result.noise,
                            m_ops=[a + a.dag()], dW_factors=[1/np.sqrt(gamma)],
                            nsubsteps=nsubsteps, store_measurement=True)
10.0%. Run time:  40.24s. Est. time left: 00:00:06:02
20.0%. Run time:  77.66s. Est. time left: 00:00:05:10
30.0%. Run time: 123.28s. Est. time left: 00:00:04:47
40.0%. Run time: 172.75s. Est. time left: 00:00:04:19
50.0%. Run time: 220.48s. Est. time left: 00:00:03:40
60.0%. Run time: 265.02s. Est. time left: 00:00:02:56
70.0%. Run time: 305.89s. Est. time left: 00:00:02:11
80.0%. Run time: 336.40s. Est. time left: 00:00:01:24
90.0%. Run time: 359.65s. Est. time left: 00:00:00:39
Total run time: 395.30s
In [23]:
plot_expectation_values([result, result_ref]);
In [24]:
for m in result.measurement:
    plt.plot(times, m[:, 0].real, 'b', alpha=0.1)

plt.plot(times,  np.array(result.measurement).mean(axis=0)[:,0].real, 'r', lw=2)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.ylim(-15, 15);

Using QuTiP built-in homodyne detection functions for $D_1$ and $D_2$

In [25]:
result = ssesolve(H, psi0, times, sc_ops, e_ops, ntraj=ntraj, nsubsteps=nsubsteps,
                  method='homodyne', store_measurement=True, dW_factors=[1])
10.0%. Run time:  65.15s. Est. time left: 00:00:09:46
20.0%. Run time: 122.95s. Est. time left: 00:00:08:11
30.0%. Run time: 166.42s. Est. time left: 00:00:06:28
40.0%. Run time: 217.26s. Est. time left: 00:00:05:25
50.0%. Run time: 267.44s. Est. time left: 00:00:04:27
60.0%. Run time: 325.35s. Est. time left: 00:00:03:36
70.0%. Run time: 372.25s. Est. time left: 00:00:02:39
80.0%. Run time: 439.13s. Est. time left: 00:00:01:49
90.0%. Run time: 490.72s. Est. time left: 00:00:00:54
Total run time: 538.31s
In [26]:
plot_expectation_values([result, result_ref]);
In [27]:
for m in result.measurement:
    plt.plot(times, m[:, 0].real, 'b', alpha=0.1)

plt.plot(times,  np.array(result.measurement).mean(axis=0)[:,0].real/np.sqrt(gamma), 'r', lw=2)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.ylim(-15, 15);
In [28]:
plt.plot(times,  np.array(result.measurement).mean(axis=0)[:,0].real/np.sqrt(gamma), 'r', lw=2)
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.plot(times, result.expect[1], 'b', lw=2)
Out[28]:
[<matplotlib.lines.Line2D at 0x7fa998eea7b8>]

Solve problem again, this time with a specified noise (from previous run)

In [29]:
result = ssesolve(H, psi0, times, sc_ops, e_ops, ntraj=ntraj, nsubsteps=nsubsteps,
                  method='homodyne', store_measurement=True, noise=result.noise)
10.0%. Run time:  44.82s. Est. time left: 00:00:06:43
20.0%. Run time: 119.79s. Est. time left: 00:00:07:59
30.0%. Run time: 174.45s. Est. time left: 00:00:06:47
40.0%. Run time: 208.24s. Est. time left: 00:00:05:12
50.0%. Run time: 250.80s. Est. time left: 00:00:04:10
60.0%. Run time: 299.31s. Est. time left: 00:00:03:19
70.0%. Run time: 348.43s. Est. time left: 00:00:02:29
80.0%. Run time: 417.13s. Est. time left: 00:00:01:44
90.0%. Run time: 451.07s. Est. time left: 00:00:00:50
Total run time: 495.12s
In [30]:
plot_expectation_values([result, result_ref]);
In [31]:
for m in result.measurement:
    plt.plot(times, m[:, 0].real, 'b', alpha=0.1)

plt.plot(times,  np.array(result.measurement).mean(axis=0)[:,0].real/np.sqrt(gamma), 'r', lw=2)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.ylim(-15, 15);

Heterodyne detection

Stochastic Schrödinger equation for heterodyne detection can be written as

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt -\frac{1}{2}\gamma\left(a^\dagger a - \langle a^\dagger \rangle a + \frac{1}{2}\langle a \rangle\langle a^\dagger \rangle)\right)|\psi(t)\rangle dt + \sqrt{\gamma/2}\left( a - \frac{1}{2} \langle a + a^\dagger \rangle_\psi \right)|\psi(t)\rangle dW_1 -i \sqrt{\gamma/2}\left( a - \frac{1}{2} \langle a - a^\dagger \rangle_\psi \right)|\psi(t)\rangle dW_2 $$

where $dW_i(t)$ is a normal distributed increment with $E[dW_i(t)] = \sqrt{dt}$.

In QuTiP format we have:

$$ d|\psi(t)\rangle = -iH|\psi(t)\rangle dt + D_{1}[c, |\psi(t)\rangle] dt + \sum_n D_{2}^{(n)}[c, |\psi(t)\rangle] dW_n $$

where $c = \sqrt{\gamma} a$, so we can identify

$$ D_1[c, |\psi(t)\rangle] = -\frac{1}{2}\left(c^\dagger c - \langle c^\dagger \rangle c + \frac{1}{2}\langle c \rangle\langle c^\dagger \rangle \right) |\psi(t)\rangle $$
In [32]:
op = sc_ops[0]
opd = (op.dag()).data
opp = (op + op.dag()).data
opm = (op + op.dag()).data
opn = (op.dag()*op).data
op = op.data
Hd = H.data * -1j
def d1_psi_func(t, psi):
    e_xd = cy.cy_expect_psi(opd, psi, 0)
    e_x = cy.cy_expect_psi(op, psi, 0)
    return  cy.spmv(Hd, psi) - 0.5 * (cy.spmv(opn, psi) - e_xd * cy.spmv(op, psi) + 0.5 * e_x * e_xd * psi)
$$D_{2}^{(1)}[c, |\psi(t)\rangle] = \sqrt{1/2} (c - \langle c + c^\dagger \rangle / 2) \psi$$$$D_{2}^{(1)}[c, |\psi(t)\rangle] = -i\sqrt{1/2} (c - \langle c - c^\dagger \rangle /2) \psi$$
In [33]:
sqrt2 = np.sqrt(0.5)
def d2_psi_func(t, psi):
    out = np.zeros((2,len(psi)), dtype=complex)
    e_p = cy.cy_expect_psi(opp, psi, 0)
    e_m = cy.cy_expect_psi(opm, psi, 0)
    out[0,:] = (cy.spmv(op,psi) - e_p * 0.5 * psi)*sqrt2
    out[1,:] = (cy.spmv(op,psi) - e_m * 0.5 * psi)*sqrt2*-1j
    return out
In [34]:
result = general_stochastic(psi0, times, d1=d1_psi_func, d2=d2_psi_func, 
                            e_ops=e_ops, ntraj=ntraj, len_d2=2,
                            m_ops=[(a + a.dag()), (-1j)*(a - a.dag())], dW_factors=[2/np.sqrt(gamma), 2/np.sqrt(gamma)],
                            nsubsteps=nsubsteps, store_measurement=True)
10.0%. Run time:  37.96s. Est. time left: 00:00:05:41
20.0%. Run time:  76.64s. Est. time left: 00:00:05:06
30.0%. Run time: 114.65s. Est. time left: 00:00:04:27
40.0%. Run time: 151.38s. Est. time left: 00:00:03:47
50.0%. Run time: 182.40s. Est. time left: 00:00:03:02
60.0%. Run time: 217.82s. Est. time left: 00:00:02:25
70.0%. Run time: 256.56s. Est. time left: 00:00:01:49
80.0%. Run time: 299.20s. Est. time left: 00:00:01:14
90.0%. Run time: 342.48s. Est. time left: 00:00:00:38
Total run time: 382.84s
In [35]:
plot_expectation_values([result, result_ref]);
In [36]:
#fig, ax = subplots()

for m in result.measurement:
    plt.plot(times, m[:, 0].real, 'r', alpha=0.025)
    plt.plot(times, m[:, 1].real, 'b', alpha=0.025)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2);
plt.plot(times, result_ref.expect[2], 'k', lw=2);
plt.ylim(-10, 10)

plt.plot(times, np.array(result.measurement).mean(axis=0)[:,0].real, 'r', lw=2);
plt.plot(times, np.array(result.measurement).mean(axis=0)[:,1].real, 'b', lw=2);

Using QuTiP built-in heterodyne detection functions for $D_1$ and $D_2$

In [37]:
result = ssesolve(H, psi0, times, sc_ops, e_ops, ntraj=ntraj, nsubsteps=nsubsteps,
                  method='heterodyne', store_measurement=True)
10.0%. Run time:  88.82s. Est. time left: 00:00:13:19
20.0%. Run time: 208.68s. Est. time left: 00:00:13:54
30.0%. Run time: 318.16s. Est. time left: 00:00:12:22
40.0%. Run time: 416.87s. Est. time left: 00:00:10:25
50.0%. Run time: 510.53s. Est. time left: 00:00:08:30
60.0%. Run time: 621.76s. Est. time left: 00:00:06:54
70.0%. Run time: 733.11s. Est. time left: 00:00:05:14
80.0%. Run time: 837.64s. Est. time left: 00:00:03:29
90.0%. Run time: 950.83s. Est. time left: 00:00:01:45
Total run time: 1053.90s
In [38]:
plot_expectation_values([result, result_ref]);
In [39]:
for m in result.measurement:
    plt.plot(times, m[:, 0, 0].real, 'r', alpha=0.025)
    plt.plot(times, m[:, 0, 1].real, 'b', alpha=0.025)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2)
plt.plot(times, result_ref.expect[2], 'k', lw=2)

plt.plot(times, np.array(result.measurement).mean(axis=0)[:,0,0].real/np.sqrt(gamma), 'r', lw=2)
plt.plot(times, np.array(result.measurement).mean(axis=0)[:,0,1].real/np.sqrt(gamma), 'b', lw=2)
Out[39]:
[<matplotlib.lines.Line2D at 0x7fa9990c29b0>]

Solve problem again, this time with a specified noise (from previous run)

In [40]:
result = ssesolve(H, psi0, times, sc_ops, e_ops, ntraj=ntraj, nsubsteps=nsubsteps,
                  method='heterodyne', store_measurement=True, noise=result.noise)
10.0%. Run time: 101.94s. Est. time left: 00:00:15:17
20.0%. Run time: 211.18s. Est. time left: 00:00:14:04
30.0%. Run time: 328.08s. Est. time left: 00:00:12:45
40.0%. Run time: 429.03s. Est. time left: 00:00:10:43
50.0%. Run time: 536.44s. Est. time left: 00:00:08:56
60.0%. Run time: 636.26s. Est. time left: 00:00:07:04
70.0%. Run time: 742.40s. Est. time left: 00:00:05:18
80.0%. Run time: 834.82s. Est. time left: 00:00:03:28
90.0%. Run time: 889.97s. Est. time left: 00:00:01:38
Total run time: 937.45s
In [41]:
plot_expectation_values([result, result_ref]);
In [42]:
for m in result.measurement:
    plt.plot(times, m[:, 0, 0].real, 'r', alpha=0.025)
    plt.plot(times, m[:, 0, 1].real, 'b', alpha=0.025)
    
plt.plot(times, result_ref.expect[1], 'k', lw=2);
plt.plot(times, result_ref.expect[2], 'k', lw=2);

plt.plot(times, np.array(result.measurement).mean(axis=0)[:,0,0].real/np.sqrt(gamma), 'r', lw=2);
plt.plot(times, np.array(result.measurement).mean(axis=0)[:,0,1].real/np.sqrt(gamma), 'b', lw=2);