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

# # Arnoldi vs Power Iteration

# 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.linspace(1., n, n)
eigvecs = np.random.randn(n, n)
#To work with symmetric matrix, orthogonalize eigvecs
eigvecs, R = la.qr(eigvecs)

print(eigvals)

A = la.solve(eigvecs, np.dot(np.diag(eigvals), eigvecs))
print(la.eig(A)[0])

# ## 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)

# Set the first column of Q to be the normalized starting vector
Q[:, k] = x0.copy()


# Make a list to save arrays of Ritz values:

# In[5]:


ritz_values = []
ritz_max = []

# ## Algorithm

# Carry out one iteration of Arnoldi iteration.
# 
# Run this cell in-place (Ctrl-Enter) until H is filled.

# In[6]:


def Arnoldi_step(A,Q,H,k):
    
    u = A @ Q[:, k]

    # Carry out Gram-Schmidt on u against Q
    # to do Lanczos change range start to k-1
    for j in range(0,k+1):
        qj = Q[:, j]
        H[j,k] = qj @ u
        u = u - H[j,k]*qj

    if k+1 < n:
        H[k+1, k] = la.norm(u)
        Q[:, k+1] = u/H[k+1, k]

# In[21]:


print(k)

Arnoldi_step(A,Q,H,k)

k += 1

pt.spy(H)


if k>1:
    D = la.eig(H)[0]
    max_ritz = D[np.argmax(np.abs(D))]
    ritz_vals = np.zeros(k)
    for i in range(k):
        ritz_vals[i] = D[np.argmax(np.abs(D))]
        D[np.argmax(np.abs(D))] = 0
    ritz_max.append(max_ritz)
    ritz_values.append(ritz_vals)

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

# In[22]:


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[23]:


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

# Check that Q is orthogonal:

# In[24]:


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

# ## Compare Max Ritz Value to Power Iteration

# In[25]:


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

# ## Plot convergence of Ritz values

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

# In[26]:


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

# In[ ]:




# In[ ]: