Preconditioned LOBPCG solver for large non-periodic problems (v0.4.0)

Project.toml

@@ -1,12 +1,13 @@
 name = "Numerov"
 uuid = "19e90f0b-5b53-4501-bfb7-be6aec03b1b3"
-version = "0.3.1"
 authors = ["Jakob Gamper <97gamjak@gmail.com>"]
+version = "0.4.0"
 
 [deps]
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
 Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
@@ -22,6 +23,7 @@ Aqua = "0.8"
 Arpack = "0.5"
 Comonicon = "1"
 DelimitedFiles = "1"
+IterativeSolvers = "0.9"
 KrylovKit = "0.10"
 LinearAlgebra = "1.10"
 PhysicalConstants = "0.2"

README.md

@@ -157,7 +157,7 @@ The only required keyword is `potential-file`. All others have defaults:
 | `k-points` | Number of k-points sampled per direction between the Gamma point and the Brillouin-zone boundary; if omitted, a single calculation at k = 0 is performed | integer > 1 | not set |
 | `datapoints` | Number of grid points per dimension, comma- or space-separated (e.g. `datapoints = 20, 30`) | integers | required for 2D/3D; in 1D taken from the potential file |
 | `band-structure` | Compute the band structure along the path through the high-symmetry points of the Brillouin zone (requires `k-points`) | `on`, `true`, `off`, `false` | `off` |
-| `solver` | Eigensolver backend | `arpack`, `krylov`, `lu` | `arpack` |
+| `solver` | Eigensolver backend (`lobpcg` is recommended for large non-periodic 3D problems) | `arpack`, `krylov`, `lobpcg`, `lu` | `arpack` |
 | `output-file` | Name of the log file | file path | `Numerov.out` |
 | `timings-file` | Name of the timings file | file path | `timings.out` |
 | `read-k-points` | Read the k-points from a file instead of generating them [^3] | `true`, `false` | `false` |

benchmark/benchmarks.jl

@@ -52,6 +52,14 @@ function write_3d_harmonic(dir::String, n::Int)
     end
 end
 
+function run_3d_harmonic_lobpcg(n::Int)
+    xs = range(-5.5, 5.5; length = n)
+    V = [0.5 * (x^2 + y^2 + z^2) for x in xs, y in xs, z in xs]
+    redirect_stdout(devnull) do
+        solve_schrodinger(V, (xs, xs, xs); n_eigenvalues = 4, solver = :lobpcg)
+    end
+end
+
 function run_3d_harmonic(n::Int)
     mktempdir() do tmp
         write_3d_harmonic(tmp, n)
@@ -94,6 +102,7 @@ SUITE["solve"] = BenchmarkGroup()
 SUITE["solve"]["1D_harmonic_201"] = @benchmarkable run_case("1DHarmonicOscillator") seconds = 30 samples = 5
 SUITE["solve"]["2D_water"]        = @benchmarkable run_case("2DWater") seconds = 60 samples = 3
 SUITE["solve"]["3D_harmonic_15"]  = @benchmarkable run_3d_harmonic(15) seconds = 120 samples = 3
+SUITE["solve"]["3D_harmonic_25_lobpcg"] = @benchmarkable run_3d_harmonic_lobpcg(25) seconds = 120 samples = 3
 
 SUITE["bandstructure"] = BenchmarkGroup()
 SUITE["bandstructure"]["1D_kronigpenney_10k"] =

docs/src/input.md

@@ -40,7 +40,7 @@ The only required keyword is `potential-file`. All others have defaults:
 | `k-points` | Number of k-points sampled per direction between the Gamma point and the Brillouin-zone boundary; if omitted, a single calculation at k = 0 is performed | integer > 1 | not set |
 | `datapoints` | Number of grid points per dimension, comma- or space-separated (e.g. `datapoints = 20, 30`) | integers | required for 2D/3D; in 1D taken from the potential file |
 | `band-structure` | Compute the band structure along the path through the high-symmetry points of the Brillouin zone (requires `k-points`) | `on`, `true`, `off`, `false` | `off` |
-| `solver` | Eigensolver backend | `arpack`, `krylov`, `lu` | `arpack` |
+| `solver` | Eigensolver backend (`lobpcg` is recommended for large non-periodic 3D problems) | `arpack`, `krylov`, `lobpcg`, `lu` | `arpack` |
 | `output-file` | Name of the log file | file path | `Numerov.out` |
 | `timings-file` | Name of the timings file | file path | `timings.out` |
 | `read-k-points` | Read the k-points from a file instead of generating them (accepted but not yet implemented) | `true`, `false` | `false` |

src/Numerov.jl

@@ -11,6 +11,7 @@ module Numerov
     using DelimitedFiles
     using TimerOutputs
     using KrylovKit
+    using IterativeSolvers: lobpcg
     using StatsBase
 
     import PhysicalConstants.CODATA2018: h, ħ, N_A, c_0
@@ -60,6 +61,7 @@ module Numerov
     include("setupSystem.jl")
 
     include("normalize.jl")
+    include("preconditioner.jl")
     include("solve.jl")
 
     include("main.jl")

src/api.jl

@@ -54,6 +54,7 @@ end
 const SOLVER_NAMES = Dict(
     :arpack => ARPACK,
     :krylov => KRYLOV,
+    :lobpcg => LOBPCG,
     :lu     => LU,
 )
 
@@ -84,7 +85,8 @@ writing any files.
   real, non-periodic problem.
 - `stencil = 9`: finite-difference stencil size (`3`, `5`, `7`, `9`, `11` or
   `13`); `stencil_laplace` and `stencil_nabla` override it individually.
-- `solver = :arpack`: eigensolver backend (`:arpack`, `:krylov` or `:lu`).
+- `solver = :arpack`: eigensolver backend (`:arpack`, `:krylov`, `:lobpcg` or
+  `:lu`); `:lobpcg` is recommended for large non-periodic 3D problems.
 - `potential_unit = u"hartree"`, `coord_unit = u"bohr"`, `mass_unit = u"m_e"`:
   units of the inputs; energies are returned in `potential_unit`.
 
@@ -233,9 +235,12 @@ function setup_problem(V::AbstractArray{<:Real}, coords;
     all(isfinite, V) || throw(ArgumentError("the potential contains non-finite values"))
 
     haskey(SOLVER_NAMES, solver) ||
-        throw(ArgumentError("unknown solver :$solver - valid options are :arpack, :krylov and :lu"))
-    solver === :arpack && n_eigenvalues + 5 >= length(V) &&
-        throw(ArgumentError("the arpack solver needs n_eigenvalues + 5 < number of grid points ($(length(V)))"))
+        throw(ArgumentError("unknown solver :$solver - valid options are :arpack, :krylov, :lobpcg and :lu"))
+    # :lobpcg transparently falls back to solve_arpack, and :krylov uses the
+    # same nev, so all three solvers share arpack's nev < N requirement - only
+    # :lu diagonalizes the full dense matrix and has no such limit
+    solver in (:arpack, :lobpcg, :krylov) && n_eigenvalues + 5 >= length(V) &&
+        throw(ArgumentError("the $(solver) solver needs n_eigenvalues + 5 < number of grid points ($(length(V)))"))
 
     stencil         in (3, 5, 7, 9, 11, 13) || throw(ArgumentError("stencil has to be 3, 5, 7, 9, 11 or 13"))
     stencil_laplace in (3, 5, 7, 9, 11, 13) || throw(ArgumentError("stencil-laplace has to be 3, 5, 7, 9, 11 or 13"))

src/checkInput/checkSystem.jl

@@ -28,12 +28,14 @@ function checkSolver(system::System)
     isempty(solver)    && (system.solver = ARPACK ; return) #write to log file about default setting
     solver == "arpack" && (system.solver = ARPACK ; return)
     solver == "krylov" && (system.solver = KRYLOV ; return)
+    solver == "lobpcg" && (system.solver = LOBPCG; return)
     solver == "cuda"   && throw(ArgumentError("the cuda solver is not implemented -- use arpack, krylov or lu"))
     solver == "lu"     && (system.solver = LU     ; return)
 
     throw(ArgumentError("\nThe given solver $(inputDictionary["solver"]) was not recognised!\n" *
                         "Valid options are:                                               \n" *
                         "    - arpack                                                     \n" *
                         "    - krylov                                                     \n" *
+                        "    - lobpcg                                                     \n" *
                         "    - lu                                                         \n"))
 end

src/datatypes/SolverEnum.jl

@@ -1,6 +1,7 @@
 @enum SolverEnum begin
     ARPACK
     KRYLOV
+    LOBPCG
     GPU
     LU
 end

src/preconditioner.jl

@@ -0,0 +1,135 @@
+"""
+    KineticPreconditioner
+
+Tensor-product preconditioner for the shifted kinetic operator, used to
+precondition the LOBPCG eigensolver.
+
+`T̃`, the Kronecker SUM of the per-dimension 1D operators
+`t_d = -Δ_d / (2 Δq_d²)`, replaces the true (possibly non-separable) kinetic
+operator: `T̃`'s eigendecomposition factorizes into the per-dimension
+eigenpairs, so `(T̃ + σI)⁻¹ x` is applied exactly - with one small dense
+eigenbasis transform per dimension - without ever factorizing a full-size
+operator, so there is no fill-in.
+
+`T̃` equals the true production kinetic operator exactly for 1D problems
+(trivially) and for the 2D **5-point** stencil (`buildLaplace_2d`'s 5-point
+Laplacian is a pure row/column "cross" pattern, which is separable: dividing
+by `2^(dimension-1)`, as `solve()` does, recovers the Kronecker sum exactly).
+Every other case - 2D with stencil 3, 7, 9 (the package default) or 11, and
+3D with any stencil (`buildLaplace_3d` always uses a more elaborate,
+non-separable stencil) - is only an APPROXIMATION; this does not compromise
+correctness, only convergence speed, since every LOBPCG result is
+independently verified against the true Hamiltonian's residual in
+`solveWrapper` (and re-solved with Arpack if that check fails) regardless of
+how good an approximation `T̃` is.
+"""
+struct KineticPreconditioner
+    Q     ::Vector{Matrix{Float64}}   # eigenbasis per dimension, d = 1..D
+    denom ::Array{Float64}            # Σ_d λ_d + σ, shaped (n_D, ..., n_1)
+end
+
+"""
+1D kinetic matrix `-Δ/(2 Δq²)` for one dimension, built with the same stencil
+machinery as the full operator.
+"""
+function kinetic_1d(n::Int, spacing::Float64, periodic::Bool, stencil::Int)
+    pot1              = Potential()
+    pot1.dimension    = 1
+    pot1.n_datapoints = [n]
+    pot1.periodic     = [periodic]
+
+    sys1              = System()
+    sys1.n_datapoints = pot1.n_datapoints
+    sys1.periodic     = pot1.periodic
+    sys1.reciprocal   = false
+    sys1.stencil      = stencil
+    sys1.stencilΔ     = stencil
+    sys1.stencil∇     = stencil
+
+    buildΔ(sys1, pot1)
+
+    return Symmetric(0.5 .* (-Matrix(sys1.Δ) ./ spacing^2))
+end
+
+"""
+    KineticPreconditioner(potential, system; σ = 1.0)
+
+Build the preconditioner for the given problem. `σ > 0` keeps the operator
+safely positive definite; its exact value only affects convergence speed.
+"""
+function KineticPreconditioner(potential::Potential, system::System; σ::Float64 = 1.0)
+    D    = potential.dimension
+    dims = potential.n_datapoints
+
+    Q       = Vector{Matrix{Float64}}(undef, D)
+    lambdas = Vector{Vector{Float64}}(undef, D)
+    for d in 1:D
+        e = eigen(kinetic_1d(dims[d], potential.intervall[d], potential.periodic[d], system.stencilΔ))
+        Q[d], lambdas[d] = e.vectors, e.values
+    end
+
+    # flattened grid ordering: the LAST dimension varies fastest, so array
+    # axis k corresponds to dimension D + 1 - k
+    denom = zeros(reverse(Tuple(dims)))
+    for idx in CartesianIndices(denom)
+        acc = σ
+        for k in 1:D
+            acc += lambdas[D + 1 - k][idx[k]]
+        end
+        denom[idx] = acc
+    end
+
+    return KineticPreconditioner(Q, denom)
+end
+
+"""
+Apply `(T̃ + σI)⁻¹` to a flattened grid vector via per-dimension eigenbasis
+transforms.
+"""
+function apply_preconditioner!(y::AbstractVector, P::KineticPreconditioner, x::AbstractVector)
+    D    = length(P.Q)
+    dims = size(P.denom)
+    X    = reshape(copy(convert(Vector{Float64}, x)), dims)
+
+    X = transform_modes(X, P, adjoint)
+    X ./= P.denom
+    X = transform_modes(X, P, identity)
+
+    y .= vec(X)
+    return y
+end
+
+"Multiply every mode of `X` by `op(Q_d)` for its dimension's eigenbasis."
+function transform_modes(X::AbstractArray, P::KineticPreconditioner, op)
+    D = length(P.Q)
+    for k in 1:D
+        d    = D + 1 - k                     # dimension of array axis k
+        perm = (k, setdiff(1:D, k)...)
+        Xp   = permutedims(X, perm)
+        sz   = size(Xp)
+        Xm   = op(P.Q[d]) * reshape(Xp, sz[1], :)
+        X    = permutedims(reshape(Xm, sz), invperm(collect(perm)))
+    end
+    return X
+end
+
+LinearAlgebra.ldiv!(P::KineticPreconditioner, x::AbstractVector) =
+    apply_preconditioner!(x, P, copy(x))
+
+function LinearAlgebra.ldiv!(P::KineticPreconditioner, X::AbstractMatrix)
+    for j in axes(X, 2)
+        col = view(X, :, j)
+        apply_preconditioner!(col, P, copy(col))
+    end
+    return X
+end
+
+LinearAlgebra.ldiv!(y::AbstractVector, P::KineticPreconditioner, x::AbstractVector) =
+    apply_preconditioner!(y, P, x)
+
+function LinearAlgebra.ldiv!(Y::AbstractMatrix, P::KineticPreconditioner, X::AbstractMatrix)
+    for j in axes(X, 2)
+        apply_preconditioner!(view(Y, :, j), P, view(X, :, j))
+    end
+    return Y
+end

src/setupSystem.jl

@@ -13,6 +13,12 @@ function setupSystem(potential::Potential, system::System)
     system.reciprocal && !any(system.periodic) && throw(ArgumentError("You have defined a number of k-points - this option is only valid in combination with \"periodic = true\""))
     any(system.n_datapoints .< system.stencil) && throw(ArgumentError("The number of datapoints in each dimension has at least to be equal to the stencil size!"))
 
+    # solveWrapper only ever builds a real Hamiltonian for lobpcg, so a
+    # reciprocal (k-point) run has to be rejected here - before any output
+    # file is written or an existing eigenvalues.dat is removed - rather than
+    # deep inside solve() on the first k-point
+    system.reciprocal && system.solver == LOBPCG && throw(ArgumentError("the lobpcg solver supports non-periodic (real symmetric) problems - use arpack or krylov for periodic k-point runs"))
+
     ###############################################################
     #                                                             #
     # set stencil for laplace and nabla if not defined seperately #