Nonlocal term for dynamical matrix (#944)

Author: epolack

src/DFTK.jl

@@ -44,6 +44,7 @@ include("common/norm.jl")
 include("common/quadrature.jl")
 include("common/hankel.jl")
 include("common/hydrogenic.jl")
+include("common/derivatives.jl")
 
 export PspHgh
 export PspUpf

src/common/derivatives.jl

@@ -0,0 +1,9 @@
+# Helper: derivatives of a position-dependent function `f`. First-order derivatives for a
+# displacement of the atom `s` in the direction `α`.
+function derivative_wrt_αs(f, positions::AbstractVector{Vec3{T}}, α, s) where {T}
+    displacement = zero.(positions)
+    displacement[s] = setindex(displacement[s], one(T), α)
+    ForwardDiff.derivative(zero(T)) do ε
+        f(ε*displacement .+ positions)
+    end
+end

src/densities.jl

@@ -70,19 +70,19 @@ end
     end
     range = [(ik, n) for ik = 1:length(basis.kpoints) for n = mask_occ[ik]]
 
-    k_to_k_plus_q = k_to_kpq_permutation(basis, q)
+    # The variation of the orbital ψ_k defined in the basis ℬ_k is δψ_{[k+q]} in ℬ_{[k+q]},
+    # where [k+q] is equivalent to the basis k+q (see find_equivalent_kpt).
+    # The perturbation of the density
+    #   |ψ_{n,k}|² is 2 ψ_{n,k} * δψ_{n,k+q}.
+    # Hence, we first get the δψ_{[k+q]} as δψ_{k+q}…
+    δψ_plus_k = transfer_blochwave_equivalent_to_actual(basis, δψ, q)
     storages = parallel_loop_over_range(allocate_local_storage, range) do storage, kn
         (ik, n) = kn
 
         kpt = basis.kpoints[ik]
         ifft!(storage.ψnk_real, basis, kpt, ψ[ik][:, n])
-        # We return the δψk in the basis k+q which are associated to a displacement of the ψk.
-        kpt_plus_q, equivalent_kpt_plus_q = get_kpoint(basis, kpt.coordinate + q, kpt.spin)
-        δψk_plus_q = transfer_blochwave_kpt(δψ[k_to_k_plus_q[ik]], basis,
-                                            equivalent_kpt_plus_q, basis, kpt_plus_q)
-        # The perturbation of the density
-        #   |ψ_nk|² is 2 ψ_{n,k} * δψ_{n,k+q}.
-        ifft!(storage.δψnk_real, basis, kpt_plus_q, δψk_plus_q[:, n])
+        # … 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])
 
         storage.δρ[:, :, :, kpt.spin] .+= real_qzero(
             2 .* occupation[ik][n] .* basis.kweights[ik] .* conj(storage.ψnk_real)
@@ -94,9 +94,7 @@ end
     δρ = sum(getfield.(storages, :δρ))
 
     mpi_sum!(δρ, basis.comm_kpts)
-    δρ = symmetrize_ρ(basis, δρ; do_lowpass=false)
-
-    δρ
+    symmetrize_ρ(basis, δρ; do_lowpass=false)
 end
 
 @views @timing function compute_kinetic_energy_density(basis::PlaneWaveBasis, ψ, occupation)
@@ -106,7 +104,7 @@ end
     dαψnk_real = zeros(complex(eltype(basis)), basis.fft_size)
     for (ik, kpt) in enumerate(basis.kpoints)
         G_plus_k = [[p[α] for p in Gplusk_vectors_cart(basis, kpt)] for α = 1:3]
-        for n in axes(ψ[ik], 2), α = 1:3
+        for n = 1:size(ψ[ik], 2), α = 1:3
             ifft!(dαψnk_real, basis, kpt, im .* G_plus_k[α] .* ψ[ik][:, n])
             @. τ[:, :, :, kpt.spin] += occupation[ik][n] * basis.kweights[ik] / 2 * abs2(dαψnk_real)
         end

src/postprocess/phonon.jl

@@ -7,7 +7,7 @@ function dynmat_red_to_cart(model::Model, dynmat)
     # We have δr_cart = lattice * δr_red, δF_cart = lattice⁻ᵀ δF_red, so
     #   δF_cart = lattice⁻ᵀ D_red lattice⁻¹ δr_cart
     dynmat_cart = zero.(dynmat)
-    for s in axes(dynmat_cart, 2), α in axes(dynmat_cart, 4)
+    for s = 1:size(dynmat_cart, 2), α = 1:size(dynmat_cart, 4)
         dynmat_cart[:, α, :, s] = inv_lattice' * dynmat[:, α, :, s] * inv_lattice
     end
     dynmat_cart
@@ -37,7 +37,7 @@ function phonon_modes(basis::PlaneWaveBasis{T}, ψ, occupation; kwargs...) where
 
     modes = _phonon_modes(basis, dynmat_cart)
     vectors = similar(modes.vectors_cart)
-    for s in axes(vectors, 2), t in axes(vectors, 4)
+    for s = 1:size(vectors, 2), t = 1:size(vectors, 4)
         vectors[:, s, :, t] = vector_cart_to_red(basis.model, modes.vectors_cart[:, s, :, t])
     end
 

src/terms/ewald.jl

@@ -230,7 +230,7 @@ function compute_dynmat(ewald::TermEwald, basis::PlaneWaveBasis{T}, ψ, occupati
     n_dim = model.n_dim
     charges = T.(charge_ionic.(model.atoms))
 
-    dynmat = zeros(complex(T), n_dim, n_atoms, n_dim, n_atoms)
+    dynmat = zeros(complex(T), 3, n_atoms, 3, n_atoms)
     # Real part
     for s = 1:n_atoms, α = 1:n_dim
         displacement = zero.(model.positions)

src/terms/local.jl

@@ -96,7 +96,7 @@ function compute_local_potential(basis::PlaneWaveBasis{T}; positions=basis.model
     # Pre-compute the form factors at unique values of |G| to speed up
     # the potential Fourier transform (by a lot). Using a hash map gives O(1)
     # lookup.
-    form_factors = IdDict{Tuple{Int,T},T}()  # IdDict for Dual compatability
+    form_factors = IdDict{Tuple{Int,T},T}()  # IdDict for Dual compatibility
     for G in Gqs_cart
         p = norm(G)
         for (igroup, group) in enumerate(model.atom_groups)
@@ -158,62 +158,46 @@ end
     forces
 end
 
-# Phonon: Perturbation of the local potential with respect to a displacement on the
-# direction α of the atom s.
-function compute_δV_αs(::TermAtomicLocal, basis::PlaneWaveBasis{T}, α, s; q=zero(Vec3{T}),
-                       positions=basis.model.positions) where {T}
-    S = iszero(q) ? T : complex(T)
-    displacement = zero.(positions)
-    displacement[s] = setindex(displacement[s], one(T), α)
-    ForwardDiff.derivative(zero(T)) do ε
-        positions = ε*displacement .+ positions
-        compute_local_potential(basis; q, positions)
-    end
-end
-
-# Phonon: Second-order perturbation of the local potential with respect to a displacement on
-# the directions α and β of the atoms s and t.
-function compute_δ²V_βtαs(term::TermAtomicLocal, basis::PlaneWaveBasis{T}, β, t, α, s) where {T}
-    positions = basis.model.positions
-    displacement = zero.(positions)
-    displacement[t] = setindex(displacement[t], one(T), β)
-    ForwardDiff.derivative(zero(T)) do ε
-        positions = ε*displacement .+ positions
-        compute_δV_αs(term, basis, α, s; positions)
-    end
-end
-
-function compute_dynmat(term::TermAtomicLocal, basis::PlaneWaveBasis{T}, ψ, occupation; ρ,
-                        δρs, q=zero(Vec3{T}), kwargs...) where {T}
+@views function compute_dynmat(term::TermAtomicLocal, basis::PlaneWaveBasis{T}, ψ, occupation;
+                               ρ, δρs, q=zero(Vec3{T}), kwargs...) where {T}
     S = complex(T)
     model = basis.model
     positions = model.positions
     n_atoms = length(positions)
     n_dim = model.n_dim
 
-    ∫δρδV = zeros(S, 3, n_atoms, 3, n_atoms)
+    # Two contributions: dynmat_δH and dynmat_δ²H
+
+    # dynmat_δH, which is ∫δρδV.
+    dynmat_δH = zeros(S, 3, n_atoms, 3, n_atoms)
     for s = 1:n_atoms, α = 1:n_dim
-        ∫δρδV[:, :, α, s] .-= stack(forces_local(S, basis, δρs[α, s], q))
+        dynmat_δH[:, :, α, s] .-= stack(forces_local(S, basis, δρs[α, s], q))
     end
 
-    ∫ρδ²V = zeros(S, 3, n_atoms, 3, n_atoms)
+    # dynmat_δ²H, which is ∫ρδ²V.
+    dynmat_δ²H = zeros(S, 3, n_atoms, 3, n_atoms)
     ρ_fourier = fft(basis, total_density(ρ))
     δ²V_fourier = similar(ρ_fourier)
-    for s = 1:n_atoms, α = 1:n_dim
-        for t = 1:n_atoms, β = 1:n_dim
-            fft!(δ²V_fourier, basis, compute_δ²V_βtαs(term, basis, β, t, α, s))
-            ∫ρδ²V[β, t, α, s] = sum(conj(ρ_fourier) .* δ²V_fourier)
+    for s = 1:n_atoms, α = 1:n_dim, β = 1:n_dim  # zero if s ≠ t
+        δ²V = derivative_wrt_αs(basis.model.positions, β, s) do positions_βs
+            derivative_wrt_αs(positions_βs, α, s) do positions_βsαs
+                compute_local_potential(basis; positions=positions_βsαs)
+            end
         end
+        dynmat_δ²H[β, s, α, s] += sum(conj(ρ_fourier) .* fft!(δ²V_fourier, basis, δ²V))
     end
 
-    ∫δρδV + ∫ρδ²V
+    dynmat_δH + dynmat_δ²H
 end
 
 # δH is the perturbation of the local potential due to a position displacement e^{iq·r} of
 # the α coordinate of atom s.
 function compute_δHψ_αs(term::TermAtomicLocal, basis::PlaneWaveBasis, ψ, α, s, q)
     δV_αs = similar(ψ[1], basis.fft_size..., basis.model.n_spin_components)
-    # All spin components get the same potential.
-    δV_αs .= compute_δV_αs(term, basis, α, s; q)
+    # Perturbation of the local potential with respect to a displacement on the direction α
+    # of the atom s. All spin components get the same.
+    δV_αs .= derivative_wrt_αs(basis.model.positions, α, s) do positions_αs
+        compute_local_potential(basis; q, positions=positions_αs)
+    end
     multiply_ψ_by_blochwave(basis, ψ, δV_αs, q)
 end

src/terms/nonlocal.jl

@@ -55,17 +55,19 @@ end
     psp_groups = [group for group in model.atom_groups
                   if model.atoms[first(group)] isa ElementPsp]
 
-    # early return if no pseudopotential atoms
+    # Early return if no pseudopotential atoms.
     isempty(psp_groups) && return nothing
 
-    # energy terms are of the form <psi, P C P' psi>, where P(G) = form_factor(G) * structure_factor(G)
+    # Energy terms are of the form <ψ, P C P' ψ>, where
+    #   P(G) = form_factor(G) * structure_factor(G).
     forces = [zero(Vec3{T}) for _ = 1:length(model.positions)]
+
     for group in psp_groups
         element = model.atoms[first(group)]
 
         C = build_projection_coefficients(T, element.psp)
         for (ik, kpt) in enumerate(basis.kpoints)
-            # we compute the forces from the irreductible BZ; they are symmetrized later
+            # We compute the forces from the irreductible BZ; they are symmetrized later.
             G_plus_k = Gplusk_vectors(basis, kpt)
             G_plus_k_cart = to_cpu(Gplusk_vectors_cart(basis, kpt))
             form_factors = build_form_factors(element.psp, G_plus_k_cart)
@@ -77,9 +79,9 @@ end
                 forces[idx] += map(1:3) do α
                     dPdR = [-2T(π)*im*p[α] for p in G_plus_k] .* P
                     ψk = ψ[ik]
-                    dHψk = P * (C * (dPdR' * ψk))
+                    δHψk = P * (C * (dPdR' * ψk))
                     -sum(occupation[ik][iband] * basis.kweights[ik] *
-                         2real(dot(ψk[:, iband], dHψk[:, iband]))
+                             2real(dot(ψk[:, iband], δHψk[:, iband]))
                          for iband=1:size(ψk, 2))
                 end  # α
             end  # r
@@ -155,11 +157,12 @@ where ``\widehat{\rm proj}_i(p) = ∫_{ℝ^3} {\rm proj}_i(r) e^{-ip·r} dr``.
 We store ``\frac{1}{\sqrt Ω} \widehat{\rm proj}_i(k+G)`` in `proj_vectors`.
 """
 function build_projection_vectors(basis::PlaneWaveBasis{T}, kpt::Kpoint,
-                                  psps, psp_positions) where {T}
+                                  psps::AbstractVector{<: NormConservingPsp},
+                                  psp_positions) where {T}
     unit_cell_volume = basis.model.unit_cell_volume
     n_proj = count_n_proj(psps, psp_positions)
     n_G    = length(G_vectors(basis, kpt))
-    proj_vectors = zeros(Complex{T}, n_G, n_proj)
+    proj_vectors = zeros(Complex{eltype(psp_positions[1][1])}, n_G, n_proj)
     G_plus_k = to_cpu(Gplusk_vectors(basis, kpt))
 
     # Compute the columns of proj_vectors = 1/√Ω \hat proj_i(k+G)
@@ -230,3 +233,120 @@ function build_form_factors(psp, G_plus_k::AbstractVector{Vec3{TT}}) where {TT}
     end
     form_factors
 end
+
+# Helpers for phonon computations.
+function build_projection_coefficients(basis::PlaneWaveBasis{T}, psp_groups) where {T}
+    psps          = [basis.model.atoms[first(group)].psp for group in psp_groups]
+    psp_positions = [basis.model.positions[group] for group in psp_groups]
+    build_projection_coefficients(T, psps, psp_positions)
+end
+function build_projection_vectors(basis::PlaneWaveBasis, kpt::Kpoint,
+                                  psp_groups::AbstractVector{<: AbstractVector{<: Int}},
+                                  positions)
+    psps          = [basis.model.atoms[first(group)].psp for group in psp_groups]
+    psp_positions = [positions[group] for group in psp_groups]
+    build_projection_vectors(basis, kpt, psps, psp_positions)
+end
+function PDPψk(basis, positions, psp_groups, kpt, kpt_minus_q, ψk)
+    D = build_projection_coefficients(basis, psp_groups)
+    P = build_projection_vectors(basis, kpt, psp_groups, positions)
+    P_minus_q = build_projection_vectors(basis, kpt_minus_q, psp_groups, positions)
+    P * (D * P_minus_q' * ψk)
+end
+
+function compute_dynmat_δH(::TermAtomicNonlocal, basis::PlaneWaveBasis{T}, ψ, occupation,
+                           δψ, δoccupation, q) where {T}
+    S = complex(T)
+    model = basis.model
+    psp_groups = [group for group in model.atom_groups
+                  if model.atoms[first(group)] isa ElementPsp]
+
+    # Early return if no pseudopotential atoms.
+    isempty(psp_groups) && return nothing
+
+    δforces = [zero(Vec3{S}) for _ = 1:length(model.positions)]
+    for group in psp_groups
+        δψ_plus_q = transfer_blochwave_equivalent_to_actual(basis, δψ, q)
+        for (ik, kpt) in enumerate(basis.kpoints)
+            ψk = ψ[ik]
+            δψk_plus_q = δψ_plus_q[ik].ψk
+            kpt_plus_q = δψ_plus_q[ik].kpt
+
+            for idx in group
+                δforces[idx] += map(1:3) do α
+                    δHψk = derivative_wrt_αs(model.positions, α, idx) do positions_αs
+                        PDPψk(basis, positions_αs, psp_groups, kpt_plus_q, kpt, ψ[ik])
+                    end
+                    δHψk_plus_q = derivative_wrt_αs(model.positions, α, idx) do positions_αs
+                        PDPψk(basis, positions_αs, psp_groups, kpt, kpt, ψ[ik])
+                    end
+                    -sum(  2occupation[ik][iband] * basis.kweights[ik]
+                               * dot(δψk_plus_q[:, iband], δHψk[:, iband])
+                         + δoccupation[ik][iband]  * basis.kweights[ik]
+                               * 2real(dot(ψk[:, iband], δHψk_plus_q[:, iband]))
+                         for iband=1:size(ψk, 2))
+                end
+            end
+        end
+    end
+
+    mpi_sum!(δforces, basis.comm_kpts)
+end
+
+@views function compute_dynmat(term::TermAtomicNonlocal, basis::PlaneWaveBasis{T}, ψ,
+                               occupation; δψs, δoccupations, q=zero(Vec3{T}),
+                               kwargs...) where {T}
+    S = complex(T)
+    model = basis.model
+    positions = model.positions
+    n_atoms = length(positions)
+    n_dim = model.n_dim
+
+    # Two contributions: dynmat_δH and dynmat_δ²H.
+
+    # dynmat_δH
+    dynmat_δH = zeros(S, 3, n_atoms, 3, n_atoms)
+    for s = 1:n_atoms, α = 1:n_dim
+        dynmat_δH[:, :, α, s] .-= stack(
+            compute_dynmat_δH(term, basis, ψ, occupation, δψs[α, s], δoccupations[α, s], q)
+        )
+    end
+
+    psp_groups = [group for group in model.atom_groups
+                  if model.atoms[first(group)] isa ElementPsp]
+    # Early return if no pseudopotential atoms.
+    isempty(psp_groups) && return dynmat_δH
+
+    # dynmat_δ²H
+    dynmat_δ²H = zeros(S, 3, n_atoms, 3, n_atoms)
+    δ²Hψ = zero.(ψ)
+    for s = 1:n_atoms, α = 1:n_dim, β = 1:n_dim  # zero if s ≠ t
+        for (ik, kpt) in enumerate(basis.kpoints)
+            δ²Hψ[ik] = derivative_wrt_αs(basis.model.positions, β, s) do positions_βs
+                derivative_wrt_αs(positions_βs, α, s) do positions_βsαs
+                    PDPψk(basis, positions_βsαs, psp_groups, kpt, kpt, ψ[ik])
+                end
+            end
+            dynmat_δ²H[β, s, α, s] += sum(occupation[ik][n] * basis.kweights[ik] *
+                                              dot(ψ[ik][:, n], δ²Hψ[ik][:, n])
+                                          for n=1:size(ψ[ik], 2))
+        end
+    end
+
+    dynmat_δH + dynmat_δ²H
+end
+
+# δH is the Fourier transform perturbation of the nonlocal potential due to a position
+# displacement e^{iq·r} of the α coordinate of atom s.
+function compute_δHψ_αs(::TermAtomicNonlocal, basis::PlaneWaveBasis{T}, ψ, α, s, q) where {T}
+    model = basis.model
+    psp_groups = [group for group in model.atom_groups
+                  if model.atoms[first(group)] isa ElementPsp]
+
+    ψ_minus_q = transfer_blochwave_equivalent_to_actual(basis, ψ, -q)
+    map(enumerate(basis.kpoints)) do (ik, kpt)
+        derivative_wrt_αs(model.positions, α, s) do positions_αs
+            PDPψk(basis, positions_αs, psp_groups, kpt, ψ_minus_q[ik].kpt, ψ_minus_q[ik].ψk)
+        end
+    end
+end

src/terms/xc.jl

@@ -362,11 +362,16 @@ end
 
 
 function apply_kernel(term::TermXc, basis::PlaneWaveBasis{T}, δρ::AbstractArray{Tδρ};
-                      ρ, kwargs...) where {T, Tδρ}
+                      ρ, q=zero(Vec3{T}), kwargs...) where {T, Tδρ}
     n_spin = basis.model.n_spin_components
     isempty(term.functionals) && return nothing
     @assert all(family(xc) in (:lda, :gga) for xc in term.functionals)
 
+    if !iszero(q) && !isnothing(term.ρcore)
+        error("Phonon computations are not supported for models using nonlinear core \
+              correction.")
+    end
+
     # Take derivatives of the density and the perturbation if needed.
     max_ρ_derivs = maximum(max_required_derivative, term.functionals)
     density      = LibxcDensities(basis, max_ρ_derivs, ρ, nothing)