import numpy as np
import pylab as pl

a = 1
n_vals = [n for n in range(1, 7)]

x_vals = np.linspace(-a / 2, a / 2, 1000)

psi_evens = []
psi_odds = []
for n in n_vals:
    if n % 2 == 0:
        psi_evens.append(np.sqrt(2 / a) * np.sin(n * np.pi * x_vals / a))
    else:
        psi_odds.append(np.sqrt(2 / a) * np.cos(n * np.pi * x_vals / a))

pl.title("First 6 Solutions for Infinite Square Well")
for psi_even in psi_evens:
    pl.plot(x_vals, psi_even, color='blue')
for psi_odd in psi_odds:
    pl.plot(x_vals, psi_odd, color='green')
pl.vlines([-a / 2, a / 2], ymin=-2, ymax=2, color='red')
pl.grid()
pl.show()