Implement forces (#106)

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

Project.toml

@@ -17,14 +17,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
-NLSolversBase = "d41bc354-129a-5804-8e4c-c37616107c6c"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -36,7 +35,8 @@ julia = "1.2"
 [extras]
 DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
 IntervalArithmetic = "d1acc4aa-44c8-5952-acd4-ba5d80a2a253"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "DoubleFloats", "IntervalArithmetic"]
+test = ["Test", "DoubleFloats", "IntervalArithmetic", "Random"]

src/DFTK.jl

@@ -153,4 +153,7 @@ export load_density
 export load_atoms
 include("external/etsf_nanoquanta.jl")
 
+export forces
+include("forces.jl")
+
 end # module DFTK

src/energy_ewald.jl

@@ -1,4 +1,5 @@
 using SpecialFunctions: erfc
+using ForwardDiff
 
 """
     energy_ewald(lattice, [recip_lattice, ]charges, positions; η=nothing)
@@ -8,9 +9,10 @@ charges in a uniform background of compensating charge to yield net
 neutrality. the `lattice` and `recip_lattice` should contain the
 lattice and reciprocal lattice vectors as columns. `charges` and
 `positions` are the point charges and their positions (as an array of
-arrays) in fractional coordinates.
+arrays) in fractional coordinates. If `forces` is not nothing, minus the derivatives
+of the energy with respect to `positions` is computed.
 """
-function energy_ewald(lattice, charges, positions; η=nothing)
+function energy_ewald(lattice, charges, positions; η=nothing, forces=nothing)
     T = eltype(lattice)
 
     for i=1:3
@@ -21,18 +23,22 @@ function energy_ewald(lattice, charges, positions; η=nothing)
             return T(0)
         end
     end
-    energy_ewald(lattice, T(2π) * inv(lattice'), charges, positions, η=η)
+    energy_ewald(lattice, T(2π) * inv(lattice'), charges, positions; η=η, forces=forces)
 end
-function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing)
+function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing, forces=nothing)
     T = eltype(lattice)
     @assert T == eltype(recip_lattice)
-    @assert(length(charges) == length(positions),
-            "Length of charges and positions does not match")
+    @assert length(charges) == length(positions)
     if η === nothing
         # Balance between reciprocal summation and real-space summation
         # with a slight bias towards reciprocal summation
         η = sqrt(sqrt(T(1.69) * norm(recip_lattice ./ 2T(π)) / norm(lattice))) / 2
     end
+    if forces !== nothing
+        @assert size(forces) == size(positions)
+        forces_real = copy(forces)
+        forces_recip = copy(forces)
+    end
 
     #
     # Numerical cutoffs
@@ -89,11 +95,24 @@ function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing)
 
             any_term_contributes = true
             sum_recip += sum_strucfac * exp(-exponent) / Gsq
+
+            if forces !== nothing
+                for (ir, r) in enumerate(positions)
+                    Z = charges[ir]
+                    dc = -Z*2T(π)*G*sin(2T(π) * dot(r, G))
+                    ds = +Z*2T(π)*G*cos(2T(π) * dot(r, G))
+                    dsum = 2cos_strucfac*dc + 2sin_strucfac*ds
+                    forces_recip[ir] -= dsum * exp(-exponent)/Gsq
+                end
+            end
         end
         gsh += 1
     end
     # Amend sum_recip by proper scaling factors:
-    sum_recip = sum_recip * 4T(π) / abs(det(lattice))
+    sum_recip *= 4T(π) / abs(det(lattice))
+    if forces !== nothing
+        forces_recip .*= 4T(π) / abs(det(lattice))
+    end
 
     #
     # Real-space sum
@@ -109,7 +128,12 @@ function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing)
 
         # Loop over R vectors for this shell patch
         for R in shell_indices(rsh)
-            for (ti, Zi) in zip(positions, charges), (tj, Zj) in zip(positions, charges)
+            for i = 1:length(positions), j = 1:length(positions)
+                ti = positions[i]
+                Zi = charges[i]
+                tj = positions[j]
+                Zj = charges[j]
+
                 # Avoid self-interaction
                 if rsh == 0 && ti == tj
                     continue
@@ -124,9 +148,18 @@ function energy_ewald(lattice, recip_lattice, charges, positions; η=nothing)
 
                 any_term_contributes = true
                 sum_real += Zi * Zj * erfc(η * dist) / dist
-            end # iat, jat
+                if forces !== nothing
+                    # Use ForwardDiff here because I'm lazy. TODO do it properly
+                    forces_real[i] -= ForwardDiff.gradient(r -> (dist=norm(lattice * (r - tj - R)); Zi * Zj * erfc(η * dist) / dist), ti)
+                    forces_real[j] -= ForwardDiff.gradient(r -> (dist=norm(lattice * (ti - r - R)); Zi * Zj * erfc(η * dist) / dist), tj)
+                end
+            end # i,j
         end # R
         rsh += 1
     end
-    (sum_recip + sum_real) / 2  # Divide by 1/2 (because of double counting)
+    energy = (sum_recip + sum_real) / 2  # Divide by 1/2 (because of double counting)
+    if forces !== nothing
+        forces .= (forces_recip .+ forces_real) ./ 2
+    end
+    energy
 end

src/energy_nuclear.jl

@@ -32,9 +32,9 @@ negative charge following the Ewald summation procedure. The function assumes al
 are modelled as point charges. If pseudopotentials are used, one needs to additionally
 compute the `energy_nuclear_psp_correction` to get the correct energy.
 """
-energy_nuclear_ewald(model::Model; η=nothing) = energy_nuclear_ewald(model.lattice, model.atoms)
-function energy_nuclear_ewald(lattice, atoms; η=nothing)
+energy_nuclear_ewald(model::Model; kwargs...) = energy_nuclear_ewald(model.lattice, model.atoms; kwargs...)
+function energy_nuclear_ewald(lattice, atoms; kwargs...)
     charges = [charge_ionic(type) for (type, positions) in atoms for pos in positions]
     positions = [pos for (elem, positions) in atoms for pos in positions]
-    energy_ewald(lattice, charges, positions; η=η)
+    energy_ewald(lattice, charges, positions; kwargs...)
 end

src/forces.jl

@@ -0,0 +1,80 @@
+"""
+Computes minus the derivatives of the energy with respect to atomic positions.
+"""
+function forces(basis, ρ, Psi, occupation)
+    # By generalized Hellmann-Feynman, dE/dR = ∂E/∂R. The atomic
+    # positions come up explicitly only in the external and nonlocal
+    # part of the energy,
+
+    # TODO this assumes that the build_external and build_nonlocal are consistent with model.atoms
+    # Find a way to move this computation closer to each term
+
+    # minus signs here because f = -∇E
+    model = basis.model
+    T = real(eltype(Psi[1]))
+
+    forces = []
+    for (type, positions) in model.atoms
+        @assert type.psp != nothing
+        # force for this atom type
+        f = zeros(Vec3{T}, length(positions))
+
+        # Local part
+        @assert model.build_external !== nothing
+        # energy = sum of form_factor(G) * struct_factor(G) * rho(G)
+        # where struct_factor(G) = cis(-2T(π) * dot(G, r))
+        form_factors = [Complex{T}(eval_psp_local_fourier(type.psp, model.recip_lattice * G))
+                        for G in G_vectors(basis)] ./ sqrt(model.unit_cell_volume)
+        form_factors[1] = 0
+
+        for (ir, r) in enumerate(positions)
+            f[ir] -= real(sum(conj(ρ.fourier[iG]) .*
+                              form_factors[iG] .*
+                              cis(-2T(π) * dot(G, r)) .*
+                              (-2T(π)) .* G .* im
+                              for (iG, G) in enumerate(G_vectors(basis))))
+        end
+
+        # Nonlocal part
+        @assert model.build_nonlocal !== nothing
+        C = build_projection_coefficients_(type.psp)
+        for (ir, r) in enumerate(positions)
+            fr = zeros(T, 3)
+            for idir = 1:3
+                for (ik, kpt) in enumerate(basis.kpoints)
+                    # energy terms are of the form <psi, P C P' psi>, where P(G) = form_factor(G) * structure_factor(G)
+                    qs = [model.recip_lattice * (kpt.coordinate + G) for G in G_vectors(kpt)]
+                    form_factors = build_form_factors(type.psp, qs)
+                    structure_factors = [cis(-2T(π)*dot(kpt.coordinate + G, r)) for G in G_vectors(kpt)]
+                    P = structure_factors .* form_factors ./ sqrt(model.unit_cell_volume)
+                    dPdR = [-2T(π)*im*(kpt.coordinate + G)[idir] for G in G_vectors(kpt)] .* P
+
+                    # TODO BLASify this further
+                    for iband = 1:size(Psi[ik], 2)
+                        psi = Psi[ik][:, iband]
+                        fr[idir] -= basis.kweights[ik] * occupation[ik][iband] *
+                                    real(dot(psi, P*C*dPdR'*psi) + dot(psi, dPdR*C*P'*psi))
+                    end
+                end
+            end
+            f[ir] += fr
+        end
+
+        push!(forces, f)
+    end
+
+    # Add Ewald forces
+    forces_ewald = zeros(Vec3{T}, sum(length(positions) for (elem, positions) in model.atoms))
+    energy_nuclear_ewald(model; forces=forces_ewald)
+
+    count = 1
+    for i = 1:length(model.atoms)
+        for j = 1:length(model.atoms[i][2])
+            forces[i][j] += forces_ewald[count]
+            count += 1
+        end
+    end
+
+    forces
+end
+forces(scfres) = forces(scfres.ham.basis, scfres.ρ, scfres.Psi, scfres.occupation)

test/energy_ewald.jl → test/ewald.jl

@@ -46,3 +46,21 @@ end
     γ_E = energy_ewald(lattice, charges, positions)
     @test abs(γ_E - ref) < 1e-7
 end
+
+@testset "Forces" begin
+    lattice = [0.0  5.131570667152971 5.131570667152971;
+               5.131570667152971 0.0 5.131570667152971;
+               5.131570667152971 5.131570667152971  0.0]
+    # perturb positions away from equilibrium to get nonzero force
+    positions = [ones(3)/8+rand(3)/20, -ones(3)/8]
+    charges = [14, 14]
+
+    forces = zeros(Vec3{Float64}, 2)
+    γ1 = energy_ewald(lattice, charges, positions, forces=forces)
+
+    # Compare forces to finite differences
+    disp = [rand(3)/20, rand(3)/20]
+    ε = 1e-8
+    γ2 = energy_ewald(lattice, charges, positions .+ ε .* disp)
+    @test (γ2-γ1)/ε ≈ -dot(disp, forces) atol=abs(γ1*1e-6)
+end

test/forces.jl

@@ -0,0 +1,37 @@
+using DFTK
+using Test
+include("testcases.jl")
+
+@testset "Forces" begin
+    function energy(pos)
+        kgrid = [2, 1, 2]        # k-Point grid
+        Ecut = 5                # kinetic energy cutoff in Hartree
+
+        # Setup silicon lattice
+        lattice = silicon.lattice
+        Si = Element(silicon.atnum, psp=load_psp(silicon.psp))
+        atoms = [Si => pos]
+        model = model_dft(silicon.lattice, :lda_xc_teter93, atoms)
+        kcoords, ksymops = bzmesh_ir_wedge(kgrid, lattice, atoms)
+        basis = PlaneWaveBasis(model, Ecut, kcoords, ksymops)
+
+        n_bands_scf = Int(model.n_electrons / 2)
+        ham = Hamiltonian(basis, guess_density(basis))
+        scfres = self_consistent_field(ham, n_bands_scf, tol=1e-12)
+
+        sum(values(scfres.energies)), forces(scfres)
+    end
+
+    pos1 = [(ones(3)+0.1*randn(3))/8, -ones(3)/8]
+    disp = randn(3)
+    ε = 1e-8
+    pos2 = [pos1[1]+ε*disp, pos1[2]]
+
+    E1, F1 = energy(pos1)
+    E2, F2 = energy(pos2)
+
+    diff_findiff = -(E2-E1)/ε
+    diff_forces = dot(F1[1][1], disp)
+
+    @test abs(diff_findiff - diff_forces) < 1e-6
+end

test/runtests.jl

@@ -1,5 +1,6 @@
 using Test
 using DFTK
+using Random
 
 #
 # This test suite test arguments. For example:
@@ -26,6 +27,8 @@ else
     println("   Running tests (TAGS = $(join(TAGS, ", "))).")
 end
 
+# Initialise seed
+Random.seed!(0)
 
 # Wrap in an outer testset to get a full report if one test fails
 @testset "DFTK.jl" begin
@@ -65,11 +68,12 @@ end
     end
 
     if "all" in TAGS
-        include("energy_ewald.jl")
+        include("ewald.jl")
         include("energy_nuclear.jl")
         include("occupation.jl")
         include("energies_guess_density.jl")
         include("compute_density.jl")
+        include("forces.jl")
     end
 
     if "all" in TAGS