Merge pull request #19 from hspark1212/mda

Mda

README.md

@@ -4,7 +4,7 @@
 <h1> Fast Grid 🏁 </h1>
 
   <p>
-    <strong>High-speed Voxel Grid Calculations with Cython</strong>
+    <strong>High-speed Voxel Grid Calculations</strong>
   </p>
 
 </div>
@@ -66,27 +66,28 @@ calculate_grid(
 - Cutoff: 12.8 Å
 - Gas Probe Parameters: TraPPE for methane united atom model
 
-![lj_irmof-1](./images/irmof-1_lj.png)
+![lj_irmof-1](./images/lj_example.png)
 
 ### 2. Gaussian potential
 
 Calculate a voxel grid with the Gaussian function:
 
 ```python
 from fast_grid import calculate_grid
+from ase.build import bulk
 
+atoms = bulk("Cu", "fcc", a=3.6, cubic=True)
 calculate_grid(
-    structure="examples/irmof-1.cif",
-    grid_size=30,
+    structure=atoms,
+    grid_spacing=0.2,
     potential="Gaussian",
-    cutoff=12.8,
     gaussian_height=1.0,
-    gaussian_width=5.0,
+    gaussian_width=1.0,
     visualize=True,
+    pallete="atomic",
 )
 ```
 
-- Default Cutoff: 12.8 Å
-- Gaussian Parameters: Height - 1.0, Width - 5.0
+- Gaussian Parameters: Height - 1.0, Width - 1.0
 
-![gaussian_irmof-1](./images/irmof-1_gaussian.png)
+![gaussian_irmof-1](./images/gaussian_example.png)

fast_grid/calculate_grid.py

@@ -7,9 +7,10 @@
 import numpy as np
 from ase import Atoms
 from ase.io import read
+import MDAnalysis as mda
 
 from fast_grid.ff import get_mixing_epsilon_sigma
-from fast_grid.libs import lj_potential_cython, gaussian_cython
+from fast_grid.potential import lj_potential, gaussian
 from fast_grid.visualize import visualize_grid
 
 warnings.filterwarnings("ignore")
@@ -30,6 +31,7 @@ def calculate_grid(
     float16: bool = False,
     emax: float = 5000.0,
     emin: float = -5000.0,
+    pallete: str = "RdBu",
     return_dict: bool = False,
 ) -> np.array:
     """Calculate the energy grid for a given structure and force field.
@@ -54,6 +56,7 @@ def calculate_grid(
     :param float16: use float16 to save memory, defaults to False
     :param emax: clip energy values for better visualization, defaults to 5000.0
     :param emin: clip energy values for better visualization, defaults to -5000.0
+    :param pallete: color pallete for visualization, defaults to "RdBu"
     :param return_dict: return a dictionary of outputs, defaults to False
     :return: energy grid
     """
@@ -68,16 +71,24 @@ def calculate_grid(
     else:
         raise TypeError("structure must be an ase Atoms object or a cif file path")
 
-    # make supercell when distance between planes is less than cutoff * 2
-    cell_volume = atoms.get_volume()
-    cell_vectors = np.array(atoms.cell)
-    dist_a = cell_volume / np.linalg.norm(np.cross(cell_vectors[1], cell_vectors[2]))
-    dist_b = cell_volume / np.linalg.norm(np.cross(cell_vectors[2], cell_vectors[0]))
-    dist_c = cell_volume / np.linalg.norm(np.cross(cell_vectors[0], cell_vectors[1]))
-    plane_distances = np.array([dist_a, dist_b, dist_c])
-    supercell = np.ceil(2 * cutoff / plane_distances).astype(int)
-    atoms = atoms.repeat(supercell)  # make supercell
-
+    if potential.lower() == "lj":
+        # make supercell when distance between planes is less than cutoff * 2
+        cell_volume = atoms.get_volume()
+        cell_vectors = np.array(atoms.cell)
+        dist_a = cell_volume / np.linalg.norm(
+            np.cross(cell_vectors[1], cell_vectors[2])
+        )
+        dist_b = cell_volume / np.linalg.norm(
+            np.cross(cell_vectors[2], cell_vectors[0])
+        )
+        dist_c = cell_volume / np.linalg.norm(
+            np.cross(cell_vectors[0], cell_vectors[1])
+        )
+        plane_distances = np.array([dist_a, dist_b, dist_c])
+        supercell = np.ceil(2 * cutoff / plane_distances).astype(int)
+        atoms = atoms.repeat(supercell)  # make supercell
+    else:
+        supercell = np.array((1, 1, 1))
     cell_vectors = np.array(atoms.cell)  # redefine cell_vectors after supercell
 
     # get position for grid
@@ -99,6 +110,11 @@ def calculate_grid(
     # get positions for atoms
     pos_atoms = atoms.get_positions()  # (N, 3)
 
+    # distance matrix
+    dist_matrix = mda.lib.distances.distance_array(
+        pos_grid, pos_atoms, box=atoms.cell.cellpar()
+    )  # (G, N)
+
     # setting force field
     symbols = atoms.get_chemical_symbols()
     epsilon, sigma = get_mixing_epsilon_sigma(
@@ -107,23 +123,10 @@ def calculate_grid(
 
     # calculate energy
     if potential.lower() == "lj":
-        calculated_grid = lj_potential_cython(
-            pos_grid,
-            pos_atoms,
-            cell_vectors,
-            epsilon,
-            sigma,
-            cutoff,
-        )  # (G,)
+        calculated_grid = lj_potential(dist_matrix, epsilon, sigma, cutoff)  # (G,)
+
     elif potential.lower() == "gaussian":
-        calculated_grid = gaussian_cython(
-            pos_grid,
-            pos_atoms,
-            cell_vectors,
-            gaussian_height,
-            gaussian_width,
-            cutoff,
-        )  # (G,)
+        calculated_grid = gaussian(dist_matrix, gaussian_height, gaussian_width)  # (G,)
     else:
         raise NotImplementedError(f"{potential} should be one of ['LJ', 'Gaussian']")
 
@@ -138,14 +141,23 @@ def calculate_grid(
 
     if visualize:
         print(f"Visualizing energy grid | supercell {supercell}...")
-        visualize_grid(pos_grid, pos_atoms, calculated_grid, emax, emin)
+        visualize_grid(
+            pos_grid,
+            atoms,
+            calculated_grid,
+            dist_matrix,
+            emax,
+            emin,
+            pallete,
+        )
 
     if return_dict:
         return {
             "atoms": atoms,  # supercelled atoms
             "supercell": supercell,
             "pos_grid": pos_grid,
             "calculated_grid": calculated_grid,
+            "dist_matrix": dist_matrix,
         }
 
     return calculated_grid