%pylab inline

from qutip import *

def steadystate_fock_distribution(N, offset):
    """
    Solve for the steady state of a driven dissipative cavity, given
    the upper (N) and lower (offset) cut offs for the number states in
    in the basis. Calculate the average photon number in the steady
    state and plot the fock state distribution.
    """
    N = N - offset
    Omega = 1.0 * 2 * pi
    kappa = 0.9
    
    a = destroy(N, offset)
    n = num(N, offset)
    H = Omega * (a + a.dag())
    c_ops = [sqrt(kappa) * a]
    
    rho_ss = steadystate(H, c_ops)
    
    title = "Steady state avg. photon number = %.2f" % expect(rho_ss, n)    
    
    fig, ax = fock_distribution(rho_ss, offset=offset, unit_y_range=False, title=title)
    
    return fig, ax

steadystate_fock_distribution(250, 0);

steadystate_fock_distribution(250, 150);

%pylab inline

from qutip import *

N = 10
offset = 0

n = num(N, offset=offset)
psi = (basis(N, 2, offset=offset) + basis(N, 4, offset=offset) + basis(N, 6, offset=offset)).unit()

fock_distribution(psi, offset=offset, unit_y_range=False, figsize=(4,3))

expect(n, psi)

N = 5
offset = 2

n = num(N, offset=offset)
psi = (basis(N, 2, offset=offset) + basis(N, 4, offset=offset) + basis(N, 6, offset=offset)).unit()

fock_distribution(psi, offset=offset, unit_y_range=False, figsize=(4,3))

expect(n, psi)

N = 75
offset = 0

n = num(N, offset=offset)
psi = coherent(N, 6.5, offset=offset)

fock_distribution(psi, offset=offset, unit_y_range=False, figsize=(4,3))

expect(n, psi)

N = 50
offset = 20

n = num(N, offset=offset)
psi = coherent(N, 6.5, offset=offset)

fock_distribution(psi, offset=offset, unit_y_range=False, figsize=(4,3))

expect(n, psi)

from qutip.ipynbtools import version_table; version_table()