Implementation

from scipy.special import hermite as Hermite

class HarmonicOscillator:

    def __init__(self, m, omega):
        self.mass = m
        self.omega = omega
        self.force_constant = self.mass * self.omega**2
        self.alpha = (1.0 / (self.mass * self.force_constant))**(1.0 / 4.0)

    def get_potential(self, x):
        return 0.5 * self.force_constant * (x**2)

    def get_energy(self, n):
        return (n + 0.5) * self.omega

    def get_eigenfunction(self, x, n):
        norm = (1.0 / np.sqrt((2**n) * np.math.factorial(n) * np.sqrt(np.pi) * self.alpha))
        y = x / self.alpha
        herm = Hermite(n)
        return norm * herm(y) * np.exp(-0.5 * y * y)
    
    def plot(self, xvalues=np.linspace(-10,10,1000), nstate=5, scale=1.0):
        f, ax = plt.subplots()
        ax.plot(xvalues, self.get_potential(xvalues), lw=3.0, color="blue")
        plt.xlabel(r"$x$")
        plt.ylabel(r"$V(x)$")
        for state in range(nstate):
            en = self.get_energy(state)
            ax.plot([xvalues[0], xvalues[-1]], [en, en], lw=2.0, color="red")
            ax.plot(xvalues, en + scale * self.get_eigenfunction(xvalues, state),
                    color="black", alpha=0.5)
        return f, ax

mass = 1.0
omega = 1.0
ho  = HarmonicOscillator(mass, omega)
xmin, xmax = -10, 10
ngrid = 128
x = np.linspace(xmin, xmax, ngrid)
v = ho.get_potential(x)
fig_ho, ax_ho = ho.plot(np.linspace(-5, 5, 1000), nstate=10)