Collinear spin for GGA functionals (#327)

src/Model.jl

@@ -118,7 +118,7 @@ function Model(lattice::AbstractMatrix{T};
     spin_polarization in (:none, :collinear, :full, :spinless) ||
         error("Only :none, :collinear, :full and :spinless allowed for spin_polarization")
     !isempty(magnetic_moments) && !(spin_polarization in (:collinear, :full)) && @warn(
-        "Non-empty magnetic_moments on a Model without spin polarization detected"
+        "Non-empty magnetic_moments on a Model without spin polarization detected."
     )
     n_spin = length(spin_components(spin_polarization))
 

src/external/etsf_nanoquanta.jl

@@ -72,7 +72,7 @@ load_atoms(folder; kwargs...) = load_atoms(Float64, folder; kwargs...)
 Load a DFTK-compatible model object from the ETSF folder.
 Use the scalar type `T` to represent the data.
 """
-function load_model(T, folder::EtsfFolder)
+function load_model(T, folder::EtsfFolder; magnetic_moments=[])
     # Parse functional ...
     functional = Vector{Symbol}()
     ixc = Int(folder.gsr["ixc"][:])
@@ -107,37 +107,45 @@ function load_model(T, folder::EtsfFolder)
 
     n_spin = size(folder.gsr["eigenvalues"], 3)
     spin_polarization = n_spin == 2 ? :collinear : :none
+    if spin_polarization == :collinear && isempty(magnetic_moments)
+        @warn("load_basis will most likely fail for collinear spin if magnetic_moments " *
+              "not explicitly specified.")
+    end
 
     # Build model
     lattice = load_lattice(T, folder)
     atoms = load_atoms(T, folder)
     model_DFT(Array{T}(lattice), atoms, functional;
               smearing=smearing_function, temperature=Tsmear,
-              spin_polarization=spin_polarization)
+              spin_polarization=spin_polarization, magnetic_moments=magnetic_moments)
 end
 load_model(folder; kwargs...) = load_model(Float64, folder; kwargs...)
 
-
 """
 Load a DFTK-compatible basis object from the ETSF folder.
 Use the scalar type `T` to represent the data.
 """
-function load_basis(T, folder::EtsfFolder)
-    model = load_model(T, folder)
+function load_basis(T, folder::EtsfFolder; magnetic_moments=[])
+    model = load_model(T, folder, magnetic_moments=magnetic_moments)
     atoms = load_atoms(T, folder)
 
     Ecut = folder.gsr["kinetic_energy_cutoff"][:]
     kcoords = Vec3{T}.(eachcol(folder.gsr["reduced_coordinates_of_kpoints"]))
 
     # Try to determine whether this is a shifted kpoint mesh or not
-    ksmallest = sort(filter(k -> all(k .> 0), kcoords), by=norm)[1]
+    if length(kcoords) > 1
+        ksmallest = sort(filter(k -> all(k .≥ 0), kcoords), by=norm)[1]
+    else
+        ksmallest = kcoords[1]
+    end
     kshift = [(abs(k) > 10eps(T)) ? 0.5 : 0.0 for k in ksmallest]
 
     kgrid = Vector{Int}(folder.gsr["monkhorst_pack_folding"])
     kcoords_new, ksymops, _ = bzmesh_ir_wedge(kgrid, model.symmetries, kshift=kshift)
-    @assert kcoords_new ≈ kcoords
+    @assert kcoords_new ≈ normalize_kpoint_coordinate.(kcoords)
 
-    PlaneWaveBasis(model, Ecut, kcoords, ksymops)
+    fft_size = size(folder.den["density"])[2:4]
+    PlaneWaveBasis(model, Ecut, kcoords, ksymops, fft_size=fft_size)
 end
 load_basis(folder; kwargs...) = load_basis(Float64, folder; kwargs...)
 
@@ -146,18 +154,29 @@ load_basis(folder; kwargs...) = load_basis(Float64, folder; kwargs...)
 Load a DFTK-compatible density object from the ETSF folder.
 Use the scalar type `T` to represent the data.
 """
-function load_density(T, folder::EtsfFolder)
-    # So far this function is untested
-    basis = load_basis(T, folder)
+function load_density(T, folder::EtsfFolder; magnetic_moments=[])
+    basis = load_basis(T, folder, magnetic_moments=magnetic_moments)
 
-    n_comp = size(folder.den["density"], 5)
     n_real_cpx = size(folder.den["density"], 1)
-    @assert n_comp == 1
     @assert n_real_cpx == 1
 
-    ρ_real = Array{Complex{T}}(folder.den["density"][1, :, :, :, 1])
-    @assert basis.fft_size == size(ρ_real)
+    n_spin = size(folder.den["density"], 5)
+    @assert n_spin in (1, 2)
 
-    from_real(basis, ρ_real)
+    if n_spin == 1
+        @assert basis.model.spin_polarization == :none
+        ρtot_real = Array{T}(folder.den["density"][1, :, :, :, 1])
+        @assert basis.fft_size == size(ρ_real)
+
+        return from_real(basis, ρtot_real), nothing
+    else
+        @assert basis.model.spin_polarization == :collinear
+        ρtot_real = Array{T}(folder.den["density"][1, :, :, :, 1])
+        ρα_real = Array{T}(folder.den["density"][1, :, :, :, 2])
+        @assert basis.fft_size == size(ρtot_real)
+        @assert basis.fft_size == size(ρα_real)
+
+        return from_real(basis, ρtot_real), from_real(basis, 2ρα_real - ρtot_real)
+    end
 end
 load_density(folder; kwargs...) = load_density(Float64, folder; kwargs...)

src/guess_density.jl

@@ -47,6 +47,7 @@ end
     @assert length(magnetic_moments) == length(atoms)
     for (ispec, (spec, magmoms)) in enumerate(magnetic_moments)
         positions = atoms[ispec][2]
+        @assert charge_nuclear(spec) == charge_nuclear(atoms[ispec][1])
         @assert length(magmoms) == length(positions)
         for (ipos, r) in enumerate(positions)
             magmom = Vec3{T}(normalize_magnetic_moment(magmoms[ipos]))

src/terms/xc.jl

@@ -24,8 +24,9 @@ end
 
 function ene_ops(term::TermXc, ψ, occ; ρ, ρspin=nothing, kwargs...)
     basis = term.basis
-    T = eltype(basis)
+    T     = eltype(basis)
     model = basis.model
+    n_spin = model.n_spin_components
     @assert all(xc.family in (:lda, :gga) for xc in term.functionals)
 
     if isempty(term.functionals)
@@ -37,7 +38,7 @@ function ene_ops(term::TermXc, ψ, occ; ρ, ρspin=nothing, kwargs...)
     max_ρ_derivs = maximum(max_required_derivative, term.functionals)
     density = DensityDerivatives(basis, max_ρ_derivs, ρ, ρspin)
 
-    potential = [zeros(T, basis.fft_size) for _ in 1:model.n_spin_components]
+    potential = [zeros(T, basis.fft_size) for _ in 1:n_spin]
     zk = zeros(T, basis.fft_size)  # Energy per unit particle
     E = zero(T)
     for xc in term.functionals
@@ -49,16 +50,25 @@ function ene_ops(term::TermXc, ψ, occ; ρ, ρspin=nothing, kwargs...)
         E += sum(terms.zk .* ρ.real) * dVol
 
         # Add potential contributions Vρ -2 ∇⋅(Vσ ∇ρ)
-        for iσ in 1:model.n_spin_components
-            potential[iσ] .+= @view terms.vrho[iσ, :, :, :]
+        for σ in 1:n_spin
+            potential[σ] .+= @view terms.vrho[σ, :, :, :]
         end
+
         if haskey(terms, :vsigma)
-            @assert term.basis.model.spin_polarization in (:none, :spinless)
-            potential[1] .+=
-                -2 * divergence_real(
-                    α -> @view(terms.vsigma[1, :, :, :]) .* density.∇ρ_real[α],
-                    density.basis,
-                )
+            # TODO Potential to save some FFTs for spin-polarised calculations:
+            #      For n_spin == 2 this calls divergence_real 4 times, where only 2 are
+            #      needed (one for potential[1] and one for potential[2])
+            @views for σ in 1:n_spin, τ in 1:n_spin
+                iστ = density.index_spin_σ[(σ, τ)]
+                # Extra factor (1/2) for σ != τ is needed because libxc only keeps σ_{αβ}
+                # in the energy expression. See comment block below on spin-polarised XC.
+                spinfac = (σ == τ ? 2 : 1)
+                potential[σ] .+=
+                     -spinfac .* divergence_real(
+                        α -> terms.vsigma[iστ, :, :, :] .* density.∇ρ_real[τ, :, :, :, α],
+                        density.basis,
+                    )
+            end
         end
     end
     if term.scaling_factor != 1
@@ -115,27 +125,24 @@ end
 
 #=
 The total energy is
-Etot = ∫ ρ E(ρ,σ), where σ = |∇ρ|^2 and E(ρ,σ) is the energy per unit particle
+Etot = ∫ ρ E(ρ,σ), where σ = |∇ρ|^2 and E(ρ,σ) is the energy per unit particle.
 libxc provides the scalars
-Vρ = ∂(ρ E)/∂ρ
-Vσ = ∂(ρ E)/∂σ
+    Vρ = ∂(ρ E)/∂ρ
+    Vσ = ∂(ρ E)/∂σ
 
 Consider a variation dϕi of an orbital (considered real for
 simplicity), and let dEtot be the corresponding variation of the
 energy. Then the potential Vxc is defined by
-dEtot = 2 ∫ Vxc ϕi dϕi
+    dEtot = 2 ∫ Vxc ϕi dϕi
 
-dρ = 2 ϕi dϕi
-dσ = 2 ∇ρ ⋅ ∇dρ = 4 ∇ρ ⋅ ∇(ϕi dϕi)
-dEtot = ∫ Vρ dρ + Vσ dσ
-      = 2 ∫ Vρ ϕi dϕi + 4 ∫ Vσ ∇ρ ⋅ ∇(ϕi dϕi)
-      = 2 ∫ Vρ ϕi dϕi - 4 ∫ div(Vσ ∇ρ) ϕi dϕi
+    dρ = 2 ϕi dϕi
+    dσ = 2 ∇ρ ⋅ ∇dρ = 4 ∇ρ ⋅ ∇(ϕi dϕi)
+    dEtot = ∫ Vρ dρ + Vσ dσ
+          = 2 ∫ Vρ ϕi dϕi + 4 ∫ Vσ ∇ρ ⋅ ∇(ϕi dϕi)
+          = 2 ∫ Vρ ϕi dϕi - 4 ∫ div(Vσ ∇ρ) ϕi dϕi
 where we performed an integration by parts in the last equation
-(boundary terms drop by periodicity).
-
-Therefore,
-Vxc = Vρ - 2 div(Vσ ∇ρ)
-
+(boundary terms drop by periodicity). Therefore,
+    Vxc = Vρ - 2 div(Vσ ∇ρ)
 See eg Richard Martin, Electronic stucture, p. 158
 
 A more algebraic way of deriving these for GGA, which helps see why
@@ -153,26 +160,46 @@ K : C -> C be the gradient operator (multiplication by iK, assuming dimension 1
 simplify notations)
 
 Then
+    ρ = IF(Xϕ).^2
+    ρf = F ρ
+    ∇ρf = K ρf
+    ∇ρ = IF ∇ρf
+
+    dρ  = 2 IF(Xϕ) .* IF(X dϕ)
+    d∇ρ = IF K F dρ
+        = 2 IF K F (IF(Xϕ) .* IF(X dϕ))
+
+    Etot  = sum(ρ . * E(ρ, ∇ρ.^2))
+    dEtot = sum(Vρ .* dρ + 2 (Vσ .* ∇ρ) .* d∇ρ)
+          = sum(Vρ .* dρ + 2 (Vσ .* ∇ρ) .* [(IF K F) dρ)]
+          = sum(Vρ .* dρ - 2 [(IF K F) (Vσ .* ∇ρ)] .* dρ)
+where we used that (IF K F) is anti-self-adjoint and thought of sum(A .* op B )
+like ⟨A | op B⟩; this is the discrete integration by parts. Now note that
+
+    sum(V .* dρ) = 2 sum(V .* IF(Xϕ) .* IF(X dϕ))
+                 = 2 sum(R(F[ V .* IF(Xϕ) ]) .* dϕ)
+
+where we took the adjoint (IF X)^† = (R F) and the result follows.
+=#
 
-ρ = IF(Xϕ).^2
-ρf = F ρ
-∇ρf = K ρf
-∇ρ = IF ∇ρf
-
-dρ = 2 IF(Xϕ) .* IF(X dϕ)
-d∇ρ = IF K F dρ
-    = 2 IF K F (IF(Xϕ) .* IF(X dϕ))
-
-Etot = sum(ρ . * E(ρ,∇ρ.^2))
-dEtot = sum(Vρ .* dρ + 2 (Vσ .* ∇ρ) .* d∇ρ)
-      = sum(Vρ .* dρ + (IF K F) (Vσ .* ∇ρ) .* dρ)
-(IF K F is self-adjoint; this is the discrete integration by parts)
-
-Now note that
-
-sum(V .* dρ) = sum(V .* IF(Xϕ) .* IF(X dϕ))
-             = sum(R(F(V .* IF(Xϕ))) .* dϕ)
-and the result follows.
+#=  Spin-polarised calculations
+
+These expressions can be generalised for spin-polarised calculations.
+In this case for example the energy per unit particle becomes
+E(ρ_α, ρ_β, σ_αα, σ_αβ, σ_βα, σ_ββ), where σ_ij = ∇ρ_i ⋅ ∇ρ_j
+and the XC potential is analogously
+    Vxc_s = Vρ_s - 2 ∑_t div(Vσ_{st} ∇ρ_t)
+where s, t ∈ {α, β} are the spin components and we understand
+    Vρ_s     = ∂(ρ E)/∂(ρ_s)
+    Vσ_{s,t} = ∂(ρ E)/∂(σ_{s,t})
+
+Now, in contrast to this libxc explicitly uses the symmetry σ_αβ = σ_βα and sets σ
+to be a vector of the three independent components only
+    σ = [σ_αα, σ_x, σ_ββ]  where     σ_x = (σ_αβ + σ_βα)/2
+Accordingly Vσ has the compoments
+    [∂(ρ E)/∂σ_αα, ∂(ρ E)/∂σ_x, ∂(ρ E)/∂σ_ββ]
+where in particular ∂(ρ E)/∂σ_x = (1/2) ∂(ρ E)/∂σ_αβ = (1/2) ∂(ρ E)/∂σ_βα.
+This explains the extra factor (1/2) needed in the GGA term of the XC potential.
 =#
 
 function max_required_derivative(functional)
@@ -181,12 +208,13 @@ function max_required_derivative(functional)
     error("Functional family $(functional.family) not known.")
 end
 
-struct DensityDerivatives
+struct DensityDerivatives  # TODO Rename to LibxcDensity as libxc-specific, refactor
     basis
     max_derivative::Int
     ρ         # density on real-space grid
     ∇ρ_real   # density gradient on real-space grid
     σ_real    # contracted density gradient on real-space grid
+    index_spin_σ  # Dict mapping a pair of spins to the index on the spin axis in σ
 end
 
 """
@@ -199,28 +227,51 @@ function DensityDerivatives(basis, max_derivative::Integer, ρ::RealFourierArray
     model.spin_polarization == :collinear && (@assert !isnothing(ρspin))
     ifft(x) = real_checked(G_to_r(basis, clear_without_conjugate!(x)))
 
-    ρF = ρ.fourier
-    σ_real = nothing
-    ∇ρ_real = nothing
-    if max_derivative > 0
-        ∇ρ_real = [ifft(im * [G[α] for G in G_vectors_cart(basis)] .* ρF)
-                   for α in 1:3]
-        # TODO The above assumes CPU arrays
-        σ_real = sum(∇ρ_real[α] .* ∇ρ_real[α] for α in 1:3)
-    end
+    n_spin    = model.n_spin_components
+    σ_real    = nothing
+    ∇ρ_real   = nothing
     if model.spin_polarization == :collinear
-        @assert max_derivative == 0
-        ρα = (ρ.real + ρspin.real) / 2
-        ρβ = (ρ.real - ρspin.real) / 2
+        ρ_real    = similar(ρ.real,    n_spin, basis.fft_size...)
+        ρ_real[1, :, :, :] .= @. (ρ.real + ρspin.real) / 2
+        ρ_real[2, :, :, :] .= @. (ρ.real - ρspin.real) / 2
+
+        if max_derivative > 0
+            ρ_fourier = similar(ρ.fourier, n_spin, basis.fft_size...)
+            ρ_fourier[1, :, :, :] .= @. (ρ.fourier + ρspin.fourier) / 2
+            ρ_fourier[2, :, :, :] .= @. (ρ.fourier - ρspin.fourier) / 2
+        end
 
-        # TODO This is the wrong place. This is specific to the interface towards Libxc.jl ...
-        ρ_real = vcat(reshape(ρα, 1, basis.fft_size...), reshape(ρβ, 1, basis.fft_size...))
+        index_spin_σ = Dict((1, 1) => 1, (1, 2) => 2, (2, 1) => 2, (2, 2) => 3)
+        @assert n_spin == 2
     else
-        # TODO This is the wrong place. This is specific to the interface towards Libxc.jl ...
-        ρ_real = reshape(ρ.real, 1, basis.fft_size...)
+        ρ_real    = reshape(ρ.real,    1, basis.fft_size...)
+        ρ_fourier = reshape(ρ.fourier, 1, basis.fft_size...)
+        index_spin_σ = Dict((1, 1) => 1)
+        @assert n_spin == 1
     end
 
-    DensityDerivatives(basis, max_derivative, ρ_real, ∇ρ_real, σ_real)
+    if max_derivative > 0
+        n_spin_σ = div((n_spin + 1) * n_spin, 2)
+        ∇ρ_real = similar(ρ.real,   n_spin, basis.fft_size..., 3)
+        σ_real  = similar(ρ.real, n_spin_σ, basis.fft_size...)
+
+        for α = 1:3
+            iGα = [im * G[α] for G in G_vectors_cart(basis)]
+            for σ = 1:n_spin
+                ∇ρ_real[σ, :, :, :, α] .= ifft(iGα .* @view ρ_fourier[σ, :, :, :])
+            end
+        end
+
+        for σ = 1:n_spin, τ=σ:n_spin
+            iστ = index_spin_σ[(σ, τ)]
+            σ_real[iστ, :, :, :] .= 0
+            @views for α in 1:3
+                σ_real[iστ, :, :, :] .+= ∇ρ_real[σ, :, :, :, α] .* ∇ρ_real[τ, :, :, :, α]
+            end
+        end
+    end
+
+    DensityDerivatives(basis, max_derivative, ρ_real, ∇ρ_real, σ_real, index_spin_σ)
 end
 
 function input_kwargs(family, density)
@@ -252,11 +303,12 @@ function apply_kernel_term_gga(terms, density, perturbation)
     # dV(r) = f(r,r') dρ(r') = (∂V/∂ρ) dρ + (∂V/∂σ) dσ
     #
     # therefore
-    # dV(r) = (∂^2ε/∂ρ^2) dρ - 2 ∇⋅((∂^2ε/∂σ∂ρ) ∇ρ) dρ
-    #       + (∂^2ε/∂ρ∂σ) dσ - 2 ∇⋅((∂^ε/∂σ^2) ∇ρ + (∂ε/∂σ) (∂∇ρ/∂σ)) dσ
+    # dV(r) = (∂^2ε/∂ρ^2) dρ - 2 ∇⋅[(∂^2ε/∂σ∂ρ) ∇ρ + (∂ε/∂σ) (∂∇ρ/∂ρ)] dρ
+    #       + (∂^2ε/∂ρ∂σ) dσ - 2 ∇⋅[(∂^ε/∂σ^2) ∇ρ  + (∂ε/∂σ) (∂∇ρ/∂σ)] dσ
     #
-    # Note dσ = 2∇ρ ∇dρ = 2∇ρ ∇dρ, therefore
+    # Note dσ = 2∇ρ d∇ρ = 2∇ρ ∇dρ, therefore
     #    - 2 ∇⋅((∂ε/∂σ) (∂∇ρ/∂σ)) dσ = - 2 ∇⋅((∂ε/∂σ) ∇dρ)
+    # and (because assumed independent variables): (∂∇ρ/∂ρ) = 0.
     #
     # Note that below the LDA term (∂^2ε/∂ρ^2) dρ is ignored (already dealt with)
     ρ   = density.ρ
@@ -265,13 +317,13 @@ function apply_kernel_term_gga(terms, density, perturbation)
     ∇dρ = perturbation.∇ρ_real
     dσ = 2sum(∇ρ[α] .* ∇dρ[α] for α in 1:3)
 
-    (
-        @view(terms.v2rhosigma[1, :, :, :]) .* dσ + divergence_real(density.basis) do α
+    @views begin
+        terms.v2rhosigma[1, :, :, :] .* dσ + divergence_real(density.basis) do α
             @. begin
-                -2 * @view(terms.v2rhosigma[1, :, :, :]) * ∇ρ[α] * dρ
-                -2 * @view(terms.v2sigma2[1, :, :, :])   * ∇ρ[α] * dσ
-                -2 * @view(terms.vsigma[1, :, :, :])     * ∇dρ[α]
+                -2 * terms.v2rhosigma[1, :, :, :] * ∇ρ[α] * dρ
+                -2 * terms.v2sigma2[1, :, :, :]   * ∇ρ[α] * dσ
+                -2 * terms.vsigma[1, :, :, :]     * ∇dρ[α]
             end
         end
-    )
+    end
 end