Phonons docs (#953)

Author: antoine-levitt

src/DFTK.jl

@@ -222,9 +222,6 @@ include("response/hessian.jl")
 export compute_current
 include("postprocess/current.jl")
 export phonon_modes
-export phonon_modes_cart
-export compute_dynmat
-export compute_dynmat_cart
 include("postprocess/phonon.jl")
 
 # Workarounds

src/PlaneWaveBasis.jl

@@ -13,18 +13,19 @@ abstract type AbstractBasis{T <: Real} end
 """
 Discretization information for ``k``-point-dependent quantities such as orbitals.
 More generally, a ``k``-point is a block of the Hamiltonian;
-eg collinear spin is treated by doubling the number of kpoints.
+e.g. collinear spin is treated by doubling the number of ``k``-points.
 """
 struct Kpoint{T <: Real, GT <: AbstractVector{Vec3{Int}}}
     spin::Int                     # Spin component can be 1 or 2 as index into what is
     #                             # returned by the `spin_components` function
     coordinate::Vec3{T}           # Fractional coordinate of k-point
+    G_vectors::GT                 # Wave vectors in integer coordinates (vector of Vec3{Int})
+    #                             # ({G, 1/2 |k+G|^2 ≤ Ecut})
+                                  # This is not assumed to be in any particular order
     mapping::Vector{Int}          # Index of G_vectors[i] on the FFT grid:
     #                             # G_vectors(basis)[kpt.mapping[i]] == G_vectors(basis, kpt)[i]
     mapping_inv::Dict{Int, Int}   # Inverse of `mapping`:
     #                             # G_vectors(basis)[i] == G_vectors(basis, kpt)[mapping_inv[i]]
-    G_vectors::GT                 # Wave vectors in integer coordinates (vector of Vec3{Int})
-    #                             # ({G, 1/2 |k+G|^2 ≤ Ecut})
 end
 
 @doc raw"""
@@ -145,12 +146,36 @@ function Kpoint(spin::Integer, coordinate::AbstractVector{<:Real},
     Gvecs_k = to_device(architecture, Gvecs_k)
 
     mapping_inv = Dict(ifull => iball for (iball, ifull) in enumerate(mapping))
-    Kpoint(spin, k, mapping, mapping_inv, Gvecs_k)
+    Kpoint(spin, k, Gvecs_k, mapping, mapping_inv)
 end
 function Kpoint(basis::PlaneWaveBasis, coordinate::AbstractVector, spin::Int)
     Kpoint(spin, coordinate, basis.model.recip_lattice, basis.fft_size, basis.Ecut;
            basis.variational, basis.architecture)
 end
+# Construct the kpoint with coordinate equivalent_kpt.coordinate + ΔG.
+# Equivalent to (but faster than) Kpoint(equivalent_kpt.coordinate + ΔG).
+function construct_from_equivalent_kpt(basis, equivalent_kpt, coordinate, ΔG)
+    linear = LinearIndices(basis.fft_size)
+    # Mapping is the same as if created from scratch, although it is not ordered.
+    mapping = map(CartesianIndices(basis.fft_size)[equivalent_kpt.mapping]) do G
+        linear[CartesianIndex(mod1.(Tuple(G + CartesianIndex(ΔG...)), basis.fft_size))]
+    end
+    mapping_inv = Dict(ifull => iball for (iball, ifull) in enumerate(mapping))
+    Kpoint(equivalent_kpt.spin, Vec3(coordinate), equivalent_kpt.G_vectors .+ Ref(ΔG),
+           mapping, mapping_inv)
+end
+# Returns the kpoint at given coordinate. If outside the Brillouin zone, it is created
+# from an equivalent kpoint in the basis (also returned)
+function get_kpoint(basis::PlaneWaveBasis{T}, kcoord, spin) where {T}
+    index, ΔG = find_equivalent_kpt(basis, kcoord, spin)
+    equivalent_kpt = basis.kpoints[index]
+    if iszero(ΔG)
+        kpt = equivalent_kpt
+    else
+        kpt =  construct_from_equivalent_kpt(basis, equivalent_kpt, kcoord, ΔG)
+    end
+    (; kpt, equivalent_kpt)
+end
 
 @timing function build_kpoints(model::Model{T}, fft_size, kcoords, Ecut;
                                variational=true,
@@ -163,7 +188,7 @@ end
     for iσ = 1:model.n_spin_components
         for kpt in kpoints_spin_1
             push!(all_kpoints, Kpoint(iσ, kpt.coordinate,
-                                      kpt.mapping, kpt.mapping_inv, kpt.G_vectors))
+                                      kpt.G_vectors, kpt.mapping, kpt.mapping_inv))
         end
     end
     all_kpoints

src/densities.jl

@@ -58,8 +58,6 @@ end
     Tδρ = iszero(q) ? real(Tψ) : Tψ
     real_qzero = iszero(q) ? real : identity
 
-    ordering(kdata) = kdata[k_to_equivalent_kpq_permutation(basis, q)]
-
     # occupation should be on the CPU as we are going to be doing scalar indexing.
     occupation = [to_cpu(oc) for oc in occupation]
     mask_occ = [findall(occnk -> abs(occnk) ≥ occupation_threshold, occk)
@@ -72,14 +70,16 @@ 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)
     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, δψk_plus_q = kpq_equivalent_blochwave_to_kpq(basis, kpt, q,
-                                                                 ordering(δψ)[ik])
+        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])
@@ -106,7 +106,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 = 1:size(ψ[ik], 2), α = 1:3
+        for n in axes(ψ[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

@@ -2,26 +2,19 @@
 function dynmat_red_to_cart(model::Model, dynmat)
     inv_lattice = model.inv_lattice
 
-    # The dynamical matrix `D` acts on vectors `dr` and gives a covector `dF`:
-    #   dF = D · dr.
-    # Thus the transformation between reduced and Cartesian coordinates is not a comatrix.
-    # To transform `dynamical_matrix` from reduced coordinates to `dynmat_cart` in Cartesian
-    # coordinates, we write
-    #   dr_cart = lattice · dr_red,
-    #   ⇒ dr_redᵀ · D_red · dr_red = dr_cartᵀ · lattice⁻ᵀ · D_red · lattice⁻¹ · dr_cart
-    #                              = dr_cartᵀ · D_cart · dr_cart
-    #   ⇒ D_cart = lattice⁻ᵀ · D_red · lattice⁻¹.
-
+    # The dynamical matrix `D` acts on vectors `δr` and gives a covector `δF`:
+    #   δF = D δr
+    # 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 = 1:size(dynmat_cart, 2), α = 1:size(dynmat_cart, 4)
+    for s in axes(dynmat_cart, 2), α in axes(dynmat_cart, 4)
         dynmat_cart[:, α, :, s] = inv_lattice' * dynmat[:, α, :, s] * inv_lattice
     end
     dynmat_cart
 end
 
-# Create a ``3×n_{\rm atoms}×3×n_{\rm atoms}`` tensor (for consistency with the format of
-# dynamical matrices) equivalent to a diagonal matrix with the atomic masses of the atoms in
-# a.u. on the diagonal.
+# Create a ``3×n_{\rm atoms}×3×n_{\rm atoms}`` tensor equivalent to a diagonal matrix with
+# the atomic masses of the atoms in a.u. on the diagonal.
 function mass_matrix(T, atoms)
     n_atoms = length(atoms)
     atoms_mass = atomic_mass.(atoms)
@@ -34,26 +27,42 @@ function mass_matrix(T, atoms)
 end
 mass_matrix(model::Model{T}) where {T} = mass_matrix(T, model.atoms)
 
-@doc raw"""
-Solve the eigenproblem for a dynamical matrix: returns the `frequencies` and eigenvectors
-(`vectors`).
 """
-function phonon_modes(basis::PlaneWaveBasis{T}, dynmat) where {T}
+Get phonon quantities. We return the frequencies, the mass matrix and reduced and Cartesian
+eigenvectors and dynamical matrices.
+"""
+function phonon_modes(basis::PlaneWaveBasis{T}, ψ, occupation; kwargs...) where {T}
+    dynmat = compute_dynmat(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
+    dynmat_cart = dynmat_red_to_cart(basis.model, dynmat)
+
+    modes = _phonon_modes(basis, dynmat_cart)
+    vectors = similar(modes.vectors_cart)
+    for s in axes(vectors, 2), t in axes(vectors, 4)
+        vectors[:, s, :, t] = vector_cart_to_red(basis.model, modes.vectors_cart[:, s, :, t])
+    end
+
+    (; modes.mass_matrix, modes.frequencies, dynmat, dynmat_cart, vectors, modes.vectors_cart)
+end
+# Compute the frequencies and vectors. Internal because of the potential misuse:
+# the diagonalization of the phonon modes has to be done in Cartesian coordinates.
+function _phonon_modes(basis::PlaneWaveBasis{T}, dynmat_cart) where {T}
     n_atoms = length(basis.model.positions)
     M = reshape(mass_matrix(T, basis.model.atoms), 3*n_atoms, 3*n_atoms)
 
-    res = eigen(reshape(dynmat, 3*n_atoms, 3*n_atoms), M)
+    res = eigen(reshape(dynmat_cart, 3*n_atoms, 3*n_atoms), M)
     maximum(abs, imag(res.values)) > sqrt(eps(T)) &&
-        @warn "Some eigenvalues of the dynamical matrix have a large imaginary part"
+        @warn "Some eigenvalues of the dynamical matrix have a large imaginary part."
 
     signs = sign.(real(res.values))
     frequencies = signs .* sqrt.(abs.(real(res.values)))
 
-    (; frequencies, res.vectors)
+    vectors_cart = 
+    (; mass_matrix=M, frequencies, vectors_cart=reshape(res.vectors, 3, n_atoms, 3, n_atoms))
 end
-function phonon_modes_cart(basis::PlaneWaveBasis{T}, dynmat) where {T}
-    dynmat_cart = dynmat_red_to_cart(basis.model, dynmat)
-    phonon_modes(basis, dynmat_cart)
+# For convenience
+function phonon_modes(scfres::NamedTuple; kwargs...)
+    phonon_modes(scfres.basis, scfres.ψ, scfres.occupation; scfres.ρ, scfres.ham,
+                 scfres.occupation_threshold, scfres.εF, scfres.eigenvalues, kwargs...)
 end
 
 @doc raw"""
@@ -69,44 +78,27 @@ in reduced coordinates.
     δoccupations = [zero.(occupation) for _ = 1:3, _ = 1:n_atoms]
     δψs = [zero.(ψ) for _ = 1:3, _ = 1:n_atoms]
     for s = 1:n_atoms, α = 1:basis.model.n_dim
+        # Get δH ψ
         δHψs_αs = compute_δHψ_αs(basis, ψ, α, s, q)
         isnothing(δHψs_αs) && continue
+        # Response solver to get δψ
         (; δψ, δρ, δoccupation) = solve_ΩplusK_split(ham, ρ, ψ, occupation, εF, eigenvalues,
                                                      -δHψs_αs; q, kwargs...)
         δoccupations[α, s] = δoccupation
         δρs[α, s] = δρ
         δψs[α, s] = δψ
     end
-
+    # Query each energy term for their contribution to the dynamical matrix.
     dynmats_per_term = [compute_dynmat(term, basis, ψ, occupation; ρ, δψs, δρs,
                                        δoccupations, q)
                         for term in basis.terms]
     sum(filter(!isnothing, dynmats_per_term))
 end
 
 """
-Cartesian form of [`compute_dynmat`](@ref).
-"""
-function compute_dynmat_cart(basis::PlaneWaveBasis, ψ, occupation; kwargs...)
-    dynmats_reduced = compute_dynmat(basis, ψ, occupation; kwargs...)
-    dynmat_red_to_cart(basis.model, dynmats_reduced)
-end
-
-function compute_dynmat(scfres::NamedTuple; kwargs...)
-    compute_dynmat(scfres.basis, scfres.ψ, scfres.occupation; scfres.ρ, scfres.ham,
-                   scfres.occupation_threshold, scfres.εF, scfres.eigenvalues, kwargs...)
-end
-
-function compute_dynmat_cart(scfres; kwargs...)
-    compute_dynmat_cart(scfres.basis, scfres.ψ, scfres.occupation; scfres.ρ, scfres.ham,
-                        scfres.occupation_threshold, scfres.εF, scfres.eigenvalues, kwargs...)
-end
-
-# TODO: Document relation to non-local potential in future phonon PR.
-"""
-Assemble the right-hand side term for the Sternheimer equation for all relevant quantities:
-Compute the perturbation of the Hamiltonian with respect to a variation of the local
-potential produced by a displacement of the atom s in the direction α.
+Get ``δH·ψ``, with ``δH`` the perturbation of the Hamiltonian with respect to a position
+displacement ``e^{iq·r}`` of the ``α`` coordinate of atom ``s``.
+`δHψ[ik]` is ``δH·ψ_{k-q}``, expressed in `basis.kpoints[ik]`.
 """
 @timing function compute_δHψ_αs(basis::PlaneWaveBasis, ψ, α, s, q)
     δHψ_per_term = [compute_δHψ_αs(term, basis, ψ, α, s, q) for term in basis.terms]

src/response/chi0.jl

@@ -290,6 +290,7 @@ end
 Perform in-place computations of the derivatives of the wave functions by solving
 a Sternheimer equation for each `k`-points. It is assumed the passed `δψ` are initialised
 to zero.
+For phonon, `δHψ[ik]` is ``δH·ψ_{k-q}``, expressed in `basis.kpoints[ik]`.
 """
 function compute_δψ!(δψ, basis::PlaneWaveBasis{T}, H, ψ, εF, ε, δHψ, ε_minus_q=ε;
                      ψ_extra=[zeros_like(ψk, size(ψk,1), 0) for ψk in ψ],
@@ -299,7 +300,7 @@ function compute_δψ!(δψ, basis::PlaneWaveBasis{T}, H, ψ, εF, ε, δHψ, ε
     # where P_{k} is the projector on ψ_{k} and with the conventions:
     # * δψ_{k} is the variation of ψ_{k-q}, which implies (for ℬ_{k} the `basis.kpoints`)
     #     δψ_{k-q} ∈ ℬ_{k-q} and δHψ_{k-q} ∈ ℬ_{k};
-    # * δHψ[ik] = δHψ_{k-q};
+    # * δHψ[ik] = δH ψ_{k-q};
     # * ε_minus_q[ik] = ε_{·, k-q}.
     model = basis.model
     temperature = model.temperature
@@ -343,7 +344,7 @@ end
                                     occupation_threshold, q=zero(Vec3{eltype(ham.basis)}),
                                     kwargs_sternheimer...)
     basis = ham.basis
-    ordering(kdata) = kdata[k_to_equivalent_kpq_permutation(basis, -q)]
+    k_to_k_minus_q = k_to_kpq_permutation(basis, -q)
 
     # We first select orbitals with occupation number higher than
     # occupation_threshold for which we compute the associated response δψn,
@@ -358,9 +359,10 @@ end
     ψ_occ   = [ψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_occ)]
     ψ_extra = [ψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_extra)]
     ε_occ   = [eigenvalues[ik][maskk] for (ik, maskk) in enumerate(mask_occ)]
-    δHψ_minus_q_occ = [δHψ[ik][:, maskk] for (ik, maskk) in enumerate(ordering(mask_occ))]
+    δHψ_minus_q_occ = [δHψ[ik][:, mask_occ[k_to_k_minus_q[ik]]] for ik = 1:length(basis.kpoints)]
     # Only needed for phonon calculations.
-    ε_minus_q_occ  = ordering([eigenvalues[ik][maskk] for (ik, maskk) in enumerate(mask_occ)])
+    ε_minus_q_occ  = [eigenvalues[k_to_k_minus_q[ik]][mask_occ[k_to_k_minus_q[ik]]]
+                      for ik = 1:length(basis.kpoints)]
 
     # First we compute δoccupation. We only need to do this for the actually occupied
     # orbitals. So we make a fresh array padded with zeros, but only alter the elements
@@ -371,15 +373,17 @@ end
     δoccupation = zero.(occupation)
     if iszero(q)
         δocc_occ = [δoccupation[ik][maskk] for (ik, maskk) in enumerate(mask_occ)]
-        δεF = compute_δocc!(δocc_occ, basis, ψ_occ, εF, ε_occ, δHψ_minus_q_occ).δεF
+        (; δεF) = compute_δocc!(δocc_occ, basis, ψ_occ, εF, ε_occ, δHψ_minus_q_occ)
     else
+        # When δH is not periodic, δH ψnk is a Bloch wave at k+q and ψnk at k,
+        # so that δεnk = <ψnk|δH|ψnk> = 0 and there is no occupation shift
         δεF = zero(εF)
     end
 
     # Then we compute δψ (again in-place into a zero-padded array) with elements of
     # `basis.kpoints` that are equivalent to `k+q`.
-    δψ = zero.(δHψ)  # δHψ ≠ δHψ_minus_q
-    δψ_occ = [δψ[ik][:, maskk] for (ik, maskk) in enumerate(ordering(mask_occ))]
+    δψ = zero.(δHψ)
+    δψ_occ = [δψ[ik][:, maskk] for (ik, maskk) in enumerate(mask_occ[k_to_k_minus_q])]
     compute_δψ!(δψ_occ, basis, ham.blocks, ψ_occ, εF, ε_occ, δHψ_minus_q_occ, ε_minus_q_occ;
                 ψ_extra, q, kwargs_sternheimer...)
 

src/response/hessian.jl

@@ -173,7 +173,7 @@ Solve the problem `(Ω+K) δψ = rhs` using a split algorithm, where `rhs` is ty
     δVind = apply_kernel(basis, δρ; ρ, q)  # Change in potential induced by δρ
     # For phonon calculations, assemble
     #   δHψ_k = δV_{q} · ψ_{k-q}.
-    δHψ = multiply_ψ_by_blochwave(basis, ψ, δVind, q) - rhs
+    δHψ = multiply_ψ_by_blochwave(basis, ψ, δVind, q) .- rhs
 
     # Compute total change in eigenvalues
     δeigenvalues = map(ψ, δHψ) do ψk, δHψk