#!/usr/bin/env python
# coding: utf-8

# # Krylov Subspace Methods for Extremal Eigenvalue Problems

# In[1]:


import numpy as np
import numpy.linalg as la
import matplotlib.pyplot as pt


# Let us make a matrix with a defined set of eigenvalues and eigenvectors, given by `eigvals` and `eigvecs`.

# In[2]:


np.random.seed(40)

# Generate matrix with eigenvalues 1...50
n = 50
eigvals = np.geomspace(1., n, n)
eigvecs = np.random.randn(n, n)
#To work with symmetric matrix, orthogonalize eigvecs
eigvecs, R = la.qr(eigvecs)

print(eigvals)
max_eigval = eigvals[np.argmax(np.abs(eigvals))]
min_eigval = eigvals[np.argmin(np.abs(eigvals))]

A = la.solve(eigvecs, np.dot(np.diag(eigvals), eigvecs))


# ## Initialization

# Set up $Q$ and $H$:

# In[3]:


Q = np.zeros((n, n))
H = np.zeros((n, n))

k = 0


# Pick a starting vector, normalize it

# In[4]:


x0 = np.random.randn(n)
x0 = x0/la.norm(x0)

# Poke it into the first column of Q
Q[:, k] = x0.copy()


# Make a list to save arrays of Ritz values:

# In[5]:


ritz_values = []
ritz_max = []
ritz_min = []


# ## Algorithm

# Carry out one iteration of Arnoldi iteration.

# In[6]:



k = n-1

for kk in range(k-1):
    
    u = A @ Q[:, kk]
    
    # Carry out Gram-Schmidt on u against Q
    # to do Lanczos change range start to k-1
    # u = (np.eye(n,n) - Q[:,:kk+1] @ Q[:,:kk+1].T) @ u
    for j in range(0,kk+1):
        qj = Q[:, j]
        H[j,kk] = qj @ u
        u = u - H[j,kk]*qj
    
    if kk+1 < n:
        H[kk+1, kk] = la.norm(u)
        Q[:, kk+1] = u/H[kk+1, kk]
    
    if kk>1:
        D = la.eig(H[:kk,:kk])[0]
        max_ritz = D[np.argmax(np.abs(D))]
        min_ritz = D[np.argmin(np.abs(D))]
        ritz_vals = np.zeros(kk)
        for i in range(kk):
            ritz_vals[i] = D[np.argmax(np.abs(D))]
            D[np.argmax(np.abs(D))] = 0
        ritz_max.append(max_ritz)
        ritz_min.append(min_ritz)
        ritz_values.append(ritz_vals)
    
pt.spy(H)


# Check that $Q^T A Q =H$:

# In[7]:


la.norm(Q[:,:k-1].T @ A @ Q[:,:k-1] - H[:k-1,:k-1])/ la.norm(A)


# Check that $AQ-QH$ is fairly small

# In[8]:


la.norm(A @ Q[:,:k-1] - Q[:,:k-1]@H[:k-1,:k-1])/ la.norm(A)


# Check that Q is orthogonal:

# In[9]:


la.norm((Q.T.conj() @ Q)[:k-1,:k-1] - np.eye(k-1))


# ## Compare Max Ritz Value to Power Iteration

# In[10]:


#true largest eigenvalue is 50
r = 0
x = A @ x0.copy()
x = x / la.norm(x)
mrs = []
for i in range(k-1):
    y = A @ x
    r = x @ y
    x = y / la.norm(y)
    mrs.append(r)
print(r,max_ritz)
pt.plot(mrs, "x", label="Power iteration")
pt.plot(ritz_max, "o", label="Arnoldi")
pt.xlabel("iteration")
pt.ylabel("maximum eigenvalue estimate")
pt.legend()


# In[11]:


pt.plot(np.abs(max_eigval-mrs)/max_eigval, "x", label="Power iteration")
pt.plot(np.abs(max_eigval-ritz_max)/max_eigval, "o", label="Arnoldi")
pt.xlabel("iteration")
pt.ylabel("relative error")
pt.yscale("log")
pt.legend()


# ## Plot convergence of Ritz values

# Enable the Ritz value collection above to make this work.

# In[12]:


for i, rv in enumerate(ritz_values):
    pt.plot([i] * len(rv), rv, "x")


# ## Compare Min Ritz Value to Inverse Iteration

# In[13]:


#true largest eigenvalue is 50
r = 0
x = la.solve(A, x0.copy())
x = x / la.norm(x)
rs = []
for i in range(k-1):
    y = A @ x
    r = x @ y
    y = la.solve(A,x)
    x = y / la.norm(y)
    rs.append(r)
print(r,min_ritz)
pt.plot(rs, "x", label="Inverse iteration")
pt.plot(ritz_min, "o", label="Arnoldi")
pt.xlabel("iteration")
pt.ylabel("minimum eigenvalue estimate")
pt.legend()


# In[14]:


print(min_eigval-rs)
pt.plot(np.abs(min_eigval-rs)/min_eigval, "x", label="Inverse iteration")
pt.plot(np.abs(min_eigval-ritz_min)/min_eigval, "o", label="Arnoldi")
pt.xlabel("iteration")
pt.ylabel("relative error")
pt.yscale("log")
pt.legend()


# In[ ]:





# In[ ]: