Refactoring diagonalize_all_kblocks (JuliaMolSim#617)

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)
 

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

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

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]
 

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)
 

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)