From 365ec8196d5fa62bca43bc40cda577c68e7c1b43 Mon Sep 17 00:00:00 2001
From: Nuwan Yapa <yapanuwan@gmail.com>
Date: Mon, 6 Jan 2025 18:35:19 -0500
Subject: [PATCH] Gram-Schmidt improved

---
 helper.jl    | 37 +++++++++++++++++++++++--------------
 test/misc.jl |  4 ++--
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/helper.jl b/helper.jl
index 6eeee4a..40656e3 100644
--- a/helper.jl
+++ b/helper.jl
@@ -49,25 +49,34 @@ function exportCSV(file::String, data, labels=nothing)
     end
 end
 
-"c-orthonormalization via Gram-Schmidt"
-function gram_schmidt(vecs, ws, precision=1e-10, verbose=false)
+"In-place c-orthonormalization via (modified) Gram-Schmidt. Most significant vectors determined by the singular values (compared to the threshold) are returned."
+function gram_schmidt!(vecs, ws, threshold=1e-5; verbose=false)
     c_product(i, j) = sum(vecs[i] .* ws .* vecs[j])
-    proj(i, j) = (c_product(i, j) / c_product(i, i)) .* vecs[i] # component of vec[i] in vec[j]
+    norm(i) = c_product(i, i)
+    proj(i, j) = (c_product(i, j) / norm(i)) .* vecs[i] # component of vec[i] in vec[j]
 
-    mask = ones(Bool, size(vecs)) # vectors to retain
+    # initial normalization
     for i in eachindex(vecs)
-        for j in findall(mask[1 : (i - 1)])
-            vecs[i] -= proj(j, i)
-        end
-        norm = c_product(i, i)
-        verbose && println("Mag. of basis vector $i = $(abs(norm))")
-        if abs(norm) > precision
-            vecs[i] = vecs[i] ./ sqrt(norm)
-        else
-            mask[i] = false
+        vecs[i] ./= sqrt(norm(i))
+    end
+
+    target_dim = c_rank(vecs, ws, threshold)
+
+    for i in eachindex(vecs)
+        for j in (i + 1):length(vecs)
+            vecs[j] .-= proj(i, j)
         end
     end
-    return vecs[mask]
+    
+    abs_norms = abs.(norm.(eachindex(vecs)))
+    verbose && println("Absolute norms = $(round.(abs_norms; sigdigits=1))")
+    selected_vecs_i = partialsortperm(abs_norms, 1:target_dim; rev=true)
+
+    for i in selected_vecs_i
+        vecs[i] ./= sqrt(norm(i))
+    end
+
+    return vecs[selected_vecs_i]
 end
 
 "Rank of a basis set (pre-normalized) under c-product"
diff --git a/test/misc.jl b/test/misc.jl
index 19d47db..7ff7d12 100644
--- a/test/misc.jl
+++ b/test/misc.jl
@@ -23,7 +23,7 @@ evals = eigvals(H_EC, N_EC)
 println("Eigenvalues with GEVP:")
 display(evals)
 
-ortho_basis = hcat(gram_schmidt(vecs, ws)...)
+ortho_basis = hcat(gram_schmidt!(vecs, ws)...)
 N_ortho = transpose(ortho_basis) * ws_mat * ortho_basis
 println("Norm matrix after Gram-Schmidt:")
 display(round.(N_ortho; digits=2)) # should be ≈I
@@ -53,7 +53,7 @@ println("Dimensionality before Gram-Schmidt = $(length(vecs))")
 println("Absolute singular values before Gram-Schmidt:")
 display(abs.(c_singular_values(vecs, ws)))
 
-ortho_vecs = gram_schmidt(vecs, ws)
+ortho_vecs = gram_schmidt!(vecs, ws; verbose=true)
 ortho_basis = hcat(ortho_vecs...)
 println("Dimensionality after Gram-Schmidt = $(length(ortho_vecs))")