complete Lanzco's method

mogu/spectral/lanzcos.py

@@ -2,6 +2,7 @@
 # implementing Lanzco's algorithm
 
 import numpy as np
+from scipy.linalg import eigh_tridiagonal
 
 
 def lanzcos(M, k):
@@ -12,18 +13,20 @@ def lanzcos(M, k):
 
     a = np.zeros(k+1)
     b = np.zeros(k)
+    v = np.zeros((d, k+1))
 
-    v0 = np.random.uniform(size=d)
-    v0 /= np.linalg.norm(v0)
+    v[:, 0] = np.random.uniform(size=d)
+    v[:, 0] /= np.linalg.norm(v[:, 0])
 
-    a[0] = v0.T @ M @ v0
-    v1 = M @ v0 - a[0]*v0
-    v1 /= np.linalg.norm(v1)
+    a[0] = v[:, 0].T @ M @ v[:, 0]
+    v[:, 1] = M @ v[:, 0] - a[0] * v[:, 0]
+    v[:, 1] /= np.linalg.norm(v[:, 1])
 
-    for i in range(1, k+1):
-        a[i] = v1.T @ M @ v1
-        b[i-1] = v0.T @ M @ v1
-        v2 = M @ v1 - a[i]*v1 - b[i-1]*v0
-        v2 /= np.linalg.norm(v2)
+    for i in range(1, k):
+        a[i] = v[:, i].T @ M @ v[:, i]
+        b[i-1] = v[:, i-1].T @ M @ v[:, i]
+        v[:, i+1] = M @ v[:, i] - a[i] * v[:, i] - b[i-1] * v[:, i-1]
+        v[:, i+1] /= np.linalg.norm(v[:, i+1])
 
-        v0, v1 = v1, v2
+    eigvals, eigvecs = eigh_tridiagonal(a, b)
+    return eigvals, v @ eigvecs