From 2f8aeefc4d285b47dc2a6812d74a87e8bd1da752 Mon Sep 17 00:00:00 2001 From: Antoine Levitt Date: Fri, 4 Mar 2022 19:06:51 +0100 Subject: [PATCH] Refactoring diagonalize_all_kblocks (#617) --- src/eigen/diag.jl | 21 ++++++++++----------- src/fft.jl | 2 +- src/orbitals.jl | 4 ++++ src/scf/direct_minimization.jl | 3 +-- src/scf/self_consistent_field.jl | 2 +- test/compute_density.jl | 2 +- 6 files changed, 18 insertions(+), 16 deletions(-) diff --git a/src/eigen/diag.jl b/src/eigen/diag.jl index a661b21a1a..38a555f262 100644 --- a/src/eigen/diag.jl +++ b/src/eigen/diag.jl @@ -9,7 +9,7 @@ that really does the work, operating on a single ``k``-Block. `prec_type` should be a function that returns a preconditioner when called as `prec(ham, kpt)` """ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint::Int; - guess=nothing, + ψguess=nothing, prec_type=PreconditionerTPA, interpolate_kpoints=true, tol=1e-6, miniter=1, maxiter=100, n_conv_check=nothing, show_progress=false) @@ -22,31 +22,30 @@ function diagonalize_all_kblocks(eigensolver, ham::Hamiltonian, nev_per_kpoint:: progress = Progress(length(kpoints), desc="Diagonalising Hamiltonian kblocks: ") end for (ik, kpt) in enumerate(kpoints) - n_Gk = length(G_vectors(ham.basis, kpoints[ik])) + n_Gk = length(G_vectors(ham.basis, kpt)) if n_Gk < nev_per_kpoint error("The size of the plane wave basis is $n_Gk, and you are asking for " * "$nev_per_kpoint eigenvalues. Increase Ecut.") end - # Get guessk + # Get ψguessk @timing "QR orthonormalization" begin - if guess !== nothing - # guess provided - guessk = guess[ik] + if ψguess !== nothing + # ψguess provided + ψguessk = ψguess[ik] elseif interpolate_kpoints && ik > 1 # use information from previous k-point X0 = interpolate_kpoint(results[ik - 1].X, ham.basis, kpoints[ik - 1], ham.basis, kpoints[ik]) - guessk = ortho_qr(X0) # Re-orthogonalize and renormalize + ψguessk = ortho_qr(X0) # Re-orthogonalize and renormalize else - # random initial guess - guessk = ortho_qr(randn(T, n_Gk, nev_per_kpoint)) + ψguessk = random_orbitals(ham.basis, kpt, nev_per_kpoint) end end - @assert size(guessk) == (n_Gk, nev_per_kpoint) + @assert size(ψguessk) == (n_Gk, nev_per_kpoint) prec = nothing prec_type !== nothing && (prec = prec_type(ham.basis, kpt)) - results[ik] = eigensolver(ham.blocks[ik], guessk; + results[ik] = eigensolver(ham.blocks[ik], ψguessk; prec=prec, tol=tol, miniter=miniter, maxiter=maxiter, n_conv_check=n_conv_check) diff --git a/src/fft.jl b/src/fft.jl index 3433427370..2ce0a46a04 100644 --- a/src/fft.jl +++ b/src/fft.jl @@ -7,7 +7,7 @@ import FFTW # For densities and potentials defined on the cubic basis set, r_to_G/G_to_r # do a simple FFT/IFFT from the cubic basis set to the real-space grid. -# These function do not take a k-point as input +# These functions do not take a k-point as input # For orbitals, G_to_r converts the orbitals defined on a spherical # basis set to the cubic basis set using zero padding, then performs diff --git a/src/orbitals.jl b/src/orbitals.jl index 4c3fdf49ed..d9980b7907 100644 --- a/src/orbitals.jl +++ b/src/orbitals.jl @@ -52,3 +52,7 @@ function unsafe_unpack_ψ(x, sizes_ψ) end end unpack_ψ(x, sizes_ψ) = deepcopy(unsafe_unpack_ψ(x, sizes_ψ)) + +function random_orbitals(basis::PlaneWaveBasis{T}, kpt::Kpoint, howmany) where {T} + ortho_qr(randn(Complex{T}, length(G_vectors(basis, kpt)), howmany)) +end diff --git a/src/scf/direct_minimization.jl b/src/scf/direct_minimization.jl index 040c2812bb..138aa5fd95 100644 --- a/src/scf/direct_minimization.jl +++ b/src/scf/direct_minimization.jl @@ -81,8 +81,7 @@ function direct_minimization(basis::PlaneWaveBasis{T}, ψ0; Nk = length(basis.kpoints) if ψ0 === nothing - ψ0 = [ortho_qr(randn(Complex{T}, length(G_vectors(basis, kpt)), n_bands)) - for kpt in basis.kpoints] + ψ0 = [random_orbitals(basis, kpt, n_bands) for kpt in basis.kpoints] end occupation = [filled_occ * ones(T, n_bands) for ik = 1:Nk] diff --git a/src/scf/self_consistent_field.jl b/src/scf/self_consistent_field.jl index cbe9cf6469..4e03b80830 100644 --- a/src/scf/self_consistent_field.jl +++ b/src/scf/self_consistent_field.jl @@ -21,7 +21,7 @@ function next_density(ham::Hamiltonian; end # Diagonalize - eigres = diagonalize_all_kblocks(eigensolver, ham, n_bands + n_ep_extra; guess=ψ, + eigres = diagonalize_all_kblocks(eigensolver, ham, n_bands + n_ep_extra; ψguess=ψ, n_conv_check=n_bands, kwargs...) eigres.converged || (@warn "Eigensolver not converged" iterations=eigres.iterations) diff --git a/test/compute_density.jl b/test/compute_density.jl index 74ab56b276..8396134f75 100644 --- a/test/compute_density.jl +++ b/test/compute_density.jl @@ -28,7 +28,7 @@ if mpi_nprocs() == 1 # not easy to distribute for it in 1:n_rounds ham = Hamiltonian(basis; ρ=ρnew) - res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol=tol, guess=res.X) + res = diagonalize_all_kblocks(lobpcg_hyper, ham, n_bands; tol=tol, ψguess=res.X) occ, εF = DFTK.compute_occupation(basis, res.λ) ρnew = compute_density(basis, res.X, occ)