From 978195c9ff7c8beb25afb6b33e003e932cfb1ece Mon Sep 17 00:00:00 2001 From: Cloudac7 <812556867@qq.com> Date: Tue, 20 Aug 2019 19:53:24 +0800 Subject: [PATCH 1/5] Improve of function construct_graph --- gdynet/preprocess.py | 51 ++++++++++++++++++++++++-------------------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/gdynet/preprocess.py b/gdynet/preprocess.py index c97e9f1..5918767 100644 --- a/gdynet/preprocess.py +++ b/gdynet/preprocess.py @@ -1,9 +1,11 @@ from __future__ import print_function, division +import functools +from multiprocessing import Pool 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 Structure class Preprocess(object): @@ -142,34 +144,37 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): 'target_index': target_index, 'nbr_lists': nbr_lists} elif self.backend == 'direct': - nbr_lists, nbr_dists = [], [] - for coord, lattice in tqdm(zip(traj_coords, lattices), - total=len(traj_coords), - disable=not self.verbose): - crystal = IStructure(lattice=lattice, - species=atom_types, - coords=coord, - coords_are_cartesian=True) - all_nbrs = crystal.get_all_neighbors(r=self.radius, - include_index=True) - all_nbrs = [sorted(nbrs, key=lambda x: x[1]) - for nbrs in all_nbrs] - nbr_list, nbr_dist = [], [] - for nbr in all_nbrs: - assert len(nbr) >= self.n_nbrs, 'not find enough neighbors' - nbr_list.append(list(map(lambda x: x[2], - nbr[:self.n_nbrs]))) - nbr_dist.append(list(map(lambda x: x[1], - nbr[:self.n_nbrs]))) - nbr_lists.append(np.array(nbr_list, dtype='int32')) - nbr_dists.append(np.array(nbr_dist, dtype='float32')) - nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists) + stcs = [Structure(lattice=lattices[i], + species=atom_types, + coords=traj_coords[i], + coords_are_cartesian=True) for i in range(len(traj_coords))] + 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 stcs] + # with Pool(n_core) as p: + nbr_info = map(functools.partial(self._get_sort_distance, + target_index=target_index), + tqdm(stcs, disable=not self.verbose)) + nbr_lists = [s[0] for s in nbr_info] + nbr_dists = [s[1] for s in nbr_info] + 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 _get_sort_distance(self, stc, target_index): + dm = stc.distance_matrix[target_index] + idx = [np.where((d0)) for d in dm] + n_nbrs_dists = np.stack([d[idx][:self.n_nbrs] for d, idx in zip(dm, idx)]) + nbr_lists = np.argsort(n_nbrs_dists, axis=1) + nbr_dists = np.sort(n_nbrs_dists, axis=1) + return nbr_lists, nbr_dists + def preprocess(self): if self.verbose: print('Loading data files') From 5ebf5d33e5eec2ec3b8313fe738230104ad3c1d0 Mon Sep 17 00:00:00 2001 From: Cloudac7 <812556867@qq.com> Date: Wed, 21 Aug 2019 17:22:05 +0800 Subject: [PATCH 2/5] fix map --- gdynet/preprocess.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/gdynet/preprocess.py b/gdynet/preprocess.py index 5918767..4f6fecd 100644 --- a/gdynet/preprocess.py +++ b/gdynet/preprocess.py @@ -147,18 +147,22 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): stcs = [Structure(lattice=lattices[i], species=atom_types, coords=traj_coords[i], - coords_are_cartesian=True) for i in range(len(traj_coords))] + coords_are_cartesian=True) + for i in tqdm(range(len(traj_coords)), + desc='Step 1/3', 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 stcs] + [np.ceil(2*self.radius/d).astype('int') for d in stc.lattice.abc]) + for stc in tqdm(stcs, desc='Step 2/3', disable=not self.verbose)] # with Pool(n_core) as p: - nbr_info = map(functools.partial(self._get_sort_distance, - target_index=target_index), - tqdm(stcs, disable=not self.verbose)) + nbr_info = list(map(functools.partial(self._get_sort_distance, + target_index=target_index), + tqdm(stcs, desc='Step 3/3', disable=not self.verbose))) nbr_lists = [s[0] for s in nbr_info] nbr_dists = [s[1] for s in nbr_info] + 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, @@ -169,8 +173,9 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): def _get_sort_distance(self, stc, target_index): dm = stc.distance_matrix[target_index] - idx = [np.where((d0)) for d in dm] - n_nbrs_dists = np.stack([d[idx][:self.n_nbrs] for d, idx in zip(dm, idx)]) + idx = [np.where((d < self.radius) & (d > 0)) for d in dm] + n_nbrs_dists = np.stack([d[idx][:self.n_nbrs] + for d, idx in zip(dm, idx)]) nbr_lists = np.argsort(n_nbrs_dists, axis=1) nbr_dists = np.sort(n_nbrs_dists, axis=1) return nbr_lists, nbr_dists From 62b933c73f9566d2a03f832c4967980041ee4771 Mon Sep 17 00:00:00 2001 From: Cloudac7 <812556867@qq.com> Date: Sun, 22 Sep 2019 22:05:19 +0800 Subject: [PATCH 3/5] fix some bugs --- gdynet/preprocess.py | 62 ++++++++++++++++++++++++++++---------------- 1 file changed, 40 insertions(+), 22 deletions(-) diff --git a/gdynet/preprocess.py b/gdynet/preprocess.py index 4f6fecd..e81c670 100644 --- a/gdynet/preprocess.py +++ b/gdynet/preprocess.py @@ -1,7 +1,5 @@ from __future__ import print_function, division -import functools -from multiprocessing import Pool import numpy as np from .utils import PeriodicCKDTree, distance_pbc from tqdm import tqdm @@ -74,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 @@ -144,25 +142,54 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): 'target_index': target_index, 'nbr_lists': nbr_lists} elif self.backend == 'direct': + nbr_lists, nbr_dists = [], [] + for coord, lattice in tqdm(zip(traj_coords, lattices), + total=len(traj_coords), + disable=not self.verbose): + crystal = IStructure(lattice=lattice, + species=atom_types, + coords=coord, + coords_are_cartesian=True) + all_nbrs = crystal.get_all_neighbors(r=self.radius, + include_index=True) + all_nbrs = [sorted(nbrs, key=lambda x: x[1]) + for nbrs in all_nbrs] + nbr_list, nbr_dist = [], [] + for nbr in all_nbrs: + assert len(nbr) >= self.n_nbrs, 'not find enough neighbors' + nbr_list.append(list(map(lambda x: x[2], + nbr[:self.n_nbrs]))) + nbr_dist.append(list(map(lambda x: x[1], + nbr[:self.n_nbrs]))) + nbr_lists.append(np.array(nbr_list, dtype='int32')) + nbr_dists.append(np.array(nbr_dist, dtype='float32')) + nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists) + return {'traj_coords': traj_coords, + 'atom_types': atom_types, + '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='Step 1/3', disable=not self.verbose)] + desc='Step 1/2', 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='Step 2/3', disable=not self.verbose)] - # with Pool(n_core) as p: - nbr_info = list(map(functools.partial(self._get_sort_distance, - target_index=target_index), - tqdm(stcs, desc='Step 3/3', disable=not self.verbose))) - nbr_lists = [s[0] for s in nbr_info] - nbr_dists = [s[1] for s in nbr_info] - nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists) + for stc in tqdm(stcs, desc='Step 2/2', disable=not self.verbose)] + nbr_lists = np.array([stc.distance_matrix.argsort()[ + :, 1:1+self.n_nbrs] for stc in tqdm( + stcs, desc='Step 3/4', disable=not self.verbose)]) + nbr_dists = np.array([np.sort(stc.distance_matrix)[ + :, 1:1+self.n_nbrs] for stc in tqdm( + stcs, desc='Step 4/4', disable=not self.verbose)]) + nbr_lists, nbr_dists = np.stack( + nbr_lists, dtype='int32'), np.stack(nbr_dists, dtype='float32') if not np.all((nbr_dists < self.radius) & (nbr_dists > 0)): raise('not find enough neighbors') return {'traj_coords': traj_coords, @@ -171,15 +198,6 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): 'nbr_lists': nbr_lists, 'nbr_dists': nbr_dists} - def _get_sort_distance(self, stc, target_index): - dm = stc.distance_matrix[target_index] - idx = [np.where((d < self.radius) & (d > 0)) for d in dm] - n_nbrs_dists = np.stack([d[idx][:self.n_nbrs] - for d, idx in zip(dm, idx)]) - nbr_lists = np.argsort(n_nbrs_dists, axis=1) - nbr_dists = np.sort(n_nbrs_dists, axis=1) - return nbr_lists, nbr_dists - def preprocess(self): if self.verbose: print('Loading data files') From 45955683a382925d8531e6f0e0aa445e0574ad80 Mon Sep 17 00:00:00 2001 From: Cloudac7 <812556867@qq.com> Date: Sun, 22 Sep 2019 22:05:56 +0800 Subject: [PATCH 4/5] fix --- gdynet/preprocess.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gdynet/preprocess.py b/gdynet/preprocess.py index e81c670..d82c8fb 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 Structure +from pymatgen.core.structure import IStructure, Structure class Preprocess(object): From 02e4ecd71faee1c12c7854730403b575be4b844b Mon Sep 17 00:00:00 2001 From: Cloudac7 <812556867@qq.com> Date: Tue, 24 Sep 2019 17:32:16 +0800 Subject: [PATCH 5/5] add ndirect flag for preprocess --- README.md | 2 +- gdynet/parsers.py | 8 +++++--- gdynet/preprocess.py | 11 +++++------ 3 files changed, 11 insertions(+), 10 deletions(-) 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 d82c8fb..f2f0518 100644 --- a/gdynet/preprocess.py +++ b/gdynet/preprocess.py @@ -175,21 +175,20 @@ def construct_graph(self, traj_coords, lattices, atom_types, target_index): coords=traj_coords[i], coords_are_cartesian=True) for i in tqdm(range(len(traj_coords)), - desc='Step 1/2', disable=not self.verbose)] + 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='Step 2/2', disable=not self.verbose)] + 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='Step 3/4', disable=not self.verbose)]) + 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='Step 4/4', disable=not self.verbose)]) - nbr_lists, nbr_dists = np.stack( - nbr_lists, dtype='int32'), np.stack(nbr_dists, dtype='float32') + 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,