Fixes of some edge cases in checkpoint handling (#937)

ext/DFTKJLD2Ext.jl

@@ -102,6 +102,7 @@ function load_scfres_jld2(jld, basis; skip_hamiltonian, strict)
         scfdict = Dict{Symbol, Any}(
             :εF       => jld["εF"],
             :ρ        => get(jld, "ρ", nothing),
+            :ψ        => nothing,
             :energies => DFTK.Energies(e_keys, e_values),
         )
         for key in jld["scfres_extra_keys"]
@@ -135,16 +136,15 @@ function load_scfres_jld2(jld, basis; skip_hamiltonian, strict)
     has_ψ = MPI.bcast(has_ψ, 0, MPI.COMM_WORLD)
     if has_ψ
         n_G_vectors = reshape_and_scatter(jld, "kpt_n_G_vectors")
-
         basisok = all(n_G_vectors[ik] == length(DFTK.G_vectors(basis, kpt))
                       for (ik, kpt) in enumerate(basis.kpoints))
         basisok = DFTK.mpi_min(basisok, basis.comm_kpts)
-        basisok || error("Mismatch in number of G-vectors per k-point.")
-
-        ψ_padded = reshape_and_scatter(jld, "ψ")
-        scfdict[:ψ] = DFTK.unblockify_ψ(ψ_padded, n_G_vectors)
-    else
-        scfdict[:ψ] = nothing
+        if basisok
+            ψ_padded = reshape_and_scatter(jld, "ψ")
+            scfdict[:ψ] = DFTK.unblockify_ψ(ψ_padded, n_G_vectors)
+        elseif strict
+            error("Mismatch in number of G-vectors per k-point.")
+        end
     end
 
     if !skip_hamiltonian && has_ψ && !isnothing(scfdict[:ρ])

src/Energies.jl

@@ -25,7 +25,7 @@ Base.pairs(energies::Energies)          = (k => v for (k, v) in zip(energies.key
 Base.iterate(energies::Energies)        = iterate(pairs(energies))
 Base.iterate(energies::Energies, state) = iterate(pairs(energies), state)
 Base.length(energies::Energies)         = length(energies.keys)
-Base.haskey(energies::Energies, key)    = haskey(energies, key)
+Base.haskey(energies::Energies, key)    = !isnothing(findfirst(isequal(key), energies.keys))
 
 Base.propertynames(energies::Energies, ::Bool=false) = append!(Symbol.(energies.keys), :total)
 function Base.getproperty(energies::Energies{T}, x::Symbol) where {T}

src/scf/mixing.jl

@@ -163,7 +163,7 @@ Important `kwargs` passed on to [`χ0Mixing`](@ref)
 - `verbose`: Run the GMRES in verbose mode.
 - `reltol`: Relative tolerance for GMRES
 """
-function HybridMixing(; εr=1.0, kTF=0.8, localization=identity,
+function HybridMixing(; εr=10.0, kTF=0.8, localization=identity,
                       adjust_temperature=IncreaseMixingTemperature(), kwargs...)
     χ0terms = [DielectricModel(; εr, kTF, localization),
                LdosModel(;adjust_temperature)]

src/scf/scf_callbacks.jl

@@ -121,6 +121,7 @@ function (conv::ScfConvergenceEnergy)(info)
     if last(info.history_Δρ) > 10sqrt(conv.tolerance)
         return false  # The ρ change should also be small to avoid the SCF being just stuck
     end
+    length(info.history_Etot) < 2 && return false
     ΔE = (info.history_Etot[end-1] - info.history_Etot[end])
     ΔE < conv.tolerance
 end

src/scf/self_consistent_field.jl

@@ -14,29 +14,29 @@ function kwargs_scf_checkpoints(basis::AbstractBasis;
                                 ρ=guess_density(basis),
                                 ψ=nothing, save_ψ=false,
                                 kwargs...)
-    callback = callback ∘ ScfSaveCheckpoints(; filename, save_ψ)
     if isfile(filename)
         # Disable strict checking, since we can live with only the density data
         previous = load_scfres(filename, basis; skip_hamiltonian=true, strict=false)
 
         # If we can expect the guess to be good, tighten the diagtol.
-        has_ρ = !isnothing(previous.ρ)
-        consistent_kpts = hasproperty(previous, :eigenvalues)
-        if has_ρ && consistent_kpts && hasproperty(previous, :history_Δρ)
-            diagtol_first = determine_diagtol(diagtolalg, previous)
-        elseif has_ρ
-            diagtol_first = diagtolalg.diagtol_max
-        else
-            diagtol_first = diagtolalg.diagtol_first
+        if !isnothing(previous.ρ)
+            ρ = previous.ρ
+            consistent_kpts = hasproperty(previous, :eigenvalues)
+            if consistent_kpts && hasproperty(previous, :history_Δρ)
+                diagtol_first = determine_diagtol(diagtolalg, previous)
+            else
+                diagtol_first = diagtolalg.diagtol_max
+            end
+            diagtolalg = AdaptiveDiagtol(; diagtol_first,
+                                           diagtolalg.diagtol_max,
+                                           diagtolalg.diagtol_min,
+                                           diagtolalg.ratio_ρdiff)
         end
-        diagtolalg = AdaptiveDiagtol(; diagtol_first,
-                                       diagtolalg.diagtol_max,
-                                       diagtolalg.diagtol_min,
-                                       diagtolalg.ratio_ρdiff)
-        return (; callback, diagtolalg, previous.ψ, previous.ρ, kwargs...)
-    else
-        return (; callback, diagtolalg, ψ, ρ, kwargs...)
+        ψ = something(previous.ψ, Some(ψ))
     end
+
+    callback = callback ∘ ScfSaveCheckpoints(; filename, save_ψ)
+    (; callback, diagtolalg, ψ, ρ, kwargs...)
 end
 
 

test/energy_orbital_eigenvalues.jl

@@ -0,0 +1,45 @@
+@testitem "Energy from orbital eigenvalues"  setup=[TestCases] begin
+using DFTK
+
+function total_energy_from_eigenvalues(energies::Energies, ham::Hamiltonian,
+                                       ρ, eigenvalues, occupation)
+    # Orbital energies take care of electron-electron and electron-nuclear
+    # interaction energies, but suffer from double counting for the Hartree
+    # and from introducing a ∫ρ vxc contribution for the XC term, which is
+    # not the XC energy.
+    data = map(eigenvalues, occupation) do εk, occk
+        sum(εk .* occk)
+    end
+    sum_eigenvalues = DFTK.weighted_ksum(ham.basis, data)
+
+    # Keys for the nuclear-nuclear interaction and Entropy term
+    sum_energies = 0.0
+    for key in ("Ewald", "PspCorrection", "Entropy")
+        haskey(energies, key) || continue
+        sum_energies += energies[key]
+    end
+
+    # Correcting for Hartree double counting and XC energy
+    @assert ham.basis.model.n_spin_components == 1
+    ρtot = DFTK.total_density(ρ)
+    i_xc = findfirst(t -> t isa DFTK.TermXc, ham.basis.terms)
+    @assert !isnothing(i_xc)
+    pot_xc = ham[1].operators[i_xc].potential
+    energy_xcpot = sum(pot_xc .* ρtot) * ham.basis.dvol
+    energy_correction = - energies.Hartree + energies.Xc - energy_xcpot
+
+    sum_eigenvalues + sum_energies + energy_correction
+end
+function total_energy_from_eigenvalues(scfres::NamedTuple)
+    total_energy_from_eigenvalues(scfres.energies, scfres.ham, scfres.ρ,
+                                  scfres.eigenvalues, scfres.occupation)
+end
+
+silicon = TestCases.silicon
+model   = model_PBE(silicon.lattice, silicon.atoms, silicon.positions, temperature=1e-2)
+basis   = PlaneWaveBasis(model; Ecut=15, kgrid=(1, 2, 3), kshift=(0, 1/2, 0))
+scfres  = self_consistent_field(basis; tol=1e-6)
+
+etot_eigenvalues = total_energy_from_eigenvalues(scfres)
+@test abs(etot_eigenvalues - scfres.energies.total) < 1e-5
+end