Merge pull request #6 from hspark1212/develop-hs

Develop hs

.gitignore

@@ -7,5 +7,8 @@ demo*
 build/
 dist/
 eggs/
+wheelhouse/
 *.egg-info/
-*.egg
+*.egg
+*.so
+*.c

README.md

@@ -4,11 +4,27 @@
 <h1> Fast Grid 🏁 </h1>
 
   <p>
-    <strong>High-speed Voxel Grid Calculations with Numba</strong>
+    <strong>High-speed Voxel Grid Calculations with Cython</strong>
   </p>
 
 </div>
 
+## Features
+
+### Supported potentials
+- LJ Potential (Lennard-Jones)
+- Gaussian Potential
+
+### Gas Probe Model
+- TraPPE (Methane) [Default]
+
+### Performance ⏱️
+- 30 x 30 x 30 grid with 484 atoms (IRMOF-1) using LJ potential : 472 ms ± 7.64
+- 30 x 30 x 30 grid with 484 atoms (IRMOF-1) using Gaussian potential : 564 ms ± 10.1 ms
+
+
+Check out a [tutorial](tutorial.ipynb) file for more details
+
 ## Installation
 
 To install Fast Grid, run the following command in your terminal:

fast_grid/__init__.py

@@ -1,3 +1,3 @@
-from fast_grid.calculate_grids import calculate_grids
+from fast_grid.calculate_grid import calculate_grid
 
-__all__ = ["calculate_grids"]
+__all__ = ["calculate_grid"]

fast_grid/calculate_grids.py → fast_grid/calculate_grid.py

@@ -9,28 +9,30 @@
 from ase.io import read
 
 from fast_grid.ff import get_mixing_epsilon_sigma
-from fast_grid.utils import mic_distance_matrix
-from fast_grid.potential import calculate_lj_potential, calculate_gaussian
-from fast_grid.visualize import visualize_grids
+from fast_grid.utils import check_inputs_energy_grid
+from fast_grid.potential import lj_potential_cython, gaussian_cython
+from fast_grid.visualize import visualize_grid
 
 warnings.filterwarnings("ignore")
 
 
-def calculate_grids(
+def calculate_grid(
     structure: Union[Atoms, str],
     grid_size: Union[int, Iterable] = 30,
+    grid_spacing: float = None,
     ff_type: str = "UFF",
     potential: str = "LJ",
     cutoff: float = 12.8,
-    gas_epsilon: float = 148.0,
-    gas_sigma: float = 3.73,
+    gas_epsilon: float = 148.0,  # LJ
+    gas_sigma: float = 3.73,  # LJ
     visualize: bool = False,
     max_num_atoms: int = 3000,
-    gaussian_height: float = 0.1,
-    gaussian_width: float = 5.0,
+    gaussian_height: float = 0.1,  # Gaussian
+    gaussian_width: float = 5.0,  # Gaussian
     float16: bool = False,
     emax: float = 5000.0,
     emin: float = -5000.0,
+    output_shape_grid: bool = False,
 ) -> np.array:
     """Calculate the energy grid for a given structure and force field.
     It takes a structure (ase Atoms object or cif file path) and returns the energy grid.
@@ -42,6 +44,7 @@ def calculate_grids(
 
     :param structure: structure (ase Atoms object or cif file path)
     :param grid_size: grid size, for example, 30 or "(30, 30, 30)", defaults to 30
+    :param grid_spacing: grid spacing, overrides grid_size, defaults to None
     :param ff_type: force field type, defaults to "UFF"
     :param potential: potential function, gaussian or lj, defaults to "LJ"
     :param cutoff: cutoff distance, defaults to 12.8
@@ -54,8 +57,10 @@ def calculate_grids(
     :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 output_shape_grid: output shape of energy grid, defaults to False
     :return: energy grid
     """
+    # read structure
     if isinstance(structure, Atoms):
         atoms = structure
     elif isinstance(structure, str):
@@ -86,8 +91,15 @@ def calculate_grids(
         grid_size = np.array([grid_size] * 3)
     else:
         grid_size = eval(grid_size)
-        assert len(grid_size) == 3, "grid_size must be a 3-dim vector"
-    indices = np.indices(grid_size).reshape(3, -1).T
+
+    # override grid_size if grid_spacing is not None
+    if grid_spacing is not None:
+        grid_size = np.ceil(
+            np.array(atoms.get_cell_lengths_and_angles()[:3]) / grid_spacing
+        ).astype(int)
+    assert len(grid_size) == 3, "grid_size must be a 3-dim vector"
+
+    indices = np.indices(grid_size).reshape(3, -1).T  # (G, 3)
     pos_grid = indices.dot(cell_vectors / grid_size)  # (G, 3)
 
     # get positions for atoms
@@ -99,42 +111,69 @@ def calculate_grids(
         symbols, ff_type, gas_epsilon, gas_sigma
     )  # (N,) (N,)
 
-    # calculate distance matrix
-    dist_matrix = mic_distance_matrix(pos_grid, pos_atoms, cell_vectors)  # (G, N)
+    # check inputs for energy grid
+    inverse_cell = np.linalg.inv(cell_vectors)
+    energy_grid = np.zeros([grid_size[0] * grid_size[1] * grid_size[2]])
+    check_inputs_energy_grid(
+        pos1=pos_grid,
+        pos2=pos_atoms,
+        cell_vectors=cell_vectors,
+        inverse_cell=inverse_cell,
+        cutoff=cutoff,
+        energy_grid=energy_grid,
+        epsilon=epsilon,
+        sigma=sigma,
+        gaussian_height=gaussian_height,
+        gaussian_width=gaussian_width,
+    )
 
     # calculate energy
     if potential.lower() == "lj":
-        calculated_grids = calculate_lj_potential(
-            dist_matrix,
-            epsilon=epsilon,
-            sigma=sigma,
-            cutoff=cutoff,
-        )  # (G,)
+        calculated_grid = lj_potential_cython(
+            pos_grid,
+            pos_atoms,
+            cell_vectors,
+            inverse_cell,
+            epsilon,
+            sigma,
+            cutoff,
+            energy_grid,
+        )  # (G, 3)
     elif potential.lower() == "gaussian":
-        calculated_grids = calculate_gaussian(
-            dist_matrix,
-            height=gaussian_height,
-            width=gaussian_width,
-            cutoff=cutoff,
-        )  # (G,)
+        calculated_grid = gaussian_cython(
+            pos_grid,
+            pos_atoms,
+            cell_vectors,
+            inverse_cell,
+            gaussian_height,
+            gaussian_width,
+            cutoff,
+            energy_grid,
+        )  # (G, 3)
     else:
         raise NotImplementedError(f"{potential} should be one of ['LJ', 'Gaussian']")
 
+    # flatten energy grid
+    calculated_grid = calculated_grid.reshape(-1)  # (G,)
+
     # convert to float16 to save memory
     if float16:
         # clip energy values for np.float16
         min_float16 = np.finfo(np.float16).min
         max_float16 = np.finfo(np.float16).max
-        calculated_grids = np.clip(calculated_grids, min_float16, max_float16)
+        calculated_grid = np.clip(calculated_grid, min_float16, max_float16)
         # convert to float16
-        calculated_grids = calculated_grids.astype(np.float16)
+        calculated_grid = calculated_grid.astype(np.float16)
+
+    if output_shape_grid:
+        return calculated_grid.reshape(grid_size)
 
     if visualize:
-        print("Visualizing energy grids...")
-        visualize_grids(pos_grid, pos_atoms, calculated_grids, emax, emin)
+        print(f"Visualizing energy grid with {grid_size} grid points")
+        visualize_grid(pos_grid, pos_atoms, calculated_grid, emax, emin)
 
-    return calculated_grids
+    return calculated_grid
 
 
 def cli():
-    Fire(calculate_grids)
+    Fire(calculate_grid)

fast_grid/potential/__init__.py

@@ -0,0 +1,4 @@
+from fast_grid.potential.lj_potential import lj_potential_cython
+from fast_grid.potential.gaussian import gaussian_cython
+
+__all__ = ["lj_potential_cython", "gaussian_cython"]

fast_grid/potential/gaussian.pyx

@@ -0,0 +1,88 @@
+# cython: language_level=3
+# cython: boundscheck=False, wraparound=False, cdivision=True
+
+import numpy as np
+
+import cython
+cimport numpy as np
+from cython.parallel import prange
+from libc.math cimport round, exp
+
+
+def gaussian_cython(np.ndarray[np.float64_t, ndim=2] pos1,
+                            np.ndarray[np.float64_t, ndim=2] pos2,
+                            np.ndarray[np.float64_t, ndim=2] cell_vectors,
+                            np.ndarray[np.float64_t, ndim=2] inverse_cell,
+                            float height,
+                            float width,
+                            float cutoff,
+                            np.ndarray[np.float64_t, ndim=1] energy_grid,
+                            ):
+
+    cdef int G = pos1.shape[0] # grid size
+    cdef int N = pos2.shape[0] # number of atoms
+    cdef int i, j = 0
+    cdef float diff_x, diff_y, diff_z
+    cdef float diff_cell_basis_x, diff_cell_basis_y, diff_cell_basis_z
+    cdef float r2, lj6, lj12, inv_r2, inv_r6, inv_r12, e, s, s6, s12 #remove this line
+    cdef float energy
+    cdef float threshold = 1e-10
+    cdef float width_squared = width * width
+    cdef float cutoff_squared = cutoff * cutoff
+
+    for i in prange(G, nogil=True):
+        energy = 0.0
+        for j in range(N):
+            diff_x = pos1[i, 0] - pos2[j, 0]
+            diff_y = pos1[i, 1] - pos2[j, 1]
+            diff_z = pos1[i, 2] - pos2[j, 2]
+
+            # Matrix multiplication with the inverse cell matrix
+            diff_cell_basis_x = (
+                inverse_cell[0, 0] * diff_x
+                + inverse_cell[0, 1] * diff_y
+                + inverse_cell[0, 2] * diff_z
+            )
+            diff_cell_basis_y = (
+                inverse_cell[1, 0] * diff_x
+                + inverse_cell[1, 1] * diff_y
+                + inverse_cell[1, 2] * diff_z
+            )
+            diff_cell_basis_z = (
+                inverse_cell[2, 0] * diff_x
+                + inverse_cell[2, 1] * diff_y
+                + inverse_cell[2, 2] * diff_z
+            )
+
+            # Applying the minimum image convention
+            diff_cell_basis_x = diff_cell_basis_x - round(diff_cell_basis_x)
+            diff_cell_basis_y = diff_cell_basis_y - round(diff_cell_basis_y)
+            diff_cell_basis_z = diff_cell_basis_z - round(diff_cell_basis_z)
+
+            # Transforming back to the original space
+            diff_x = (
+                cell_vectors[0, 0] * diff_cell_basis_x
+                + cell_vectors[0, 1] * diff_cell_basis_y
+                + cell_vectors[0, 2] * diff_cell_basis_z
+            )
+            diff_y = (
+                cell_vectors[1, 0] * diff_cell_basis_x
+                + cell_vectors[1, 1] * diff_cell_basis_y
+                + cell_vectors[1, 2] * diff_cell_basis_z
+            )
+            diff_z = (
+                cell_vectors[2, 0] * diff_cell_basis_x
+                + cell_vectors[2, 1] * diff_cell_basis_y
+                + cell_vectors[2, 2] * diff_cell_basis_z
+            )
+
+            # Calculating the distance
+            r2 = diff_x * diff_x + diff_y * diff_y + diff_z * diff_z
+
+            if r2 < cutoff_squared and r2 > threshold:
+                # Calculate Guassian
+                energy += height * exp(r2 / width_squared)
+
+        energy_grid[i] += energy
+
+    return energy_grid