Merge pull request #2 from fxcoudert/faster2

Make neighbors and LJ calculation much faster

src/moxel/utils.py

@@ -178,6 +178,11 @@ def calculate(self, cubic_box=False, length=30, potential='lj', n_jobs=None):
             self._simulation_box = self.structure * scale
 
         if potential == 'lj':
+            # Cache LJ parameters for all atoms in the simulation box
+            self._lj_params = np.array([lj_params[atom.species_string] for atom in self._simulation_box])
+            # Cache fractional coordinates, since this is a slow function in pymatgen
+            self._frac_coords = self._simulation_box.frac_coords
+
             # Embarrassingly parallel.
             with Pool(processes=n_jobs) as p:
                 energies = p.map(
@@ -205,21 +210,24 @@ def lj_potential(self, coords):
         """
         if self.cubic_box:
             cartesian_coords = coords
-            neighbors = self._simulation_box.get_sites_in_sphere(cartesian_coords, self.cutoff)
         else:
             cartesian_coords = self._simulation_box.lattice.get_cartesian_coords(coords)
-            neighbors = self._simulation_box.get_sites_in_sphere(cartesian_coords, self.cutoff)
-
-        energy = 0
-        if len(neighbors) != 0:
-            for atom in neighbors:
-                r_ij = np.linalg.norm(cartesian_coords - atom.coords)
-                if r_ij <= 1e-3:
-                    energy += 1000
-                else:
-                    e_j, s_j = lj_params[atom.species_string]
-                    x = (0.5 * (s_j + self.sigma)) / r_ij
-                    energy += 4 * np.sqrt(e_j * self.epsilon) * (x**12 - x**6)
+
+        _, r_ij, indices, _ = self._simulation_box._lattice.get_points_in_sphere(self._frac_coords, cartesian_coords, self.cutoff, zip_results=False)
+
+        # No neighbor, zero energy
+        # Need to check for length of r_ij because of https://github.com/materialsproject/pymatgen/issues/3794
+        if len(r_ij) == 0:
+            return 1.
+
+        # Close contact, infinite energy
+        if np.any(r_ij < 1e-3):
+            return 0.
+
+        es_j = self._lj_params[indices]
+        x = (0.5 * (es_j[:,1] + self.sigma)) / r_ij
+        e = 4 * np.sqrt(es_j[:,0] * self.epsilon)
+        energy = sum(e * (x**12 - x**6))
 
         # This should be changed with clipping in future versions.
         return np.exp(-(1 / 298) * energy)  # For numerical stability.