doc added

descriptor.py

@@ -6,12 +6,29 @@ from numba import jit
 def gauss(x, m, s):
     return np.exp(-(x-m)*(x-m)/2/s/s)/np.sqrt(2*np.pi)/s
 
-def get_lmbtr(xyz: np.array, z: np.array, centers: list, species: list, grid: dict) -> np.array:
+DEFAULT_GRID = {'min': 0.8, 'max': 6.0, 'n': 40, 'sigma': 0.6}
+
+def get_lmbtr(xyz: np.array, z: np.array, centers: list, species: list, grid: dict = DEFAULT_GRID) -> np.array:
+    """Calculates the smoothed by Gaussian distibution of interatomic distances for each atomic type.
+
+    Args:
+        xyz (np.array): coordinates of the atoms (in Angstrom), L x 3
+        z (np.array): charges/atomic numbers, L
+        centers (list): indices of the atoms to calculate LMBTR, len(centers) -> M
+        species (list): atomic numbers to calculate. E.g. [1, 8, 16] refer to Hydrogen, Oxygen and Sulpfur. len(species) -> N. 
+        grid (dict): parameters of the grid to calculate gaussians {'min', 'max', 'n', 'sigma'}
+            n points between min and max (including), and smoothed with a gaussian function of sigma.
+
+    Returns:
+        np.array: all gaussians expanded into 1-dimensional vector. Overall N x M x n values.
+    """
     xmin = grid['min']
     xmax = grid['max']
     n = grid['n']
     sigma = grid['sigma']
     x = np.linspace(xmin, xmax, n)[:, None]
+    if type(species) == int:
+        species = [species]
     N = len(species)
     M = len(centers)
     out = np.zeros(N*M*n)
@@ -24,13 +41,23 @@ def get_lmbtr(xyz: np.array, z: np.array, centers: list, species: list, grid: di
             offset += n
     return out
 
-def get_lmbtr_h2so4(xyz: np.array, z: np.array, center: int, species: list, grid: dict):
+def get_lmbtr_h2so4(xyz: np.array, z: np.array, center: int, species: list, grid: dict = DEFAULT_GRID):
+    """Calculates LMBTR for a sample H2SO4 system: one sulfur atom and 4 closest oxygens, i.e. 5 atoms in total.
+
+    Args:
+        xyz (np.array): coordinates of the atoms (in Angstrom), L x 3
+        z (np.array): charges/atomic numbers, L
+        center (int): index of the target S atom. 
+        species (list): atomic numbers to caalculate LMBTR. E.g. [1, 8, 16] refer to Hydrogen, Oxygen and Sulpfur. len(species) -> N.
+        grid (dict): parameters of the grid to calculate gaussians {'min', 'max', 'n', 'sigma'}
+            n points between min and max (including), and smoothed with a gaussian function of sigma.
+
+    Returns:
+        _type_: all gaussians expanded into 1-dimensional vector. Overall N x 5 x n values.
+    """
     n = grid['n']
     N = len(species)
-    L = n * N
     n_o = np.where(z == 8)[0]
     d_o = cdist(xyz[[center]], xyz[n_o])[0]
-    idx_o = n_o[np.argsort(d_o)[:4]]  # indices of 5 closest O atoms
+    idx_o = n_o[np.argsort(d_o)[:4]]  # indices of 4 closest O atoms
     return get_lmbtr(xyz, z, [center] + list(idx_o), species, grid)
-
-DEFAULT_GRID = {'min': 0.8, 'max': 6.0, 'n': 40, 'sigma': 0.6}
\ No newline at end of file