Skip to content

Commit

Permalink
Add Lanczos
Browse files Browse the repository at this point in the history
  • Loading branch information
cohensbw committed Sep 19, 2024
1 parent ec5d2ec commit b7a274c
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 3 deletions.
4 changes: 3 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand Down
8 changes: 6 additions & 2 deletions docs/src/api.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -29,6 +31,8 @@ kpm_mul
kpm_mul!
kpm_eval
kpm_eval!
apply_jackson_kernel!
apply_jackson_kernel
apply_jackson_kernel!
lanczos
lanczos!
```
4 changes: 4 additions & 0 deletions src/SmoQyKPMCore.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module SmoQyKPMCore

using LinearAlgebra
using Random
using FFTW

# functional api implementation of KPM algorithm
Expand All @@ -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)
Expand Down
113 changes: 113 additions & 0 deletions src/lanczos.jl
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit b7a274c

Please sign in to comment.