Optimize Ω+K a little (#1266)

Author: Technici4n

src/densities.jl

@@ -85,11 +85,13 @@ end
         (ik, n) = kn
 
         kpt = basis.kpoints[ik]
-        ifft!(storage.ψnk_real, basis, kpt, ψ[ik][:, n])
+        ifft!(storage.ψnk_real, basis, kpt, ψ[ik][:, n]; normalize=false)
         # … and then we compute the real Fourier transform in the adequate basis.
-        ifft!(storage.δψnk_real, basis, δψ_plus_k[ik].kpt, δψ_plus_k[ik].ψk[:, n])
+        ifft!(storage.δψnk_real, basis, δψ_plus_k[ik].kpt, δψ_plus_k[ik].ψk[:, n]; normalize=false)
+        # use unnormalized plans for extra speed, normalize at the end
+        ifft_normalization = basis.fft_grid.ifft_normalization
 
-        storage.δρ[:, :, :, kpt.spin] .+= real_qzero.(
+        storage.δρ[:, :, :, kpt.spin] .+= ifft_normalization^2 .* real_qzero.(
             2 .* occupation[ik][n]  .* basis.kweights[ik] .* conj.(storage.ψnk_real)
                                                           .* storage.δψnk_real
               .+ δoccupation[ik][n] .* basis.kweights[ik] .* abs2.(storage.ψnk_real))

src/response/hessian.jl

@@ -29,8 +29,12 @@ from ψ and Λ is the set of Rayleigh coefficients ψk' * Hk * ψk at each k-poi
 """
 @timing function apply_Ω(δψ, ψ, H::Hamiltonian, Λ)
     δψ = proj_tangent(δψ, ψ)
-    Ωδψ = [H.blocks[ik] * δψk - δψk * Λ[ik] for (ik, δψk) in enumerate(δψ)]
-    proj_tangent!(Ωδψ, ψ)
+    map(enumerate(δψ)) do (ik, δψk)
+        Ωδψ = H.blocks[ik] * δψk
+        mul!(Ωδψ, δψk, Λ[ik], -1, 1)
+        proj_tangent_kpt!(Ωδψ, ψ[ik])
+        Ωδψ
+    end
 end
 
 """
@@ -48,16 +52,18 @@ Compute the application of K defined at ψ to δψ. ρ is the density issued fro
     δψ = proj_tangent(δψ, ψ)
     δρ = compute_δρ(basis, ψ, δψ, occupation)
     δV = apply_kernel(basis, δρ; ρ)
+    # normalize here so we can use unnormalized FFTs for extra speed
+    δV .*= basis.fft_grid.ifft_normalization * basis.fft_grid.fft_normalization
 
     ψnk_real = similar(G_vectors(basis), promote_type(T, eltype(ψ[1])))
     Kδψ = map(enumerate(ψ)) do (ik, ψk)
         kpt = basis.kpoints[ik]
         δVψk = similar(ψk)
 
         for n = 1:size(ψk, 2)
-            ifft!(ψnk_real, basis, kpt, ψk[:, n])
+            ifft!(ψnk_real, basis, kpt, ψk[:, n]; normalize=false)
             ψnk_real .*= δV[:, :, :, kpt.spin]
-            fft!(δVψk[:, n], basis, kpt, ψnk_real)
+            fft!(δVψk[:, n], basis, kpt, ψnk_real; normalize=false)
         end
         δVψk
     end