import numpy as np
import scipy as sp
import scipy.sparse as sparse
import scipy.sparse.linalg as sla
import matplotlib.pyplot as plt
%matplotlib inline

n = 64
h = 1.0 / n
A = sparse.diags([-1, 2, -1], [-1, 0, 1], shape=(n,n), format='csr')
b = np.zeros((n,))

plt.figure(figsize=(6,6))

K = np.arange(1,n)

omega = 1.0 /2
lmbda = 1 - (omega / 2) * 4.0 * np.sin(np.pi * K / (2*(n+1)))**2
plt.plot(lmbda,'-k',label='residual',linewidth=4, clip_on=False)
plt.text(n+2, lmbda[-1], r'$\omega=1/2$', fontsize=24)

omega = 1.0 / 3.0
lmbda = 1 - (omega / 2) * 4.0 * np.sin(np.pi * K / (2*(n+1)))**2
plt.plot(lmbda,'-k',label='residual',linewidth=4, clip_on=False)
plt.text(n+2, lmbda[-1], r'$\omega=1/3$', fontsize=24)

omega = 2.0 / 3.0
lmbda = 1 - (omega / 2) * 4.0 * np.sin(np.pi * K / (2*(n+1)))**2
plt.plot(lmbda,'-k',label='residual',linewidth=4, clip_on=False)
plt.text(n+2, lmbda[-1], r'$\omega=2/3$', fontsize=24)

omega = 1
lmbda = 1 - (omega / 2) * 4.0 * np.sin(np.pi * K/ (2*(n+1)))**2
plt.plot(lmbda,'-k',label='residual',linewidth=4, clip_on=False)
plt.text(n+2, lmbda[-1], r'$\omega=1$', fontsize=24)

####### fancy plotting vvvvvvvvvv
ax = plt.gca()
ax.plot([0, n], [0, 0], '--b')
plt.axis([0,n,-1,1])
plt.ylabel(r'$\lambda^{\omega_J}_k$', fontsize=32)
plt.xlabel(r'wave number, $k$', fontsize=32)
ax.spines["right"].set_visible(False)
ax.spines["top"].set_visible(False)
ax.get_xaxis().tick_bottom()   # remove unneeded ticks 
ax.get_yaxis().tick_left()
plt.show()