import qutip
import numpy as np
import matplotlib.pyplot as plt
import scipy

N = 5
D = qutip.displace(N, 3)
C = qutip.dag(D)
thermal  = C * (qutip.states.thermal_dm(N,1)) * D
thermal2 = C * (qutip.states.thermal_dm(N,2)) * D
thermal3 = C * (qutip.states.thermal_dm(N,5)) * D

print thermal



fig, graph1 = plt.subplots(1, 3, figsize =(20, 5))

bar1 = graph1[0].bar(np.arange(0,N),thermal.diag())
label1 = graph1[0].set_title("N = 100, n = 1, alpha = 1")
lim1 = graph1[0].set_xlim([-.5, 30])

bar2 = graph1[1].bar(np.arange(0,N),thermal2.diag())
label2 = graph1[1].set_title("N = 100, n = 2, alpha = 1")
lim2 = graph1[1].set_xlim([-.5, 40])

bar3 = graph1[2].bar(np.arange(0,N),thermal3.diag())
label3 = graph1[2].set_title("N = 100, n = 5, alpha = 1")
lim3 = graph1[2].set_xlim([-.5, 60])



t1diag = np.matrix(thermal, dtype=object)
a = np.arange,thermal.diag()
b = np.array(a)
x = str(b[1])

print x