Uniformize memory access

Author: antoinewdg

src/core/src/energy/global/ewald.rs

@@ -378,12 +378,9 @@ impl Ewald {
         return energy;
     }
 
-    /// Loop over each k_space force.
-    ///
-    /// Abstracts over the very similar computations in `kspace_forces`
-    /// and `kspace_virial`.
-    fn for_each_kspace_force<F>(&mut self, system: &System, mut callback: F)
-    where F : FnMut(usize, usize, Vector3D){
+    /// k-space contribution to the forces
+    fn kspace_forces(&mut self, system: &System, forces: &mut [Vector3D]) {
+        assert_eq!(forces.len(), system.size());
         self.density_fft(system);
 
         let factor = 4.0 * PI / (system.cell().volume() * ELCC);
@@ -394,62 +391,93 @@ impl Ewald {
                 for ikz in 0..self.kmax {
                     // The k = 0 and the cutoff in k-space are already handled
                     // in `expfactors`.
-                    let expfactor = self.expfactors[(ikx, iky, ikz)].abs() ;
-                    if expfactor < f64::EPSILON {continue}
+                    let expfactor = self.expfactors[(ikx, iky, ikz)].abs();
+                    if expfactor < f64::EPSILON { continue }
 
                     let f = expfactor * factor;
                     let k = (ikx as f64) * rec_kx + (iky as f64) * rec_ky + (ikz as f64) * rec_kz;
 
                     for i in 0..system.size() {
                         let qi = system[i].charge;
-                        let fourier_i = self.fourier_phases[(ikx, i, 0)].imag_mul(
-                            self.fourier_phases[(iky, i, 1)] * self.fourier_phases[(ikz, i, 2)]
-                        );
+
+                        let fourier_i = self.fourier_phases[(ikx, i, 0)] *
+                                        self.fourier_phases[(iky, i, 1)] *
+                                        self.fourier_phases[(ikz, i, 2)];
+                        let fourier_i = fourier_i.imag();
+
+                        let mut force_i = Vector3D::zero();
+
                         for j in (i + 1)..system.size() {
                             let qj = system[j].charge;
                             let force = f * self.kspace_force_factor(j, ikx, iky, ikz, qi, qj, fourier_i) * k;
-                            callback(i, j, force);
+                            force_i -= force;
+                            forces[j] += force;
                         }
+
+                        forces[i] += force_i;
                     }
                 }
             }
         }
     }
 
-    /// k-space contribution to the forces
-    fn kspace_forces(&mut self, system: &System, forces: &mut [Vector3D]) {
-        assert_eq!(forces.len(), system.size());
-
-        self.for_each_kspace_force(system, |i, j, force|{
-            forces[i] -= force;
-            forces[j] += force;
-        });
-    }
-
     /// Get the force factor for particles `i` and `j` with charges `qi` and
     /// `qj`, at k point  `(ikx, iky, ikz)`
     #[inline]
     #[allow(too_many_arguments)]
     fn kspace_force_factor(&self, j: usize, ikx: usize, iky: usize, ikz: usize,
                            qi: f64, qj: f64, fourier_i: f64) -> f64 {
 
-        let fourier_j = self.fourier_phases[(ikx, j, 0)].imag_mul(
-            self.fourier_phases[(iky, j, 1)] * self.fourier_phases[(ikz, j, 2)]
-        );
+        // Here the compiler is smart enough to optimize away
+        // the useless computation of the last real part.
+        let fourier_j = self.fourier_phases[(ikx, j, 0)] *
+                        self.fourier_phases[(iky, j, 1)] *
+                        self.fourier_phases[(ikz, j, 2)];
+        let fourier_j = fourier_j.imag();
 
         return qi * qj * (fourier_i - fourier_j);
     }
 
     /// k-space contribution to the virial
     fn kspace_virial(&mut self, system: &System) -> Matrix3 {
+        self.density_fft(system);
         let mut virial = Matrix3::zero();
 
-        self.for_each_kspace_force(system, |i, j, force|{
-            let rij = system.nearest_image(i, j);
-            virial += force.tensorial(&rij);
-        });
+        let factor = 4.0 * PI / (system.cell().volume() * ELCC);
+        let (rec_kx, rec_ky, rec_kz) = system.cell().reciprocal_vectors();
 
-        return virial;
+        for ikx in 0..self.kmax {
+            for iky in 0..self.kmax {
+                for ikz in 0..self.kmax {
+                    // The k = 0 and the cutoff in k-space are already handled
+                    // in `expfactors`.
+                    let expfactor = self.expfactors[(ikx, iky, ikz)].abs();
+                    if expfactor < f64::EPSILON { continue }
+
+                    let f = expfactor * factor;
+                    let k = (ikx as f64) * rec_kx + (iky as f64) * rec_ky + (ikz as f64) * rec_kz;
+
+                    for i in 0..system.size() {
+                        let qi = system[i].charge;
+
+                        let fourier_i = self.fourier_phases[(ikx, i, 0)] *
+                                        self.fourier_phases[(iky, i, 1)] *
+                                        self.fourier_phases[(ikz, i, 2)];
+                        let fourier_i = fourier_i.imag();
+
+                        for j in (i + 1)..system.size() {
+                            let qj = system[j].charge;
+                            let force = f * self.kspace_force_factor(j, ikx, iky, ikz, qi, qj, fourier_i) * k;
+                            let rij = system.nearest_image(i, j);
+                            virial += force.tensorial(&rij);
+                        }
+
+                    }
+                }
+            }
+        }
+
+        virial
     }
 
     fn compute_delta_rho_move_particles(&mut self, system: &System, idxes: &[usize], newpos: &[Vector3D]) {