"""
MATH3476 Example 3.3
Theoretically determined omega^*
"""
import scipy as sp
from scipy import linalg

A = sp.array([[-4.,  0., 1.,  1.],
              [ 0., -4., 1.,  1.],
              [ 1., 1., -4.,  0.],
              [ 1., 1.,  0., -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

# Jocabi iteration matrix
H_J = L + U
print('H_J ='); print(H_J, '\n')

lam_HJ = linalg.eigvals(H_J)   # eigenvalues

rho_HJ = sp.amax(abs(lam_HJ))  # spectral radius
print('spectral radius of H_J =', rho_HJ, '\n')

# Gauss-Seidel iteration matrix
H_GS = sp.dot(linalg.inv(I-L),U)
print('H_GS ='); print(H_GS, '\n')

lam_HGS = linalg.eigvals(H_GS)

rho_HGS = sp.amax(abs(lam_HGS))
print('spectral radius of H_GS =', rho_HGS, '\n')

# Theoretical optimal omega for H_omega
omega1 = 2./(1+sp.sqrt(1-rho_HJ**2.))   # omega^*
print('optimal omega =',  omega1, '\n')

# optimal SOR iteration matrix
H_omega1 = sp.dot(linalg.inv(I-omega1*L),(1-omega1)*I + omega1*U)
print('optimal H_omega ='); print(H_omega1, '\n')

lam_Homega1 = linalg.eigvals(H_omega1)

rho_Homega1 = sp.amax(abs(lam_Homega1))
print('spectral radius of optimal H_omega =', rho_Homega1)