"""
MATH3476 Example 3.4
Experiments with numerically determined omega^*
"""
import scipy as sp
from scipy import linalg
import matplotlib.pyplot as plt
matplotlib.rcParams.update({'font.size': 18})

A = sp.array([[-4.,  1.,  1.,  1.],
              [ 1., -4.,  1.,  1.],
              [ 1.,  1., -4.,  1.],
              [ 1.,  1.,  1., -4.]])

A1 = A/(-4.)          # scaled diagonal to 1

I = sp.identity(4)    # 4x4 identity matrix

L = I - sp.tril(A1)   # extract L from A

U = I - sp.triu(A1)   # extract U from A

N = 51                # no. of data points to scan
omega_list = sp.linspace(1,2,N)        # scan omega in [1,2]
rho = sp.zeros(N)     # initialise array for spectral radius

for i in range(0,N):
    omega = omega_list[i]
    H_SOR = sp.dot(linalg.inv(I-omega*L),(1-omega)*I + omega*U)
    lam = sp.real(linalg.eigvals(H_SOR))   # eigenvalues
    rho[i] = sp.amax(abs(lam))             # spectral radius

#--- Plotting
plt.figure(1); plt.clf()
plt.plot(omega_list,rho,'-o',fillstyle='none')
plt.xlim((1,2))
plt.xlabel('$\omega$')
plt.ylabel('$\\rho(H_{SOR})$')
plt.show()
plt.savefig('rho_HSOR.pdf',dpi=600,bbox_inches='tight')

# Jocabi iteration matrix
H_J = L + U
lam_HJ = linalg.eigvals(H_J)   # eigenvalues
rho_HJ = sp.amax(lam_HJ)       # spectral radius
print('spectral radius of H_J =', rho_HJ)

# Gauss-Seidel iteration matrix
H_GS = sp.dot(linalg.inv(I-L),U)
lam_HGS = linalg.eigvals(H_GS)
rho_HGS = sp.amax(lam_HGS)
print('spectral radius of H_GS =', rho_HGS)