A new "ndirect" mode for preprocessing

README.md

@@ -96,7 +96,7 @@ Then, you can use the `preprocess.py` to preprocess the `traj.npz`. It will crea
 python preprocess.py traj.npz graph.npz
 ```
 
-Note that the graph construction is slow especially for large MD trajectories. There two different graph construction algorithms implemented. The default `--backend kdtree` has a linear scaling but only works for orthogonal simulation box. For non-orthogonal simulation, use flag `--backend direct` which has a quadratic scaling. You can also take advantage of the multiprocessing with flag `--n-workers`. For other flags, checkout the help information with `python preprocess.py -h`.
+Note that the graph construction is slow especially for large MD trajectories. There two different graph construction algorithms implemented. The default `--backend kdtree` has a linear scaling but only works for orthogonal simulation box. For non-orthogonal simulation, use flag `--backend direct` or `--backend ndirect` which has a quadratic scaling (for the two choices, the latter is specially efficient for large cells, while the former could be quick for small ones). You can also take advantage of the multiprocessing with flag `--n-workers`. For other flags, checkout the help information with `python preprocess.py -h`.
 
 ### Training the model
 

gdynet/parsers.py

@@ -56,10 +56,12 @@
 prep_parser.add_argument('--radius', type=float, default=7.,
                          help='search radius for finding nearest neighbors '
                          '(default: 7.)')
-prep_parser.add_argument('--backend', choices=['kdtree', 'direct'],
-                         default='kdtree', help='either "kdtree" or "direct", '
+prep_parser.add_argument('--backend', choices=['kdtree', 'direct', 'ndirect'],
+                         default='kdtree', help='"kdtree", "direct" or "ndirect" available, '
                          'the backend used to search for nearest neighbors. '
                          '"kdtree" has linear scaling but only works for '
                          'orthogonal lattices. "direct" works for trigonal '
-                         'lattices but has quadratic scaling. '
+                         'lattices but has quadratic scaling. "ndirect" is '
+                         'an enhanced method for "direct" which could '
+                         'accelarate the process ofdealing with large lattices.'
                          '(default: "kdtree")')

gdynet/preprocess.py

@@ -3,7 +3,7 @@
 import numpy as np
 from .utils import PeriodicCKDTree, distance_pbc
 from tqdm import tqdm
-from pymatgen.core.structure import IStructure
+from pymatgen.core.structure import IStructure, Structure
 
 
 class Preprocess(object):
@@ -72,8 +72,8 @@ def __init__(self, input_file, output_file, n_nbrs=20, radius=7.,
         self.output_file = output_file
         self.n_nbrs = n_nbrs
         self.radius = radius
-        if backend not in ['kdtree', 'direct']:
-            raise ValueError('backend should be either "kdtree" or "direct", '
+        if backend not in ['kdtree', 'direct', 'ndirect']:
+            raise ValueError('backend should be "kdtree", "ndirect" or "direct", '
                              'but got {}'.format(backend))
         self.backend = backend
         self.verbose = verbose
@@ -169,6 +169,33 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index):
                     'target_index': target_index,
                     'nbr_lists': nbr_lists,
                     'nbr_dists': nbr_dists}
+        elif self.backend == 'ndirect':
+            stcs = [Structure(lattice=lattices[i],
+                              species=atom_types,
+                              coords=traj_coords[i],
+                              coords_are_cartesian=True)
+                    for i in tqdm(range(len(traj_coords)),
+                                  desc='Generating structure...', disable=not self.verbose)]
+            a, b, c = [np.ceil(2*self.radius/d).astype('int')
+                       for d in stcs[0].lattice.abc]
+            if [a, b, c] != [1, 1, 1]:
+                _ = [stc.make_supercell(
+                    [np.ceil(2*self.radius/d).astype('int') for d in stc.lattice.abc])
+                    for stc in tqdm(stcs, desc='Building supercell...', disable=not self.verbose)]
+            nbr_lists = np.array([stc.distance_matrix.argsort()[
+                                 :, 1:1+self.n_nbrs] for stc in tqdm(
+                stcs, desc='Generating neighbor index...', disable=not self.verbose)], dtype='int32')
+            nbr_dists = np.array([np.sort(stc.distance_matrix)[
+                                 :, 1:1+self.n_nbrs] for stc in tqdm(
+                stcs, desc='Generating neighbor distance...', disable=not self.verbose)], dtype='float32')
+            nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists)
+            if not np.all((nbr_dists < self.radius) & (nbr_dists > 0)):
+                raise('not find enough neighbors')
+            return {'traj_coords': traj_coords,
+                    'atom_types': atom_types,
+                    'target_index': target_index,
+                    'nbr_lists': nbr_lists,
+                    'nbr_dists': nbr_dists}
 
     def preprocess(self):
         if self.verbose: