From b7a274cca1515e43e2878b78e77d282a1f3e302b Mon Sep 17 00:00:00 2001
From: Benjamin Cohen-Stead <benwcs@gmail.com>
Date: Thu, 19 Sep 2024 14:34:17 -0400
Subject: [PATCH] Add Lanczos

---
 Project.toml        |   4 +-
 docs/src/api.md     |   8 +++-
 src/SmoQyKPMCore.jl |   4 ++
 src/lanczos.jl      | 113 ++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 126 insertions(+), 3 deletions(-)
 create mode 100644 src/lanczos.jl

diff --git a/Project.toml b/Project.toml
index e37ec86..59a5c7f 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "1.0.0-DEV"
 [deps]
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -14,10 +15,11 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 SmoQyKPMCoreCUDA = "CUDA"
 
 [compat]
+CUDA = "5.5"
 FFTW = "1.8"
 LinearAlgebra = "1.10"
+Random = "1.10"
 julia = "1.10"
-CUDA = "5.5"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/docs/src/api.md b/docs/src/api.md
index 74e6679..1266ec9 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -14,8 +14,10 @@ level one that does not.
 - [`kpm_mul!`](@ref)
 - [`kpm_eval`](@ref)
 - [`kpm_eval!`](@ref)
-- [`apply_jackson_kernel!`](@ref)
 - [`apply_jackson_kernel`](@ref)
+- [`apply_jackson_kernel!`](@ref)
+- [`lanczos`](@ref)
+- [`lanczos!`](@ref)
 
 ```@docs
 KPMExpansion
@@ -29,6 +31,8 @@ kpm_mul
 kpm_mul!
 kpm_eval
 kpm_eval!
-apply_jackson_kernel!
 apply_jackson_kernel
+apply_jackson_kernel!
+lanczos
+lanczos!
 ```
\ No newline at end of file
diff --git a/src/SmoQyKPMCore.jl b/src/SmoQyKPMCore.jl
index 4bc5fff..1188e7c 100644
--- a/src/SmoQyKPMCore.jl
+++ b/src/SmoQyKPMCore.jl
@@ -1,6 +1,7 @@
 module SmoQyKPMCore
 
 using LinearAlgebra
+using Random
 using FFTW
 
 # functional api implementation of KPM algorithm
@@ -14,6 +15,9 @@ include("KPMExpansion.jl")
 export KPMExpansion
 export update_kpmexpansion!, update_kpmexpansion_bounds!, update_kpmexpansion_order!
 
+include("lanczos.jl")
+export lanczos, lanczos!
+
 # ensure FFTW uses only a single thread
 function __init__()
     FFTW.set_num_threads(1)
diff --git a/src/lanczos.jl b/src/lanczos.jl
new file mode 100644
index 0000000..9334df2
--- /dev/null
+++ b/src/lanczos.jl
@@ -0,0 +1,113 @@
+
+# Use Lanczos iterations to find a truncated tridiagonal representation of A S,
+# up to similarity transformation. Here, A is any Hermitian matrix, while S is
+# both Hermitian and positive definite. Traditional Lanczos uses the identity
+# matrix, S = I. The extension to non-identity matrices S is as follows: Each
+# matrix-vector product A v becomes A S v, and each vector inner product w† v
+# becomes w† S v. The implementation below follows Wikipedia, and is the most
+# stable of the four variants considered by Paige [1]. This implementation
+# introduces additional vector storage so that each Lanczos iteration requires
+# only one matrix-vector multiplication for A and S, respectively.
+
+# Similar generalizations of Lanczos have been considered in [2] and [3].
+#
+# 1. C. C. Paige, IMA J. Appl. Math., 373-381 (1972),
+# https://doi.org/10.1093%2Fimamat%2F10.3.373.
+# 2. H. A. van der Vorst, Math. Comp. 39, 559-561 (1982),
+# https://doi.org/10.1090/s0025-5718-1982-0669648-0
+# 3. M. Grüning, A. Marini, X. Gonze, Comput. Mater. Sci. 50, 2148-2156 (2011),
+# https://doi.org/10.1016/j.commatsci.2011.02.021.
+
+@doc raw"""
+    lanczos(niters, v, A, S = I, rng = Random.default_rng())
+
+Use `niters` Lanczos iterations to find a truncated tridiagonal representation of ``A\cdot S``, up to similarity transformation.
+Here, ``A`` is any Hermitian matrix, while ``S`` is both Hermitian and positive definite.
+Traditional Lanczos uses the identity matrix, ``S = I``.
+The extension to non-identity matrices ``S`` is as follows:
+Each matrix-vector product ``A\cdot v`` becomes ``(A S) \cdot v``, and each vector inner product ``w^\dagger \cdot v`` becomes ``w^\dagger \cdot S \cdot v``.
+The implementation below follows Wikipedia, and is the most stable of the four variants considered by Paige [1].
+This implementation introduces additional vector storage so that each Lanczos iteration requires only one matrix-vector multiplication for ``A`` and ``S``, respectively.
+
+This function returns a [`SymTridiagonal`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.SymTridiagonal) matrix.
+Note that the [`eigmin`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.eigmin)
+and [`eigmax`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.eigmax) routines have
+[specialized implementations](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#Elementary-operations) for a
+[`SymTridiagonal`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.SymTridiagonal) matrix type.
+"""
+function lanczos(niters, v, A, S = I, rng = Random.default_rng())
+
+    αs = zeros( real(eltype(v)) , niters)
+    βs = zeros( real(eltype(v)) , niters-1)
+    tmp = zeros(size(v)..., 5)
+    symtridiag = lanczos!(αs, βs, v, A, S, tmp, rng)
+
+    return symtridiag
+end
+
+@doc raw"""
+    lanczos!(
+        αs::AbstractVector, βs::AbstractVector, v::AbstractVector,
+        A, S = I,
+        tmp::AbstractMatrix = zeros(length(v), 5),
+        rng = Random.default_rng()
+    )
+
+Use Lanczos iterations to find a truncated tridiagonal representation of ``A\cdot S``, up to similarity transformation.
+Here, ``A`` is any Hermitian matrix, while ``S`` is both Hermitian and positive definite.
+Traditional Lanczos uses the identity matrix, ``S = I``.
+The extension to non-identity matrices ``S`` is as follows:
+Each matrix-vector product ``A\cdot v`` becomes ``(A S) \cdot v``, and each vector inner product ``w^\dagger \cdot v`` becomes ``w^\dagger \cdot S \cdot v``.
+The implementation below follows Wikipedia, and is the most stable of the four variants considered by Paige [1].
+This implementation introduces additional vector storage so that each Lanczos iteration requires only one matrix-vector multiplication for ``A`` and ``S``, respectively.
+
+The number of Lanczos iterations performed equals `niters = length(αs)`, and `niters - 1 == length(βs)`.
+This function returns a [`SymTridiagonal`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.SymTridiagonal) matrix
+based on the contents of the vectors `αs` and `βs`. Note that the [`eigmin`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.eigmin)
+and [`eigmax`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.eigmax) routines have
+[specialized implementations](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#Elementary-operations) for a
+[`SymTridiagonal`](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.SymTridiagonal) matrix type.
+"""
+function lanczos!(
+    αs::AbstractVector, βs::AbstractVector, v::AbstractVector,
+    A, S = I,
+    tmp::AbstractMatrix = zeros(length(v), 5),
+    rng = Random.default_rng()
+)
+
+    (vp, Sv, Svp, w, Sw) = eachcol(tmp)
+
+    niters = length(αs)
+    @assert niters - 1 == length(βs)
+    
+    if iszero(norm(v))
+        randn!(rng, v)
+    end
+
+    _matrix_multiply!(Sv, S, v)
+    c = sqrt(dot(v, Sv))
+    @. v /= c
+    @. Sv /= c
+
+    _matrix_multiply!(w, A, Sv)
+    α = dot(w, Sv)
+    @. w = w - α * v
+    _matrix_multiply!(Sw, S, w)
+    αs[1] = α
+
+    for n in 2:niters
+        β = sqrt(dot(Sw, w))
+        iszero(β) && break
+        @. vp = w / β
+        @. Svp = Sw / β
+        _matrix_multiply!(w, A, Svp)
+        α = dot(w, Svp)
+        @. w = w - α * vp - β * v
+        _matrix_multiply!(Sw, S, w)
+        @. v = vp
+        αs[n] = α
+        βs[n-1] = β
+    end
+
+    return SymTridiagonal(αs, βs)
+end
\ No newline at end of file