Factorisation of phonon test cases (#822)

Author: epolack

src/postprocess/phonon.jl

@@ -63,25 +63,19 @@ in reduced coordinates.
 @timing function compute_dynmat(basis::PlaneWaveBasis{T}, ψ, occupation; q=zero(Vec3{T}),
                                 ρ=nothing, ham=nothing, εF=nothing, eigenvalues=nothing,
                                 kwargs...) where {T}
-    model = basis.model
-    positions = model.positions
-    n_atoms = length(positions)
-    n_dim = model.n_dim
-
+    n_atoms = length(basis.model.positions)
     δρs = [zeros(complex(T), basis.fft_size..., basis.model.n_spin_components)
            for _ = 1:3, _ = 1:n_atoms]
     δoccupations = [zero.(occupation) for _ = 1:3, _ = 1:n_atoms]
     δψs = [zero.(ψ) for _ = 1:3, _ = 1:n_atoms]
-    if !isempty(ψ)
-        for s = 1:n_atoms, α = 1:n_dim
-            δHψs_αs = compute_δHψ_αs(basis, ψ, α, s, q)
-            (; δψ, δρ, δoccupation) = solve_ΩplusK_split(ham, ρ, ψ, occupation, εF,
-                                                         eigenvalues, -δHψs_αs; q,
-                                                         kwargs...)
-            δoccupations[α, s] = δoccupation
-            δρs[α, s] = δρ
-            δψs[α, s] = δψ
-        end
+    for s = 1:n_atoms, α = 1:basis.model.n_dim
+        δHψs_αs = compute_δHψ_αs(basis, ψ, α, s, q)
+        isnothing(δHψs_αs) && continue
+        (; δψ, δρ, δoccupation) = solve_ΩplusK_split(ham, ρ, ψ, occupation, εF, eigenvalues,
+                                                     -δHψs_αs; q, kwargs...)
+        δoccupations[α, s] = δoccupation
+        δρs[α, s] = δρ
+        δψs[α, s] = δψ
     end
 
     dynmats_per_term = [compute_dynmat(term, basis, ψ, occupation; ρ, δψs, δρs,
@@ -116,5 +110,7 @@ potential produced by a displacement of the atom s in the direction α.
 """
 @timing function compute_δHψ_αs(basis::PlaneWaveBasis, ψ, α, s, q)
     δHψ_per_term = [compute_δHψ_αs(term, basis, ψ, α, s, q) for term in basis.terms]
-    sum(filter(!isnothing, δHψ_per_term))
+    filter!(!isnothing, δHψ_per_term)
+    isempty(δHψ_per_term) && return nothing
+    sum(δHψ_per_term)
 end

src/response/cg.jl

@@ -28,14 +28,21 @@ function cg!(x::AbstractVector{T}, A::LinearMap{T}, b::AbstractVector{T};
     # p is the descent direction
     p = copy(c)
     n_iter = 0
-    residual_norm = zero(real(T))
+    residual_norm = norm(r)
 
     # convergence history
     converged = false
 
     # preconditioned conjugate gradient
     while n_iter < maxiter
+        # output
+        info = (; A, b, n_iter, x, r, residual_norm, converged, stage=:iterate)
+        callback(info)
         n_iter += 1
+        if (n_iter ≥ miniter) && residual_norm ≤ tol
+            converged = true
+            break
+        end
         mul!(c, A, p)
         α = γ / dot(p, c)
 
@@ -44,14 +51,6 @@ function cg!(x::AbstractVector{T}, A::LinearMap{T}, b::AbstractVector{T};
         r .= proj(r .- α .* c)
         residual_norm = norm(r)
 
-        # output
-        info = (; A, b, n_iter, x, r, residual_norm, converged, stage=:iterate)
-        callback(info)
-        if (n_iter > miniter) && residual_norm <= tol
-            converged = true
-            break
-        end
-
         # apply preconditioner and prepare next iteration
         ldiv!(c, precon, r)
         γ_prev, γ = γ, dot(r, c)

src/response/chi0.jl

@@ -349,7 +349,7 @@ end
     # occupation_threshold for which we compute the associated response δψn,
     # the others being discarded to ψ_extra.
     # We then use the extra information we have from these additional bands,
-    # non-necessarily converged, to split the sternheimer_solver with a Schur
+    # non-necessarily converged, to split the Sternheimer_solver with a Schur
     # complement.
     occ_thresh = occupation_threshold
     mask_occ   = map(occk -> findall(occnk -> abs(occnk) ≥ occ_thresh, occk), occupation)

test/phonon/ewald.jl

@@ -1,53 +1,67 @@
+@testsetup module PhononEwald
+using DFTK
+
+function model_tested(lattice::AbstractMatrix, atoms::Vector{<:DFTK.Element},
+                      positions::Vector{<:AbstractVector}; kwargs...)
+    terms = [Kinetic(),
+             Ewald()]
+    if :temperature in keys(kwargs) && kwargs[:temperature] != 0
+        terms = [terms..., Entropy()]
+    end
+    Model(lattice, atoms, positions; model_name="atomic", terms, kwargs...)
+end
+end
+
 @testitem "Phonon: Ewald: comparison to ref testcase" #=
-    =#    tags=[:phonon, :dont_test_mpi] setup=[Phonon, TestCases] begin
+    =#    tags=[:phonon, :dont_test_mpi] setup=[Phonon, PhononEwald, TestCases] begin
     using DFTK
+    using .Phonon: test_frequencies
+    using .PhononEwald: model_tested
 
-    testcase = TestCases.silicon
-    terms = [Ewald()]
-    ω_ref = [ -0.2442311083805831
-              -0.24423110838058237
-              -0.23442208238107232
-              -0.23442208238107184
-              -0.1322944535508822
-              -0.13229445355088176
-              -0.10658869539441493
-              -0.10658869539441468
-              -0.10658869539441346
-              -0.10658869539441335
-              -4.891274318712944e-16
-              -3.773447798738169e-17
-               1.659776058962626e-15
-               0.09553958285993536
-               0.18062696253387409
-               0.18062696253387464
-               0.4959725605665635
-               0.4959725605665648
-               0.49597256056656597
-               0.5498259359834827
-               0.5498259359834833
-               0.6536453595829087
-               0.6536453595829091
-               0.6536453595829103
-               0.6536453595829105
-               0.6961890494198791
-               0.6961890494198807
-               0.7251587593311752
-               0.7251587593311782
-               0.9261195383192374
-               0.9261195383192381
-               1.2533843205271504
-               1.2533843205271538
-               1.7010950724885228
-               1.7010950724885254
-               1.752506588801463 ]
-    Phonon.test_frequencies(testcase, terms, ω_ref)
+    ω_ref = [ -3.720615299046614e-12
+               1.969314371029982e-11
+               1.9739956911274832e-11
+               0.00029302379784864935
+               0.0002930237978486494
+               0.000293023797851601
+               0.0002930237978516018
+               0.0005105451353059533
+               0.0005105451353059533
+               0.000510545135311239
+               0.0005105451353112397
+               0.0005676024288436319
+               0.000591265950289604
+               0.0005912659502958081
+               0.0007328535013566558
+               0.0007328535013566561
+               0.0008109743140779055
+               0.0008109743140779056
+               0.000938673782810113
+               0.000987619635716976
+               0.0009876196357169761
+               0.0010949497272589232
+               0.0011998186659486743
+               0.0011998186659486745
+               0.001523238357971607
+               0.0019593679918607546
+               0.0022394777249719524
+               0.0022394777249719524
+               0.0024681196094789985
+               0.0024681196094789993
+               0.0024809296524054506
+               0.0025805236057819345
+               0.002614761988704579
+               0.002614761988704579
+               0.0026807773193116675
+               0.0026807773193116675 ]
+    test_frequencies(model_tested, TestCases.magnesium; ω_ref)
 end
 
 @testitem "Phonon: Ewald: comparison to automatic differentiation" #=
-    =#    tags=[:phonon, :slow, :dont_test_mpi] setup=[Phonon, TestCases] begin
+    =#    tags=[:phonon, :slow, :dont_test_mpi] setup=[Phonon, PhononEwald, TestCases] begin
     using DFTK
-    testcase = TestCases.silicon
+    using .Phonon: test_frequencies
+    using .PhononEwald: model_tested
 
-    terms = [Ewald()]
-    Phonon.test_rand_frequencies(testcase, terms)
+    test_frequencies(model_tested, TestCases.magnesium)
 end

test/phonon/helpers.jl

@@ -2,41 +2,10 @@
 @testsetup module Phonon
 using Test
 using DFTK
-using DFTK: TermAtomicLocal, TermAtomicNonlocal
-using DFTK: compute_dynmat_cart, setindex, dynmat_red_to_cart, normalize_kpoint_coordinate
+using DFTK: normalize_kpoint_coordinate
 using LinearAlgebra
 using ForwardDiff
-
-# We do not take the square root to compare eigenvalues with machine precision.
-function squared_frequencies(matrix)
-    n, m = size(matrix, 1), size(matrix, 2)
-    Ω = eigvals(reshape(matrix, n*m, n*m))
-    real(Ω)
-end
-
-# Reference against automatic differentiation.
-function reference_squared_frequencies(basis; kwargs...)
-    model = basis.model
-    n_atoms = length(model.positions)
-    n_dim = model.n_dim
-    T = eltype(model.lattice)
-    dynmat_ad = zeros(T, 3, n_atoms, 3, n_atoms)
-    for s = 1:n_atoms, α = 1:n_dim
-        displacement = zero.(model.positions)
-        displacement[s] = setindex(displacement[s], one(T), α)
-        dynmat_ad[:, :, α, s] = -ForwardDiff.derivative(zero(T)) do ε
-            lattice = convert(Matrix{eltype(ε)}, model.lattice)
-            positions = ε*displacement .+ model.positions
-            model_disp = Model(convert(Model{eltype(ε)}, model); lattice, positions)
-            # TODO: Would be cleaner with PR #675.
-            basis_disp_bs = PlaneWaveBasis(model_disp; Ecut=5)
-            forces = compute_forces(basis_disp_bs, nothing, nothing)
-            stack(forces)
-        end
-    end
-    hessian_ad = DFTK.dynmat_red_to_cart(model, dynmat_ad)
-    sort(squared_frequencies(hessian_ad))
-end
+using FiniteDifferences
 
 function generate_random_supercell(; max_length=6)
     n_max = min(max_length, 5)
@@ -57,47 +26,97 @@ function generate_supercell_qpoints(; supercell_size=generate_random_supercell()
     (; supercell_size, qpoints)
 end
 
-# Test against a reference array.
-function test_frequencies(testcase, terms, ω_ref; tol=1e-9, supercell_size=[2, 1, 3])
-    model = Model(testcase.lattice, testcase.atoms, testcase.positions; terms)
-    basis_bs = PlaneWaveBasis(model; Ecut=5)
+function test_approx_frequencies(ω_uc, ω_ref; tol=1e-10)
+    # Because three eigenvalues should be close to zero and the square root near
+    # zero decrease machine accuracy, we expect at least ``3×2×2 - 3 = 9``
+    # eigenvalues to have norm related to the accuracy of the SCF convergence
+    # parameter and the rest to be larger.
+    n_dim = 3
+    n_atoms = length(ω_uc) ÷ 3
+
+    @test count(abs.(ω_uc - ω_ref) .< sqrt(tol)) ≥ n_dim*n_atoms - n_dim
+    @test count(sqrt(tol) .< abs.(ω_uc - ω_ref) .< tol) ≤ n_dim
+end
+
+function test_frequencies(model_tested, testcase; ω_ref=nothing, Ecut=7, kgrid=[2, 1, 3],
+                          tol=1e-12, randomize=false, compute_ref=nothing)
+    supercell_size = randomize ? generate_random_supercell() : kgrid
+    qpoints = generate_supercell_qpoints(; supercell_size).qpoints
+    scf_tol = tol
+    χ0_tol  = scf_tol/10
+    scf_kwargs = (; is_converged=ScfConvergenceDensity(scf_tol),
+                  diagtolalg=AdaptiveDiagtol(; diagtol_max=scf_tol))
 
-    phonon = (; supercell_size, generate_supercell_qpoints(; supercell_size).qpoints)
+    model = model_tested(testcase.lattice, testcase.atoms, testcase.positions;
+                         symmetries=false, testcase.temperature)
+    nbandsalg = AdaptiveBands(model; occupation_threshold=1e-10)
+    scf_kwargs = merge(scf_kwargs, (; nbandsalg))
+    basis = PlaneWaveBasis(model; Ecut, kgrid)
+    scfres = self_consistent_field(basis; scf_kwargs...)
 
-    ω_uc = sort!(reduce(vcat, map(phonon.qpoints) do q
-        hessian = compute_dynmat_cart(basis_bs, [], []; q)
-        squared_frequencies(hessian)
+    ω_uc = sort!(reduce(vcat, map(qpoints) do q
+        dynamical_matrix = compute_dynmat(scfres; q, tol=χ0_tol)
+        phonon_modes_cart(basis, dynamical_matrix).frequencies
     end))
 
-    @test norm(ω_uc - ω_ref) < tol
-end
+    !isnothing(ω_ref) && return test_approx_frequencies(ω_uc, ω_ref; tol=10scf_tol)
 
-# Random test. Slow but more robust than against some reference.
-# TODO: Will need rework for local term in future PR.
-function test_rand_frequencies(testcase, terms; tol=1e-9)
-    model = Model(testcase.lattice, testcase.atoms, testcase.positions; terms)
-    basis_bs = PlaneWaveBasis(model; Ecut=5)
+    supercell = create_supercell(testcase.lattice, testcase.atoms, testcase.positions,
+                                 supercell_size)
+    model_supercell = model_tested(supercell.lattice, supercell.atoms, supercell.positions;
+                               symmetries=false, testcase.temperature)
+    nbandsalg = AdaptiveBands(model_supercell; occupation_threshold=1e-10)
+    scf_kwargs = merge(scf_kwargs, (; nbandsalg))
+    basis_supercell = PlaneWaveBasis(model_supercell; Ecut, kgrid=[1, 1, 1])
+    scfres_supercell = self_consistent_field(basis_supercell; scf_kwargs...)
 
-    supercell_size = supercell_size=generate_random_supercell()
-    phonon = (; supercell_size, generate_supercell_qpoints(; supercell_size).qpoints)
+    dynamical_matrix_sc = compute_dynmat(scfres_supercell; tol=χ0_tol)
+    ω_sc = sort(phonon_modes_cart(basis_supercell, dynamical_matrix_sc).frequencies)
+    test_approx_frequencies(ω_uc, ω_sc; tol=10scf_tol)
 
-    ω_uc = []
-    for q in phonon.qpoints
-        hessian = compute_dynmat_cart(basis_bs, [], []; q)
-        push!(ω_uc, squared_frequencies(hessian))
-    end
-    ω_uc = sort!(collect(Iterators.flatten(ω_uc)))
+    isnothing(compute_ref) && return
 
-    supercell = create_supercell(testcase.lattice, testcase.atoms, testcase.positions,
-                                 phonon.supercell_size)
-    model_supercell = Model(supercell.lattice, supercell.atoms, supercell.positions; terms)
-    basis_supercell_bs = PlaneWaveBasis(model_supercell; Ecut=5)
-    hessian_supercell = compute_dynmat_cart(basis_supercell_bs, [], [])
-    ω_supercell = sort(squared_frequencies(hessian_supercell))
-    @test norm(ω_uc - ω_supercell) < tol
+    dynamical_matrix_ref = compute_dynmat_ref(scfres_supercell.basis, model_tested; Ecut,
+                                              kgrid=[1, 1, 1], scf_tol, method=compute_ref)
+    ω_ref = sort(phonon_modes_cart(basis_supercell, dynamical_matrix_ref).frequencies)
+
+    test_approx_frequencies(ω_uc, ω_ref; tol=10scf_tol)
+end
 
-    ω_ad = reference_squared_frequencies(basis_supercell_bs)
+# Reference results using finite differences or automatic differentiation.
+# This should be run by hand to obtain the reference values of the quick computations of the
+# tests, as they are too slow for CI runs.
+function compute_dynmat_ref(basis, model_tested; Ecut=5, kgrid=[1,1,1], scf_tol, method=:ad)
+    # TODO: Cannot use symmetries: https://github.com/JuliaMolSim/DFTK.jl/issues/817
+    @assert isone(only(basis.model.symmetries))
+    @assert method ∈ [:ad, :fd]
 
-    @test norm(ω_ad - ω_supercell) < tol
+    model = basis.model
+    n_atoms = length(model.positions)
+    n_dim = model.n_dim
+    T = eltype(model.lattice)
+    dynmat = zeros(T, 3, n_atoms, 3, n_atoms)
+    scf_kwargs = (; is_converged=ScfConvergenceDensity(scf_tol),
+                  diagtolalg=AdaptiveDiagtol(; diagtol_max=scf_tol))
+
+    diff_fn = method == :ad ? ForwardDiff.derivative : FiniteDifferences.central_fdm(5, 1)
+    for s = 1:n_atoms, α = 1:n_dim
+        displacement = zero.(model.positions)
+        displacement[s] = DFTK.setindex(displacement[s], one(T), α)
+        dynmat[:, :, α, s] = -diff_fn(zero(T)) do ε
+            lattice = convert(Matrix{eltype(ε)}, model.lattice)
+            positions = ε*displacement .+ model.positions
+            model_disp = model_tested(lattice, model.atoms, positions; symmetries=false,
+                                      model.temperature)
+            # TODO: Would be cleaner with PR #675.
+            basis_disp = PlaneWaveBasis(model_disp; Ecut, kgrid)
+            nbandsalg = AdaptiveBands(model_disp; occupation_threshold=1e-10)
+            scf_kwargs = merge(scf_kwargs, (; nbandsalg))
+            scfres_disp = self_consistent_field(basis_disp; scf_kwargs...)
+            forces = compute_forces(scfres_disp)
+            stack(forces)
+        end
+    end
+    dynmat
 end
 end