ForwardDiff for response derivatives with PspUpf pseudopotentials (#1013

)

Co-authored-by: Michael F. Herbst <info@michael-herbst.com>

src/terms/xc.jl

@@ -337,13 +337,18 @@ end
 
 
 function compute_kernel(term::TermXc, basis::PlaneWaveBasis; ρ, kwargs...)
-    density = LibxcDensities(basis, 0, ρ, nothing)
     n_spin  = basis.model.n_spin_components
     @assert 1 ≤ n_spin ≤ 2
     if !all(family(xc) == :lda for xc in term.functionals)
         error("compute_kernel only implemented for LDA")
     end
 
+    # Add the model core charge density (non-linear core correction)
+    if !isnothing(term.ρcore)
+        ρ = ρ + term.ρcore
+    end
+
+    density = LibxcDensities(basis, 0, ρ, nothing)
     kernel = kernel_terms(term.functionals, density).Vρρ
     fac = term.scaling_factor
     if n_spin == 1
@@ -372,6 +377,11 @@ function apply_kernel(term::TermXc, basis::PlaneWaveBasis{T}, δρ::AbstractArra
               correction.")
     end
 
+    # Add the model core charge density (non-linear core correction)
+    if !isnothing(term.ρcore)
+        ρ = ρ + term.ρcore
+    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)

src/workarounds/forwarddiff_rules.jl

@@ -126,7 +126,11 @@ function construct_value(psp::PspHgh{T}) where {T <: Dual}
            psp.identifier,
            psp.description)
 end
-
+function construct_value(psp::PspUpf{T,I}) where {T <: AbstractFloat, I <: AbstractArray{<:AbstractFloat}}
+    # NOTE: This permits non-Dual UPF pseudos to be used in ForwardDiff computations,
+    #       but does not yet permit response derivatives w.r.t. UPF parameters.
+    psp
+end
 
 function construct_value(basis::PlaneWaveBasis{T}) where {T <: Dual}
     # NOTE: This is a pretty slow function as it *recomputes* basically

test/forwarddiff.jl

@@ -5,15 +5,15 @@
     using LinearAlgebra
     silicon = TestCases.silicon
 
-    function compute_force(ε1, ε2; metal=false, tol=1e-10)
+    function compute_force(ε1, ε2; metal=false, tol=1e-10, atoms=silicon.atoms)
         T = promote_type(typeof(ε1), typeof(ε2))
         pos = [[1.01, 1.02, 1.03] / 8, -ones(3) / 8 + ε1 * [1., 0, 0] + ε2 * [0, 1., 0]]
         if metal
             # Silicon reduced HF is metallic
-            model = model_DFT(Matrix{T}(silicon.lattice), silicon.atoms, pos;
+            model = model_DFT(Matrix{T}(silicon.lattice), atoms, pos;
                               functionals=[], temperature=1e-3)
         else
-            model = model_DFT(Matrix{T}(silicon.lattice), silicon.atoms, pos;
+            model = model_DFT(Matrix{T}(silicon.lattice), atoms, pos;
                               functionals=LDA())
         end
         basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2], kshift=[0, 0, 0])
@@ -54,6 +54,17 @@
         derivative_ε1 = ForwardDiff.derivative(ε1 -> compute_force(ε1, 0.0; metal), 0.0)
         @test norm(derivative_ε1 - derivative_ε1_fd) < 1e-4
     end
+
+    @testset "Using PspUpf" begin
+        Si = ElementPsp(:Si; psp=load_psp(silicon.psp_upf))
+        atoms = [Si, Si]
+
+        derivative_ε1_fd = let ε1 = 1e-5
+            (compute_force(ε1, 0.0; atoms) - compute_force(-ε1, 0.0; atoms)) / 2ε1
+        end
+        derivative_ε1 = ForwardDiff.derivative(ε1 -> compute_force(ε1, 0.0; atoms), 0.0)
+        @test norm(derivative_ε1 - derivative_ε1_fd) < 1e-4
+    end
 end
 
 @testitem "scfres PSP sensitivity using ForwardDiff" #=

test/kernel.jl

@@ -1,24 +1,27 @@
 @testitem "Kernels" setup=[TestCases] begin
     using DFTK
     using LinearAlgebra
+    using LazyArtifacts
 
-    function test_kernel(spin_polarization, termtype; test_compute=true)
+    function test_kernel(spin_polarization, termtype; test_compute=true, psp=TestCases.silicon.psp_hgh)
         kgrid  = [2, 2, 2]
         kshift = ones(3) / 2
         testcase = TestCases.silicon
+        Si = ElementPsp(TestCases.silicon.atnum; psp=load_psp(psp))
+        atoms = [Si, Si]
         ε   = 1e-8
         tol = 1e-5
 
         xcsym = (termtype isa Xc) ? join(string.(termtype.functionals), " ") : ""
-        @testset "Kernel $(typeof(termtype)) $xcsym ($spin_polarization)" begin
+        @testset "Kernel $(typeof(termtype)) $xcsym ($spin_polarization) $(psp)" begin
             magnetic_moments = []
             n_spin = 1
             if spin_polarization == :collinear
                 magnetic_moments = 2rand(2)
                 n_spin = 2
             end
 
-            model = Model(testcase.lattice, testcase.atoms, testcase.positions;
+            model = Model(testcase.lattice, atoms, testcase.positions;
                           terms=[termtype], magnetic_moments, spin_polarization)
             @test model.n_spin_components == n_spin
             basis = PlaneWaveBasis(model; Ecut=2, kgrid, kshift)
@@ -50,7 +53,7 @@
 
 
     function test_kernel_collinear_vs_noncollinear(termtype)
-        Ecut=2
+        Ecut = 2
         kgrid = [2, 2, 2]
         kshift = ones(3) / 2
         testcase = TestCases.silicon
@@ -97,4 +100,10 @@
     test_kernel(:collinear, Xc([:gga_c_pbe]), test_compute=false)
     test_kernel(:collinear, Xc([:gga_x_pbe]), test_compute=false)
     test_kernel(:collinear, Xc([:gga_x_pbe, :gga_c_pbe]), test_compute=false)
+
+    @testset "Non-linear core correction (NLCC)" begin
+        psp = TestCases.silicon.psp_upf  # PseudoDojo v0.4.1 Si includes NLCC
+        @test DFTK.has_core_density(load_psp(psp))
+        test_kernel(:none, Xc([:lda_xc_teter93]); psp)
+    end
 end