Skip to content

Commit 2f6020b

Browse files
amcadmusamcadmus
authored
Merge pull request #1126 from deepmodeling/devel
Merge devel into master
2 parents 1a25414 + 95a0340 commit 2f6020b

File tree

83 files changed

+1231
-1022
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

83 files changed

+1231
-1022
lines changed

‎.github/workflows/test_python.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,4 +71,9 @@ jobs:
7171
CXX: g++-${{ matrix.gcc }}
7272
TENSORFLOW_VERSION: ${{ matrix.tf }}
7373
- run: dp --version
74+
- name: Prepare parallel runtime
75+
if: ${{ matrix.tf == '' }}
76+
run: |
77+
sudo apt install libopenmpi-dev openmpi-bin
78+
HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 pip install horovod mpi4py
7479
- run: pytest --cov=deepmd source/tests && codecov

‎deepmd/descriptor/descriptor.py‎

Lines changed: 344 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,344 @@
1+
from abc import ABC, abstractmethod
2+
from typing import Any, Dict, List, Tuple
3+
4+
import numpy as np
5+
from deepmd.env import tf
6+
7+
8+
class Descriptor(ABC):
9+
r"""The abstract class for descriptors. All specific descriptors should
10+
be based on this class.
11+
12+
The descriptor :math:`\mathcal{D}` describes the environment of an atom,
13+
which should be a function of coordinates and types of its neighbour atoms.
14+
15+
Notes
16+
-----
17+
Only methods and attributes defined in this class are generally public,
18+
that can be called by other classes.
19+
"""
20+
21+
@abstractmethod
22+
def get_rcut(self) -> float:
23+
"""
24+
Returns the cut-off radius.
25+
26+
Returns
27+
-------
28+
float
29+
the cut-off radius
30+
31+
Notes
32+
-----
33+
This method must be implemented, as it's called by other classes.
34+
"""
35+
36+
@abstractmethod
37+
def get_ntypes(self) -> int:
38+
"""
39+
Returns the number of atom types.
40+
41+
Returns
42+
-------
43+
int
44+
the number of atom types
45+
46+
Notes
47+
-----
48+
This method must be implemented, as it's called by other classes.
49+
"""
50+
51+
@abstractmethod
52+
def get_dim_out(self) -> int:
53+
"""
54+
Returns the output dimension of this descriptor.
55+
56+
Returns
57+
-------
58+
int
59+
the output dimension of this descriptor
60+
61+
Notes
62+
-----
63+
This method must be implemented, as it's called by other classes.
64+
"""
65+
66+
def get_dim_rot_mat_1(self) -> int:
67+
"""
68+
Returns the first dimension of the rotation matrix. The rotation is of shape
69+
dim_1 x 3
70+
71+
Returns
72+
-------
73+
int
74+
the first dimension of the rotation matrix
75+
"""
76+
# TODO: I think this method should be implemented as it's called by dipole and
77+
# polar fitting network. However, currently not all descriptors have this
78+
# method.
79+
raise NotImplementedError
80+
81+
def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]:
82+
"""
83+
Returns neighbor information.
84+
85+
Returns
86+
-------
87+
nlist : tf.Tensor
88+
Neighbor list
89+
rij : tf.Tensor
90+
The relative distance between the neighbor and the center atom.
91+
sel_a : list[int]
92+
The number of neighbors with full information
93+
sel_r : list[int]
94+
The number of neighbors with only radial information
95+
"""
96+
# TODO: I think this method should be implemented as it's called by energy
97+
# model. However, se_ar and hybrid doesn't have this method.
98+
raise NotImplementedError
99+
100+
@abstractmethod
101+
def compute_input_stats(self,
102+
data_coord: List[np.ndarray],
103+
data_box: List[np.ndarray],
104+
data_atype: List[np.ndarray],
105+
natoms_vec: List[np.ndarray],
106+
mesh: List[np.ndarray],
107+
input_dict: Dict[str, List[np.ndarray]]
108+
) -> None:
109+
"""
110+
Compute the statisitcs (avg and std) of the training data. The input will be
111+
normalized by the statistics.
112+
113+
Parameters
114+
----------
115+
data_coord : list[np.ndarray]
116+
The coordinates. Can be generated by
117+
:meth:`deepmd.model.model_stat.make_stat_input`
118+
data_box : list[np.ndarray]
119+
The box. Can be generated by
120+
:meth:`deepmd.model.model_stat.make_stat_input`
121+
data_atype : list[np.ndarray]
122+
The atom types. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input`
123+
natoms_vec : list[np.ndarray]
124+
The vector for the number of atoms of the system and different types of
125+
atoms. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input`
126+
mesh : list[np.ndarray]
127+
The mesh for neighbor searching. Can be generated by
128+
:meth:`deepmd.model.model_stat.make_stat_input`
129+
input_dict : dict[str, list[np.ndarray]]
130+
Dictionary for additional input
131+
132+
Notes
133+
-----
134+
This method must be implemented, as it's called by other classes.
135+
"""
136+
137+
@abstractmethod
138+
def build(self,
139+
coord_: tf.Tensor,
140+
atype_: tf.Tensor,
141+
natoms: tf.Tensor,
142+
box_: tf.Tensor,
143+
mesh: tf.Tensor,
144+
input_dict: Dict[str, Any],
145+
reuse: bool = None,
146+
suffix: str = '',
147+
) -> tf.Tensor:
148+
"""
149+
Build the computational graph for the descriptor.
150+
151+
Parameters
152+
----------
153+
coord_ : tf.Tensor
154+
The coordinate of atoms
155+
atype_ : tf.Tensor
156+
The type of atoms
157+
natoms : tf.Tensor
158+
The number of atoms. This tensor has the length of Ntypes + 2
159+
natoms[0]: number of local atoms
160+
natoms[1]: total number of atoms held by this processor
161+
natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
162+
box : tf.Tensor
163+
The box of frames
164+
mesh : tf.Tensor
165+
For historical reasons, only the length of the Tensor matters.
166+
if size of mesh == 6, pbc is assumed.
167+
if size of mesh == 0, no-pbc is assumed.
168+
input_dict : dict[str, Any]
169+
Dictionary for additional inputs
170+
reuse : bool, optional
171+
The weights in the networks should be reused when get the variable.
172+
suffix : str, optional
173+
Name suffix to identify this descriptor
174+
175+
Returns
176+
-------
177+
descriptor: tf.Tensor
178+
The output descriptor
179+
180+
Notes
181+
-----
182+
This method must be implemented, as it's called by other classes.
183+
"""
184+
185+
def enable_compression(self,
186+
min_nbor_dist: float,
187+
model_file: str = 'frozon_model.pb',
188+
table_extrapolate: float = 5.,
189+
table_stride_1: float = 0.01,
190+
table_stride_2: float = 0.1,
191+
check_frequency: int = -1,
192+
suffix: str = "",
193+
) -> None:
194+
"""
195+
Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the
196+
training data.
197+
198+
Parameters
199+
----------
200+
min_nbor_dist : float
201+
The nearest distance between atoms
202+
model_file : str, default: 'frozon_model.pb'
203+
The original frozen model, which will be compressed by the program
204+
table_extrapolate : float, default: 5.
205+
The scale of model extrapolation
206+
table_stride_1 : float, default: 0.01
207+
The uniform stride of the first table
208+
table_stride_2 : float, default: 0.1
209+
The uniform stride of the second table
210+
check_frequency : int, default: -1
211+
The overflow check frequency
212+
suffix : str, optional
213+
The suffix of the scope
214+
215+
Notes
216+
-----
217+
This method is called by others when the descriptor supported compression.
218+
"""
219+
raise NotImplementedError(
220+
"Descriptor %s doesn't support compression!" % type(self).__name__)
221+
222+
@abstractmethod
223+
def prod_force_virial(self,
224+
atom_ener: tf.Tensor,
225+
natoms: tf.Tensor
226+
) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
227+
"""
228+
Compute force and virial.
229+
230+
Parameters
231+
----------
232+
atom_ener : tf.Tensor
233+
The atomic energy
234+
natoms : tf.Tensor
235+
The number of atoms. This tensor has the length of Ntypes + 2
236+
natoms[0]: number of local atoms
237+
natoms[1]: total number of atoms held by this processor
238+
natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
239+
240+
Returns
241+
-------
242+
force : tf.Tensor
243+
The force on atoms
244+
virial : tf.Tensor
245+
The total virial
246+
atom_virial : tf.Tensor
247+
The atomic virial
248+
"""
249+
250+
def get_feed_dict(self,
251+
coord_: tf.Tensor,
252+
atype_: tf.Tensor,
253+
natoms: tf.Tensor,
254+
box: tf.Tensor,
255+
mesh: tf.Tensor
256+
) -> Dict[str, tf.Tensor]:
257+
"""
258+
Generate the feed_dict for current descriptor
259+
260+
Parameters
261+
----------
262+
coord_ : tf.Tensor
263+
The coordinate of atoms
264+
atype_ : tf.Tensor
265+
The type of atoms
266+
natoms : tf.Tensor
267+
The number of atoms. This tensor has the length of Ntypes + 2
268+
natoms[0]: number of local atoms
269+
natoms[1]: total number of atoms held by this processor
270+
natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
271+
box : tf.Tensor
272+
The box. Can be generated by deepmd.model.make_stat_input
273+
mesh : tf.Tensor
274+
For historical reasons, only the length of the Tensor matters.
275+
if size of mesh == 6, pbc is assumed.
276+
if size of mesh == 0, no-pbc is assumed.
277+
278+
Returns
279+
-------
280+
feed_dict : dict[str, tf.Tensor]
281+
The output feed_dict of current descriptor
282+
"""
283+
feed_dict = {
284+
't_coord:0' :coord_,
285+
't_type:0' :atype_,
286+
't_natoms:0' :natoms,
287+
't_box:0' :box,
288+
't_mesh:0' :mesh
289+
}
290+
return feed_dict
291+
292+
def init_variables(self,
293+
model_file: str,
294+
suffix : str = "",
295+
) -> None:
296+
"""
297+
Init the embedding net variables with the given dict
298+
299+
Parameters
300+
----------
301+
model_file : str
302+
The input model file
303+
suffix : str, optional
304+
The suffix of the scope
305+
306+
Notes
307+
-----
308+
This method is called by others when the descriptor supported initialization from the given variables.
309+
"""
310+
raise NotImplementedError(
311+
"Descriptor %s doesn't support initialization from the given variables!" % type(self).__name__)
312+
313+
def get_tensor_names(self, suffix : str = "") -> Tuple[str]:
314+
"""Get names of tensors.
315+
316+
Parameters
317+
----------
318+
suffix : str
319+
The suffix of the scope
320+
321+
Returns
322+
-------
323+
Tuple[str]
324+
Names of tensors
325+
"""
326+
raise NotImplementedError("Descriptor %s doesn't support this property!" % type(self).__name__)
327+
328+
def pass_tensors_from_frz_model(self,
329+
*tensors : tf.Tensor,
330+
) -> None:
331+
"""
332+
Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def
333+
334+
Parameters
335+
----------
336+
*tensors : tf.Tensor
337+
passed tensors
338+
339+
Notes
340+
-----
341+
The number of parameters in the method must be equal to the numbers of returns in
342+
:meth:`get_tensor_names`.
343+
"""
344+
raise NotImplementedError("Descriptor %s doesn't support this method!" % type(self).__name__)

0 commit comments

Comments
 (0)