diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 8517560fc9..27f8e866d6 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -71,4 +71,9 @@ jobs: CXX: g++-${{ matrix.gcc }} TENSORFLOW_VERSION: ${{ matrix.tf }} - run: dp --version + - name: Prepare parallel runtime + if: ${{ matrix.tf == '' }} + run: | + sudo apt install libopenmpi-dev openmpi-bin + HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 pip install horovod mpi4py - run: pytest --cov=deepmd source/tests && codecov diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py new file mode 100644 index 0000000000..893493a201 --- /dev/null +++ b/deepmd/descriptor/descriptor.py @@ -0,0 +1,344 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Tuple + +import numpy as np +from deepmd.env import tf + + +class Descriptor(ABC): + r"""The abstract class for descriptors. All specific descriptors should + be based on this class. + + The descriptor :math:`\mathcal{D}` describes the environment of an atom, + which should be a function of coordinates and types of its neighbour atoms. + + Notes + ----- + Only methods and attributes defined in this class are generally public, + that can be called by other classes. + """ + + @abstractmethod + def get_rcut(self) -> float: + """ + Returns the cut-off radius. + + Returns + ------- + float + the cut-off radius + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def get_ntypes(self) -> int: + """ + Returns the number of atom types. + + Returns + ------- + int + the number of atom types + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def get_dim_out(self) -> int: + """ + Returns the output dimension of this descriptor. + + Returns + ------- + int + the output dimension of this descriptor + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + def get_dim_rot_mat_1(self) -> int: + """ + Returns the first dimension of the rotation matrix. The rotation is of shape + dim_1 x 3 + + Returns + ------- + int + the first dimension of the rotation matrix + """ + # TODO: I think this method should be implemented as it's called by dipole and + # polar fitting network. However, currently not all descriptors have this + # method. + raise NotImplementedError + + def get_nlist(self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns neighbor information. + + Returns + ------- + nlist : tf.Tensor + Neighbor list + rij : tf.Tensor + The relative distance between the neighbor and the center atom. + sel_a : list[int] + The number of neighbors with full information + sel_r : list[int] + The number of neighbors with only radial information + """ + # TODO: I think this method should be implemented as it's called by energy + # model. However, se_ar and hybrid doesn't have this method. + raise NotImplementedError + + @abstractmethod + def compute_input_stats(self, + data_coord: List[np.ndarray], + data_box: List[np.ndarray], + data_atype: List[np.ndarray], + natoms_vec: List[np.ndarray], + mesh: List[np.ndarray], + input_dict: Dict[str, List[np.ndarray]] + ) -> None: + """ + Compute the statisitcs (avg and std) of the training data. The input will be + normalized by the statistics. + + Parameters + ---------- + data_coord : list[np.ndarray] + The coordinates. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` + data_box : list[np.ndarray] + The box. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` + data_atype : list[np.ndarray] + The atom types. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input` + natoms_vec : list[np.ndarray] + The vector for the number of atoms of the system and different types of + atoms. Can be generated by :meth:`deepmd.model.model_stat.make_stat_input` + mesh : list[np.ndarray] + The mesh for neighbor searching. Can be generated by + :meth:`deepmd.model.model_stat.make_stat_input` + input_dict : dict[str, list[np.ndarray]] + Dictionary for additional input + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + @abstractmethod + def build(self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box_: tf.Tensor, + mesh: tf.Tensor, + input_dict: Dict[str, Any], + reuse: bool = None, + suffix: str = '', + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor. + + Parameters + ---------- + coord_ : tf.Tensor + The coordinate of atoms + atype_ : tf.Tensor + The type of atoms + natoms : tf.Tensor + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + box : tf.Tensor + The box of frames + mesh : tf.Tensor + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict : dict[str, Any] + Dictionary for additional inputs + reuse : bool, optional + The weights in the networks should be reused when get the variable. + suffix : str, optional + Name suffix to identify this descriptor + + Returns + ------- + descriptor: tf.Tensor + The output descriptor + + Notes + ----- + This method must be implemented, as it's called by other classes. + """ + + def enable_compression(self, + min_nbor_dist: float, + model_file: str = 'frozon_model.pb', + table_extrapolate: float = 5., + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + suffix: str = "", + ) -> None: + """ + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the + training data. + + Parameters + ---------- + min_nbor_dist : float + The nearest distance between atoms + model_file : str, default: 'frozon_model.pb' + The original frozen model, which will be compressed by the program + table_extrapolate : float, default: 5. + The scale of model extrapolation + table_stride_1 : float, default: 0.01 + The uniform stride of the first table + table_stride_2 : float, default: 0.1 + The uniform stride of the second table + check_frequency : int, default: -1 + The overflow check frequency + suffix : str, optional + The suffix of the scope + + Notes + ----- + This method is called by others when the descriptor supported compression. + """ + raise NotImplementedError( + "Descriptor %s doesn't support compression!" % type(self).__name__) + + @abstractmethod + def prod_force_virial(self, + atom_ener: tf.Tensor, + natoms: tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial. + + Parameters + ---------- + atom_ener : tf.Tensor + The atomic energy + natoms : tf.Tensor + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + + Returns + ------- + force : tf.Tensor + The force on atoms + virial : tf.Tensor + The total virial + atom_virial : tf.Tensor + The atomic virial + """ + + def get_feed_dict(self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box: tf.Tensor, + mesh: tf.Tensor + ) -> Dict[str, tf.Tensor]: + """ + Generate the feed_dict for current descriptor + + Parameters + ---------- + coord_ : tf.Tensor + The coordinate of atoms + atype_ : tf.Tensor + The type of atoms + natoms : tf.Tensor + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + box : tf.Tensor + The box. Can be generated by deepmd.model.make_stat_input + mesh : tf.Tensor + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + + Returns + ------- + feed_dict : dict[str, tf.Tensor] + The output feed_dict of current descriptor + """ + feed_dict = { + 't_coord:0' :coord_, + 't_type:0' :atype_, + 't_natoms:0' :natoms, + 't_box:0' :box, + 't_mesh:0' :mesh + } + return feed_dict + + def init_variables(self, + model_file: str, + suffix : str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + model_file : str + The input model file + suffix : str, optional + The suffix of the scope + + Notes + ----- + This method is called by others when the descriptor supported initialization from the given variables. + """ + raise NotImplementedError( + "Descriptor %s doesn't support initialization from the given variables!" % type(self).__name__) + + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + raise NotImplementedError("Descriptor %s doesn't support this property!" % type(self).__name__) + + def pass_tensors_from_frz_model(self, + *tensors : tf.Tensor, + ) -> None: + """ + Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def + + Parameters + ---------- + *tensors : tf.Tensor + passed tensors + + Notes + ----- + The number of parameters in the method must be equal to the numbers of returns in + :meth:`get_tensor_names`. + """ + raise NotImplementedError("Descriptor %s doesn't support this method!" % type(self).__name__) diff --git a/deepmd/descriptor/hybrid.py b/deepmd/descriptor/hybrid.py index 013ee9f753..37b9578b4e 100644 --- a/deepmd/descriptor/hybrid.py +++ b/deepmd/descriptor/hybrid.py @@ -12,6 +12,7 @@ # from deepmd.descriptor import DescrptSeAEbd # from deepmd.descriptor import DescrptSeAEf # from deepmd.descriptor import DescrptSeR +from .descriptor import Descriptor from .se_a import DescrptSeA from .se_r import DescrptSeR from .se_ar import DescrptSeAR @@ -20,7 +21,7 @@ from .se_a_ef import DescrptSeAEf from .loc_frame import DescrptLocFrame -class DescrptHybrid (): +class DescrptHybrid (Descriptor): """Concate a list of descriptors to form a new descriptor. Parameters @@ -219,3 +220,88 @@ def prod_force_virial(self, virial += vv atom_virial += av return force, virial, atom_virial + + def enable_compression(self, + min_nbor_dist: float, + model_file: str = 'frozon_model.pb', + table_extrapolate: float = 5., + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + suffix: str = "" + ) -> None: + """ + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the + training data. + + Parameters + ---------- + min_nbor_dist : float + The nearest distance between atoms + model_file : str, default: 'frozon_model.pb' + The original frozen model, which will be compressed by the program + table_extrapolate : float, default: 5. + The scale of model extrapolation + table_stride_1 : float, default: 0.01 + The uniform stride of the first table + table_stride_2 : float, default: 0.1 + The uniform stride of the second table + check_frequency : int, default: -1 + The overflow check frequency + suffix : str, optional + The suffix of the scope + """ + for idx, ii in enumerate(self.descrpt_list): + ii.enable_compression(min_nbor_dist, model_file, table_extrapolate, table_stride_1, table_stride_2, check_frequency, suffix=f"{suffix}_{idx}") + + def init_variables(self, + model_file : str, + suffix : str = "", + ) -> None: + """ + Init the embedding net variables with the given dict + + Parameters + ---------- + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope + """ + for idx, ii in enumerate(self.descrpt_list): + ii.init_variables(model_file, suffix=f"{suffix}_{idx}") + + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + tensor_names = [] + for idx, ii in enumerate(self.descrpt_list): + tensor_names.extend(ii.get_tensor_names(suffix=f"{suffix}_{idx}")) + return tuple(tensor_names) + + def pass_tensors_from_frz_model(self, + *tensors : tf.Tensor, + ) -> None: + """ + Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def + + Parameters + ---------- + *tensors : tf.Tensor + passed tensors + """ + jj = 0 + for ii in self.descrpt_list: + n_tensors = len(ii.get_tensor_names()) + ii.pass_tensors_from_frz_model(*tensors[jj:jj+n_tensors]) + jj += n_tensors diff --git a/deepmd/descriptor/loc_frame.py b/deepmd/descriptor/loc_frame.py index be57403d60..a664384ba6 100644 --- a/deepmd/descriptor/loc_frame.py +++ b/deepmd/descriptor/loc_frame.py @@ -7,8 +7,9 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptLocFrame () : +class DescrptLocFrame (Descriptor) : """Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame. diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index fbc9a77b56..39485463a9 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -13,9 +13,10 @@ from deepmd.utils.tabulate import DPTabulate from deepmd.utils.type_embed import embed_atom_type from deepmd.utils.sess import run_sess -from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph +from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph, get_embedding_net_variables +from .descriptor import Descriptor -class DescrptSeA (): +class DescrptSeA (Descriptor): r"""DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input. @@ -156,6 +157,7 @@ def __init__ (self, self.dstd = None self.davg = None self.compress = False + self.embedding_net_variables = None self.place_holders = {} nei_type = np.array([]) for ii in range(self.ntypes): @@ -296,7 +298,8 @@ def enable_compression(self, table_extrapolate : float = 5, table_stride_1 : float = 0.01, table_stride_2 : float = 0.1, - check_frequency : int = -1 + check_frequency : int = -1, + suffix : str = "", ) -> None: """ Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. @@ -315,10 +318,15 @@ def enable_compression(self, The uniform stride of the second table check_frequency The overflow check frequency + suffix : str, optional + The suffix of the scope """ + assert ( + not self.filter_resnet_dt + ), "Model compression error: descriptor resnet_dt must be false!" self.compress = True self.table = DPTabulate( - model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn) + model_file, self.type_one_side, self.exclude_types, self.compress_activation_fn, suffix=suffix) self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency] self.lower, self.upper \ = self.table.build(min_nbor_dist, @@ -327,8 +335,8 @@ def enable_compression(self, table_stride_2) graph, _ = load_graph_def(model_file) - self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr/t_avg') - self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr/t_std') + self.davg = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_avg' % suffix) + self.dstd = get_tensor_by_name_from_graph(graph, 'descrpt_attr%s/t_std' % suffix) @@ -425,10 +433,10 @@ def build (self, tf.summary.histogram('nlist', self.nlist) self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) - self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') - self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') - self.rij = tf.identity(self.rij, name = 'o_rij') - self.nlist = tf.identity(self.nlist, name = 'o_nlist') + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix) + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix) + self.rij = tf.identity(self.rij, name = 'o_rij' + suffix) + self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix) self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, @@ -448,6 +456,21 @@ def get_rot_mat(self) -> tf.Tensor: """ return self.qmat + def get_tensor_names(self, suffix : str = "") -> Tuple[str]: + """Get names of tensors. + + Parameters + ---------- + suffix : str + The suffix of the scope + + Returns + ------- + Tuple[str] + Names of tensors + """ + return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0') + def pass_tensors_from_frz_model(self, descrpt_reshape : tf.Tensor, descrpt_deriv : tf.Tensor, @@ -473,46 +496,22 @@ def pass_tensors_from_frz_model(self, self.descrpt_deriv = descrpt_deriv self.descrpt_reshape = descrpt_reshape - def get_feed_dict(self, - coord_, - atype_, - natoms, - box, - mesh): + def init_variables(self, + model_file : str, + suffix : str = "", + ) -> None: """ - generate the deed_dict for current descriptor + Init the embedding net variables with the given dict Parameters ---------- - coord_ - The coordinate of atoms - atype_ - The type of atoms - natoms - The number of atoms. This tensor has the length of Ntypes + 2 - natoms[0]: number of local atoms - natoms[1]: total number of atoms held by this processor - natoms[i]: 2 <= i < Ntypes+2, number of type i atoms - box - The box. Can be generated by deepmd.model.make_stat_input - mesh - For historical reasons, only the length of the Tensor matters. - if size of mesh == 6, pbc is assumed. - if size of mesh == 0, no-pbc is assumed. - - Returns - ------- - feed_dict - The output feed_dict of current descriptor + model_file : str + The input frozen model file + suffix : str, optional + The suffix of the scope """ - feed_dict = { - 't_coord:0' :coord_, - 't_type:0' :atype_, - 't_natoms:0' :natoms, - 't_box:0' :box, - 't_mesh:0' :mesh - } - return feed_dict + self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix) + def prod_force_virial(self, atom_ener : tf.Tensor, @@ -667,17 +666,36 @@ def _concat_type_embedding( nframes, natoms, type_embedding, - ): + ): + '''Concatenate `type_embedding` of neighbors and `xyz_scatter`. + If not self.type_one_side, concatenate `type_embedding` of center atoms as well. + + Parameters + ---------- + xyz_scatter: + shape is [nframes*natoms[0]*self.nnei, 1] + nframes: + shape is [] + natoms: + shape is [1+1+self.ntypes] + type_embedding: + shape is [self.ntypes, Y] where Y=jdata['type_embedding']['neuron'][-1] + + Returns + ------- + embedding: + environment of each atom represented by embedding. + ''' te_out_dim = type_embedding.get_shape().as_list()[-1] - nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) #nnei*nchnl - nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) # shape is [self.nnei, 1+te_out_dim] + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) # shape is [nframes*natoms[0]*self.nnei, te_out_dim] nei_embed = tf.reshape(nei_embed,[-1,te_out_dim]) - embedding_input = tf.concat([xyz_scatter,nei_embed],1) + embedding_input = tf.concat([xyz_scatter,nei_embed],1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim] if not self.type_one_side: - atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) - atm_embed = tf.tile(atm_embed,(1,self.nnei)) - atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) - embedding_input = tf.concat([embedding_input,atm_embed],1) + atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) # shape is [natoms[0], te_out_dim] + atm_embed = tf.tile(atm_embed,(nframes,self.nnei)) # shape is [nframes*natoms[0], self.nnei*te_out_dim] + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) # shape is [nframes*natoms[0]*self.nnei, te_out_dim] + embedding_input = tf.concat([embedding_input,atm_embed],1) # shape is [nframes*natoms[0]*self.nnei, 1+te_out_dim+te_out_dim] return embedding_input @@ -740,7 +758,8 @@ def _filter_lower( bavg = bavg, seed = self.seed, trainable = trainable, - uniform_seed = self.uniform_seed) + uniform_seed = self.uniform_seed, + initial_variables = self.embedding_net_variables) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: # we can safely return the final xyz_scatter filled with zero directly diff --git a/deepmd/descriptor/se_a_ef.py b/deepmd/descriptor/se_a_ef.py index 996de6ff2d..e2424aeae4 100644 --- a/deepmd/descriptor/se_a_ef.py +++ b/deepmd/descriptor/se_a_ef.py @@ -10,8 +10,9 @@ from deepmd.env import op_module from deepmd.env import default_tf_session_config from .se_a import DescrptSeA +from .descriptor import Descriptor -class DescrptSeAEf (): +class DescrptSeAEf (Descriptor): """ Parameters diff --git a/deepmd/descriptor/se_ar.py b/deepmd/descriptor/se_ar.py index 8ded2c5849..518cf9315c 100644 --- a/deepmd/descriptor/se_ar.py +++ b/deepmd/descriptor/se_ar.py @@ -5,8 +5,9 @@ from .se_a import DescrptSeA from .se_r import DescrptSeR from deepmd.env import op_module +from .descriptor import Descriptor -class DescrptSeAR (): +class DescrptSeAR (Descriptor): def __init__ (self, jdata): args = ClassArg()\ .add('a', dict, must = True) \ diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index fafa03a350..b6de76be76 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -10,8 +10,9 @@ from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptSeR (): +class DescrptSeR (Descriptor): """DeepPot-SE constructed from radial information of atomic configurations. The embedding takes the distance between atoms as input. diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index 94298c6eb4..2ab7a732be 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -10,8 +10,9 @@ from deepmd.env import default_tf_session_config from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift from deepmd.utils.sess import run_sess +from .descriptor import Descriptor -class DescrptSeT (): +class DescrptSeT (Descriptor): """DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. @@ -468,6 +469,7 @@ def _filter(self, inputs_i = tf.slice (inputs, [ 0, start_index_i *4], [-1, self.sel_a[type_i] *4] ) + start_index_j = start_index_i start_index_i += self.sel_a[type_i] nei_type_i = self.sel_a[type_i] shape_i = inputs_i.get_shape().as_list() @@ -476,7 +478,6 @@ def _filter(self, env_i = tf.reshape(inputs_i, [-1, nei_type_i, 4]) # with natom x nei_type_i x 3 env_i = tf.slice(env_i, [0, 0, 1], [-1, -1, -1]) - start_index_j = 0 for type_j in range(type_i, self.ntypes): # with natom x (nei_type_j x 4) inputs_j = tf.slice (inputs, diff --git a/deepmd/entrypoints/compress.py b/deepmd/entrypoints/compress.py index c689481b96..58c3c2a075 100644 --- a/deepmd/entrypoints/compress.py +++ b/deepmd/entrypoints/compress.py @@ -5,8 +5,8 @@ import logging from typing import Optional -from deepmd.env import tf -from deepmd.common import j_loader, GLOBAL_TF_FLOAT_PRECISION +from deepmd.common import j_loader +from deepmd.env import tf, GLOBAL_ENER_FLOAT_PRECISION from deepmd.utils.argcheck import normalize from deepmd.utils.compat import updata_deepmd_input from deepmd.utils.errors import GraphTooLargeError, GraphWithoutTensorError @@ -89,7 +89,7 @@ def compress( tf.constant(t_min_nbor_dist, name = 'train_attr/min_nbor_dist', - dtype = GLOBAL_TF_FLOAT_PRECISION) + dtype = GLOBAL_ENER_FLOAT_PRECISION) jdata["model"]["compress"] = {} jdata["model"]["compress"]["type"] = 'se_e2_a' jdata["model"]["compress"]["compress"] = True @@ -101,15 +101,11 @@ def compress( 10 * step, int(frequency), ] + jdata["training"]["save_ckpt"] = "model-compression/model.ckpt" jdata = normalize(jdata) # check the descriptor info of the input file - assert ( - jdata["model"]["descriptor"]["type"] == "se_a" or jdata["model"]["descriptor"]["type"] == "se_e2_a" - ), "Model compression error: descriptor type must be se_a or se_e2_a!" - assert ( - jdata["model"]["descriptor"]["resnet_dt"] is False - ), "Model compression error: descriptor resnet_dt must be false!" + # move to the specific Descriptor class # stage 1: training or refining the model with tabulation log.info("\n\n") diff --git a/deepmd/entrypoints/freeze.py b/deepmd/entrypoints/freeze.py index c6bb89fb8f..511f58598d 100755 --- a/deepmd/entrypoints/freeze.py +++ b/deepmd/entrypoints/freeze.py @@ -6,6 +6,7 @@ https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc """ +import logging from deepmd.env import tf from deepmd.env import op_module from deepmd.utils.sess import run_sess @@ -18,6 +19,8 @@ __all__ = ["freeze"] +log = logging.getLogger(__name__) + def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> List[str]: """Get node names based on model type. @@ -108,10 +111,12 @@ def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> Li "modifier_attr/sys_charge_map", "modifier_attr/ewald_h", "modifier_attr/ewald_beta", + "dipole_charge/model_type", "dipole_charge/descrpt_attr/rcut", "dipole_charge/descrpt_attr/ntypes", "dipole_charge/model_attr/tmap", "dipole_charge/model_attr/model_type", + "dipole_charge/model_attr/model_version", "o_dm_force", "dipole_charge/model_attr/sel_type", "dipole_charge/o_dipole", @@ -175,9 +180,18 @@ def freeze( modifier_type = None if node_names is None: output_node_list = _make_node_names(model_type, modifier_type) + different_set = set(output_node_list) - set(nodes) + if different_set: + log.warning( + "The following nodes are not in the graph: %s. " + "Skip freezeing these nodes. You may be freezing " + "a checkpoint generated by an old version." % different_set + ) + # use intersection as output list + output_node_list = list(set(output_node_list) & set(nodes)) else: output_node_list = node_names.split(",") - print(f"The following nodes will be frozen: {output_node_list}") + log.info(f"The following nodes will be frozen: {output_node_list}") # We use a built-in TF helper to export variables to constants output_graph_def = tf.graph_util.convert_variables_to_constants( @@ -189,4 +203,4 @@ def freeze( # Finally we serialize and dump the output graph to the filesystem with tf.gfile.GFile(output_graph, "wb") as f: f.write(output_graph_def.SerializeToString()) - print(f"{len(output_graph_def.node):d} ops in the final graph.") + log.info(f"{len(output_graph_def.node):d} ops in the final graph.") diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index bef0d2bf58..721eed357c 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -311,7 +311,7 @@ def parse_args(args: Optional[List[str]] = None): "-c", "--checkpoint-folder", type=str, - default=".", + default="model-compression", help="path to checkpoint folder", ) parser_compress.add_argument( diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index bbb1e55bd2..817d603f3c 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -10,7 +10,7 @@ from typing import Dict, List, Optional, Any from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have -from deepmd.env import tf, reset_default_tf_session_config +from deepmd.env import tf, reset_default_tf_session_config, GLOBAL_ENER_FLOAT_PRECISION from deepmd.infer.data_modifier import DipoleChargeModifier from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions from deepmd.train.trainer import DPTrainer @@ -262,6 +262,16 @@ def get_nbor_stat(jdata, rcut): neistat = NeighborStat(ntypes, rcut) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) + + # moved from traier.py as duplicated + # TODO: this is a simple fix but we should have a clear + # architecture to call neighbor stat + tf.constant(min_nbor_dist, + name = 'train_attr/min_nbor_dist', + dtype = GLOBAL_ENER_FLOAT_PRECISION) + tf.constant(max_nbor_size, + name = 'train_attr/max_nbor_size', + dtype = tf.int32) return min_nbor_dist, max_nbor_size def get_sel(jdata, rcut): diff --git a/deepmd/entrypoints/transfer.py b/deepmd/entrypoints/transfer.py index 576df74c80..3a755fd61a 100644 --- a/deepmd/entrypoints/transfer.py +++ b/deepmd/entrypoints/transfer.py @@ -130,7 +130,7 @@ def transform_graph(raw_graph: tf.Graph, old_graph: tf.Graph) -> tf.Graph: if raw_graph_dtype == np.float16: if old_graph_dtype == np.float64 or old_graph_dtype == np.float32: if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): - tensor = np.frombuffer(old_node.tensor_content).astype(raw_graph_dtype) + tensor = np.frombuffer(old_node.tensor_content, dtype = raw_graph_dtype) cp_attr.from_array(tensor, tf.float16, shape = tensor_shape) else: tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype) @@ -143,7 +143,7 @@ def transform_graph(raw_graph: tf.Graph, old_graph: tf.Graph) -> tf.Graph: elif raw_graph_dtype == np.float64 or raw_graph_dtype == np.float32: if old_graph_dtype == np.float64 or old_graph_dtype == np.float32: if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): - tensor = np.frombuffer(old_node.tensor_content).astype(raw_graph_dtype) + tensor = np.frombuffer(old_node.tensor_content, dtype = raw_graph_dtype) cp_attr.from_str(tensor) else: tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype) diff --git a/deepmd/env.py b/deepmd/env.py index aaa148b357..92287d8aa5 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -2,7 +2,7 @@ import logging import os -import distutils.ccompiler +import platform from configparser import ConfigParser from imp import reload from pathlib import Path @@ -119,6 +119,7 @@ def get_tf_session_config() -> Any: set_tf_default_nthreads() intra, inter = get_tf_default_nthreads() config = tf.ConfigProto( + gpu_options=tf.GPUOptions(allow_growth=True), intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter ) return config @@ -156,8 +157,12 @@ def get_module(module_name: str) -> "ModuleType": FileNotFoundError if module is not found in directory """ - # https://discuss.python.org/t/how-to-get-the-file-extension-of-dynamic-libraries-for-current-os/3916/5 - ext = distutils.ccompiler.new_compiler().shared_lib_extension + if platform.system() == "Windows": + ext = ".dll" + elif platform.system() == "Darwin": + ext = ".dylib" + else: + ext = ".so" module_file = ( (Path(__file__).parent / SHARED_LIB_MODULE / module_name) diff --git a/deepmd/infer/data_modifier.py b/deepmd/infer/data_modifier.py index 6b473e93e6..224b8fbd06 100644 --- a/deepmd/infer/data_modifier.py +++ b/deepmd/infer/data_modifier.py @@ -358,7 +358,16 @@ def _extend_system(self, coord, box, atype, charge): ref_coord = coord3[:,sel_idx_map,:] ref_coord = np.reshape(ref_coord, [nframes, nsel * 3]) - dipole = DeepDipole.eval(self, coord, box, atype) + batch_size = 8 + all_dipole = [] + for ii in range(0,nframes,batch_size): + dipole = DeepDipole.eval(self, + coord[ii:ii+batch_size], + box[ii:ii+batch_size], + atype) + all_dipole.append(dipole) + dipole = np.concatenate(all_dipole, axis = 0) + assert(dipole.shape[0] == nframes) dipole = np.reshape(dipole, [nframes, nsel * 3]) wfcc_coord = ref_coord + dipole diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 846efc9bde..d594be32f8 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -84,7 +84,6 @@ def __init__( DeepTensor.__init__( self, model_file, - 9, load_prefix=load_prefix, default_tf_graph=default_tf_graph, ) diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 3c5fb28f6e..63625905e8 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -137,6 +137,7 @@ def __init__( t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0") t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0") [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = run_sess(self.sess, [t_mdl_name, t_mdl_charge_map, t_sys_charge_map, t_ewald_h, t_ewald_beta]) + mdl_name = mdl_name.decode("UTF-8") mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()] sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()] self.dm = DipoleChargeModifier(mdl_name, mdl_charge_map, sys_charge_map, ewald_h = ewald_h, ewald_beta = ewald_beta) diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index dec6ca66f4..441709caff 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -173,10 +173,11 @@ def build (self, name = 'descrpt_attr/ntypes', dtype = tf.int32) feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh) - return_elements = ['o_rmat:0', 'o_rmat_deriv:0', 'o_rij:0', 'o_nlist:0', 'o_descriptor:0'] - descrpt_reshape, descrpt_deriv, rij, nlist, dout \ + return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0'] + imported_tensors \ = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements) - self.descrpt.pass_tensors_from_frz_model(descrpt_reshape, descrpt_deriv, rij, nlist) + dout = imported_tensors[-1] + self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1]) if self.srtab is not None : diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py index 5c996ec38d..2a63eda4d6 100644 --- a/deepmd/model/tensor.py +++ b/deepmd/model/tensor.py @@ -3,7 +3,7 @@ from deepmd.env import tf from deepmd.common import ClassArg -from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION +from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION from deepmd.env import op_module from deepmd.utils.graph import load_graph_def from .model_stat import make_stat_input, merge_sys_stat @@ -138,10 +138,11 @@ def build (self, name = 'descrpt_attr/ntypes', dtype = tf.int32) feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh) - return_elements = ['o_rmat:0', 'o_rmat_deriv:0', 'o_rij:0', 'o_nlist:0', 'o_descriptor:0'] - descrpt_reshape, descrpt_deriv, rij, nlist, dout \ + return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0'] + imported_tensors \ = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements) - self.descrpt.pass_tensors_from_frz_model(descrpt_reshape, descrpt_deriv, rij, nlist) + dout = imported_tensors[-1] + self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1]) rot_mat = self.descrpt.get_rot_mat() rot_mat = tf.identity(rot_mat, name = 'o_rot_mat'+suffix) diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index e301a0753b..dc888ad3e0 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -26,7 +26,7 @@ from deepmd.utils.neighbor_stat import NeighborStat from deepmd.utils.sess import run_sess from deepmd.utils.type_embed import TypeEmbedNet -from deepmd.utils.graph import get_tensor_by_name, get_fitting_net_variables +from deepmd.utils.graph import get_tensor_by_name, get_embedding_net_variables, get_fitting_net_variables from tensorflow.python.client import timeline from deepmd.env import op_module @@ -278,7 +278,6 @@ def _init_param(self, jdata): # if init the graph with the frozen model self.frz_model = None self.model_type = None - self.init_from_frz_model = False def build (self, @@ -314,18 +313,10 @@ def build (self, if self.run_opt.init_mode == 'init_from_frz_model': self._init_from_frz_model() - self.neighbor_stat \ - = NeighborStat(self.ntypes, self.descrpt_param['rcut']) - self.min_nbor_dist, self.max_nbor_size \ - = self.neighbor_stat.get_stat(data) - tf.constant(self.min_nbor_dist, - name = 'train_attr/min_nbor_dist', - dtype = GLOBAL_TF_FLOAT_PRECISION) - tf.constant(self.max_nbor_size, - name = 'train_attr/max_nbor_size', - dtype = GLOBAL_TF_FLOAT_PRECISION) + # neighbor_stat is moved to train.py as duplicated + # TODO: this is a simple fix but we should have a clear + # architecture to call neighbor stat else : - assert 'rcut' in self.descrpt_param, "Error: descriptor must have attr rcut!" self.descrpt.enable_compression(self.model_param['compress']["min_nbor_dist"], self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3]) self.fitting.init_variables(get_fitting_net_variables(self.model_param['compress']['model_file'])) @@ -385,10 +376,10 @@ def _build_training(self): optimizer = self.run_opt._HVD.DistributedOptimizer(optimizer) else: optimizer = tf.train.AdamOptimizer(learning_rate = self.learning_rate) - grads = tf.gradients(self.l2_l, trainable_variables) - apply_op = optimizer.apply_gradients (zip (grads, trainable_variables), - global_step=self.global_step, - name='train_step') + apply_op = optimizer.minimize(loss=self.l2_l, + global_step=self.global_step, + var_list=trainable_variables, + name='train_step') train_ops = [apply_op] + self._extra_train_ops self.train_op = tf.group(*train_ops) log.info("built training") @@ -397,7 +388,6 @@ def _init_session(self): config = get_tf_session_config() device, idx = self.run_opt.my_device.split(":", 1) if device == "gpu": - config.gpu_options.allow_growth = True config.gpu_options.visible_device_list = idx self.sess = tf.Session(config=config) @@ -408,8 +398,9 @@ def _init_session(self): if self.run_opt.init_mode == 'init_from_scratch' : log.info("initialize model from scratch") run_sess(self.sess, init_op) - fp = open(self.disp_file, "w") - fp.close () + if not self.is_compress: + fp = open(self.disp_file, "w") + fp.close () elif self.run_opt.init_mode == 'init_from_model' : log.info("initialize from model %s" % self.run_opt.init_model) run_sess(self.sess, init_op) @@ -695,14 +686,17 @@ def _init_from_frz_model(self): "which is not supported by the 'dp train init-frz-model' interface. " % self.run_opt.init_frz_model ) from e + if self.fitting_type != 'ener': + raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!") # self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch... # initialize fitting net with the given compressed frozen model - if self.model_type == 'compressed_model' and self.fitting_type == 'ener': - self.init_from_frz_model = True + if self.model_type == 'original_model': + self.descrpt.init_variables(self.run_opt.init_frz_model) + self.fitting.init_variables(get_fitting_net_variables(self.run_opt.init_frz_model)) + tf.constant("original_model", name = 'model_type', dtype = tf.string) + elif self.model_type == 'compressed_model': self.frz_model = self.run_opt.init_frz_model self.fitting.init_variables(get_fitting_net_variables(self.frz_model)) tf.constant("compressed_model", name = 'model_type', dtype = tf.string) - elif self.fitting_type != 'ener': - raise RuntimeError("The 'dp train init-frz-model' command only supports the 'ener' type fitting net currently!") else: - raise RuntimeError("The 'dp train init-frz-model' command only supports the compressed model currently!") + raise RuntimeError("Unknown model type %s" % self.model_type) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fcaba80908..6279b3e088 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -581,7 +581,7 @@ def training_args(): # ! modified by Ziyao: data configuration isolated. arg_validation_data, Argument("numb_steps", int, optional=False, doc=doc_numb_steps, alias=["stop_batch"]), Argument("seed", [int,None], optional=True, doc=doc_seed), - Argument("disp_file", str, optional=True, default='lcueve.out', doc=doc_disp_file), + Argument("disp_file", str, optional=True, default='lcurve.out', doc=doc_disp_file), Argument("disp_freq", int, optional=True, default=1000, doc=doc_disp_freq), Argument("numb_test", [list,int,str], optional=True, default=1, doc=doc_numb_test), Argument("save_freq", int, optional=True, default=1000, doc=doc_save_freq), diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py index 31452be7f5..53fd05cc67 100644 --- a/deepmd/utils/graph.py +++ b/deepmd/utils/graph.py @@ -108,11 +108,11 @@ def get_tensor_by_type(node, elif data_type == np.float32: tensor = np.array(node.float_val) else: - raise RunTimeError('model compression does not support the half precision') + raise RuntimeError('model compression does not support the half precision') return tensor -def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict: +def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef, suffix: str = "") -> Dict: """ Get the embedding net nodes with the given tf.GraphDef object @@ -120,6 +120,8 @@ def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict: ---------- graph_def The input tf.GraphDef object + suffix : str, optional + The scope suffix Returns ---------- @@ -127,7 +129,7 @@ def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict: The embedding net nodes within the given tf.GraphDef object """ embedding_net_nodes = {} - embedding_net_pattern = "filter_type_\d+/matrix_\d+_\d+|filter_type_\d+/bias_\d+_\d+|filter_type_\d+/idt_\d+_\d+|filter_type_all/matrix_\d+_\d+|filter_type_all/bias_\d+_\d+|filter_type_all/idt_\d+_\d" + embedding_net_pattern = f"filter_type_\d+{suffix}/matrix_\d+_\d+|filter_type_\d+{suffix}/bias_\d+_\d+|filter_type_\d+{suffix}/idt_\d+_\d+|filter_type_all{suffix}/matrix_\d+_\d+|filter_type_all{suffix}/bias_\d+_\d+|filter_type_all{suffix}/idt_\d+_\d" for node in graph_def.node: if re.fullmatch(embedding_net_pattern, node.name) != None: embedding_net_nodes[node.name] = node.attr["value"].tensor @@ -137,7 +139,7 @@ def get_embedding_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict: return embedding_net_nodes -def get_embedding_net_nodes(model_file: str) -> Dict: +def get_embedding_net_nodes(model_file: str, suffix: str = "") -> Dict: """ Get the embedding net nodes with the given frozen model(model_file) @@ -145,6 +147,8 @@ def get_embedding_net_nodes(model_file: str) -> Dict: ---------- model_file The input frozen model path + suffix : str, optional + The suffix of the scope Returns ---------- @@ -152,10 +156,10 @@ def get_embedding_net_nodes(model_file: str) -> Dict: The embedding net nodes with the given frozen model """ _, graph_def = load_graph_def(model_file) - return get_embedding_net_nodes_from_graph_def(graph_def) + return get_embedding_net_nodes_from_graph_def(graph_def, suffix=suffix) -def get_embedding_net_variables_from_graph_def(graph_def : tf.GraphDef) -> Dict: +def get_embedding_net_variables_from_graph_def(graph_def : tf.GraphDef, suffix: str = "") -> Dict: """ Get the embedding net variables with the given tf.GraphDef object @@ -163,6 +167,8 @@ def get_embedding_net_variables_from_graph_def(graph_def : tf.GraphDef) -> Dict: ---------- graph_def The input tf.GraphDef object + suffix : str, optional + The suffix of the scope Returns ---------- @@ -170,19 +176,19 @@ def get_embedding_net_variables_from_graph_def(graph_def : tf.GraphDef) -> Dict: The embedding net variables within the given tf.GraphDef object """ embedding_net_variables = {} - embedding_net_nodes = get_embedding_net_nodes_from_graph_def(graph_def) + embedding_net_nodes = get_embedding_net_nodes_from_graph_def(graph_def, suffix=suffix) for item in embedding_net_nodes: node = embedding_net_nodes[item] dtype = tf.as_dtype(node.dtype).as_numpy_dtype tensor_shape = tf.TensorShape(node.tensor_shape).as_list() if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): - tensor_value = np.frombuffer(node.tensor_content) + tensor_value = np.frombuffer(node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) else: tensor_value = get_tensor_by_type(node, dtype) embedding_net_variables[item] = np.reshape(tensor_value, tensor_shape) return embedding_net_variables -def get_embedding_net_variables(model_file : str) -> Dict: +def get_embedding_net_variables(model_file : str, suffix: str = "") -> Dict: """ Get the embedding net variables with the given frozen model(model_file) @@ -190,6 +196,8 @@ def get_embedding_net_variables(model_file : str) -> Dict: ---------- model_file The input frozen model path + suffix : str, optional + The suffix of the scope Returns ---------- @@ -197,7 +205,7 @@ def get_embedding_net_variables(model_file : str) -> Dict: The embedding net variables within the given frozen model """ _, graph_def = load_graph_def(model_file) - return get_embedding_net_variables_from_graph_def(graph_def) + return get_embedding_net_variables_from_graph_def(graph_def, suffix=suffix) def get_fitting_net_nodes_from_graph_def(graph_def: tf.GraphDef) -> Dict: @@ -264,7 +272,7 @@ def get_fitting_net_variables_from_graph_def(graph_def : tf.GraphDef) -> Dict: dtype= tf.as_dtype(node.dtype).as_numpy_dtype tensor_shape = tf.TensorShape(node.tensor_shape).as_list() if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): - tensor_value = np.frombuffer(node.tensor_content) + tensor_value = np.frombuffer(node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) else: tensor_value = get_tensor_by_type(node, dtype) fitting_net_variables[item] = np.reshape(tensor_value, tensor_shape) diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index d42f3a5c48..5c78031167 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -92,7 +92,8 @@ def embedding_net(xx, bavg = 0.0, seed = None, trainable = True, - uniform_seed = False): + uniform_seed = False, + initial_variables = None): r"""The embedding network. The embedding network function :math:`\mathcal{N}` is constructed by is the @@ -141,6 +142,11 @@ def embedding_net(xx, Random seed for initializing network parameters trainable: boolean If the network is trainable + uniform_seed : boolean + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + initial_variables : dict + The input dict which stores the embedding net variables + References ---------- @@ -152,37 +158,47 @@ def embedding_net(xx, outputs_size = [input_shape[1]] + network_size for ii in range(1, len(outputs_size)): - w = tf.get_variable('matrix_'+str(ii)+name_suffix, + w_initializer = tf.random_normal_initializer( + stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), + seed = seed if (seed is None or uniform_seed) else seed + ii*3+0 + ) + b_initializer = tf.random_normal_initializer( + stddev=stddev, + mean = bavg, + seed = seed if (seed is None or uniform_seed) else seed + 3*ii+1 + ) + if initial_variables is not None: + scope = tf.get_variable_scope().name + w_initializer = tf.constant_initializer(initial_variables[scope+'/matrix_'+str(ii)+name_suffix]) + b_initializer = tf.constant_initializer(initial_variables[scope+'/bias_'+str(ii)+name_suffix]) + w = tf.get_variable('matrix_'+str(ii)+name_suffix, [outputs_size[ii - 1], outputs_size[ii]], precision, - tf.random_normal_initializer( - stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), - seed = seed if (seed is None or uniform_seed) else seed + ii*3+0 - ), + w_initializer, trainable = trainable) variable_summaries(w, 'matrix_'+str(ii)+name_suffix) b = tf.get_variable('bias_'+str(ii)+name_suffix, [1, outputs_size[ii]], precision, - tf.random_normal_initializer( - stddev=stddev, - mean = bavg, - seed = seed if (seed is None or uniform_seed) else seed + 3*ii+1 - ), + b_initializer, trainable = trainable) variable_summaries(b, 'bias_'+str(ii)+name_suffix) hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]]) if resnet_dt : + idt_initializer = tf.random_normal_initializer( + stddev=0.001, + mean = 1.0, + seed = seed if (seed is None or uniform_seed) else seed + 3*ii+2 + ) + if initial_variables is not None: + scope = tf.get_variable_scope().name + idt_initializer = tf.constant_initializer(initial_variables[scope+'/idt_'+str(ii)+name_suffix]) idt = tf.get_variable('idt_'+str(ii)+name_suffix, [1, outputs_size[ii]], precision, - tf.random_normal_initializer( - stddev=0.001, - mean = 1.0, - seed = seed if (seed is None or uniform_seed) else seed + 3*ii+2 - ), + idt_initializer, trainable = trainable) variable_summaries(idt, 'idt_'+str(ii)+name_suffix) diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index f1057b38f2..fe95173890 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -34,12 +34,16 @@ class DPTabulate(): For example, `[[0, 1]]` means no interaction between type 0 and type 1. activation_function The activation function in the embedding net. Supported options are {"tanh","gelu"} in common.ACTIVATION_FN_DICT. + suffix : str, optional + The suffix of the scope """ def __init__(self, model_file : str, type_one_side : bool = False, exclude_types : List[List[int]] = [], - activation_fn : Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh) -> None: + activation_fn : Callable[[tf.Tensor], tf.Tensor] = tf.nn.tanh, + suffix : str = "", + ) -> None: """ Constructor """ @@ -47,8 +51,9 @@ def __init__(self, self.model_file = model_file self.type_one_side = type_one_side self.exclude_types = exclude_types + self.suffix = suffix if self.type_one_side and len(self.exclude_types) != 0: - raise RunTimeError('"type_one_side" is not compatible with "exclude_types"') + raise RuntimeError('"type_one_side" is not compatible with "exclude_types"') # functype if activation_fn == ACTIVATION_FN_DICT["tanh"]: @@ -56,7 +61,7 @@ def __init__(self, elif activation_fn == ACTIVATION_FN_DICT["gelu"]: self.functype = 2 else: - raise RunTimeError("Unknown actication function type!") + raise RuntimeError("Unknown actication function type!") self.activation_fn = activation_fn self.graph, self.graph_def = load_graph_def(self.model_file) @@ -72,15 +77,15 @@ def __init__(self, self.sel_a = self.graph.get_operation_by_name('DescrptSeA').get_attr('sel_a') self.descrpt = self.graph.get_operation_by_name ('DescrptSeA') - self.davg = get_tensor_by_name_from_graph(self.graph, 'descrpt_attr/t_avg') - self.dstd = get_tensor_by_name_from_graph(self.graph, 'descrpt_attr/t_std') + self.davg = get_tensor_by_name_from_graph(self.graph, f'descrpt_attr{self.suffix}/t_avg') + self.dstd = get_tensor_by_name_from_graph(self.graph, f'descrpt_attr{self.suffix}/t_std') self.ntypes = get_tensor_by_name_from_graph(self.graph, 'descrpt_attr/ntypes') self.rcut = self.descrpt.get_attr('rcut_r') self.rcut_smth = self.descrpt.get_attr('rcut_r_smth') - self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(self.graph_def) + self.embedding_net_nodes = get_embedding_net_nodes_from_graph_def(self.graph_def, suffix=self.suffix) for tt in self.exclude_types: if (tt[0] not in range(self.ntypes)) or (tt[1] not in range(self.ntypes)): @@ -174,14 +179,16 @@ def _get_bias(self): bias["layer_" + str(layer)] = [] if self.type_one_side: for ii in range(0, self.ntypes): - tensor_value = np.frombuffer (self.embedding_net_nodes["filter_type_all/bias_" + str(layer) + "_" + str(ii)].tensor_content) - tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_all/bias_" + str(layer) + "_" + str(ii)].tensor_shape).as_list() + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/bias_{layer}_{ii}"] + tensor_value = np.frombuffer (node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types: - tensor_value = np.frombuffer(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content) - tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list() + node = self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/bias_{layer}_{ii % self.ntypes}"] + tensor_value = np.frombuffer(node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) else: bias["layer_" + str(layer)].append(np.array([])) @@ -193,14 +200,16 @@ def _get_matrix(self): matrix["layer_" + str(layer)] = [] if self.type_one_side: for ii in range(0, self.ntypes): - tensor_value = np.frombuffer (self.embedding_net_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(ii)].tensor_content) - tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(ii)].tensor_shape).as_list() + node = self.embedding_net_nodes[f"filter_type_all{self.suffix}/matrix_{layer}_{ii}"] + tensor_value = np.frombuffer (node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) else: for ii in range(0, self.ntypes * self.ntypes): if (ii // self.ntypes, int(ii % self.ntypes)) not in self.exclude_types: - tensor_value = np.frombuffer(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content) - tensor_shape = tf.TensorShape(self.embedding_net_nodes["filter_type_" + str(ii // self.ntypes) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list() + node = self.embedding_net_nodes[f"filter_type_{ii // self.ntypes}{self.suffix}/matrix_{layer}_{ii % self.ntypes}"] + tensor_value = np.frombuffer(node.tensor_content, dtype = tf.as_dtype(node.dtype).as_numpy_dtype) + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) else: matrix["layer_" + str(layer)].append(np.array([])) diff --git a/doc/conf.py b/doc/conf.py index 02610d172b..dd946fa2e1 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -220,6 +220,13 @@ def setup(app): ), "ase": ("https://wiki.fysik.dtu.dk/ase/", None), } +numpydoc_xref_param_type = True + + +numpydoc_xref_aliases = {} +import typing +for typing_type in typing.__all__: + numpydoc_xref_aliases[typing_type] = "typing.%s" % typing_type # -- Options for HTML output ------------------------------------------------- diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 805361120b..7218b75090 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -94,5 +94,53 @@ kspace_modify gewald 0.45 ``` Please notice that the DeePMD does nothing to the direct space part of the electrostatic interaction, because this part is assumed to be fitted in the DeePMD model (the direct space cut-off is thus the cut-off of the DeePMD model). The splitting parameter `gewald` is modified by the `kspace_modify` command. +## Use of the centroid/stress/atom to get the full 3x3 "atomic-virial" + +The [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit) allows also the computation of per-atom stress tensor defined as: + + + +Where is the atomic position of nth atom, velocity of atom and the derivative of the atomic energy. + +In LAMMPS one can get the per-atom stress using the command `centroid/stress/atom`: +```bash +compute ID group-ID centroid/stress/atom NULL virial +``` +see [LAMMPS doc page](https://docs.lammps.org/compute_stress_atom.html#thompson2) for more detailes on the meaning of the keywords. +### Examples +In order of computing the 9-component per-atom stress +```bash +compute stress all centroid/stress/atom NULL virial +``` +Thus `c_stress` is an array with 9 component in the order `xx,yy,zz,xy,xz,yz,yx,zx,zy`. + +If you use this feature please cite [D. Tisi, L. Zhang, R. Bertossa, H. Wang, R. Car, S. Baroni - arXiv preprint arXiv:2108.10850, 2021](https://arxiv.org/abs/2108.10850) + +## Computation of heat flux +Using per-atom stress tensor one can, for example, compute the heat flux defined as: + + + +to compute the heat flux with LAMMPS: +```bash +compute ke_ID all ke/atom +compute pe_ID all pe/atom +compute stress_ID group-ID centroid/stress/atom NULL virial +compute flux_ID all heat/flux ke_ID pe_ID stress_ID +``` + +### Examples + +```bash +compute ke all ke/atom +compute pe all pe/atom +compute stress all centroid/stress/atom NULL virial +compute flux all heat/flux ke pe stress +``` +`c_flux` is a global vector of length 6. The first three components are the `x`, `y` and `z` components of the full heat flux vector. The others are the components of the so-called convective portion, see [LAMMPS doc page](https://docs.lammps.org/compute_heat_flux.html) for more detailes. + +If you use these features please cite [D. Tisi, L. Zhang, R. Bertossa, H. Wang, R. Car, S. Baroni - arXiv preprint arXiv:2108.10850, 2021](https://arxiv.org/abs/2108.10850) + + [DP]:https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.143001 -[DP-SE]:https://dl.acm.org/doi/10.5555/3327345.3327356 \ No newline at end of file +[DP-SE]:https://dl.acm.org/doi/10.5555/3327345.3327356 diff --git a/doc/train-input-auto.rst b/doc/train-input-auto.rst index ea8da06a61..9201809549 100644 --- a/doc/train-input-auto.rst +++ b/doc/train-input-auto.rst @@ -898,7 +898,7 @@ model: .. _`model/fitting_net[polar]/scale`: scale: - | type: ``list`` | ``float``, optional, default: ``1.0`` + | type: ``float`` | ``list``, optional, default: ``1.0`` | argument path: ``model/fitting_net[polar]/scale`` The output of the fitting net (polarizability matrix) will be scaled by ``scale`` @@ -1102,7 +1102,7 @@ loss: .. _`loss[ener]/start_pref_e`: start_pref_e: - | type: ``int`` | ``float``, optional, default: ``0.02`` + | type: ``float`` | ``int``, optional, default: ``0.02`` | argument path: ``loss[ener]/start_pref_e`` The prefactor of energy loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the energy label should be provided by file energy.npy in each data system. If both start_pref_energy and limit_pref_energy are set to 0, then the energy will be ignored. @@ -1110,7 +1110,7 @@ loss: .. _`loss[ener]/limit_pref_e`: limit_pref_e: - | type: ``int`` | ``float``, optional, default: ``1.0`` + | type: ``float`` | ``int``, optional, default: ``1.0`` | argument path: ``loss[ener]/limit_pref_e`` The prefactor of energy loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1118,7 +1118,7 @@ loss: .. _`loss[ener]/start_pref_f`: start_pref_f: - | type: ``int`` | ``float``, optional, default: ``1000`` + | type: ``float`` | ``int``, optional, default: ``1000`` | argument path: ``loss[ener]/start_pref_f`` The prefactor of force loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the force label should be provided by file force.npy in each data system. If both start_pref_force and limit_pref_force are set to 0, then the force will be ignored. @@ -1126,7 +1126,7 @@ loss: .. _`loss[ener]/limit_pref_f`: limit_pref_f: - | type: ``int`` | ``float``, optional, default: ``1.0`` + | type: ``float`` | ``int``, optional, default: ``1.0`` | argument path: ``loss[ener]/limit_pref_f`` The prefactor of force loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1134,7 +1134,7 @@ loss: .. _`loss[ener]/start_pref_v`: start_pref_v: - | type: ``int`` | ``float``, optional, default: ``0.0`` + | type: ``float`` | ``int``, optional, default: ``0.0`` | argument path: ``loss[ener]/start_pref_v`` The prefactor of virial loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the virial label should be provided by file virial.npy in each data system. If both start_pref_virial and limit_pref_virial are set to 0, then the virial will be ignored. @@ -1142,7 +1142,7 @@ loss: .. _`loss[ener]/limit_pref_v`: limit_pref_v: - | type: ``int`` | ``float``, optional, default: ``0.0`` + | type: ``float`` | ``int``, optional, default: ``0.0`` | argument path: ``loss[ener]/limit_pref_v`` The prefactor of virial loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1150,7 +1150,7 @@ loss: .. _`loss[ener]/start_pref_ae`: start_pref_ae: - | type: ``int`` | ``float``, optional, default: ``0.0`` + | type: ``float`` | ``int``, optional, default: ``0.0`` | argument path: ``loss[ener]/start_pref_ae`` The prefactor of atom_ener loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the atom_ener label should be provided by file atom_ener.npy in each data system. If both start_pref_atom_ener and limit_pref_atom_ener are set to 0, then the atom_ener will be ignored. @@ -1158,7 +1158,7 @@ loss: .. _`loss[ener]/limit_pref_ae`: limit_pref_ae: - | type: ``int`` | ``float``, optional, default: ``0.0`` + | type: ``float`` | ``int``, optional, default: ``0.0`` | argument path: ``loss[ener]/limit_pref_ae`` The prefactor of atom_ener loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity. @@ -1166,7 +1166,7 @@ loss: .. _`loss[ener]/relative_f`: relative_f: - | type: ``NoneType`` | ``float``, optional + | type: ``float`` | ``NoneType``, optional | argument path: ``loss[ener]/relative_f`` If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label. @@ -1179,7 +1179,7 @@ loss: .. _`loss[tensor]/pref`: pref: - | type: ``int`` | ``float`` + | type: ``float`` | ``int`` | argument path: ``loss[tensor]/pref`` The prefactor of the weight of global loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to global label, i.e. 'polarizability.npy` or `dipole.npy`, whose shape should be #frames x [9 or 3]. If it's larger than 0.0, this npy should be included. @@ -1187,7 +1187,7 @@ loss: .. _`loss[tensor]/pref_atomic`: pref_atomic: - | type: ``int`` | ``float`` + | type: ``float`` | ``int`` | argument path: ``loss[tensor]/pref_atomic`` The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to atomic label, i.e. `atomic_polarizability.npy` or `atomic_dipole.npy`, whose shape should be #frames x ([9 or 3] x #selected atoms). If it's larger than 0.0, this npy should be included. Both `pref` and `pref_atomic` should be provided, and either can be set to 0.0. @@ -1408,7 +1408,7 @@ training: .. _`training/disp_file`: disp_file: - | type: ``str``, optional, default: ``lcueve.out`` + | type: ``str``, optional, default: ``lcurve.out`` | argument path: ``training/disp_file`` The file for printing learning curve. diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index 609dc8721d..d619569c8d 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -5,13 +5,19 @@ Currently, parallel training is enabled in a sychoronized way with help of [Horo Testing `examples/water/se_e2_a` on a 8-GPU host, linear acceleration can be observed with increasing number of cards. | Num of GPU cards | Seconds every 100 samples | Samples per second | Speed up | | -- | -- | -- | -- | -| 1 | 1.6116 | 62.05 | 1.00 | -| 2 | 1.6310 | 61.31 | 1.98 | -| 4 | 1.6168 | 61.85 | 3.99 | -| 8 | 1.6212 | 61.68 | 7.95 | +| 1 | 1.4515 | 68.89 | 1.00 | +| 2 | 1.5962 | 62.65*2 | 1.82 | +| 4 | 1.7635 | 56.71*4 | 3.29 | +| 8 | 1.7267 | 57.91*8 | 6.72 | To experience this powerful feature, please intall Horovod and [mpi4py](https://github.com/mpi4py/mpi4py) first. For better performance on GPU, please follow tuning steps in [Horovod on GPU](https://github.com/horovod/horovod/blob/master/docs/gpus.rst). ```bash +# With GPU, prefer NCCL as communicator. +HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_NCCL_HOME=/path/to/nccl pip3 install horovod mpi4py +``` + +If your work in CPU environment, please prepare runtime as below: +```bash # By default, MPI is used as communicator. HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 pip install horovod mpi4py ``` diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 32deaca906..ea9e1e8075 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -109,7 +109,7 @@ optional arguments: **`--restart model.ckpt`**, continues the training from the checkpoint `model.ckpt`. -**`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`. Note that the `init-frz-model` command only supports the compressed model currently! +**`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`. On some resources limited machines, one may want to control the number of threads used by DeePMD-kit. This is achieved by three environmental variables: `OMP_NUM_THREADS`, `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`. `OMP_NUM_THREADS` controls the multithreading of DeePMD-kit implemented operations. `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS` controls `intra_op_parallelism_threads` and `inter_op_parallelism_threads`, which are Tensorflow configurations for multithreading. An explanation is found [here](https://stackoverflow.com/questions/41233635/meaning-of-inter-op-parallelism-threads-and-intra-op-parallelism-threads). diff --git a/examples/water/se_e3/input.json b/examples/water/se_e3/input.json index f91ccad83c..4b62f4435b 100644 --- a/examples/water/se_e3/input.json +++ b/examples/water/se_e3/input.json @@ -4,9 +4,9 @@ "type_map": ["O", "H"], "descriptor" :{ "type": "se_e3", - "sel": [40, 80], + "sel": "auto", "rcut_smth": 0.50, - "rcut": 6.00, + "rcut": 5.80, "neuron": [2, 4, 8], "resnet_dt": false, "seed": 1, diff --git a/source/lib/src/neighbor_list.cc b/source/lib/src/neighbor_list.cc index c3cd376fbe..cae7630430 100644 --- a/source/lib/src/neighbor_list.cc +++ b/source/lib/src/neighbor_list.cc @@ -1,6 +1,7 @@ #include "neighbor_list.h" #include "device.h" #include +#include // #include // using namespace std; diff --git a/source/lib/src/region.cc b/source/lib/src/region.cc index 90704016c2..d8333e4929 100644 --- a/source/lib/src/region.cc +++ b/source/lib/src/region.cc @@ -22,8 +22,8 @@ Region:: delete [] rec_boxt; } -template struct Region; -template struct Region; +template struct deepmd::Region; +template struct deepmd::Region; template inline FPTYPE diff --git a/source/lmp/fix_dplr.cpp b/source/lmp/fix_dplr.cpp index bfb2b9f543..1d2323a027 100644 --- a/source/lmp/fix_dplr.cpp +++ b/source/lmp/fix_dplr.cpp @@ -41,7 +41,7 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg) efield_fsum_all(4, 0.0), efield_force_flag(0) { -#if LAMMPS_VERSION_NUMBER>=20210702 +#if LAMMPS_VERSION_NUMBER>=20210210 // lammps/lammps#2560 energy_global_flag = 1; virial_global_flag = 1; @@ -123,7 +123,7 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg) int FixDPLR::setmask() { int mask = 0; -#if LAMMPS_VERSION_NUMBER<20210702 +#if LAMMPS_VERSION_NUMBER<20210210 // THERMO_ENERGY removed in lammps/lammps#2560 mask |= THERMO_ENERGY; #endif diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 09e6d55f0a..80b64eefb9 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -238,6 +238,7 @@ PairDeepMD::PairDeepMD(LAMMPS *lmp) error->all(FLERR,"Pair deepmd requires metal unit, please set it by \"units metal\""); } restartinfo = 1; + centroidstressflag = 2 ; // set centroidstressflag = 2 to allow the use of the centroid/stress/atom. Added by Davide Tisi pppmflag = 1; respa_enable = 0; writedata = 0; @@ -366,7 +367,8 @@ void PairDeepMD::compute(int eflag, int vflag) if (do_ghost) { deepmd::InputNlist lmp_list (list->inum, list->ilist, list->numneigh, list->firstneigh); if (single_model || multi_models_no_mod_devi) { - if ( ! (eflag_atom || vflag_atom) ) { + //cvflag_atom is the right flag for the cvatom matrix + if ( ! (eflag_atom || cvflag_atom) ) { #ifdef HIGH_PREC deep_pot.compute (dener, dforce, dvirial, dcoord, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); #else @@ -409,14 +411,26 @@ void PairDeepMD::compute(int eflag, int vflag) if (eflag_atom) { for (int ii = 0; ii < nlocal; ++ii) eatom[ii] += deatom[ii]; } - if (vflag_atom) { + // Added by Davide Tisi 2020 + // interface the atomic virial computed by DeepMD + // with the one used in centroid atoms + if (cvflag_atom) { for (int ii = 0; ii < nall; ++ii){ - vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - vatom[ii][5] += 1.0 * dvatom[9*ii+7]; + //vatom[ii][0] += 1.0 * dvatom[9*ii+0]; + //vatom[ii][1] += 1.0 * dvatom[9*ii+4]; + //vatom[ii][2] += 1.0 * dvatom[9*ii+8]; + //vatom[ii][3] += 1.0 * dvatom[9*ii+3]; + //vatom[ii][4] += 1.0 * dvatom[9*ii+6]; + //vatom[ii][5] += 1.0 * dvatom[9*ii+7]; + cvatom[ii][0] += -1.0 * dvatom[9*ii+0]; // xx + cvatom[ii][1] += -1.0 * dvatom[9*ii+4]; // yy + cvatom[ii][2] += -1.0 * dvatom[9*ii+8]; // zz + cvatom[ii][3] += -1.0 * dvatom[9*ii+3]; // xy + cvatom[ii][4] += -1.0 * dvatom[9*ii+6]; // xz + cvatom[ii][5] += -1.0 * dvatom[9*ii+7]; // yz + cvatom[ii][6] += -1.0 * dvatom[9*ii+1]; // yx + cvatom[ii][7] += -1.0 * dvatom[9*ii+2]; // zx + cvatom[ii][8] += -1.0 * dvatom[9*ii+5]; // zy } } } @@ -478,18 +492,37 @@ void PairDeepMD::compute(int eflag, int vflag) all_force[ii][jj] = all_force_[ii][jj]; } } + all_virial.resize(all_virial_.size()); + for (unsigned ii = 0; ii < all_virial_.size(); ++ii){ + all_virial[ii].resize(all_virial_[ii].size()); + for (unsigned jj = 0; jj < all_virial_[ii].size(); ++jj){ + all_virial[ii][jj] = all_virial_[ii][jj]; + } + } #endif if (eflag_atom) { for (int ii = 0; ii < nlocal; ++ii) eatom[ii] += deatom[ii]; } - if (vflag_atom) { + // Added by Davide Tisi 2020 + // interface the atomic virial computed by DeepMD + // with the one used in centroid atoms + if (cvflag_atom) { for (int ii = 0; ii < nall; ++ii){ - vatom[ii][0] += 1.0 * dvatom[9*ii+0]; - vatom[ii][1] += 1.0 * dvatom[9*ii+4]; - vatom[ii][2] += 1.0 * dvatom[9*ii+8]; - vatom[ii][3] += 1.0 * dvatom[9*ii+3]; - vatom[ii][4] += 1.0 * dvatom[9*ii+6]; - vatom[ii][5] += 1.0 * dvatom[9*ii+7]; + //vatom[ii][0] += 1.0 * dvatom[9*ii+0]; + //vatom[ii][1] += 1.0 * dvatom[9*ii+4]; + //vatom[ii][2] += 1.0 * dvatom[9*ii+8]; + //vatom[ii][3] += 1.0 * dvatom[9*ii+3]; + //vatom[ii][4] += 1.0 * dvatom[9*ii+6]; + //vatom[ii][5] += 1.0 * dvatom[9*ii+7]; + cvatom[ii][0] += -1.0 * dvatom[9*ii+0]; // xx + cvatom[ii][1] += -1.0 * dvatom[9*ii+4]; // yy + cvatom[ii][2] += -1.0 * dvatom[9*ii+8]; // zz + cvatom[ii][3] += -1.0 * dvatom[9*ii+3]; // xy + cvatom[ii][4] += -1.0 * dvatom[9*ii+6]; // xz + cvatom[ii][5] += -1.0 * dvatom[9*ii+7]; // yz + cvatom[ii][6] += -1.0 * dvatom[9*ii+1]; // yx + cvatom[ii][7] += -1.0 * dvatom[9*ii+2]; // zx + cvatom[ii][8] += -1.0 * dvatom[9*ii+5]; // zy } } if (out_freq > 0 && update->ntimestep % out_freq == 0) { diff --git a/source/lmp/pppm_dplr.cpp b/source/lmp/pppm_dplr.cpp index cd4e7da2ba..0fda38e4f3 100644 --- a/source/lmp/pppm_dplr.cpp +++ b/source/lmp/pppm_dplr.cpp @@ -28,7 +28,7 @@ enum{FORWARD_IK,FORWARD_AD,FORWARD_IK_PERATOM,FORWARD_AD_PERATOM}; /* ---------------------------------------------------------------------- */ -#if LAMMPS_VERSION_NUMBER<20190201 +#if LAMMPS_VERSION_NUMBER<20181109 // See lammps/lammps#1165 PPPMDPLR::PPPMDPLR(LAMMPS *lmp, int narg, char **arg) : PPPM(lmp, narg, arg) diff --git a/source/lmp/pppm_dplr.h b/source/lmp/pppm_dplr.h index df436170f9..20e4107385 100644 --- a/source/lmp/pppm_dplr.h +++ b/source/lmp/pppm_dplr.h @@ -21,7 +21,7 @@ namespace LAMMPS_NS { class PPPMDPLR : public PPPM { public: -#if LAMMPS_VERSION_NUMBER<20190201 +#if LAMMPS_VERSION_NUMBER<20181109 // See lammps/lammps#1165 PPPMDPLR(class LAMMPS *, int, char **); #else diff --git a/source/op/CMakeLists.txt b/source/op/CMakeLists.txt index 07d030ef01..96c02ec71c 100644 --- a/source/op/CMakeLists.txt +++ b/source/op/CMakeLists.txt @@ -32,6 +32,20 @@ if (BUILD_PY_IF) ) target_include_directories(op_abi PUBLIC ${TensorFlow_INCLUDE_DIRS}) target_include_directories(op_grads PUBLIC ${TensorFlow_INCLUDE_DIRS}) + if (APPLE) + set_target_properties( + op_abi + PROPERTIES + COMPILE_FLAGS ${OP_CXX_FLAG} + INSTALL_RPATH @loader_path + ) + set_target_properties( + op_grads + PROPERTIES + COMPILE_FLAGS ${OP_CXX_FLAG} + INSTALL_RPATH @loader_path + ) + else() set_target_properties( op_abi PROPERTIES @@ -44,6 +58,7 @@ if (BUILD_PY_IF) COMPILE_FLAGS ${OP_CXX_FLAG} INSTALL_RPATH $ORIGIN ) + endif () endif (BUILD_PY_IF) if (BUILD_CPP_IF) diff --git a/source/tests/common.py b/source/tests/common.py index 5a63666825..e767f8a3eb 100644 --- a/source/tests/common.py +++ b/source/tests/common.py @@ -3,10 +3,9 @@ import pathlib from deepmd.env import tf -from deepmd.env import GLOBAL_TF_FLOAT_PRECISION from deepmd.env import GLOBAL_NP_FLOAT_PRECISION -from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION from deepmd.common import j_loader as dp_j_loader +from deepmd.utils import random as dp_random if GLOBAL_NP_FLOAT_PRECISION == np.float32 : global_default_fv_hh = 1e-2 @@ -26,8 +25,8 @@ def del_data(): if os.path.isdir('system'): shutil.rmtree('system') -def gen_data() : - tmpdata = Data(rand_pert = 0.1, seed = 1) +def gen_data(nframes = 1) : + tmpdata = Data(rand_pert = 0.1, seed = 1, nframes = nframes) sys = dpdata.LabeledSystem() sys.data['atom_names'] = ['foo', 'bar'] sys.data['coords'] = tmpdata.coord @@ -47,14 +46,15 @@ class Data(): def __init__ (self, rand_pert = 0.1, seed = 1, - box_scale = 20) : + box_scale = 20, + nframes = 1): coord = [[0.0, 0.0, 0.1], [1.1, 0.0, 0.1], [0.0, 1.1, 0.1], [4.0, 0.0, 0.0], [5.1, 0.0, 0.0], [4.0, 1.1, 0.0]] - self.nframes = 1 + self.nframes = nframes self.coord = np.array(coord) self.coord = self._copy_nframes(self.coord) - np.random.seed(seed) - self.coord += rand_pert * np.random.random(self.coord.shape) + dp_random.seed(seed) + self.coord += rand_pert * dp_random.random(self.coord.shape) self.fparam = np.array([[0.1, 0.2]]) self.aparam = np.tile(self.fparam, [1, 6]) self.fparam = self._copy_nframes(self.fparam) @@ -69,7 +69,7 @@ def __init__ (self, self.coord = self.coord.reshape([self.nframes, -1, 3]) self.coord = self.coord[:,self.idx_map,:] self.coord = self.coord.reshape([self.nframes, -1]) - self.efield = np.random.random(self.coord.shape) + self.efield = dp_random.random(self.coord.shape) self.atype = self.atype[self.idx_map] self.datype = self._copy_nframes(self.atype) @@ -128,7 +128,7 @@ def get_test_box_data (self, coord0_, box0_, type0_ = self.get_data() coord = coord0_[0] box = box0_[0] - box += rand_pert * np.random.random(box.shape) + box += rand_pert * dp_random.random(box.shape) atype = type0_[0] nframes = 1 natoms = coord.size // 3 @@ -258,11 +258,9 @@ def virial_test (inter, num_vir = np.transpose(num_vir, [1,0]) box3 = dbox[0].reshape([3,3]) num_vir = np.matmul(num_vir, box3) - for ii in range(3): - for jj in range(3): - testCase.assertAlmostEqual(ana_vir[ii][jj], num_vir[ii][jj], - places=places, - msg = 'virial component %d %d ' % (ii,jj)) + np.testing.assert_almost_equal(ana_vir, num_vir, + places, + err_msg = 'virial component') diff --git a/source/tests/model_compression/data/set.000/box.npy b/source/tests/model_compression/data/set.000/box.npy index 1f5de3f920..aa092a9aad 100644 Binary files a/source/tests/model_compression/data/set.000/box.npy and b/source/tests/model_compression/data/set.000/box.npy differ diff --git a/source/tests/model_compression/data/set.000/coord.npy b/source/tests/model_compression/data/set.000/coord.npy index 3ab4571007..2205bfd452 100644 Binary files a/source/tests/model_compression/data/set.000/coord.npy and b/source/tests/model_compression/data/set.000/coord.npy differ diff --git a/source/tests/model_compression/data/set.000/energy.npy b/source/tests/model_compression/data/set.000/energy.npy index 35899bf0f1..6c1e5ce145 100644 Binary files a/source/tests/model_compression/data/set.000/energy.npy and b/source/tests/model_compression/data/set.000/energy.npy differ diff --git a/source/tests/model_compression/data/set.000/force.npy b/source/tests/model_compression/data/set.000/force.npy index a1f696c186..10698ebb59 100644 Binary files a/source/tests/model_compression/data/set.000/force.npy and b/source/tests/model_compression/data/set.000/force.npy differ diff --git a/source/tests/model_compression/data/type.raw b/source/tests/model_compression/data/type.raw index 97e8fdfcf8..e329bb5191 100644 --- a/source/tests/model_compression/data/type.raw +++ b/source/tests/model_compression/data/type.raw @@ -1,192 +1,6 @@ +0 +1 +1 0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -0 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 +1 +1 \ No newline at end of file diff --git a/source/tests/test_data_modifier.py b/source/tests/test_data_modifier.py index 977df9a2b6..d791ca2844 100644 --- a/source/tests/test_data_modifier.py +++ b/source/tests/test_data_modifier.py @@ -119,10 +119,9 @@ def _test_fv (self): ep, _, __ = dcm.eval(coordp, box, atype, eval_fv = False) em, _, __ = dcm.eval(coordm, box, atype, eval_fv = False) num_f = -(ep - em) / (2.*hh) - for ff in range(nframes): - self.assertAlmostEqual(vf[ff,ii], num_f[ff], - places = places, - msg = 'frame %d dof %d does not match' % (ff, ii)) + np.testing.assert_almost_equal(vf[:,ii].ravel(), num_f.ravel(), + places, + err_msg = 'dof %d does not match' % (ii)) box3 = np.reshape(box, [nframes, 3,3]) rbox3 = np.linalg.inv(box3) @@ -150,10 +149,7 @@ def _test_fv (self): t_esti = np.matmul(num_deriv, box3) # print(t_esti, '\n', vv.reshape([-1, 3, 3])) - for ff in range(nframes): - for ii in range(3): - for jj in range(3): - self.assertAlmostEqual(t_esti[ff][ii][jj], vv[ff,ii*3+jj], - places = places, - msg = "frame %d virial component [%d,%d] failed" % (ff, ii, jj)) + np.testing.assert_almost_equal(t_esti.ravel(), vv.ravel(), + places, + err_msg = "virial component failed") diff --git a/source/tests/test_data_modifier_shuffle.py b/source/tests/test_data_modifier_shuffle.py index c14b6dd105..194a80e9c4 100644 --- a/source/tests/test_data_modifier_shuffle.py +++ b/source/tests/test_data_modifier_shuffle.py @@ -185,11 +185,9 @@ def test_z_dipole(self): dv01 = dv01.reshape([self.nframes, -1]) dv1 = dv1.reshape([self.nframes, -1]) - for ii in range(self.nframes): - for jj in range(self.nsel): - self.assertAlmostEqual( - dv01[ii][jj], dv1[ii][jj], - msg = "dipole [%d,%d] dose not match" % (ii, jj)) + np.testing.assert_almost_equal( + dv01, dv1, + err_msg = "dipole dose not match") def test_modify(self): @@ -202,18 +200,12 @@ def test_modify(self): ve1, vf1, vv1 = dcm.eval(self.coords1, self.box1, self.atom_types1) vf01 = vf0[:,self.idx_map, :] - for ii in range(self.nframes): - self.assertAlmostEqual(ve0[ii], ve1[ii], - msg = 'energy %d should match' % ii) - for ii in range(self.nframes): - for jj in range(9): - self.assertAlmostEqual(vv0[ii][jj], vv1[ii][jj], - msg = 'virial [%d,%d] should match' % (ii,jj)) - for ii in range(self.nframes): - for jj in range(self.natoms): - for dd in range(3): - self.assertAlmostEqual( - vf01[ii][jj][dd], vf1[ii][jj][dd], - msg = "force [%d,%d,%d] dose not match" % (ii,jj,dd)) + np.testing.assert_almost_equal(ve0, ve1, + err_msg = 'energy should match') + np.testing.assert_almost_equal(vv0, vv1, + err_msg = 'virial should match') + np.testing.assert_almost_equal( + vf01, vf1, + err_msg = "force dose not match") diff --git a/source/tests/test_deepdipole.py b/source/tests/test_deepdipole.py index 84049a9e52..7b23e3ae00 100644 --- a/source/tests/test_deepdipole.py +++ b/source/tests/test_deepdipole.py @@ -44,8 +44,7 @@ def test_1frame_atm(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,3)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -58,8 +57,7 @@ def test_2frame_atm(self): self.assertEqual(dd.shape, (nframes,nsel,3)) # check values expected_d = np.concatenate((self.expected_d, self.expected_d)) - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), expected_d, default_places) class TestDeepDipoleNoPBC(unittest.TestCase) : @@ -87,8 +85,7 @@ def test_1frame_atm(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,3)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_1frame_atm_large_box(self): dd = self.dp.eval(self.coords, self.box, self.atype) @@ -98,8 +95,7 @@ def test_1frame_atm_large_box(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,3)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) @unittest.skipIf(parse_version(tf.__version__) < parse_version("1.15"), @@ -138,8 +134,7 @@ def test_1frame_old(self): nframes = 1 self.assertEqual(gt.shape, (nframes,self.nout)) # check values - for ii in range(gt.size): - self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) def test_1frame_old_atm(self): at = self.dp.eval(self.coords, self.box, self.atype) @@ -149,8 +144,7 @@ def test_1frame_old_atm(self): nsel = 2 self.assertEqual(at.shape, (nframes,nsel,self.nout)) # check values - for ii in range(at.size): - self.assertAlmostEqual(at.reshape([-1])[ii], self.expected_t.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places) def test_2frame_old_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -163,8 +157,7 @@ def test_2frame_old_atm(self): self.assertEqual(at.shape, (nframes,nsel,self.nout)) # check values expected_d = np.concatenate((self.expected_t, self.expected_t)) - for ii in range(at.size): - self.assertAlmostEqual(at.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(at.ravel(), expected_d, default_places) def test_1frame_full(self): gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic = False) @@ -175,12 +168,9 @@ def test_1frame_full(self): self.assertEqual(ff.shape, (nframes,self.nout,natoms,3)) self.assertEqual(vv.shape, (nframes,self.nout,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(gt.size): - self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places) - for ii in range(vv.size): - self.assertAlmostEqual(vv.reshape([-1])[ii], self.expected_gv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f, default_places) + np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) + np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places) def test_1frame_full_atm(self): gt, ff, vv, at, av = self.dp.eval_full(self.coords, self.box, self.atype, atomic = True) diff --git a/source/tests/test_deepmd_data.py b/source/tests/test_deepmd_data.py index 78d0a36cf2..8532d8d9a4 100644 --- a/source/tests/test_deepmd_data.py +++ b/source/tests/test_deepmd_data.py @@ -49,10 +49,7 @@ def test_load_set_1(self) : .add('value_1', 1, atomic=True, must=True, type_sel = [0]) data = dd._load_set(os.path.join(self.data_name, 'set.foo')) self.assertEqual(data['value_1'].shape, (self.nframes, 2)) - for ii in range(self.nframes): - for jj in range(2): - self.assertAlmostEqual(data['value_1'][ii][jj], - self.value_1[ii][jj]) + np.testing.assert_almost_equal(data['value_1'], self.value_1) def test_load_set_2(self) : @@ -60,10 +57,7 @@ def test_load_set_2(self) : .add('value_2', 1, atomic=True, must=True, type_sel = [1]) data = dd._load_set(os.path.join(self.data_name, 'set.foo')) self.assertEqual(data['value_2'].shape, (self.nframes, 4)) - for ii in range(self.nframes): - for jj in range(4): - self.assertAlmostEqual(data['value_2'][ii][jj], - self.value_2[ii][jj]) + np.testing.assert_almost_equal(data['value_2'], self.value_2) class TestData (unittest.TestCase) : @@ -217,8 +211,7 @@ def test_avg(self) : .add('test_frame', 5, atomic=False, must=True) favg = dd.avg('test_frame') fcmp = np.average(np.concatenate((self.test_frame, self.test_frame_bar), axis = 0), axis = 0) - for ii in range(favg.size) : - self.assertAlmostEqual((favg[ii]), (fcmp[ii]), places = places) + np.testing.assert_almost_equal(favg, fcmp, places) def test_check_batch_size(self) : dd = DeepmdData(self.data_name) @@ -263,8 +256,4 @@ def test_get_nbatch(self): self.assertEqual(nb, 2) def _comp_np_mat2(self, first, second) : - for ii in range(first.shape[0]) : - for jj in range(first.shape[1]) : - self.assertAlmostEqual(first[ii][jj], second[ii][jj], - msg = 'item [%d][%d] does not match' % (ii,jj), - places = places) + np.testing.assert_almost_equal(first, second, places) diff --git a/source/tests/test_deeppolar.py b/source/tests/test_deeppolar.py index 64dbcbf597..b60a1abf79 100644 --- a/source/tests/test_deeppolar.py +++ b/source/tests/test_deeppolar.py @@ -44,8 +44,7 @@ def test_1frame_atm(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,9)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -58,8 +57,8 @@ def test_2frame_atm(self): self.assertEqual(dd.shape, (nframes,nsel,9)) # check values expected_d = np.concatenate((self.expected_d, self.expected_d)) - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), expected_d, default_places) + class TestDeepPolarNoPBC(unittest.TestCase) : @@ -87,8 +86,7 @@ def test_1frame_atm(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,9)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_1frame_atm_large_box(self): dd = self.dp.eval(self.coords, self.box, self.atype) @@ -98,8 +96,7 @@ def test_1frame_atm_large_box(self): nsel = 2 self.assertEqual(dd.shape, (nframes,nsel,9)) # check values - for ii in range(dd.size): - self.assertAlmostEqual(dd.reshape([-1])[ii], self.expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) @unittest.skipIf(parse_version(tf.__version__) < parse_version("1.15"), @@ -138,8 +135,7 @@ def test_1frame_old(self): nframes = 1 self.assertEqual(gt.shape, (nframes,self.nout)) # check values - for ii in range(gt.size): - self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) def test_1frame_old_atm(self): at = self.dp.eval(self.coords, self.box, self.atype) @@ -149,8 +145,7 @@ def test_1frame_old_atm(self): nsel = 2 self.assertEqual(at.shape, (nframes,nsel,self.nout)) # check values - for ii in range(at.size): - self.assertAlmostEqual(at.reshape([-1])[ii], self.expected_t.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(at.ravel(), self.expected_t, default_places) def test_2frame_old_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -163,8 +158,7 @@ def test_2frame_old_atm(self): self.assertEqual(at.shape, (nframes,nsel,self.nout)) # check values expected_d = np.concatenate((self.expected_t, self.expected_t)) - for ii in range(at.size): - self.assertAlmostEqual(at.reshape([-1])[ii], expected_d.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(at.ravel(), expected_d, default_places) def test_1frame_full(self): gt, ff, vv = self.dp.eval_full(self.coords, self.box, self.atype, atomic = False) @@ -175,12 +169,9 @@ def test_1frame_full(self): self.assertEqual(ff.shape, (nframes,self.nout,natoms,3)) self.assertEqual(vv.shape, (nframes,self.nout,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(gt.size): - self.assertAlmostEqual(gt.reshape([-1])[ii], self.expected_gt.reshape([-1])[ii], places = default_places) - for ii in range(vv.size): - self.assertAlmostEqual(vv.reshape([-1])[ii], self.expected_gv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f, default_places) + np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) + np.testing.assert_almost_equal(vv.ravel(), self.expected_gv, default_places) def test_1frame_full_atm(self): gt, ff, vv, at, av = self.dp.eval_full(self.coords, self.box, self.atype, atomic = True) diff --git a/source/tests/test_deeppot_a.py b/source/tests/test_deeppot_a.py index 3726299fb5..2ce4ba1193 100644 --- a/source/tests/test_deeppot_a.py +++ b/source/tests/test_deeppot_a.py @@ -108,14 +108,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -128,18 +125,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_2frame_atm(self): @@ -158,18 +150,13 @@ def test_2frame_atm(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis = 0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis = 0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis = 0) - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) class TestDeepPotANoPBC(unittest.TestCase) : @@ -200,14 +187,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -220,18 +204,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_2frame_atm(self): @@ -249,18 +228,13 @@ def test_2frame_atm(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis = 0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis = 0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis = 0) - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) class TestDeepPotALargeBoxNoPBC(unittest.TestCase) : @@ -291,14 +265,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -311,18 +282,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_ase(self): from ase import Atoms @@ -334,9 +300,6 @@ def test_ase(self): ee = water.get_potential_energy() ff = water.get_forces() nframes = 1 - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) - + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) diff --git a/source/tests/test_deeppot_r.py b/source/tests/test_deeppot_r.py index c43a71677a..2a4ad62a83 100644 --- a/source/tests/test_deeppot_r.py +++ b/source/tests/test_deeppot_r.py @@ -59,14 +59,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -79,18 +76,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_2frame_atm(self): @@ -109,18 +101,13 @@ def test_2frame_atm(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis = 0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis = 0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis = 0) - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) class TestDeepPotRNoPBC(unittest.TestCase) : @@ -157,14 +144,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -177,18 +161,13 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_2frame_atm(self): @@ -206,18 +185,13 @@ def test_2frame_atm(self): expected_f = np.concatenate((self.expected_f, self.expected_f), axis = 0) expected_e = np.concatenate((self.expected_e, self.expected_e), axis = 0) expected_v = np.concatenate((self.expected_v, self.expected_v), axis = 0) - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) expected_se = np.sum(expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) class TestDeepPotRLargeBoxNoPBC(unittest.TestCase) : @@ -254,14 +228,11 @@ def test_1frame(self): self.assertEqual(ff.shape, (nframes,natoms,3)) self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) def test_1frame_atm(self): ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) @@ -274,17 +245,12 @@ def test_1frame_atm(self): self.assertEqual(ae.shape, (nframes,natoms,1)) self.assertEqual(av.shape, (nframes,natoms,9)) # check values - for ii in range(ff.size): - self.assertAlmostEqual(ff.reshape([-1])[ii], self.expected_f.reshape([-1])[ii], places = default_places) - for ii in range(ae.size): - self.assertAlmostEqual(ae.reshape([-1])[ii], self.expected_e.reshape([-1])[ii], places = default_places) - for ii in range(av.size): - self.assertAlmostEqual(av.reshape([-1])[ii], self.expected_v.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee.reshape([-1])[ii], expected_se.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv.reshape([-1])[ii], expected_sv.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) diff --git a/source/tests/test_descrpt_nonsmth.py b/source/tests/test_descrpt_nonsmth.py index decb18e053..01fdbcfe85 100644 --- a/source/tests/test_descrpt_nonsmth.py +++ b/source/tests/test_descrpt_nonsmth.py @@ -216,12 +216,8 @@ def test_pbc(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) def test_pbc_small_box(self): data0 = Data() @@ -257,12 +253,8 @@ def test_pbc_small_box(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) if __name__ == '__main__': diff --git a/source/tests/test_descrpt_se_a_type.py b/source/tests/test_descrpt_se_a_type.py index c6f3cb5a19..f65338d0a4 100644 --- a/source/tests/test_descrpt_se_a_type.py +++ b/source/tests/test_descrpt_se_a_type.py @@ -17,8 +17,8 @@ GLOBAL_NP_FLOAT_PRECISION = np.float64 class TestModel(tf.test.TestCase): - def setUp(self) : - gen_data() + def setUp(self): + gen_data(nframes=2) def test_descriptor_two_sides(self): jfile = 'water_se_a_type.json' @@ -28,7 +28,7 @@ def test_descriptor_two_sides(self): set_pfx = j_must_have(jdata, 'set_prefix') batch_size = j_must_have(jdata, 'batch_size') test_size = j_must_have(jdata, 'numb_test') - batch_size = 1 + batch_size = 2 test_size = 1 stop_batch = j_must_have(jdata, 'stop_batch') rcut = j_must_have (jdata['model']['descriptor'], 'rcut') @@ -124,8 +124,7 @@ def test_descriptor_two_sides(self): 0.0004482204892297628,-0.00016887749501640607,-0.00016887749501640607,6.364643102775375e-05,-0.001181345877677835,0.0004452029242063362,-0.0008941636427724908,0.0003369586197174627,0.0004677878512312651,-0.00017625260641095753] places = 10 - for ii in range(model_dout.size) : - self.assertAlmostEqual(model_dout[ii], ref_dout[ii], places = places) + np.testing.assert_almost_equal(model_dout, ref_dout, places) def test_descriptor_one_side(self): @@ -234,8 +233,7 @@ def test_descriptor_one_side(self): -0.0015760302086346792,-0.0012475134925387294,-0.001041074331192672,-0.0008239586048523492,0.0015319673563669856,0.0012124704278707746] places = 10 - for ii in range(model_dout.size) : - self.assertAlmostEqual(model_dout[ii], ref_dout[ii], places = places) + np.testing.assert_almost_equal(model_dout, ref_dout, places) diff --git a/source/tests/test_descrpt_se_r.py b/source/tests/test_descrpt_se_r.py index bf002736b0..e2cc6360be 100644 --- a/source/tests/test_descrpt_se_r.py +++ b/source/tests/test_descrpt_se_r.py @@ -196,12 +196,8 @@ def test_pbc(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) def test_pbc_small_box(self): @@ -238,12 +234,8 @@ def test_pbc_small_box(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) if __name__ == '__main__': diff --git a/source/tests/test_descrpt_sea_ef_rot.py b/source/tests/test_descrpt_sea_ef_rot.py index f172ec2771..e96951b275 100644 --- a/source/tests/test_descrpt_sea_ef_rot.py +++ b/source/tests/test_descrpt_sea_ef_rot.py @@ -353,10 +353,9 @@ def test_rot_field_corot(self, suffix=''): self.efield: new_defield, self.tnatoms: self.natoms}) - for ii in range(0, self.natoms[0]): - self.assertAlmostEqual(p_ae0[ii], p_ae1[ii]) - self.assertAlmostEqual(v_ae0[ii], v_ae1[ii]) - self.assertAlmostEqual(ae0[ii], ae1[ii]) + np.testing.assert_almost_equal(p_ae0, p_ae1) + np.testing.assert_almost_equal(v_ae0, v_ae1) + np.testing.assert_almost_equal(ae0, ae1) diff --git a/source/tests/test_descrpt_smooth.py b/source/tests/test_descrpt_smooth.py index adb5beb354..00dbbd61d0 100644 --- a/source/tests/test_descrpt_smooth.py +++ b/source/tests/test_descrpt_smooth.py @@ -207,12 +207,8 @@ def test_pbc(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) def test_pbc_small_box(self): data0 = Data() @@ -248,12 +244,8 @@ def test_pbc_small_box(self): inter1.tnatoms: inter1.natoms}) self.assertAlmostEqual(e0[0], e1[0]) - for ii in range(f0[0].size): - # print(ii) - self.assertAlmostEqual(f0[0][ii], f1[0][ii]) - for ii in range(v0[0].size): - # print(ii) - self.assertAlmostEqual(v0[0][ii], v1[0][ii]) + np.testing.assert_almost_equal(f0[0], f1[0]) + np.testing.assert_almost_equal(v0[0], v1[0]) if __name__ == '__main__': diff --git a/source/tests/test_dipole_se_a.py b/source/tests/test_dipole_se_a.py index 3687ba47f2..42a615300a 100644 --- a/source/tests/test_dipole_se_a.py +++ b/source/tests/test_dipole_se_a.py @@ -98,15 +98,13 @@ def test_model(self): refp = [1.616802262298876514e+01,9.809535439521079425e+00,3.572312180768947854e-01,1.336308874095981203e+00,1.057908563208963848e+01,-5.999602350098874881e-01] places = 10 - for ii in range(p.size) : - self.assertAlmostEqual(p[ii], refp[ii], places = places) + np.testing.assert_almost_equal(p, refp, places) gp = gp.reshape([-1]) refgp = np.array(refp).reshape(-1, 3).sum(0) places = 9 - for ii in range(gp.size) : - self.assertAlmostEqual(gp[ii], refgp[ii], places = places) + np.testing.assert_almost_equal(gp, refgp, places) # make sure only one frame is used feed_dict_single = {t_prop_c: test_data['prop_c'], @@ -140,11 +138,8 @@ def test_model(self): @ box0.reshape(3,3)).reshape(-1) delta = 1e-5 - for ii in range(pf.size) : - self.assertAlmostEqual(pf[ii], fdf[ii], delta = delta) - for ii in range(pv.size) : - self.assertAlmostEqual(pv[ii], fdv[ii], delta = delta) + np.testing.assert_allclose(pf, fdf, delta) + np.testing.assert_allclose(pv, fdv, delta) # make sure atomic virial sum to virial places = 10 - for ii in range(pv.size) : - self.assertAlmostEqual(pv[ii], spv[ii], places = places) \ No newline at end of file + np.testing.assert_almost_equal(pv, spv, places) \ No newline at end of file diff --git a/source/tests/test_dipolecharge.py b/source/tests/test_dipolecharge.py index ccda444a32..a36919baec 100644 --- a/source/tests/test_dipolecharge.py +++ b/source/tests/test_dipolecharge.py @@ -76,12 +76,9 @@ def test_1frame(self): ee = ee.reshape([-1]) ff = ff.reshape([-1]) vv = vv.reshape([-1]) - for ii in range(ee.size): - self.assertAlmostEqual(ee[ii], self.expected_e[ii]) - for ii in range(ff.size): - self.assertAlmostEqual(ff[ii], self.expected_f[ii]) - for ii in range(vv.size): - self.assertAlmostEqual(vv[ii], self.expected_v[ii]) + np.testing.assert_almost_equal(ee, self.expected_e) + np.testing.assert_almost_equal(ff, self.expected_f) + np.testing.assert_almost_equal(vv, self.expected_v) def test_2frame(self): nframes = 2 @@ -103,10 +100,7 @@ def test_2frame(self): ee = ee.reshape([-1]) ff = ff.reshape([-1]) vv = vv.reshape([-1]) - for ii in range(ee.size): - self.assertAlmostEqual(ee[ii], self.expected_e[ii]) - for ii in range(ff.size): - self.assertAlmostEqual(ff[ii], self.expected_f[ii]) - for ii in range(vv.size): - self.assertAlmostEqual(vv[ii], self.expected_v[ii]) + np.testing.assert_almost_equal(ee, self.expected_e) + np.testing.assert_almost_equal(ff, self.expected_f) + np.testing.assert_almost_equal(vv, self.expected_v) diff --git a/source/tests/test_embedding_net.py b/source/tests/test_embedding_net.py index 03084a8267..d1930524f7 100644 --- a/source/tests/test_embedding_net.py +++ b/source/tests/test_embedding_net.py @@ -32,9 +32,7 @@ def test_enlarger_net(self): refout = [[-0.1482171, -0.14177827, -0.76181204, 0.21266767], [-0.27800543, -0.08974353, -0.78784335, 0.3485518 ], [-0.36744368, -0.06285603, -0.80749876, 0.4347974 ]] - for ii in range(self.ndata): - for jj in range(network_size[-1]): - self.assertAlmostEqual(refout[ii][jj], myout[ii][jj], places = self.places) + np.testing.assert_almost_equal(refout, myout, self.places) def test_enlarger_net_1(self): @@ -50,9 +48,7 @@ def test_enlarger_net_1(self): refout = [[ 0.10842905, -0.61623145, -1.46738788, -0.01921788], [ 0.09376136, -0.75526936, -1.64995884, 0.01076112], [ 0.1033177, -0.8911794, -1.75530172, 0.00653156]] - for ii in range(self.ndata): - for jj in range(network_size[-1]): - self.assertAlmostEqual(refout[ii][jj], myout[ii][jj], places = self.places) + np.testing.assert_almost_equal(refout, myout, self.places) def test_enlarger_net_1_idt(self): network_size = [4, 4] @@ -68,9 +64,7 @@ def test_enlarger_net_1_idt(self): refout = [[ 0.10839754, -0.6161336, -1.46673253, -0.01927138], [ 0.09370214, -0.75516888, -1.64927868, 0.01067603], [ 0.10323835, -0.89107102, -1.75460243, 0.00642493]] - for ii in range(self.ndata): - for jj in range(network_size[-1]): - self.assertAlmostEqual(refout[ii][jj], myout[ii][jj], places = self.places) + np.testing.assert_almost_equal(refout, myout, self.places) def test_enlarger_net_2(self): network_size = [2, 4] @@ -85,9 +79,7 @@ def test_enlarger_net_2(self): refout = [[ 0.24023149, -0.66311811, -0.50951819, -0.36873654], [ 2.00858313, -0.05971232, 0.52272395, -0.12604478], [ 3.39365063, 0.63492697, 1.5780069, 0.46445682]] - for ii in range(self.ndata): - for jj in range(network_size[-1]): - self.assertAlmostEqual(refout[ii][jj], myout[ii][jj], places = self.places) + np.testing.assert_almost_equal(refout, myout, self.places) def test_enlarger_net_2(self): @@ -104,9 +96,7 @@ def test_enlarger_net_2(self): refout = [[ 0.2403889, -0.66290763, -0.50883586, -0.36869913], [ 2.00891479, -0.05936574, 0.52351633, -0.12579749], [ 3.3940202, 0.63538459, 1.57887697, 0.46486689]] - for ii in range(self.ndata): - for jj in range(network_size[-1]): - self.assertAlmostEqual(refout[ii][jj], myout[ii][jj], places = self.places) + np.testing.assert_almost_equal(refout, myout, self.places) diff --git a/source/tests/test_ewald.py b/source/tests/test_ewald.py index f4913db36f..236c295b91 100644 --- a/source/tests/test_ewald.py +++ b/source/tests/test_ewald.py @@ -76,20 +76,15 @@ def test_py_interface(self): }) er = EwaldRecp(self.ewald_h, self.ewald_beta) e1, f1, v1 = er.eval(self.dcoord, self.dcharge, self.dbox) - for ff in range(self.nframes): - self.assertAlmostEqual(e[ff], e1[ff], - places = places, - msg = "frame %d energy failed" % (ff)) - for idx in range(self.natoms): - for dd in range(3): - self.assertAlmostEqual(f[ff, idx*3+dd], f1[ff,idx*3+dd], - places = places, - msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd)) - for d0 in range(3): - for d1 in range(3): - self.assertAlmostEqual(v[ff, d0*3+d1], v[ff,d0*3+d1], - places = places, - msg = "frame %d virial component [%d,%d] failed" % (ff, d0, d1)) + np.testing.assert_almost_equal(e, e1, + places, + err_msg = "energy failed") + np.testing.assert_almost_equal(f, f1, + places, + err_msg = "force component failed") + np.testing.assert_almost_equal(v, v, + places, + err_msg = "virial component failed") @@ -129,10 +124,9 @@ def test_force(self): self.nloc: [self.natoms], }) c_force = -(energyp[0] - energym[0]) / (2*hh) - for ff in range(self.nframes): - self.assertAlmostEqual(c_force[ff], force[ff,idx*3+dd], - places = places, - msg = "frame %d force component [%d,%d] failed" % (ff, idx, dd)) + np.testing.assert_almost_equal(c_force, force[:,idx*3+dd], + places, + err_msg = "force component [%d,%d] failed" % (idx, dd)) def test_virial(self): @@ -204,12 +198,9 @@ def test_virial(self): # # print(0.5 * (t_esti[0] + t_esti[0].T) - virial[0].reshape([3,3])) # print(0.5 * (t_esti[0] + t_esti[0]) - virial[0].reshape([3,3])) # print(0.5 * (t_esti[0] + t_esti[0].T) - virial[0].reshape([3,3])) - for ff in range(self.nframes): - for ii in range(3): - for jj in range(3): - self.assertAlmostEqual(t_esti[ff][ii][jj], virial[ff,ii*3+jj], - places = places, - msg = "frame %d virial component [%d,%d] failed" % (ff, ii, jj)) + np.testing.assert_almost_equal(t_esti.ravel(), virial.ravel(), + places, + err_msg = "virial component failed") diff --git a/source/tests/test_fitting_ener_type.py b/source/tests/test_fitting_ener_type.py index 25b85391e2..1eb1002147 100644 --- a/source/tests/test_fitting_ener_type.py +++ b/source/tests/test_fitting_ener_type.py @@ -109,7 +109,6 @@ def test_fitting(self): places = 10 - for ii in range(pred_atom_ener.size) : - self.assertAlmostEqual(pred_atom_ener[ii], ref_atom_ener[ii], places = places) + np.testing.assert_almost_equal(pred_atom_ener, ref_atom_ener, places) diff --git a/source/tests/test_fitting_stat.py b/source/tests/test_fitting_stat.py index a17bef6453..970e30da4b 100644 --- a/source/tests/test_fitting_stat.py +++ b/source/tests/test_fitting_stat.py @@ -81,8 +81,7 @@ def test (self) : arefa, arefs = _brute_aparam(all_data, len(avgs)) fitting.compute_input_stats(all_data, protection = 1e-2) # print(frefa, frefs) - for ii in range(len(avgs)): - self.assertAlmostEqual(frefa[ii], fitting.fparam_avg[ii]) - self.assertAlmostEqual(frefs[ii], fitting.fparam_std[ii]) - self.assertAlmostEqual(arefa[ii], fitting.aparam_avg[ii]) - self.assertAlmostEqual(arefs[ii], fitting.aparam_std[ii]) + np.testing.assert_almost_equal(frefa, fitting.fparam_avg) + np.testing.assert_almost_equal(frefs, fitting.fparam_std) + np.testing.assert_almost_equal(arefa, fitting.aparam_avg) + np.testing.assert_almost_equal(arefs, fitting.aparam_std) diff --git a/source/tests/test_gen_stat_data.py b/source/tests/test_gen_stat_data.py index b7391f7070..f8c990b756 100644 --- a/source/tests/test_gen_stat_data.py +++ b/source/tests/test_gen_stat_data.py @@ -41,10 +41,7 @@ def tearDown(self): shutil.rmtree('system_1') def _comp_data(self, d0, d1) : - for ii in range(d0.shape[0]): - for jj in range(d0.shape[1]): - for kk in range(d0.shape[2]): - self.assertAlmostEqual(d0[ii][jj][kk], d1[ii][jj][kk]) + np.testing.assert_almost_equal(d0, d1) def test_merge_all_stat(self): dp_random.seed(0) @@ -58,13 +55,13 @@ def test_merge_all_stat(self): 5, 10, 1.0) - data1.add('force', 3, atomic = True, must = True) + data1.add('energy', 1, must = True) dp_random.seed(0) data2 = DeepmdDataSystem(['system_0', 'system_1'], 5, 10, 1.0) - data2.add('force', 3, atomic = True, must = True) + data2.add('energy', 1, must = True) dp_random.seed(0) all_stat_0 = make_stat_input(data0, 10, merge_sys = False) @@ -119,5 +116,4 @@ def test_ener_shift(self): ener_shift0 = data.compute_energy_shift(rcond = 1) all_stat = make_stat_input(data, 4, merge_sys = False) ener_shift1 = EnerFitting._compute_output_stats(all_stat, rcond = 1) - for ii in range(len(ener_shift0)): - self.assertAlmostEqual(ener_shift0[ii], ener_shift1[ii]) + np.testing.assert_almost_equal(ener_shift0, ener_shift1) diff --git a/source/tests/test_model_compression.py b/source/tests/test_model_compression.py index 8465c5a05e..2f42134409 100644 --- a/source/tests/test_model_compression.py +++ b/source/tests/test_model_compression.py @@ -1,6 +1,7 @@ import os,sys,platform,shutil,dpdata,json import numpy as np import unittest +import subprocess as sp from deepmd.infer import DeepPot from deepmd.env import MODEL_VERSION @@ -14,9 +15,21 @@ default_places = 10 def _file_delete(file) : - if os.path.exists(file): + if os.path.isdir(file): + os.rmdir(file) + elif os.path.isfile(file): os.remove(file) +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode + def _init_models(): data_file = str(tests_path / os.path.join("model_compression", "data")) frozen_model = str(tests_path / "dp-original.pb") @@ -28,12 +41,12 @@ def _init_models(): with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) - ret = os.system("dp train " + INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + frozen_model) - assert(ret == 0), "DP freeze error!" - ret = os.system("dp compress " + " -i " + frozen_model + " -o " + compressed_model) - assert(ret == 0), "DP model compression error!" + ret = _subprocess_run("dp train " + INPUT) + np.testing.assert_equal(ret, 0, 'DP train failed!') + ret = _subprocess_run("dp freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, 'DP freeze failed!') + ret = _subprocess_run("dp compress " + " -i " + frozen_model + " -o " + compressed_model) + np.testing.assert_equal(ret, 0, 'DP model compression failed!') return INPUT, frozen_model, compressed_model INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() @@ -78,12 +91,9 @@ def test_1frame(self): self.assertEqual(ff1.shape, (nframes,natoms,3)) self.assertEqual(vv1.shape, (nframes,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_1frame_atm(self): ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) @@ -102,16 +112,11 @@ def test_1frame_atm(self): self.assertEqual(ae1.shape, (nframes,natoms,1)) self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -133,16 +138,11 @@ def test_2frame_atm(self): self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) class TestDeepPotANoPBC(unittest.TestCase) : @@ -172,12 +172,9 @@ def test_1frame(self): self.assertEqual(ff1.shape, (nframes,natoms,3)) self.assertEqual(vv1.shape, (nframes,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_1frame_atm(self): ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) @@ -196,16 +193,11 @@ def test_1frame_atm(self): self.assertEqual(ae1.shape, (nframes,natoms,1)) self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -226,16 +218,11 @@ def test_2frame_atm(self): self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) class TestDeepPotALargeBoxNoPBC(unittest.TestCase) : @@ -265,12 +252,9 @@ def test_1frame(self): self.assertEqual(ff1.shape, (nframes,natoms,3)) self.assertEqual(vv1.shape, (nframes,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_1frame_atm(self): ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) @@ -289,16 +273,11 @@ def test_1frame_atm(self): self.assertEqual(ae1.shape, (nframes,natoms,1)) self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_ase(self): from ase import Atoms @@ -316,10 +295,8 @@ def test_ase(self): ee1 = water1.get_potential_energy() ff1 = water1.get_forces() nframes = 1 - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) class TestDeepPotAPBCExcludeTypes(unittest.TestCase) : @classmethod @@ -343,10 +320,17 @@ def tearDownClass(self): _file_delete("out.json") _file_delete("compress.json") _file_delete("checkpoint") - _file_delete("lcurve.out") _file_delete("model.ckpt.meta") _file_delete("model.ckpt.index") _file_delete("model.ckpt.data-00000-of-00001") + _file_delete("model.ckpt-100.meta") + _file_delete("model.ckpt-100.index") + _file_delete("model.ckpt-100.data-00000-of-00001") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression/model.ckpt.meta") + _file_delete("model-compression/model.ckpt.index") + _file_delete("model-compression/model.ckpt.data-00000-of-00001") + _file_delete("model-compression") def test_attrs(self): self.assertEqual(self.dp_original.get_ntypes(), 2) @@ -374,12 +358,9 @@ def test_1frame(self): self.assertEqual(ff1.shape, (nframes,natoms,3)) self.assertEqual(vv1.shape, (nframes,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_1frame_atm(self): ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) @@ -398,16 +379,11 @@ def test_1frame_atm(self): self.assertEqual(ae1.shape, (nframes,natoms,1)) self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) @@ -429,13 +405,8 @@ def test_2frame_atm(self): self.assertEqual(av1.shape, (nframes,natoms,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) \ No newline at end of file + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) \ No newline at end of file diff --git a/source/tests/test_model_devi.py b/source/tests/test_model_devi.py index 4d5414cea8..6843874d57 100644 --- a/source/tests/test_model_devi.py +++ b/source/tests/test_model_devi.py @@ -38,9 +38,8 @@ def test_calc_model_devi(self): fname=self.output) self.assertAlmostEqual(model_devi[0][0], 0) self.assertAlmostEqual(model_devi[1][0], self.freq) - for ii in range(1, 7): - self.assertAlmostEqual(model_devi[0][ii], self.expect[ii]) - self.assertAlmostEqual(model_devi[0][ii], model_devi[1][ii]) + np.testing.assert_almost_equal(model_devi[0][1:7], self.expect[1:7], 6) + np.testing.assert_almost_equal(model_devi[0][1:7], model_devi[1][1:7], 6) self.assertTrue(os.path.isfile(self.output)) def tearDown(self): diff --git a/source/tests/test_model_loc_frame.py b/source/tests/test_model_loc_frame.py index 3ffce1bf41..5e0399fbb7 100644 --- a/source/tests/test_model_loc_frame.py +++ b/source/tests/test_model_loc_frame.py @@ -112,10 +112,7 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_a.py b/source/tests/test_model_se_a.py index e8b12c8553..0dce8f1b28 100644 --- a/source/tests/test_model_se_a.py +++ b/source/tests/test_model_se_a.py @@ -205,9 +205,6 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_a_aparam.py b/source/tests/test_model_se_a_aparam.py index 9c9847c242..f9a6817092 100644 --- a/source/tests/test_model_se_a_aparam.py +++ b/source/tests/test_model_se_a_aparam.py @@ -111,9 +111,6 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_a_fparam.py b/source/tests/test_model_se_a_fparam.py index 23b0da569b..469a5c429b 100644 --- a/source/tests/test_model_se_a_fparam.py +++ b/source/tests/test_model_se_a_fparam.py @@ -112,9 +112,6 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_a_srtab.py b/source/tests/test_model_se_a_srtab.py index 4b87fc0bb4..acfe62ce1e 100644 --- a/source/tests/test_model_se_a_srtab.py +++ b/source/tests/test_model_se_a_srtab.py @@ -134,9 +134,6 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_a_type.py b/source/tests/test_model_se_a_type.py index 11979108b4..6e80839d75 100644 --- a/source/tests/test_model_se_a_type.py +++ b/source/tests/test_model_se_a_type.py @@ -126,9 +126,6 @@ def test_model(self): refv = np.reshape(refv, [-1]) places = 10 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_model_se_t.py b/source/tests/test_model_se_t.py index 48fffd76b0..dead21c2d0 100644 --- a/source/tests/test_model_se_t.py +++ b/source/tests/test_model_se_t.py @@ -100,17 +100,14 @@ def test_model(self): np.savetxt('e.out', e.reshape([1, -1])) np.savetxt('f.out', f.reshape([1, -1]), delimiter = ',') np.savetxt('v.out', v.reshape([1, -1]), delimiter = ',') - refe = [4.826771866004193612e+01] - reff = [5.355088169393570574e+00,5.606772412401632266e+00,2.703270748296462966e-01,5.381408138049708967e+00,5.261355614357515975e+00,-4.079549918988090162e-01,-5.182324474551911919e+00,3.695481388907447262e-01,-5.238474288082559799e-02,1.665564584447352670e-01,-5.955401876564963892e+00,-2.217626865156164251e-01,-5.967343479332643419e+00,9.073821102416884665e-02,3.703103995504785639e-01,2.466151879965444438e-01,-5.373012500109097367e+00,4.146494691512622732e-02] - refv = [-1.336768232407933077e+01,4.818050125305787801e-01,3.589284283410607568e-01,4.818050125305786691e-01,-1.225345559839458964e+01,-1.701405121682751653e-01,3.589284283410607568e-01,-1.701405121682752486e-01,-3.428455515842296353e-02] + refe = [4.8436558582194039e+01] + reff = [5.2896335066946598e+00,5.5778402259211131e+00,2.6839994229557251e-01,5.3528786387686784e+00,5.2477755362164968e+00,-4.0486366542657343e-01,-5.1297084055340498e+00,3.4607112287117253e-01,-5.1800783428369482e-02,1.5557068351407846e-01,-5.9071343228741506e+00,-2.2012359669589748e-01,-5.9156735320857488e+00,8.8397615509389127e-02,3.6701215949753935e-01,2.4729910864238122e-01,-5.3529501776440211e+00,4.1375943757728552e-02] + refv = [-1.3159448660141607e+01,4.6952048725161544e-01,3.5482003698976106e-01,4.6952048725161577e-01,-1.2178990983673918e+01,-1.6867277410496895e-01,3.5482003698976106e-01,-1.6867277410496900e-01,-3.3986741457321945e-02] refe = np.reshape(refe, [-1]) reff = np.reshape(reff, [-1]) refv = np.reshape(refv, [-1]) places = 6 - for ii in range(e.size) : - self.assertAlmostEqual(e[ii], refe[ii], places = places) - for ii in range(f.size) : - self.assertAlmostEqual(f[ii], reff[ii], places = places) - for ii in range(v.size) : - self.assertAlmostEqual(v[ii], refv[ii], places = places) + np.testing.assert_almost_equal(e, refe, places) + np.testing.assert_almost_equal(f, reff, places) + np.testing.assert_almost_equal(v, refv, places) diff --git a/source/tests/test_parallel_training.py b/source/tests/test_parallel_training.py index e2310e8267..320d8ca22c 100644 --- a/source/tests/test_parallel_training.py +++ b/source/tests/test_parallel_training.py @@ -18,9 +18,12 @@ def setUp(self): def test_two_workers(self): command = 'horovodrun -np 2 dp train -m workers ' + self.input_file penv = os.environ.copy() - if len(get_gpus() or []) > 1: + num_gpus = len(get_gpus() or []) + if num_gpus > 1: penv['CUDA_VISIBLE_DEVICES'] = '0,1' - popen = sp.Popen(command, shell=True, env=penv, stdout=sp.PIPE, stderr=sp.STDOUT) + elif num_gpus == 1: + raise unittest.SkipTest("At least 2 GPU cards are needed for parallel-training tests.") + popen = sp.Popen(command, shell=True, cwd=str(tests_path), env=penv, stdout=sp.PIPE, stderr=sp.STDOUT) for line in iter(popen.stdout.readline, b''): if hasattr(line, 'decode'): line = line.decode('utf-8') diff --git a/source/tests/test_polar_se_a.py b/source/tests/test_polar_se_a.py index 78da154864..ee298bfb8b 100644 --- a/source/tests/test_polar_se_a.py +++ b/source/tests/test_polar_se_a.py @@ -97,15 +97,13 @@ def test_model(self): refp = [3.39695248e+01, 2.16564043e+01, 8.18501479e-01, 2.16564043e+01, 1.38211789e+01, 5.22775159e-01, 8.18501479e-01, 5.22775159e-01, 1.97847218e-02, 8.08467431e-01, 3.42081126e+00, -2.01072261e-01, 3.42081126e+00, 1.54924596e+01, -9.06153697e-01, -2.01072261e-01, -9.06153697e-01, 5.30193262e-02] places = 6 - for ii in range(p.size) : - self.assertAlmostEqual(p[ii], refp[ii], places = places) + np.testing.assert_almost_equal(p, refp, places) gp = gp.reshape([-1]) refgp = np.array(refp).reshape(-1, 9).sum(0) places = 5 - for ii in range(gp.size) : - self.assertAlmostEqual(gp[ii], refgp[ii], places = places) + np.testing.assert_almost_equal(gp, refgp, places) # make sure only one frame is used feed_dict_single = {t_prop_c: test_data['prop_c'], @@ -139,11 +137,8 @@ def test_model(self): @ box0.reshape(3,3)).reshape(-1) delta = 1e-4 - for ii in range(pf.size) : - self.assertAlmostEqual(pf[ii], fdf[ii], delta = delta) - for ii in range(pv.size) : - self.assertAlmostEqual(pv[ii], fdv[ii], delta = delta) + np.testing.assert_allclose(pf, fdf, delta) + np.testing.assert_allclose(pv, fdv, delta) # make sure atomic virial sum to virial places = 10 - for ii in range(pv.size) : - self.assertAlmostEqual(pv[ii], spv[ii], places = places) + np.testing.assert_almost_equal(pv, spv, places) diff --git a/source/tests/test_prod_env_mat.py b/source/tests/test_prod_env_mat.py index 4b60da39d9..82fe3e18eb 100644 --- a/source/tests/test_prod_env_mat.py +++ b/source/tests/test_prod_env_mat.py @@ -91,8 +91,7 @@ def test_pbc_self_built_nlist(self): self.assertEqual(drij.shape, (self.nframes, self.nloc*self.nnei*3)) self.assertEqual(dnlist.shape, (self.nframes, self.nloc*self.nnei)) for ff in range(self.nframes): - for ii in range(self.nloc*self.ndescrpt): - self.assertAlmostEqual(dem[ff][ii], self.pbc_expected_output[ii], places=5) + np.testing.assert_almost_equal(dem[ff], self.pbc_expected_output, 5) def test_pbc_self_built_nlist_deriv(self): hh = 1e-4 @@ -142,8 +141,7 @@ def test_nopbc_self_built_nlist(self): self.assertEqual(drij.shape, (self.nframes, self.nloc*self.nnei*3)) self.assertEqual(dnlist.shape, (self.nframes, self.nloc*self.nnei)) for ff in range(self.nframes): - for ii in range(self.nloc*self.ndescrpt): - self.assertAlmostEqual(dem[ff][ii], self.nopbc_expected_output[ii], places=5) + np.testing.assert_almost_equal(dem[ff], self.nopbc_expected_output, 5) def test_nopbc_self_built_nlist_deriv(self): diff --git a/source/tests/test_prod_force.py b/source/tests/test_prod_force.py index 4135ee0e2e..c0aaa29ea2 100644 --- a/source/tests/test_prod_force.py +++ b/source/tests/test_prod_force.py @@ -103,5 +103,4 @@ def test_prod_force(self): ) self.assertEqual(dforce.shape, (self.nframes, self.nall*3)) for ff in range(self.nframes): - for ii in range(self.nall*3): - self.assertAlmostEqual(dforce[ff][ii], self.expected_force[ii], places=5) + np.testing.assert_almost_equal(dforce[ff], self.expected_force, 5) diff --git a/source/tests/test_prod_force_grad.py b/source/tests/test_prod_force_grad.py index 8b8ad74ed6..71a95380f0 100644 --- a/source/tests/test_prod_force_grad.py +++ b/source/tests/test_prod_force_grad.py @@ -105,5 +105,4 @@ def test_prod_force_grad(self): ) self.assertEqual(dgrad_net.shape, (self.nframes, self.nloc*self.ndescrpt)) for ff in range(self.nframes): - for ii in range(self.nloc*self.ndescrpt): - self.assertAlmostEqual(dgrad_net[ff][ii], self.expected_grad_net[ii], places=5) + np.testing.assert_almost_equal(dgrad_net[ff], self.expected_grad_net, 5) diff --git a/source/tests/test_prod_virial.py b/source/tests/test_prod_virial.py index 2a11f0acfd..691a8a652f 100644 --- a/source/tests/test_prod_virial.py +++ b/source/tests/test_prod_virial.py @@ -106,7 +106,5 @@ def test_prod_virial(self): self.assertEqual(dvirial.shape, (self.nframes, 9)) self.assertEqual(datom_virial.shape, (self.nframes, self.nall*9)) for ff in range(self.nframes): - for ii in range(9): - self.assertAlmostEqual(dvirial[ff][ii], self.expected_virial[ii], places=5) - for ii in range(self.nall*9): - self.assertAlmostEqual(datom_virial[ff][ii], self.expected_atom_virial[ii], places=5) + np.testing.assert_almost_equal(dvirial[ff], self.expected_virial, 5) + np.testing.assert_almost_equal(datom_virial[ff], self.expected_atom_virial, 5) diff --git a/source/tests/test_prod_virial_grad.py b/source/tests/test_prod_virial_grad.py index 2ca9071a56..6217fa8600 100644 --- a/source/tests/test_prod_virial_grad.py +++ b/source/tests/test_prod_virial_grad.py @@ -113,5 +113,4 @@ def test_prod_virial_grad(self): ) self.assertEqual(dgrad_net.shape, (self.nframes, self.nloc*self.ndescrpt)) for ff in range(self.nframes): - for ii in range(self.nloc*self.ndescrpt): - self.assertAlmostEqual(dgrad_net[ff][ii], self.expected_grad_net[ii], places=5) + np.testing.assert_almost_equal(dgrad_net[ff], self.expected_grad_net, 5) diff --git a/source/tests/test_tabulate.py b/source/tests/test_tabulate.py index ce26c4e3e6..341850e68c 100644 --- a/source/tests/test_tabulate.py +++ b/source/tests/test_tabulate.py @@ -22,9 +22,7 @@ def test_op_tanh(self): [5.492686739421748753e-03, 5.754985286040992763e-03, 4.493113544969218158e-03, 3.107638130764600777e-03]]) places = 18 - for ii in range(dy_array.shape[0]): - for jj in range(dy_array.shape[1]): - self.assertAlmostEqual(dy_array[ii,jj], answer[ii,jj], places=places) + np.testing.assert_almost_equal(dy_array, answer, places) def test_op_gelu(self): w = tf.constant([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], [ @@ -42,9 +40,7 @@ def test_op_gelu(self): [1.072173273655498138e-01, 2.082159073100979807e-01, 3.059816075270163083e-01, 4.032981557798429595e-01]]) places = 18 - for ii in range(dy_array.shape[0]): - for jj in range(dy_array.shape[1]): - self.assertAlmostEqual(dy_array[ii, jj], answer[ii, jj], places=places) + np.testing.assert_almost_equal(dy_array, answer, places) diff --git a/source/tests/test_transfer.py b/source/tests/test_transfer.py index 317afada63..7b5884b2b4 100644 --- a/source/tests/test_transfer.py +++ b/source/tests/test_transfer.py @@ -1,6 +1,7 @@ import os,sys,platform,shutil,dpdata,json import numpy as np import unittest +import subprocess as sp from deepmd.env import tf from deepmd.infer import DeepPot @@ -18,48 +19,28 @@ def _file_delete(file) : if os.path.exists(file): os.remove(file) -class TestTransform(unittest.TestCase) : - def setUp(self): - self.data_file = str(tests_path / os.path.join("model_compression", "data")) - self.original_model = str(tests_path / "dp-original.pb") - self.frozen_model = str(tests_path / "dp-frozen.pb") - self.transferred_model = str(tests_path / "dp-transferred.pb") - self.INPUT = str(tests_path / "input.json") - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - jdata["model"]["descriptor"]["seed"] = 1 - jdata["model"]["fitting_net"]["seed"] = 1 - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - - # generate the original input model - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.original_model) - assert(ret == 0), "DP freeze error!" +def _subprocess_run(command): + popen = sp.Popen(command.split(), shell=False, stdout=sp.PIPE, stderr=sp.STDOUT) + for line in iter(popen.stdout.readline, b''): + if hasattr(line, 'decode'): + line = line.decode('utf-8') + line = line.rstrip() + print(line) + popen.wait() + return popen.returncode - # generate the frozen raw model - jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) - jdata["training"]["training_data"]["systems"] = self.data_file - jdata["training"]["validation_data"]["systems"] = self.data_file - jdata["model"]["descriptor"]["seed"] = 2 - jdata["model"]["fitting_net"]["seed"] = 2 - _file_delete(self.INPUT) - with open(self.INPUT, "w") as fp: - json.dump(jdata, fp, indent=4) - ret = os.system("dp train " + self.INPUT) - assert(ret == 0), "DP train error!" - ret = os.system("dp freeze -o " + self.frozen_model) - assert(ret == 0), "DP freeze error!" - - # generate the transferred output model - ret = os.system("dp transfer -O " + self.original_model + " -r " + self.frozen_model + " -o " + self.transferred_model) - assert(ret == 0), "DP transfer error!" - - self.dp_original = DeepPot(self.original_model) - self.dp_transferred = DeepPot(self.transferred_model) +class TestTransform(unittest.TestCase) : + @classmethod + def setUpClass(self): + self.old_model = str(tests_path / "dp-old.pb") + self.raw_model = str(tests_path / "dp-raw.pb") + self.new_model = str(tests_path / "dp-new.pb") + convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot.pbtxt")), self.old_model) + convert_pbtxt_to_pb(str(tests_path / os.path.join("infer","deeppot-1.pbtxt")), self.raw_model) + ret = _subprocess_run("dp transfer -O " + self.old_model + " -r " + self.raw_model + " -o " + self.new_model) + np.testing.assert_equal(ret, 0, 'DP transfer failed!') + self.dp = DeepPot(self.new_model) self.coords = np.array([12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 00.25, 3.32, 1.68, @@ -71,47 +52,76 @@ def setUp(self): self.expected_e = np.array([-9.275780747115504710e+01,-1.863501786584258468e+02,-1.863392472863538103e+02,-9.279281325486221021e+01,-1.863671545232153903e+02,-1.863619822847602165e+02]) self.expected_f = np.array([-3.034045420701179663e-01,8.405844663871177014e-01,7.696947487118485642e-02,7.662001266663505117e-01,-1.880601391333554251e-01,-6.183333871091722944e-01,-5.036172391059643427e-01,-6.529525836149027151e-01,5.432962643022043459e-01,6.382357912332115024e-01,-1.748518296794561167e-01,3.457363524891907125e-01,1.286482986991941552e-03,3.757251165286925043e-01,-5.972588700887541124e-01,-5.987006197104716154e-01,-2.004450304880958100e-01,2.495901655353461868e-01]) self.expected_v = np.array([-2.912234126853306959e-01,-3.800610846612756388e-02,2.776624987489437202e-01,-5.053761003913598976e-02,-3.152373041953385746e-01,1.060894290092162379e-01,2.826389131596073745e-01,1.039129970665329250e-01,-2.584378792325942586e-01,-3.121722367954994914e-01,8.483275876786681990e-02,2.524662342344257682e-01,4.142176771106586414e-02,-3.820285230785245428e-02,-2.727311173065460545e-02,2.668859789777112135e-01,-6.448243569420382404e-02,-2.121731470426218846e-01,-8.624335220278558922e-02,-1.809695356746038597e-01,1.529875294531883312e-01,-1.283658185172031341e-01,-1.992682279795223999e-01,1.409924999632362341e-01,1.398322735274434292e-01,1.804318474574856390e-01,-1.470309318999652726e-01,-2.593983661598450730e-01,-4.236536279233147489e-02,3.386387920184946720e-02,-4.174017537818433543e-02,-1.003500282164128260e-01,1.525690815194478966e-01,3.398976109910181037e-02,1.522253908435125536e-01,-2.349125581341701963e-01,9.515545977581392825e-04,-1.643218849228543846e-02,1.993234765412972564e-02,6.027265332209678569e-04,-9.563256398907417355e-02,1.510815124001868293e-01,-7.738094816888557714e-03,1.502832772532304295e-01,-2.380965783745832010e-01,-2.309456719810296654e-01,-6.666961081213038098e-02,7.955566551234216632e-02,-8.099093777937517447e-02,-3.386641099800401927e-02,4.447884755740908608e-02,1.008593228579038742e-01,4.556718179228393811e-02,-6.078081273849572641e-02]) - - def tearDown(self): - _file_delete(self.INPUT) - _file_delete(self.original_model) - _file_delete(self.frozen_model) - _file_delete(self.transferred_model) - _file_delete("out.json") - _file_delete("compress.json") - _file_delete("checkpoint") - _file_delete("lcurve.out") - _file_delete("model.ckpt.meta") - _file_delete("model.ckpt.index") - _file_delete("model.ckpt.data-00000-of-00001") - def test(self): - ee0, ff0, vv0, ae0, av0 = self.dp_original.eval(self.coords, self.box, self.atype, atomic = True) - ee1, ff1, vv1, ae1, av1 = self.dp_transferred.eval(self.coords, self.box, self.atype, atomic = True) + @classmethod + def tearDownClass(self): + _file_delete(self.old_model) + _file_delete(self.raw_model) + _file_delete(self.new_model) + + def test_attrs(self): + self.assertEqual(self.dp.get_ntypes(), 2) + self.assertAlmostEqual(self.dp.get_rcut(), 6.0, places = default_places) + self.assertEqual(self.dp.get_type_map(), ['O', 'H']) + self.assertEqual(self.dp.get_dim_fparam(), 0) + self.assertEqual(self.dp.get_dim_aparam(), 0) + + def test_1frame(self): + ee, ff, vv = self.dp.eval(self.coords, self.box, self.atype, atomic = False) # check shape of the returns nframes = 1 natoms = len(self.atype) - self.assertEqual(ee0.shape, (nframes,1)) - self.assertEqual(ff0.shape, (nframes,natoms,3)) - self.assertEqual(vv0.shape, (nframes,9)) - self.assertEqual(ae0.shape, (nframes,natoms,1)) - self.assertEqual(av0.shape, (nframes,natoms,9)) - self.assertEqual(ee1.shape, (nframes,1)) - self.assertEqual(ff1.shape, (nframes,natoms,3)) - self.assertEqual(vv1.shape, (nframes,9)) - self.assertEqual(ae1.shape, (nframes,natoms,1)) - self.assertEqual(av1.shape, (nframes,natoms,9)) + self.assertEqual(ee.shape, (nframes,1)) + self.assertEqual(ff.shape, (nframes,natoms,3)) + self.assertEqual(vv.shape, (nframes,9)) # check values - for ii in range(ff0.size): - self.assertAlmostEqual(ff0.reshape([-1])[ii], ff1.reshape([-1])[ii], places = default_places) - for ii in range(ae0.size): - self.assertAlmostEqual(ae0.reshape([-1])[ii], ae1.reshape([-1])[ii], places = default_places) - for ii in range(av0.size): - self.assertAlmostEqual(av0.reshape([-1])[ii], av1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) - for ii in range(nframes): - self.assertAlmostEqual(ee0.reshape([-1])[ii], ee1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) - for ii in range(nframes, 9): - self.assertAlmostEqual(vv0.reshape([-1])[ii], vv1.reshape([-1])[ii], places = default_places) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + + def test_1frame_atm(self): + ee, ff, vv, ae, av = self.dp.eval(self.coords, self.box, self.atype, atomic = True) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee.shape, (nframes,1)) + self.assertEqual(ff.shape, (nframes,natoms,3)) + self.assertEqual(vv.shape, (nframes,9)) + self.assertEqual(ae.shape, (nframes,natoms,1)) + self.assertEqual(av.shape, (nframes,natoms,9)) + # check values + np.testing.assert_almost_equal(ff.ravel(), self.expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), self.expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), self.expected_v.ravel(), default_places) + expected_se = np.sum(self.expected_e.reshape([nframes, -1]), axis = 1) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + expected_sv = np.sum(self.expected_v.reshape([nframes, -1, 9]), axis = 1) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) + + + def test_2frame_atm(self): + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee, ff, vv, ae, av = self.dp.eval(coords2, box2, self.atype, atomic = True) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee.shape, (nframes,1)) + self.assertEqual(ff.shape, (nframes,natoms,3)) + self.assertEqual(vv.shape, (nframes,9)) + self.assertEqual(ae.shape, (nframes,natoms,1)) + self.assertEqual(av.shape, (nframes,natoms,9)) + # check values + expected_f = np.concatenate((self.expected_f, self.expected_f), axis = 0) + expected_e = np.concatenate((self.expected_e, self.expected_e), axis = 0) + expected_v = np.concatenate((self.expected_v, self.expected_v), axis = 0) + np.testing.assert_almost_equal(ff.ravel(), expected_f.ravel(), default_places) + np.testing.assert_almost_equal(ae.ravel(), expected_e.ravel(), default_places) + np.testing.assert_almost_equal(av.ravel(), expected_v.ravel(), default_places) + expected_se = np.sum(expected_e.reshape([nframes, -1]), axis = 1) + np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) + expected_sv = np.sum(expected_v.reshape([nframes, -1, 9]), axis = 1) + np.testing.assert_almost_equal(vv.ravel(), expected_sv.ravel(), default_places) diff --git a/source/tests/test_type_embed.py b/source/tests/test_type_embed.py index 42d1c86063..57b153c24d 100644 --- a/source/tests/test_type_embed.py +++ b/source/tests/test_type_embed.py @@ -21,10 +21,8 @@ def test_embed_atom_type(self): atom_embed = embed_atom_type(ntypes, natoms, type_embedding) sess = self.test_session().__enter__() atom_embed = sess.run(atom_embed) - for ii in range(5): - for jj in range(3): - self.assertAlmostEqual( - atom_embed[ii][jj], expected_out[ii][jj], places=10) + np.testing.assert_almost_equal( + atom_embed, expected_out, 10) def test_type_embed_net(self): ten = TypeEmbedNet([2, 4, 8], seed = 1, uniform_seed = True) @@ -44,7 +42,5 @@ def test_type_embed_net(self): # size of embedded vec 8 self.assertEqual(type_embedding.shape[1], 8) # check value - for ii in range(2): - for jj in range(8): - self.assertAlmostEqual( - type_embedding[ii][jj], expected_out[ii][jj], places=10) + np.testing.assert_almost_equal( + type_embedding, expected_out, 10) diff --git a/source/tests/test_wfc.py b/source/tests/test_wfc.py index 97af7486bd..80f9d48126 100644 --- a/source/tests/test_wfc.py +++ b/source/tests/test_wfc.py @@ -91,8 +91,7 @@ def test_model(self): refp = [-9.105016838228578990e-01,7.196284362034099935e-01,-9.548516928185298014e-02,2.764615027095288724e+00,2.661319598995644520e-01,7.579512949131941846e-02,-2.107409067376114997e+00,-1.299080016614967414e-01,-5.962778584850070285e-01,2.913899917663253514e-01,-1.226917174638697094e+00,1.829523069930876655e+00,1.015704024959750873e+00,-1.792333611099589386e-01,5.032898080485321834e-01,1.808561721292949453e-01,2.468863482075112081e+00,-2.566442546384765100e-01,-1.467453783795173994e-01,-1.822963931552128658e+00,5.843600156865462747e-01,-1.493875280832117403e+00,1.693322352814763398e-01,-1.877325443995481624e+00] places = 6 - for ii in range(p.size) : - self.assertAlmostEqual(p[ii], refp[ii], places = places) + np.testing.assert_almost_equal(p, refp, places)