From 7dae2ab6f86db2a099adb847020460be63c96bae Mon Sep 17 00:00:00 2001
From: Antoine Levitt <antoine.levitt@gmail.com>
Date: Sun, 19 Jan 2020 13:35:45 +0100
Subject: [PATCH] style

---
 src/energy_nuclear.jl |  6 +++---
 src/forces.jl         | 38 +++++++++++++++++++++-----------------
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/energy_nuclear.jl b/src/energy_nuclear.jl
index 85e21dd7c8..f2153aa95e 100644
--- a/src/energy_nuclear.jl
+++ b/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
diff --git a/src/forces.jl b/src/forces.jl
index f99b39eaab..823eb06f69 100644
--- a/src/forces.jl
+++ b/src/forces.jl
@@ -25,13 +25,15 @@ function forces(scfres)
         # Local part
         # energy = sum of form_factor(G) * struct_factor(G) * rho(G)
         # where struct_factor(G) = cis(-2T(π) * dot(G, r))
-        form_factors = [Complex{T}(1/sqrt(model.unit_cell_volume) *
-                                   eval_psp_local_fourier(type.psp, model.recip_lattice * G))
-                        for G in G_vectors(basis)]
+        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(pos)
-            f[ir] -= real(sum(conj(ρ.fourier[iG]) .* form_factors[iG] .* cis(-2T(π) * dot(G, r)) .* (-2T(π)) .* G .* im
+            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
 
@@ -41,16 +43,19 @@ function forces(scfres)
             fr = zeros(T, 3)
             for idir = 1:3
                 for (ik, kpt) in enumerate(basis.kpoints)
-                    form_factors = build_form_factors(type.psp, [model.recip_lattice * (kpt.coordinate + G) for G in kpt.basis])
+                    # 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 kpt.basis]
+                    form_factors = build_form_factors(type.psp, qs)
                     structure_factors = [cis(-2T(π)*dot(kpt.coordinate + G, r)) for G in kpt.basis]
                     P = structure_factors .* form_factors ./ sqrt(model.unit_cell_volume)
-                    dPdR = [-2T(π)*im*(kpt.coordinate + G)[idir] * cis(-2T(π)*dot(kpt.coordinate + G, r)) for G in kpt.basis] .*
-                           form_factors ./
-                           sqrt(model.unit_cell_volume)
+                    dPdR = [-2T(π)*im*(kpt.coordinate + G)[idir] for G in kpt.basis] .* P
+
                     # TODO BLASify this further
                     for iband = 1:size(Psi[ik], 2)
                         psi = Psi[ik][:, iband]
-                        fr[idir] -= basis.kweights[ik] * scfres.occupation[ik][iband] * real(dot(psi, P*C*dPdR'*psi) + dot(psi, dPdR*C*P'*psi))
+                        fr[idir] -= basis.kweights[ik] *
+                                    scfres.occupation[ik][iband] *
+                                    real(dot(psi, P*C*dPdR'*psi) + dot(psi, dPdR*C*P'*psi))
                     end
                 end
             end
@@ -60,18 +65,17 @@ function forces(scfres)
         push!(forces, f)
     end
 
-    # Add the ewald term
-    charges_all = [charge_ionic(type) for (type, positions) in model.atoms for pos in positions]
-    positions_all = [pos for (elem, positions) in model.atoms for pos in positions]
-    forces_ewald = zeros(Vec3{T}, length(positions_all))
-    energy_ewald(model.lattice, charges_all, positions_all; forces=forces_ewald)
+    # Add Ewald forces
+    forces_ewald = zeros(Vec3{T}, sum(length(positions) for (elem, positions) in model.atoms))
+    energy_nuclear_ewald(model; forces=forces_ewald)
 
-    offset = 1
+    count = 1
     for i = 1:length(model.atoms)
         for j = 1:length(model.atoms[i][2])
-            forces[i][j] += forces_ewald[offset]
-            offset += 1
+            forces[i][j] += forces_ewald[count]
+            count += 1
         end
     end
+
     forces
 end