Add large matrix tracking demo notebook

src/eigenpairflow/correction.py

@@ -3,6 +3,22 @@
 from scipy.optimize import linear_sum_assignment
 
 
+def _find_clusters(eigvals, delta):
+    """固有値を昇順に並べて差が ``delta`` 未満となるグループを検出する。"""
+    sorted_idx = np.argsort(eigvals)
+    clusters: list[list[int]] = []
+    current = [sorted_idx[0]] if sorted_idx.size > 0 else []
+    for idx in sorted_idx[1:]:
+        if abs(eigvals[idx] - eigvals[current[-1]]) < delta:
+            current.append(idx)
+        else:
+            clusters.append(current)
+            current = [idx]
+    if current:
+        clusters.append(current)
+    return clusters
+
+
 def match_decompositions(
     predicted_eigvals, predicted_eigvecs, exact_eigvals, exact_eigvecs
 ):
@@ -108,20 +124,20 @@ def correct_trajectory(A_func, t_eval, Qs_ode, Lambdas_ode, method="matching"):
 
 def ogita_aishima_refinement(A, X_hat, max_iter=10, tol=1e-12, rho=1.0):
     """
-    Refines a given approximate eigenvector matrix using the Ogita-Aishima method.
-    This implementation is vectorized for performance.
+    荻田・相島の方法で近似固有ベクトルを改良する。
+    固有値が近接している場合には小さな部分空間で再対角化を行う。
 
     Args:
-        A (np.ndarray): The target real symmetric matrix.
-        X_hat (np.ndarray): The approximate eigenvector matrix.
-        max_iter (int): Maximum number of iterations.
-        tol (float): Tolerance for convergence.
-        rho (float): Parameter to determine if eigenvalues are clustered.
+        A (np.ndarray): 対称行列。
+        X_hat (np.ndarray): 初期固有ベクトル近似。
+        max_iter (int): 最大反復回数。
+        tol (float): 収束判定の許容誤差。
+        rho (float): 固有値クラスタ判定に用いる係数。
 
     Returns:
         tuple[np.ndarray, np.ndarray]:
-            - X_new: The refined eigenvector matrix.
-            - D_new: A diagonal matrix containing the refined eigenvalues.
+            - X_new: 改良された固有ベクトル。
+            - D_new: 改良された固有値を対角に持つ行列。
     """
     n = A.shape[0]
     Id = np.eye(n)
@@ -135,7 +151,7 @@ def ogita_aishima_refinement(A, X_hat, max_iter=10, tol=1e-12, rho=1.0):
         # 2. Calculate approximate eigenvalues
         r_ii = np.diag(R)
         s_ii = np.diag(S)
-        lambda_tilde = s_ii / (1 - r_ii + 1e-15)
+        lambda_tilde = s_ii / (1 - r_ii)
 
         # 3. Calculate the correction matrix E_tilde (vectorized)
         s_off_diag = S - np.diag(s_ii)
@@ -163,15 +179,26 @@ def ogita_aishima_refinement(A, X_hat, max_iter=10, tol=1e-12, rho=1.0):
         # 4. Update the solution
         X_new = X_new @ (Id + E_tilde)
 
-        # 5. Check for convergence
+        # 5. Additional refinement for clustered eigenvalues
+        clusters = _find_clusters(lambda_tilde, delta)
+        for cluster in clusters:
+            if len(cluster) > 1:
+                X_block = X_new[:, cluster]
+                X_block, _ = np.linalg.qr(X_block)
+                B = X_block.T @ A @ X_block
+                w, V = np.linalg.eigh(B)
+                X_new[:, cluster] = X_block @ V
+
+        # 6. Check for convergence
         if np.linalg.norm(E_tilde) < tol:
             break
 
-    # Final calculation of eigenvalues and sorting
+    # Final calculation of eigenvalues without reordering
     S_final = X_new.T @ A @ X_new
     final_eigvals = np.diag(S_final)
-    sort_indices = np.argsort(final_eigvals)
-    D_new = np.diag(final_eigvals[sort_indices])
-    X_new = X_new[:, sort_indices]
+    dots = np.sum(X_hat * X_new, axis=0)
+    signs = np.where(dots >= 0, 1.0, -1.0)
+    X_new = X_new * signs
+    D_new = np.diag(final_eigvals)
 
     return X_new, D_new

test/test_correction.py

@@ -3,49 +3,44 @@
 from eigenpairflow.correction import ogita_aishima_refinement
 
 
+def _residual_norm(A: np.ndarray, X: np.ndarray) -> float:
+    """固有値分解の残差ノルムを計算する。"""
+    D = np.diag(np.diag(X.T @ A @ X))
+    return np.linalg.norm(A @ X - X @ D)
+
+
 def test_ogita_aishima_refinement_simple():
-    # Create a random symmetric matrix
+    rng = np.random.default_rng(0)
     n = 5
-    A = np.random.rand(n, n)
+    A = rng.standard_normal((n, n))
     A = (A + A.T) / 2
 
-    # Get the exact solution
-    exact_eigvals, exact_eigvecs = scipy.linalg.eigh(A)
-
-    # Create a perturbed initial guess
-    X_hat = exact_eigvecs + np.random.rand(n, n) * 0.1
+    # Exact eigenvectors for constructing a perturbed initial guess
+    _, exact_eigvecs = scipy.linalg.eigh(A)
+    X_hat = exact_eigvecs + rng.standard_normal((n, n)) * 0.1
 
-    # Refine the solution
     X_refined, D_refined = ogita_aishima_refinement(A, X_hat)
 
-    # Check if the refined solution is closer to the exact solution
-    initial_error = np.linalg.norm(X_hat - exact_eigvecs)
-    refined_error = np.linalg.norm(X_refined - exact_eigvecs)
+    initial_residual = _residual_norm(A, X_hat)
+    refined_residual = np.linalg.norm(A @ X_refined - X_refined @ D_refined)
 
-    assert refined_error < initial_error
+    assert refined_residual < initial_residual
 
 
 def test_ogita_aishima_refinement_clustered():
-    # Create a symmetric matrix with clustered eigenvalues
     n = 5
-    A = np.diag([1.0, 1.0 + 1e-9, 3.0, 4.0, 5.0])
-
-    # Add some off-diagonal noise
-    noise = np.random.rand(n, n) * 1e-12
-    noise = (noise + noise.T) / 2
-    A = A + noise
-
-    # Get the exact solution
-    exact_eigvals, exact_eigvecs = scipy.linalg.eigh(A)
+    A = np.diag([1.0, 1.0 + 1e-9, 3.0, 4.0, 5.0]).astype(float)
+    # deterministic tiny coupling for the first two eigenvalues
+    A[0, 1] = A[1, 0] = 1e-12
 
-    # Create a perturbed initial guess
-    X_hat = exact_eigvecs + np.random.rand(n, n) * 0.1
+    # Exact eigenvectors and a perturbed initial guess
+    _, exact_eigvecs = scipy.linalg.eigh(A)
+    rng = np.random.default_rng(0)
+    X_hat = exact_eigvecs + rng.standard_normal((n, n)) * 0.1
 
-    # Refine the solution
     X_refined, D_refined = ogita_aishima_refinement(A, X_hat)
 
-    # Check if the refined solution is closer to the exact solution
-    initial_error = np.linalg.norm(X_hat - exact_eigvecs)
-    refined_error = np.linalg.norm(X_refined - exact_eigvecs)
+    initial_residual = _residual_norm(A, X_hat)
+    refined_residual = np.linalg.norm(A @ X_refined - X_refined @ D_refined)
 
-    assert refined_error < initial_error
+    assert refined_residual < initial_residual