import scqubits as sc

tmon = sc.Transmon(
    EJ = 31,
    EC = .8,
    ng = 0,
    ncut = 110
)
evals = tmon.eigenvals()[:5]
omega = evals[1]-evals[0]
osc = sc.Oscillator(
    E_osc = omega
)

import numpy as np
from matplotlib import pyplot as plt

xlist = np.linspace(-np.pi,np.pi,100)
mineng = min([tmon.potential(x) for x in xlist])

plt.plot(xlist, [(tmon.potential(x)-mineng)/omega for x in xlist], c = "blue", label = "qubit")
plt.plot(xlist, [(x**2 * 31/2)/omega for x in xlist], c = "red", label = "qumode")
plt.hlines([(n+.5) for n in range(5)], -1, 1, color="red")
plt.hlines((evals-mineng)/omega, -1, 1, color="blue")
plt.xlabel(r"$\varphi$")
plt.ylabel(r"E/$\omega$")
plt.ylim([0,6])
plt.legend()

