diff --git a/README.md b/README.md index 3694661..5044c7f 100644 --- a/README.md +++ b/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 diff --git a/gdynet/parsers.py b/gdynet/parsers.py index d2e84d8..c3eaf19 100644 --- a/gdynet/parsers.py +++ b/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")') diff --git a/gdynet/preprocess.py b/gdynet/preprocess.py index c97e9f1..f2f0518 100644 --- a/gdynet/preprocess.py +++ b/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: