diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py index ca4b163734..4057de726d 100644 --- a/deepmd/dpmodel/descriptor/__init__.py +++ b/deepmd/dpmodel/descriptor/__init__.py @@ -18,10 +18,14 @@ from .se_r import ( DescrptSeR, ) +from .se_t import ( + DescrptSeT, +) __all__ = [ "DescrptSeA", "DescrptSeR", + "DescrptSeT", "DescrptDPA1", "DescrptSeAttenV2", "DescrptDPA2", diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 144071aa95..319fb2262a 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -14,6 +14,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -524,18 +527,36 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist @DescriptorBlock.register("se_atten") diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 78bf174685..9243883ebe 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -15,6 +15,9 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -743,30 +746,46 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy["repinit"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repinit"], + min_nbor_dist, repinit_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repinit"]["rcut"], + local_jdata_cpy["repinit"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - local_jdata_cpy["repformer"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repformer"], + local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0] + min_nbor_dist, repformer_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repformer"]["rcut"], + local_jdata_cpy["repformer"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - return local_jdata_cpy + local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 15825ecc10..d359bf6fb7 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -19,6 +20,9 @@ from deepmd.dpmodel.utils.nlist import ( nlist_distinguish_types, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -222,22 +226,42 @@ def call( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - BaseDescriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self) -> dict: return { diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index cba9eebe4b..a4fc8bddf9 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -7,12 +7,16 @@ Callable, List, Optional, + Tuple, Union, ) from deepmd.common import ( j_get_type, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -149,19 +153,33 @@ def deserialize(cls, data: dict) -> "BD": @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # call subprocess cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) setattr(BD, fwd_method_name, BD.fwd) delattr(BD, "fwd") diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 193383ac4f..2b6c3a843e 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -9,6 +9,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -422,15 +425,32 @@ def deserialize(cls, data: dict) -> "DescrptSeA": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 5175b91ae1..e4a0a80657 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -4,6 +4,9 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -21,6 +24,7 @@ Any, List, Optional, + Tuple, ) from deepmd.dpmodel import ( @@ -345,15 +349,32 @@ def deserialize(cls, data: dict) -> "DescrptSeR": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py new file mode 100644 index 0000000000..eac6a9640e --- /dev/null +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools + +import numpy as np + +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +try: + from deepmd._version import version as __version__ +except ImportError: + __version__ = "unknown" + +import copy +from typing import ( + List, + Optional, + Tuple, +) + +from deepmd.dpmodel import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils import ( + EmbeddingNet, + EnvMat, + NetworkCollection, + PairExcludeMask, +) + +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("se_e3") +@BaseDescriptor.register("se_at") +@BaseDescriptor.register("se_a_3be") +class DescrptSeT(NativeOP, BaseDescriptor): + r"""DeepPot-SE constructed from all information (both angular and radial) of atomic + configurations. + + The embedding takes angles between two neighboring atoms as input. + + Parameters + ---------- + rcut : float + The cut-off radius + rcut_smth : float + From where the environment matrix should be smoothed + sel : list[int] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + set_davg_zero : bool + Set the shift of embedding net input to zero. + activation_function : str + The activation function in the embedding net. Supported options are |ACTIVATION_FN| + env_protection : float + Protection parameter to prevent division by zero errors during environment matrix calculations. + exclude_types : List[List[int]] + The excluded pairs of types which have no interaction with each other. + For example, `[[0, 1]]` means no interaction between type 0 and type 1. + precision : str + The precision of the embedding net parameters. Supported options are |PRECISION| + trainable : bool + If the weights of embedding net are trainable. + seed : int, Optional + Random seed for initializing the network parameters. + """ + + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[int], + neuron: List[int] = [24, 48, 96], + resnet_dt: bool = False, + set_davg_zero: bool = False, + activation_function: str = "tanh", + env_protection: float = 0.0, + exclude_types: List[Tuple[int, int]] = [], + precision: str = DEFAULT_PRECISION, + trainable: bool = True, + seed: Optional[int] = None, + ) -> None: + self.rcut = rcut + self.rcut_smth = rcut_smth + self.sel = sel + self.neuron = neuron + self.filter_neuron = self.neuron + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.resnet_dt = resnet_dt + self.env_protection = env_protection + self.ntypes = len(sel) + self.seed = seed + # order matters, placed after the assignment of self.ntypes + self.reinit_exclude(exclude_types) + self.trainable = trainable + + in_dim = 1 # not considiering type embedding + self.embeddings = NetworkCollection( + ntypes=self.ntypes, + ndim=2, + network_type="embedding_network", + ) + for embedding_idx in itertools.product( + range(self.ntypes), repeat=self.embeddings.ndim + ): + self.embeddings[embedding_idx] = EmbeddingNet( + in_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + ) + self.env_mat = EnvMat(self.rcut, self.rcut_smth, protection=self.env_protection) + self.nnei = np.sum(self.sel) + self.davg = np.zeros( + [self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision] + ) + self.dstd = np.ones( + [self.ntypes, self.nnei, 4], dtype=PRECISION_DICT[self.precision] + ) + self.orig_sel = self.sel + + def __setitem__(self, key, value): + if key in ("avg", "data_avg", "davg"): + self.davg = value + elif key in ("std", "data_std", "dstd"): + self.dstd = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.davg + elif key in ("std", "data_std", "dstd"): + return self.dstd + else: + raise KeyError(key) + + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.get_dim_out() + + def get_dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.neuron[-1] + + def get_dim_emb(self): + """Returns the embedding (g2) dimension of this descriptor.""" + return self.neuron[-1] + + def get_rcut(self): + """Returns cutoff radius.""" + return self.rcut + + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.rcut_smth + + def get_sel(self): + """Returns cutoff radius.""" + return self.sel + + def mixed_types(self): + """Returns if the descriptor requires a neighbor list that distinguish different + atomic types or not. + """ + return False + + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.env_protection + + def share_params(self, base_class, shared_level, resume=False): + """ + Share the parameters of self to the base_class with shared_level during multitask training. + If not start from checkpoint (resume is False), + some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes. + """ + raise NotImplementedError + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.ntypes + + def compute_input_stats(self, merged: List[dict], path: Optional[DPPath] = None): + """Update mean and stddev for descriptor elements.""" + raise NotImplementedError + + def reinit_exclude( + self, + exclude_types: List[Tuple[int, int]] = [], + ): + self.exclude_types = exclude_types + self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + + def call( + self, + coord_ext, + atype_ext, + nlist, + mapping: Optional[np.ndarray] = None, + ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping from extended to lcoal region. not used by this descriptor. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x ng + gr + The rotationally equivariant and permutationally invariant single particle + representation. + This descriptor returns None. + g2 + The rotationally invariant pair-partical representation. + This descriptor returns None. + h2 + The rotationally equivariant pair-partical representation. + This descriptor returns None. + sw + The smooth switch function. + """ + del mapping + # nf x nloc x nnei x 4 + rr, diff, ww = self.env_mat.call( + coord_ext, atype_ext, nlist, self.davg, self.dstd + ) + nf, nloc, nnei, _ = rr.shape + sec = np.append([0], np.cumsum(self.sel)) + + ng = self.neuron[-1] + result = np.zeros([nf * nloc, ng], dtype=PRECISION_DICT[self.precision]) + exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) + # merge nf and nloc axis, so for type_one_side == False, + # we don't require atype is the same in all frames + exclude_mask = exclude_mask.reshape(nf * nloc, nnei) + rr = rr.reshape(nf * nloc, nnei, 4) + + for embedding_idx in itertools.product( + range(self.ntypes), repeat=self.embeddings.ndim + ): + ti, tj = embedding_idx + nei_type_i = self.sel[ti] + nei_type_j = self.sel[tj] + if ti <= tj: + # avoid repeat calculation + # nfnl x nt_i x 3 + rr_i = rr[:, sec[ti] : sec[ti + 1], 1:] + mm_i = exclude_mask[:, sec[ti] : sec[ti + 1]] + rr_i = rr_i * mm_i[:, :, None] + # nfnl x nt_j x 3 + rr_j = rr[:, sec[tj] : sec[tj + 1], 1:] + mm_j = exclude_mask[:, sec[tj] : sec[tj + 1]] + rr_j = rr_j * mm_j[:, :, None] + # nfnl x nt_i x nt_j + env_ij = np.einsum("ijm,ikm->ijk", rr_i, rr_j) + # nfnl x nt_i x nt_j x 1 + env_ij_reshape = env_ij[:, :, :, None] + # nfnl x nt_i x nt_j x ng + gg = self.embeddings[embedding_idx].call(env_ij_reshape) + # nfnl x nt_i x nt_j x ng + res_ij = np.einsum("ijk,ijkm->im", env_ij, gg) + res_ij = res_ij * (1.0 / float(nei_type_i) / float(nei_type_j)) + result += res_ij + # nf x nloc x ng + result = result.reshape(nf, nloc, ng).astype(GLOBAL_NP_FLOAT_PRECISION) + return result, None, None, None, ww + + def serialize(self) -> dict: + """Serialize the descriptor to dict.""" + for embedding_idx in itertools.product(range(self.ntypes), repeat=2): + # not actually used; to match serilization data from TF to pass the test + ti, tj = embedding_idx + if (self.exclude_types and embedding_idx in self.emask) or tj < ti: + self.embeddings[embedding_idx].clear() + + return { + "@class": "Descriptor", + "type": "se_e3", + "@version": 1, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + "neuron": self.neuron, + "resnet_dt": self.resnet_dt, + "set_davg_zero": self.set_davg_zero, + "activation_function": self.activation_function, + "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "embeddings": self.embeddings.serialize(), + "env_mat": self.env_mat.serialize(), + "exclude_types": self.exclude_types, + "env_protection": self.env_protection, + "@variables": { + "davg": self.davg, + "dstd": self.dstd, + }, + "trainable": self.trainable, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeT": + """Deserialize from dict.""" + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + data.pop("type", None) + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + env_mat = data.pop("env_mat") + obj = cls(**data) + + obj["davg"] = variables["davg"] + obj["dstd"] = variables["dstd"] + obj.embeddings = NetworkCollection.deserialize(embeddings) + return obj + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms + local_jdata : dict + The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms + """ + local_jdata_cpy = local_jdata.copy() + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index b416e57fb8..354fa58348 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -7,9 +7,14 @@ from typing import ( Any, List, + Optional, + Tuple, Type, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.plugin import ( PluginVariant, make_plugin_registry, @@ -129,12 +134,21 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": raise NotImplementedError(f"Not implemented in class {cls.__name__}") model_def_script: str + """The model definition script.""" + min_nbor_dist: Optional[float] + """The minimum distance between two atoms. Used for model compression. + None when skipping neighbor statistics. + """ @abstractmethod def get_model_def_script(self) -> str: """Get the model definition script.""" pass + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum distance between two atoms.""" + return self.min_nbor_dist + @abstractmethod def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" @@ -148,22 +162,36 @@ def get_nsel(self) -> int: @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # getting model type based on fitting type model_type = local_jdata.get("type", "standard") if model_type == "standard": model_type = local_jdata.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) return BaseBaseModel @@ -186,6 +214,7 @@ class BaseModel(make_base_model()): def __init__(self) -> None: self.model_def_script = "" + self.min_nbor_dist = None def get_model_def_script(self) -> str: """Get the model definition script.""" diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 37cb426ab7..1597ba0b14 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,26 +1,49 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, + Optional, + Tuple, +) + from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) # use "class" to resolve "Variable not allowed in type expression" class DPModelCommon: @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 90e2bb3fb4..c2cea35d27 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -227,6 +227,10 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum neighbor distance.""" + return self.backbone_model.get_min_nbor_dist() + def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" # for C++ interface diff --git a/deepmd/dpmodel/utils/update_sel.py b/deepmd/dpmodel/utils/update_sel.py index 48463b5743..dc38a6a041 100644 --- a/deepmd/dpmodel/utils/update_sel.py +++ b/deepmd/dpmodel/utils/update_sel.py @@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # TODO: save to the model in UpdateSel.hook - pass diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 5a00ba616d..879455b942 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -76,6 +76,9 @@ class DeepEvalBackend(ABC): "dos_redu": "dos", "mask_mag": "mask_mag", "mask": "mask", + # old models in v1 + "global_polar": "global_polar", + "wfc": "wfc", } @abstractmethod diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index c2089b278d..6650c349a2 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -7,8 +7,14 @@ import numpy as np +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) from deepmd.infer.deep_tensor import ( DeepTensor, + OldDeepTensor, ) @@ -36,7 +42,7 @@ def output_tensor_name(self) -> str: return "polar" -class DeepGlobalPolar(DeepTensor): +class DeepGlobalPolar(OldDeepTensor): @property def output_tensor_name(self) -> str: return "global_polar" @@ -95,3 +101,22 @@ def eval( mixed_type=mixed_type, **kwargs, ) + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + # no atomic or differentiable output is defined + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + self.output_tensor_name, + shape=[-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + atomic=False, + ), + ] + ) + ) diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 14e13e7f84..106bc3156c 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -234,3 +234,24 @@ def output_def(self) -> ModelOutputDef: ] ) ) + + +class OldDeepTensor(DeepTensor): + """Old tensor models from v1, which has no gradient output.""" + + # See https://github.com/deepmodeling/deepmd-kit/blob/1d1b251a2c5f05d1401aa89be792f9ed18b8f096/source/train/Model.py#L264 + def eval_full( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: np.ndarray, + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: dict, + ) -> Tuple[np.ndarray, ...]: + """Unsupported method.""" + raise RuntimeError( + "This model does not support eval_full method. Use eval instead." + ) diff --git a/deepmd/infer/deep_wfc.py b/deepmd/infer/deep_wfc.py index deed938e04..d92af28f5a 100644 --- a/deepmd/infer/deep_wfc.py +++ b/deepmd/infer/deep_wfc.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) from deepmd.infer.deep_tensor import ( - DeepTensor, + OldDeepTensor, ) -class DeepWFC(DeepTensor): +class DeepWFC(OldDeepTensor): """Deep WFC model. Parameters @@ -26,3 +31,22 @@ class DeepWFC(DeepTensor): @property def output_tensor_name(self) -> str: return "wfc" + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + # no reduciable or differentiable output is defined + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + self.output_tensor_name, + shape=[-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + atomic=True, + ), + ] + ) + ) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index fba22e6d24..991278f0aa 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -3,9 +3,6 @@ import json import logging import os -from copy import ( - deepcopy, -) from pathlib import ( Path, ) @@ -66,6 +63,7 @@ update_deepmd_input, ) from deepmd.utils.data_system import ( + get_data, process_systems, ) from deepmd.utils.path import ( @@ -251,24 +249,33 @@ def train(FLAGS): config = normalize(config) # do neighbor stat + min_nbor_dist = None if not FLAGS.skip_neighbor_stat: log.info( "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) + + type_map = config["model"].get("type_map") if not multi_task: - config["model"] = BaseModel.update_sel(config, config["model"]) + train_data = get_data( + config["training"]["training_data"], 0, type_map, None + ) + config["model"], min_nbor_dist = BaseModel.update_sel( + train_data, type_map, config["model"] + ) else: - training_jdata = deepcopy(config["training"]) - training_jdata.pop("data_dict", {}) - training_jdata.pop("model_prob", {}) + min_nbor_dist = {} for model_item in config["model"]["model_dict"]: - fake_global_jdata = { - "model": deepcopy(config["model"]["model_dict"][model_item]), - "training": deepcopy(config["training"]["data_dict"][model_item]), - } - fake_global_jdata["training"].update(training_jdata) - config["model"]["model_dict"][model_item] = BaseModel.update_sel( - fake_global_jdata, config["model"]["model_dict"][model_item] + train_data = get_data( + config["training"]["data_dict"][model_item]["training_data"], + 0, + type_map, + None, + ) + config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( + BaseModel.update_sel( + train_data, type_map, config["model"]["model_dict"][model_item] + ) ) with open(FLAGS.output, "w") as fp: @@ -284,6 +291,13 @@ def train(FLAGS): FLAGS.init_frz_model, shared_links=shared_links, ) + # save min_nbor_dist + if min_nbor_dist is not None: + if not multi_task: + trainer.model.min_nbor_dist = min_nbor_dist + else: + for model_item in min_nbor_dist: + trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item] trainer.run() diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py index 434bba871b..b42aa98380 100644 --- a/deepmd/pt/model/descriptor/__init__.py +++ b/deepmd/pt/model/descriptor/__init__.py @@ -35,6 +35,9 @@ from .se_r import ( DescrptSeR, ) +from .se_t import ( + DescrptSeT, +) __all__ = [ "BaseDescriptor", @@ -45,6 +48,7 @@ "DescrptSeAttenV2", "DescrptSeA", "DescrptSeR", + "DescrptSeT", "DescrptDPA1", "DescrptDPA2", "DescrptHybrid", diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 4ab39465dc..92f5cf2e15 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -27,6 +27,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -549,15 +552,33 @@ def forward( return g1, rot_mat, g2, h2, sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 678b797e6c..b33a528721 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -37,6 +37,9 @@ from deepmd.pt.utils.utils import ( to_numpy_array, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -573,30 +576,46 @@ def forward( return g1, rot_mat, g2, h2, sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy["repinit"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repinit"], + min_nbor_dist, repinit_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repinit"]["rcut"], + local_jdata_cpy["repinit"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - local_jdata_cpy["repformer"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repformer"], + local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0] + min_nbor_dist, repformer_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repformer"]["rcut"], + local_jdata_cpy["repformer"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - return local_jdata_cpy + local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index bce52fd8ca..e202005f4c 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -20,6 +21,9 @@ from deepmd.pt.utils.utils import ( to_torch_tensor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -242,22 +246,42 @@ def forward( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - BaseDescriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self) -> dict: return { diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 0035eddba6..350fceae2d 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -30,6 +30,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -82,9 +85,15 @@ def __init__( env_protection: float = 0.0, old_impl: bool = False, type_one_side: bool = True, + trainable: bool = True, seed: Optional[int] = None, - **kwargs, + ntypes: Optional[int] = None, # to be compat with input + # not implemented + spin=None, ): + del ntypes + if spin is not None: + raise NotImplementedError("old implementation of spin is not supported.") super().__init__() self.sea = DescrptBlockSeA( rcut, @@ -100,8 +109,8 @@ def __init__( env_protection=env_protection, old_impl=old_impl, type_one_side=type_one_side, + trainable=trainable, seed=seed, - **kwargs, ) def get_rcut(self) -> float: @@ -295,18 +304,35 @@ def t_cvt(xx): return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist @DescriptorBlock.register("se_e2_a") diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 340fd0c02a..b0a739f5e6 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -35,6 +35,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -413,15 +416,32 @@ def t_cvt(xx): return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py new file mode 100644 index 0000000000..2c8f52709f --- /dev/null +++ b/deepmd/pt/model/descriptor/se_t.py @@ -0,0 +1,689 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +from typing import ( + Callable, + ClassVar, + Dict, + List, + Optional, + Tuple, + Union, +) + +import numpy as np +import torch + +from deepmd.pt.model.descriptor import ( + DescriptorBlock, + prod_env_mat, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, + RESERVED_PRECISON_DICT, +) +from deepmd.pt.utils.env_mat_stat import ( + EnvMatStatSe, +) +from deepmd.pt.utils.update_sel import ( + UpdateSel, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) +from deepmd.utils.env_mat_stat import ( + StatItem, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +try: + from typing import ( + Final, + ) +except ImportError: + from torch.jit import Final + +from deepmd.dpmodel.utils import EnvMat as DPEnvMat +from deepmd.pt.model.network.mlp import ( + EmbeddingNet, + NetworkCollection, +) +from deepmd.pt.utils.exclude_mask import ( + PairExcludeMask, +) + +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("se_e3") +@BaseDescriptor.register("se_at") +@BaseDescriptor.register("se_a_3be") +class DescrptSeT(BaseDescriptor, torch.nn.Module): + r"""DeepPot-SE constructed from all information (both angular and radial) of atomic + configurations. + + The embedding takes angles between two neighboring atoms as input. + + Parameters + ---------- + rcut : float + The cut-off radius + rcut_smth : float + From where the environment matrix should be smoothed + sel : list[int] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + set_davg_zero : bool + Set the shift of embedding net input to zero. + activation_function : str + The activation function in the embedding net. Supported options are |ACTIVATION_FN| + env_protection : float + Protection parameter to prevent division by zero errors during environment matrix calculations. + exclude_types : List[List[int]] + The excluded pairs of types which have no interaction with each other. + For example, `[[0, 1]]` means no interaction between type 0 and type 1. + precision : str + The precision of the embedding net parameters. Supported options are |PRECISION| + trainable : bool + If the weights of embedding net are trainable. + seed : int, Optional + Random seed for initializing the network parameters. + """ + + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[int], + neuron: List[int] = [24, 48, 96], + resnet_dt: bool = False, + set_davg_zero: bool = False, + activation_function: str = "tanh", + env_protection: float = 0.0, + exclude_types: List[Tuple[int, int]] = [], + precision: str = "float64", + trainable: bool = True, + seed: Optional[int] = None, + ntypes: Optional[int] = None, # to be compat with input + # not implemented + spin=None, + ): + del ntypes + if spin is not None: + raise NotImplementedError("old implementation of spin is not supported.") + super().__init__() + self.seat = DescrptBlockSeT( + rcut, + rcut_smth, + sel, + neuron=neuron, + resnet_dt=resnet_dt, + set_davg_zero=set_davg_zero, + activation_function=activation_function, + env_protection=env_protection, + exclude_types=exclude_types, + precision=precision, + trainable=trainable, + seed=seed, + ) + + def get_rcut(self) -> float: + """Returns the cut-off radius.""" + return self.seat.get_rcut() + + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.seat.get_rcut_smth() + + def get_nsel(self) -> int: + """Returns the number of selected atoms in the cut-off radius.""" + return self.seat.get_nsel() + + def get_sel(self) -> List[int]: + """Returns the number of selected atoms for each type.""" + return self.seat.get_sel() + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.seat.get_ntypes() + + def get_dim_out(self) -> int: + """Returns the output dimension.""" + return self.seat.get_dim_out() + + def get_dim_emb(self) -> int: + """Returns the output dimension.""" + return self.seat.get_dim_emb() + + def mixed_types(self): + """Returns if the descriptor requires a neighbor list that distinguish different + atomic types or not. + """ + return self.seat.mixed_types() + + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.seat.get_env_protection() + + def share_params(self, base_class, shared_level, resume=False): + """ + Share the parameters of self to the base_class with shared_level during multitask training. + If not start from checkpoint (resume is False), + some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes. + """ + assert ( + self.__class__ == base_class.__class__ + ), "Only descriptors of the same type can share params!" + # For SeT descriptors, the user-defined share-level + # shared_level: 0 + # share all parameters in sea + if shared_level == 0: + self.seat.share_params(base_class.seat, 0, resume=resume) + # Other shared levels + else: + raise NotImplementedError + + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.seat.dim_out + + def compute_input_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + path: Optional[DPPath] = None, + ): + """ + Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + path : Optional[DPPath] + The path to the stat file. + + """ + return self.seat.compute_input_stats(merged, path) + + def reinit_exclude( + self, + exclude_types: List[Tuple[int, int]] = [], + ): + """Update the type exclusions.""" + self.seat.reinit_exclude(exclude_types) + + def forward( + self, + coord_ext: torch.Tensor, + atype_ext: torch.Tensor, + nlist: torch.Tensor, + mapping: Optional[torch.Tensor] = None, + comm_dict: Optional[Dict[str, torch.Tensor]] = None, + ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, not required by this descriptor. + comm_dict + The data needed for communication for parallel inference. + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x ng + gr + The rotationally equivariant and permutationally invariant single particle + representation. + This descriptor returns None. + g2 + The rotationally invariant pair-partical representation. + This descriptor returns None. + h2 + The rotationally equivariant pair-partical representation. + This descriptor returns None. + sw + The smooth switch function. + + """ + return self.seat.forward(nlist, coord_ext, atype_ext, None, mapping) + + def set_stat_mean_and_stddev( + self, + mean: torch.Tensor, + stddev: torch.Tensor, + ) -> None: + self.seat.mean = mean + self.seat.stddev = stddev + + def serialize(self) -> dict: + obj = self.seat + return { + "@class": "Descriptor", + "type": "se_e3", + "@version": 1, + "rcut": obj.rcut, + "rcut_smth": obj.rcut_smth, + "sel": obj.sel, + "neuron": obj.neuron, + "resnet_dt": obj.resnet_dt, + "set_davg_zero": obj.set_davg_zero, + "activation_function": obj.activation_function, + "precision": RESERVED_PRECISON_DICT[obj.prec], + "embeddings": obj.filter_layers.serialize(), + "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), + "exclude_types": obj.exclude_types, + "env_protection": obj.env_protection, + "@variables": { + "davg": obj["davg"].detach().cpu().numpy(), + "dstd": obj["dstd"].detach().cpu().numpy(), + }, + "trainable": obj.trainable, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptSeT": + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + data.pop("type", None) + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + env_mat = data.pop("env_mat") + obj = cls(**data) + + def t_cvt(xx): + return torch.tensor(xx, dtype=obj.seat.prec, device=env.DEVICE) + + obj.seat["davg"] = t_cvt(variables["davg"]) + obj.seat["dstd"] = t_cvt(variables["dstd"]) + obj.seat.filter_layers = NetworkCollection.deserialize(embeddings) + return obj + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms + local_jdata : dict + The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms + """ + local_jdata_cpy = local_jdata.copy() + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist + + +@DescriptorBlock.register("se_e3") +class DescrptBlockSeT(DescriptorBlock): + ndescrpt: Final[int] + __constants__: ClassVar[list] = ["ndescrpt"] + + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[int], + neuron: List[int] = [24, 48, 96], + resnet_dt: bool = False, + set_davg_zero: bool = False, + activation_function: str = "tanh", + env_protection: float = 0.0, + exclude_types: List[Tuple[int, int]] = [], + precision: str = "float64", + trainable: bool = True, + seed: Optional[int] = None, + ): + r"""Construct an embedding net of type `se_e3`. + + The embedding takes angles between two neighboring atoms as input. + + Parameters + ---------- + rcut : float + The cut-off radius + rcut_smth : float + From where the environment matrix should be smoothed + sel : list[int] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + set_davg_zero : bool + Set the shift of embedding net input to zero. + activation_function : str + The activation function in the embedding net. Supported options are |ACTIVATION_FN| + env_protection : float + Protection parameter to prevent division by zero errors during environment matrix calculations. + exclude_types : List[List[int]] + The excluded pairs of types which have no interaction with each other. + For example, `[[0, 1]]` means no interaction between type 0 and type 1. + precision : str + The precision of the embedding net parameters. Supported options are |PRECISION| + trainable : bool + If the weights of embedding net are trainable. + seed : int, Optional + Random seed for initializing the network parameters. + """ + super().__init__() + self.rcut = rcut + self.rcut_smth = rcut_smth + self.neuron = neuron + self.filter_neuron = self.neuron + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.resnet_dt = resnet_dt + self.env_protection = env_protection + self.ntypes = len(sel) + self.seed = seed + # order matters, placed after the assignment of self.ntypes + self.reinit_exclude(exclude_types) + + self.sel = sel + # should be on CPU to avoid D2H, as it is used as slice index + self.sec = [0, *np.cumsum(self.sel).tolist()] + self.split_sel = self.sel + self.nnei = sum(sel) + self.ndescrpt = self.nnei * 4 + + wanted_shape = (self.ntypes, self.nnei, 4) + mean = torch.zeros(wanted_shape, dtype=self.prec, device=env.DEVICE) + stddev = torch.ones(wanted_shape, dtype=self.prec, device=env.DEVICE) + self.register_buffer("mean", mean) + self.register_buffer("stddev", stddev) + + ndim = 2 + filter_layers = NetworkCollection( + ndim=ndim, ntypes=len(sel), network_type="embedding_network" + ) + for embedding_idx in itertools.product(range(self.ntypes), repeat=ndim): + filter_layers[embedding_idx] = EmbeddingNet( + 1, + self.filter_neuron, + activation_function=self.activation_function, + precision=self.precision, + resnet_dt=self.resnet_dt, + seed=self.seed, + ) + self.filter_layers = filter_layers + self.stats = None + # set trainable + self.trainable = trainable + for param in self.parameters(): + param.requires_grad = trainable + + def get_rcut(self) -> float: + """Returns the cut-off radius.""" + return self.rcut + + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.rcut_smth + + def get_nsel(self) -> int: + """Returns the number of selected atoms in the cut-off radius.""" + return sum(self.sel) + + def get_sel(self) -> List[int]: + """Returns the number of selected atoms for each type.""" + return self.sel + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.ntypes + + def get_dim_out(self) -> int: + """Returns the output dimension.""" + return self.dim_out + + def get_dim_emb(self) -> int: + """Returns the output dimension.""" + return self.neuron[-1] + + def get_dim_in(self) -> int: + """Returns the input dimension.""" + return self.dim_in + + def mixed_types(self) -> bool: + """If true, the discriptor + 1. assumes total number of atoms aligned across frames; + 2. requires a neighbor list that does not distinguish different atomic types. + + If false, the discriptor + 1. assumes total number of atoms of each atom type aligned across frames; + 2. requires a neighbor list that distinguishes different atomic types. + + """ + return False + + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.env_protection + + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.filter_neuron[-1] + + @property + def dim_in(self): + """Returns the atomic input dimension of this descriptor.""" + return 0 + + def __setitem__(self, key, value): + if key in ("avg", "data_avg", "davg"): + self.mean = value + elif key in ("std", "data_std", "dstd"): + self.stddev = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.mean + elif key in ("std", "data_std", "dstd"): + return self.stddev + else: + raise KeyError(key) + + def compute_input_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + path: Optional[DPPath] = None, + ): + """ + Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + path : Optional[DPPath] + The path to the stat file. + + """ + env_mat_stat = EnvMatStatSe(self) + if path is not None: + path = path / env_mat_stat.get_hash() + if path is None or not path.is_dir(): + if callable(merged): + # only get data for once + sampled = merged() + else: + sampled = merged + else: + sampled = [] + env_mat_stat.load_or_compute_stats(sampled, path) + self.stats = env_mat_stat.stats + mean, stddev = env_mat_stat() + if not self.set_davg_zero: + self.mean.copy_(torch.tensor(mean, device=env.DEVICE)) + self.stddev.copy_(torch.tensor(stddev, device=env.DEVICE)) + + def get_stats(self) -> Dict[str, StatItem]: + """Get the statistics of the descriptor.""" + if self.stats is None: + raise RuntimeError( + "The statistics of the descriptor has not been computed." + ) + return self.stats + + def reinit_exclude( + self, + exclude_types: List[Tuple[int, int]] = [], + ): + self.exclude_types = exclude_types + self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + + def forward( + self, + nlist: torch.Tensor, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + extended_atype_embd: Optional[torch.Tensor] = None, + mapping: Optional[torch.Tensor] = None, + ): + """Compute the descriptor. + + Parameters + ---------- + nlist + The neighbor list. shape: nf x nloc x nnei + extended_coord + The extended coordinates of atoms. shape: nf x (nallx3) + extended_atype + The extended aotm types. shape: nf x nall x nt + extended_atype_embd + The extended type embedding of atoms. shape: nf x nall + mapping + The index mapping, not required by this descriptor. + + Returns + ------- + result + The descriptor. shape: nf x nloc x ng + gr + The rotationally equivariant and permutationally invariant single particle + representation. + This descriptor returns None. + g2 + The rotationally invariant pair-partical representation. + This descriptor returns None. + h2 + The rotationally equivariant pair-partical representation. + This descriptor returns None. + sw + The smooth switch function. shape: nf x nloc x nnei + + """ + del extended_atype_embd, mapping + nloc = nlist.shape[1] + atype = extended_atype[:, :nloc] + dmatrix, diff, sw = prod_env_mat( + extended_coord, + nlist, + atype, + self.mean, + self.stddev, + self.rcut, + self.rcut_smth, + protection=self.env_protection, + ) + dmatrix = dmatrix.view(-1, self.nnei, 4) + dmatrix = dmatrix.to(dtype=self.prec) + nfnl = dmatrix.shape[0] + # pre-allocate a shape to pass jit + result = torch.zeros( + [nfnl, self.filter_neuron[-1]], + dtype=self.prec, + device=extended_coord.device, + ) + # nfnl x nnei + exclude_mask = self.emask(nlist, extended_atype).view(nfnl, -1) + for embedding_idx, ll in enumerate(self.filter_layers.networks): + ti = embedding_idx % self.ntypes + nei_type_j = self.sel[ti] + tj = embedding_idx // self.ntypes + nei_type_i = self.sel[tj] + if ti <= tj: + # avoid repeat calculation + # nfnl x nt_i x 3 + rr_i = dmatrix[:, self.sec[ti] : self.sec[ti + 1], 1:] + mm_i = exclude_mask[:, self.sec[ti] : self.sec[ti + 1]] + rr_i = rr_i * mm_i[:, :, None] + # nfnl x nt_j x 3 + rr_j = dmatrix[:, self.sec[tj] : self.sec[tj + 1], 1:] + mm_j = exclude_mask[:, self.sec[tj] : self.sec[tj + 1]] + rr_j = rr_j * mm_j[:, :, None] + # nfnl x nt_i x nt_j + env_ij = torch.einsum("ijm,ikm->ijk", rr_i, rr_j) + # nfnl x nt_i x nt_j x 1 + env_ij_reshape = env_ij.unsqueeze(-1) + # nfnl x nt_i x nt_j x ng + gg = ll.forward(env_ij_reshape) + # nfnl x nt_i x nt_j x ng + res_ij = torch.einsum("ijk,ijkm->im", env_ij, gg) + res_ij = res_ij * (1.0 / float(nei_type_i) / float(nei_type_j)) + result += res_ij + # xyz_scatter /= (self.nnei * self.nnei) + result = result.view(-1, nloc, self.filter_neuron[-1]) + return ( + result.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + None, + None, + None, + sw, + ) diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index fab1ff580f..d3a65db287 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,28 +1,51 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, + Optional, + Tuple, +) + from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) class DPModelCommon: """A base class to implement common methods for all the Models.""" @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist def get_fitting_net(self): """Get the fitting network.""" diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index f18e1d097f..45b76d1da6 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Dict, + List, Optional, + Tuple, ) import torch @@ -12,6 +14,9 @@ from deepmd.pt.model.model.model import ( BaseModel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .dp_model import ( DPModelCommon, @@ -105,18 +110,32 @@ def forward_lower( return model_predict @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["dpmodel"] = DPModelCommon.update_sel( - global_jdata, local_jdata["dpmodel"] + local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.update_sel( + train_data, type_map, local_jdata["dpmodel"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index e3dcd389bb..148ffaa703 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, ) import torch @@ -18,6 +19,9 @@ from deepmd.pt.model.model.model import ( BaseModel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @BaseModel.register("frozen") @@ -130,6 +134,11 @@ def get_model_def_script(self) -> str: # be a problem return self.model.get_model_def_script() + @torch.jit.export + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum neighbor distance.""" + return self.model.get_min_nbor_dist() + def serialize(self) -> dict: from deepmd.pt.model.model import ( get_model, @@ -156,17 +165,31 @@ def get_nsel(self) -> int: return self.model.get_nsel() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None @torch.jit.export def model_output_type(self) -> str: diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index bf97472e33..d3670737ba 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -18,6 +18,7 @@ def __init__(self, *args, **kwargs): """Construct a basic model for different tasks.""" torch.nn.Module.__init__(self) self.model_def_script = "" + self.min_nbor_dist = None def compute_or_load_stat( self, @@ -46,6 +47,11 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script + @torch.jit.export + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum distance between two atoms.""" + return self.min_nbor_dist + @torch.jit.export def get_ntypes(self): """Returns the number of element types.""" diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index ff932bc402..369a413f50 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -291,6 +291,11 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() + @torch.jit.export + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum neighbor distance.""" + return self.backbone_model.get_min_nbor_dist() + @torch.jit.export def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" diff --git a/deepmd/pt/utils/auto_batch_size.py b/deepmd/pt/utils/auto_batch_size.py index 13264a336c..0af7cdcc47 100644 --- a/deepmd/pt/utils/auto_batch_size.py +++ b/deepmd/pt/utils/auto_batch_size.py @@ -52,7 +52,17 @@ def is_oom_error(self, e: Exception) -> bool: e : Exception Exception """ - return isinstance(e, RuntimeError) and "CUDA out of memory." in e.args[0] + # several sources think CUSOLVER_STATUS_INTERNAL_ERROR is another out-of-memory error, + # such as https://github.com/JuliaGPU/CUDA.jl/issues/1924 + # (the meaningless error message should be considered as a bug in cusolver) + if isinstance(e, RuntimeError) and ( + "CUDA out of memory." in e.args[0] + or "cusolver error: CUSOLVER_STATUS_INTERNAL_ERROR" in e.args[0] + ): + # Release all unoccupied cached memory + torch.cuda.empty_cache() + return True + return False def execute_all( self, callable: Callable, total_size: int, natoms: int, *args, **kwargs diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index c99ddbb3c6..21a2a3fbda 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -50,9 +50,10 @@ def serialize_from_file(model_file: str) -> dict: "pt_version": torch.__version__, "model": model_dict, "model_def_script": model_def_script, - # TODO "@variables": {}, } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() return data @@ -72,4 +73,6 @@ def deserialize_to_file(model_file: str, data: dict) -> None: # JIT will happy in this way... model.model_def_script = json.dumps(data["model_def_script"]) model = torch.jit.script(model) + if "min_nbor_dist" in data.get("@variables", {}): + model.min_nbor_dist = data["@variables"]["min_nbor_dist"] torch.jit.save(model, model_file) diff --git a/deepmd/pt/utils/update_sel.py b/deepmd/pt/utils/update_sel.py index 8c2d0699f2..7f42a9f91c 100644 --- a/deepmd/pt/utils/update_sel.py +++ b/deepmd/pt/utils/update_sel.py @@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # TODO: save to the model in UpdateSel.hook - pass diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index fabaf78c85..2f813da731 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -26,6 +26,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -462,19 +465,33 @@ def explicit_ntypes(self) -> bool: @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # call subprocess cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py index 7c69ea5202..fe4fc2ae6a 100644 --- a/deepmd/tf/descriptor/hybrid.py +++ b/deepmd/tf/descriptor/hybrid.py @@ -17,6 +17,9 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -423,22 +426,42 @@ def explicit_ntypes(self) -> bool: return any(ii.explicit_ntypes for ii in self.descrpt_list) @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - Descriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = Descriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self, suffix: str = "") -> dict: return { diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 963e9bf607..e6247290c0 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -20,6 +20,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -431,14 +434,28 @@ def init_variables( self.dstd = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_std") @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index a14dc0a5ce..f5f54550f2 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -2,6 +2,7 @@ import re from typing import ( List, + Optional, Set, Tuple, ) @@ -21,6 +22,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -150,19 +154,36 @@ def precision(self) -> tf.DType: return self.filter_precision @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # default behavior is to update sel which is a list local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist def serialize_network( self, diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py index e78dfba461..b79e806fca 100644 --- a/deepmd/tf/descriptor/se_a_mask.py +++ b/deepmd/tf/descriptor/se_a_mask.py @@ -24,6 +24,9 @@ from deepmd.tf.utils.network import ( embedding_net_rand_seed_shift, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -420,14 +423,28 @@ def prod_force_virial( return force, virial, atom_virial @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 43c38b0955..6d3cfeaa6e 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -78,6 +78,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -1470,18 +1473,36 @@ def explicit_ntypes(self) -> bool: return True @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist def serialize_attention_layers( self, diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index cd0a9c0a19..b1a278703a 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -1,18 +1,28 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import re from typing import ( List, Optional, + Set, Tuple, ) import numpy as np +from deepmd.dpmodel.utils.env_mat import ( + EnvMat, +) +from deepmd.dpmodel.utils.network import ( + EmbeddingNet, + NetworkCollection, +) from deepmd.tf.common import ( cast_precision, get_activation_func, get_precision, ) from deepmd.tf.env import ( + EMBEDDING_NET_PATTERN, GLOBAL_NP_FLOAT_PRECISION, GLOBAL_TF_FLOAT_PRECISION, default_tf_session_config, @@ -32,6 +42,9 @@ from deepmd.tf.utils.tabulate import ( DPTabulate, ) +from deepmd.utils.version import ( + check_version_compatibility, +) from .descriptor import ( Descriptor, @@ -75,6 +88,8 @@ class DescrptSeT(DescrptSe): The precision of the embedding net parameters. Supported options are |PRECISION| uniform_seed Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + env_protection: float + Protection parameter to prevent division by zero errors during environment matrix calculations. """ def __init__( @@ -86,10 +101,12 @@ def __init__( resnet_dt: bool = False, trainable: bool = True, seed: Optional[int] = None, + exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: """Constructor.""" @@ -97,6 +114,10 @@ def __init__( raise RuntimeError( f"rcut_smth ({rcut_smth:f}) should be no more than rcut ({rcut:f})!" ) + if exclude_types: + raise NotImplementedError("exclude_types != [] is not supported.") + if env_protection != 0.0: + raise NotImplementedError("env_protection != 0.0 is not supported.") self.sel_a = sel self.rcut_r = rcut self.rcut_r_smth = rcut_smth @@ -107,12 +128,15 @@ def __init__( self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron) self.trainable = trainable self.filter_activation_fn = get_activation_func(activation_function) + self.activation_function_name = activation_function self.filter_precision = get_precision(precision) - # self.exclude_types = set() - # for tt in exclude_types: - # assert(len(tt) == 2) - # self.exclude_types.add((tt[0], tt[1])) - # self.exclude_types.add((tt[1], tt[0])) + self.env_protection = env_protection + self.orig_exclude_types = exclude_types + self.exclude_types = set() + for tt in exclude_types: + assert len(tt) == 2 + self.exclude_types.add((tt[0], tt[1])) + self.exclude_types.add((tt[1], tt[0])) self.set_davg_zero = set_davg_zero # descrpt config @@ -690,3 +714,240 @@ def _filter( else: result += res_ij return result, None + + def serialize_network( + self, + ntypes: int, + ndim: int, + in_dim: int, + neuron: List[int], + activation_function: str, + resnet_dt: bool, + variables: dict, + excluded_types: Set[Tuple[int, int]] = set(), + suffix: str = "", + ) -> dict: + """Serialize network. + + Parameters + ---------- + ntypes : int + The number of types + ndim : int + The dimension of elements + in_dim : int + The input dimension + neuron : List[int] + The neuron list + activation_function : str + The activation function + resnet_dt : bool + Whether to use resnet + variables : dict + The input variables + excluded_types : Set[Tuple[int, int]], optional + The excluded types + suffix : str, optional + The suffix of the scope + + Returns + ------- + dict + The converted network data + """ + assert ndim == 2, "Embeddings in descriptor 'se_e3' must have two dimensions." + embeddings = NetworkCollection( + ntypes=ntypes, + ndim=ndim, + network_type="embedding_network", + ) + + def clear_ij(type_i, type_j): + # initialize an empty network + embeddings[(type_i, type_j)] = EmbeddingNet( + in_dim=in_dim, + neuron=neuron, + activation_function=activation_function, + resnet_dt=resnet_dt, + precision=self.precision.name, + ) + embeddings[(type_i, type_j)].clear() + + for i, j in excluded_types: + clear_ij(i, j) + clear_ij(j, i) + for i in range(ntypes): + for j in range(0, i): + clear_ij(i, j) + + if suffix != "": + embedding_net_pattern = ( + EMBEDDING_NET_PATTERN.replace("/(idt)", suffix + "/(idt)") + .replace("/(bias)", suffix + "/(bias)") + .replace("/(matrix)", suffix + "/(matrix)") + ) + else: + embedding_net_pattern = EMBEDDING_NET_PATTERN + for key, value in variables.items(): + m = re.search(embedding_net_pattern, key) + m = [mm for mm in m.groups() if mm is not None] + typei = m[3] + typej = m[4] + layer_idx = int(m[2]) - 1 + weight_name = m[1] + network_idx = (int(typei), int(typej)) + if embeddings[network_idx] is None: + # initialize the network if it is not initialized + embeddings[network_idx] = EmbeddingNet( + in_dim=in_dim, + neuron=neuron, + activation_function=activation_function, + resnet_dt=resnet_dt, + precision=self.precision.name, + ) + assert embeddings[network_idx] is not None + if weight_name == "idt": + value = value.ravel() + embeddings[network_idx][layer_idx][weight_name] = value + return embeddings.serialize() + + @classmethod + def deserialize_network(cls, data: dict, suffix: str = "") -> dict: + """Deserialize network. + + Parameters + ---------- + data : dict + The input network data + suffix : str, optional + The suffix of the scope + + Returns + ------- + variables : dict + The input variables + """ + embedding_net_variables = {} + embeddings = NetworkCollection.deserialize(data) + assert ( + embeddings.ndim == 2 + ), "Embeddings in descriptor 'se_e3' must have two dimensions." + for ii in range(embeddings.ntypes**embeddings.ndim): + net_idx = [] + rest_ii = ii + for _ in range(embeddings.ndim): + net_idx.append(rest_ii % embeddings.ntypes) + rest_ii //= embeddings.ntypes + net_idx = tuple(net_idx) + key0 = "all" + key1 = f"_{net_idx[0]}" + key2 = f"_{net_idx[1]}" + network = embeddings[net_idx] + assert network is not None + for layer_idx, layer in enumerate(network.layers): + embedding_net_variables[ + f"filter_type_{key0}{suffix}/matrix_{layer_idx + 1}{key1}{key2}" + ] = layer.w + embedding_net_variables[ + f"filter_type_{key0}{suffix}/bias_{layer_idx + 1}{key1}{key2}" + ] = layer.b + if layer.idt is not None: + embedding_net_variables[ + f"filter_type_{key0}{suffix}/idt_{layer_idx + 1}{key1}{key2}" + ] = layer.idt.reshape(1, -1) + else: + # prevent keyError + embedding_net_variables[ + f"filter_type_{key0}{suffix}/idt_{layer_idx + 1}{key1}{key2}" + ] = 0.0 + return embedding_net_variables + + @classmethod + def deserialize(cls, data: dict, suffix: str = ""): + """Deserialize the model. + + Parameters + ---------- + data : dict + The serialized data + + Returns + ------- + Model + The deserialized model + """ + if cls is not DescrptSeT: + raise NotImplementedError(f"Not implemented in class {cls.__name__}") + data = data.copy() + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("@class", None) + data.pop("type", None) + embedding_net_variables = cls.deserialize_network( + data.pop("embeddings"), suffix=suffix + ) + data.pop("env_mat") + variables = data.pop("@variables") + descriptor = cls(**data) + descriptor.embedding_net_variables = embedding_net_variables + descriptor.davg = variables["davg"].reshape( + descriptor.ntypes, descriptor.ndescrpt + ) + descriptor.dstd = variables["dstd"].reshape( + descriptor.ntypes, descriptor.ndescrpt + ) + return descriptor + + def serialize(self, suffix: str = "") -> dict: + """Serialize the model. + + Parameters + ---------- + suffix : str, optional + The suffix of the scope + + Returns + ------- + dict + The serialized data + """ + if type(self) is not DescrptSeT: + raise NotImplementedError( + f"Not implemented in class {self.__class__.__name__}" + ) + if self.embedding_net_variables is None: + raise RuntimeError("init_variables must be called before serialize") + assert self.davg is not None + assert self.dstd is not None + + return { + "@class": "Descriptor", + "type": "se_e3", + "@version": 1, + "rcut": self.rcut_r, + "rcut_smth": self.rcut_r_smth, + "sel": self.sel_a, + "neuron": self.filter_neuron, + "resnet_dt": self.filter_resnet_dt, + "set_davg_zero": self.set_davg_zero, + "activation_function": self.activation_function_name, + "precision": self.filter_precision.name, + "embeddings": self.serialize_network( + ntypes=self.ntypes, + ndim=2, + in_dim=1, + neuron=self.filter_neuron, + activation_function=self.activation_function_name, + resnet_dt=self.filter_resnet_dt, + variables=self.embedding_net_variables, + excluded_types=self.exclude_types, + suffix=suffix, + ), + "env_mat": EnvMat(self.rcut_r, self.rcut_r_smth).serialize(), + "exclude_types": list(self.orig_exclude_types), + "env_protection": self.env_protection, + "@variables": { + "davg": self.davg.reshape(self.ntypes, self.nnei_a, 4), + "dstd": self.dstd.reshape(self.ntypes, self.nnei_a, 4), + }, + "trainable": self.trainable, + } diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py index 1e50d88b63..6f3ff13e3b 100644 --- a/deepmd/tf/entrypoints/compress.py +++ b/deepmd/tf/entrypoints/compress.py @@ -32,6 +32,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + get_data, +) from .freeze import ( freeze, @@ -115,9 +118,17 @@ def compress( log.info("stage 0: compute the min_nbor_dist") jdata = j_loader(training_script) jdata = update_deepmd_input(jdata) + + type_map = jdata["model"].get("type_map", None) + train_data = get_data( + jdata["training"]["training_data"], + 0, # not used + type_map, + None, + ) update_sel = UpdateSel() t_min_nbor_dist = update_sel.get_min_nbor_dist( - jdata, update_sel.get_rcut(jdata) + train_data, ) _check_compress_type(graph) diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 3c4decbe8c..d394773cf2 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -17,6 +17,7 @@ j_loader, ) from deepmd.tf.env import ( + GLOBAL_ENER_FLOAT_PRECISION, reset_default_tf_session_config, tf, ) @@ -255,5 +256,21 @@ def update_sel(jdata): "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) jdata_cpy = jdata.copy() - jdata_cpy["model"] = Model.update_sel(jdata, jdata["model"]) + type_map = jdata["model"].get("type_map") + train_data = get_data( + jdata["training"]["training_data"], + 0, # not used + type_map, + None, # not used + ) + jdata_cpy["model"], min_nbor_dist = Model.update_sel( + train_data, type_map, jdata["model"] + ) + + if min_nbor_dist is not None: + tf.constant( + min_nbor_dist, + name="train_attr/min_nbor_dist", + dtype=GLOBAL_ENER_FLOAT_PRECISION, + ) return jdata_cpy diff --git a/deepmd/tf/env.py b/deepmd/tf/env.py index cdb4feadc3..03f36fb675 100644 --- a/deepmd/tf/env.py +++ b/deepmd/tf/env.py @@ -138,6 +138,7 @@ def dlopen_library(module: str, filename: str): r"filter_type_(all)/(bias)_(\d+)_(\d+)_(\d+)|" r"filter_type_(all)/(bias)_(\d+)_(\d+)|" r"filter_type_(all)/(bias)_(\d+)|" + r"filter_type_(all)/(idt)_(\d+)_(\d+)_(\d+)|" r"filter_type_(all)/(idt)_(\d+)_(\d+)|" r"filter_type_(all)/(idt)_(\d+)|" )[:-1] diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index fa28b2cc58..3e296c00f2 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -8,6 +8,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -38,6 +39,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .model import ( Model, @@ -237,18 +241,32 @@ def get_type_map(self) -> list: return self.model.get_type_map() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # we don't know how to compress it, so no neighbor statistics here - return local_jdata + return local_jdata, None def serialize(self, suffix: str = "") -> dict: # try to recover the original model diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 26bc382569..1bd1644e54 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -10,6 +10,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -27,6 +28,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .model import ( Model, @@ -133,22 +137,42 @@ def get_type_map(self) -> list: return self.models[0].get_type_map() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["models"] = [ - Model.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["models"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["models"]: + new_sub_jdata, min_nbor_dist_ = Model.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["models"] = new_list + return local_jdata_cpy, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 194750a9d7..a1baf85dbc 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -11,6 +11,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -511,7 +512,12 @@ def get_feed_dict( @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Notes @@ -520,8 +526,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class @@ -529,9 +537,11 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: ------- dict The updated local data + float + The minimum distance between two atoms """ cls = cls.get_class_by_type(local_jdata.get("type", "standard")) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Model": @@ -744,21 +754,35 @@ def get_ntypes(self) -> int: return self.ntypes @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = Descriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = Descriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index 4f71dcd76e..8c79dedea0 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -5,6 +5,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -35,6 +36,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @Model.register("pairtab") @@ -268,7 +272,12 @@ def enable_compression(self, suffix: str = "") -> None: # nothing needs to do @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Notes @@ -277,8 +286,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class @@ -286,9 +297,15 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: ------- dict The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 3d61dfd339..44e3943e12 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -3,6 +3,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -36,6 +37,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @Model.register("pairwise_dprc") @@ -410,20 +414,33 @@ def get_feed_dict( return feed_dict @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # do not update sel; only find min distance - # rcut is not important here - UpdateSel().get_min_nbor_dist(global_jdata, 6.0) - return local_jdata + min_nbor_dist = UpdateSel().get_min_nbor_dist(train_data) + return local_jdata, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/utils/update_sel.py b/deepmd/tf/utils/update_sel.py index db0420dde8..726aec4d41 100644 --- a/deepmd/tf/utils/update_sel.py +++ b/deepmd/tf/utils/update_sel.py @@ -3,12 +3,6 @@ Type, ) -from deepmd.env import ( - GLOBAL_ENER_FLOAT_PRECISION, -) -from deepmd.tf.env import ( - tf, -) from deepmd.tf.utils.neighbor_stat import ( NeighborStat, ) @@ -21,12 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # moved from traier.py as duplicated - 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) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 31228cdb69..d18bb2fde9 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -309,6 +309,8 @@ def descrpt_se_t_args(): doc_trainable = "If the parameters in the embedding net are trainable" doc_seed = "Random seed for parameter initialization" doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used" + doc_exclude_types = "The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1." + doc_env_protection = "Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection." return [ Argument("sel", [List[int], str], optional=True, default="auto", doc=doc_sel), @@ -331,6 +333,20 @@ def descrpt_se_t_args(): Argument( "set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero ), + Argument( + "exclude_types", + List[List[int]], + optional=True, + default=[], + doc=doc_exclude_types, + ), + Argument( + "env_protection", + float, + optional=True, + default=0.0, + doc=doc_only_pt_supported + doc_env_protection, + ), ] diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index d1be8e8138..6feed525e5 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -1,14 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from abc import ( + ABC, abstractmethod, ) from typing import ( + List, + Optional, + Tuple, Type, + Union, ) from deepmd.utils.data_system import ( - get_data, + DeepmdDataSystem, ) from deepmd.utils.neighbor_stat import ( NeighborStat, @@ -17,32 +22,29 @@ log = logging.getLogger(__name__) -class BaseUpdateSel: +class BaseUpdateSel(ABC): """Update the sel field in the descriptor.""" def update_one_sel( self, - jdata, - descriptor, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + rcut: float, + sel: Union[int, List[int], str], mixed_type: bool = False, - rcut_key="rcut", - sel_key="sel", - ): - rcut = descriptor[rcut_key] - tmp_sel = self.get_sel( - jdata, + ) -> Tuple[float, List[int]]: + min_nbor_dist, tmp_sel = self.get_nbor_stat( + train_data, + type_map, rcut, mixed_type=mixed_type, ) - sel = descriptor[sel_key] if isinstance(sel, int): # convert to list and finnally convert back to int sel = [sel] - if self.parse_auto_sel(descriptor[sel_key]): - ratio = self.parse_auto_sel_ratio(descriptor[sel_key]) - descriptor[sel_key] = sel = [ - int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel - ] + if self.parse_auto_sel(sel): + ratio = self.parse_auto_sel_ratio(sel) + sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)): @@ -54,9 +56,7 @@ def update_one_sel( "not less than %d, but you set it to %d. The accuracy" " of your model may get worse." % (ii, tt, dd) ) - if mixed_type: - descriptor[sel_key] = sum(sel) - return descriptor + return min_nbor_dist, sel def parse_auto_sel(self, sel): if not isinstance(sel, str): @@ -83,65 +83,36 @@ def parse_auto_sel_ratio(self, sel): def wrap_up_4(self, xx): return 4 * ((int(xx) + 3) // 4) - def get_sel(self, jdata, rcut, mixed_type: bool = False): - _, max_nbor_size = self.get_nbor_stat(jdata, rcut, mixed_type=mixed_type) - return max_nbor_size - - def get_rcut(self, jdata): - if jdata["model"].get("type") == "pairwise_dprc": - return max( - jdata["model"]["qm_model"]["descriptor"]["rcut"], - jdata["model"]["qmmm_model"]["descriptor"]["rcut"], - ) - descrpt_data = jdata["model"]["descriptor"] - rcut_list = [] - if descrpt_data["type"] == "hybrid": - for ii in descrpt_data["list"]: - rcut_list.append(ii["rcut"]) - else: - rcut_list.append(descrpt_data["rcut"]) - return max(rcut_list) - - def get_type_map(self, jdata): - return jdata["model"].get("type_map", None) - - def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): - # it seems that DeepmdDataSystem does not need rcut - # it's not clear why there is an argument... - # max_rcut = get_rcut(jdata) - max_rcut = rcut - type_map = self.get_type_map(jdata) - + def get_nbor_stat( + self, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + rcut: float, + mixed_type: bool = False, + ) -> Tuple[float, Union[int, List[int]]]: + """Get the neighbor statistics of the data. + + Parameters + ---------- + train_data : DeepmdDataSystem + The training data. + type_map : Optional[List[str]] + The type map. + rcut : float + The cutoff radius. + mixed_type : bool, optional + Whether to mix the types. + + Returns + ------- + min_nbor_dist : float + The minimum neighbor distance. + max_nbor_size : List[int] + The maximum neighbor size. + """ if type_map and len(type_map) == 0: type_map = None - multi_task_mode = "data_dict" in jdata["training"] - if not multi_task_mode: - train_data = get_data( - jdata["training"]["training_data"], max_rcut, type_map, None - ) - train_data.get_batch() - else: - assert ( - type_map is not None - ), "Data stat in multi-task mode must have available type_map! " - train_data = None - for systems in jdata["training"]["data_dict"]: - tmp_data = get_data( - jdata["training"]["data_dict"][systems]["training_data"], - max_rcut, - type_map, - None, - ) - tmp_data.get_batch() - assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " - if train_data is None: - train_data = tmp_data - else: - train_data.system_dirs += tmp_data.system_dirs - train_data.data_systems += tmp_data.data_systems - train_data.natoms += tmp_data.natoms - train_data.natoms_vec += tmp_data.natoms_vec - train_data.default_mesh += tmp_data.default_mesh + train_data.get_batch() data_ntypes = train_data.get_ntypes() if type_map is not None: map_ntypes = len(type_map) @@ -152,7 +123,6 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): neistat = self.neighbor_stat(ntypes, rcut, mixed_type=mixed_type) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) - self.hook(min_nbor_dist, max_nbor_size) return min_nbor_dist, max_nbor_size @@ -161,10 +131,14 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): def neighbor_stat(self) -> Type[NeighborStat]: pass - @abstractmethod - def hook(self, min_nbor_dist, max_nbor_size): - pass - - def get_min_nbor_dist(self, jdata, rcut): - min_nbor_dist, _ = self.get_nbor_stat(jdata, rcut) + def get_min_nbor_dist( + self, + train_data: DeepmdDataSystem, + ): + min_nbor_dist, _ = self.get_nbor_stat( + train_data, + None, # type_map doesn't affect min_nbor_dist + 1e-6, # we don't need the max_nbor_size + mixed_type=True, # mixed_types doesn't affect min_nbor_dist + ) return min_nbor_dist diff --git a/doc/data/system.md b/doc/data/system.md index 6ca044f1c9..b5db516726 100644 --- a/doc/data/system.md +++ b/doc/data/system.md @@ -45,3 +45,40 @@ In general, we always use the following convention of units: | Force | eV/Å | | Virial | eV | | Pressure | Bar | + +## Mixed type + +:::{note} +Only the [DPA-1](../model/train-se-atten.md) and [DPA-2](../model/dpa2.md) descriptors support this format. +::: + +In the standard data format, only those frames with the same fingerprint (i.e. the number of atoms of different elements) can be put together as a unified system. +This may lead to sparse frame numbers in those rare systems. + +An ideal way is to put systems with the same total number of atoms together, which is the way we trained DPA-1 on [OC2M](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md). +This system format, which is called `mixed_type`, is proper to put frame-sparse systems together and is slightly different from the standard one. +Take an example, a `mixed_type` may contain the following files: + +``` +type.raw +type_map.raw +set.*/box.npy +set.*/coord.npy +set.*/energy.npy +set.*/force.npy +set.*/real_atom_types.npy +``` + +This system contains `Nframes` frames with the same atom number `Natoms`, the total number of element types contained in all frames is `Ntypes`. Most files are the same as those in [standard formats](../data/system.md), here we only list the distinct ones: + +| ID | Property | File | Required/Optional | Shape | Description | +| -------- | -------------------------------- | ------------------- | ----------------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------ | +| / | Atom type indexes (place holder) | type.raw | Required | Natoms | All zeros to fake the type input | +| type_map | Atom type names | type_map.raw | Required | Ntypes | Atom names that map to atom type contained in all the frames, which is unnecessart to be contained in the periodic table | +| type | Atom type indexes of each frame | real_atom_types.npy | Required | Nframes \* Natoms | Integers that describe atom types in each frame, corresponding to indexes in type_map. `-1` means virtual atoms. | + +With these edited files, one can put together frames with the same `Natoms`, instead of the same formula (like `H2O`). + +To put frames with different `Natoms` into the same system, one can pad systems by adding virtual atoms whose type is `-1`. Virtual atoms do not contribute to any fitting property, so the atomic property of virtual atoms (e.g. forces) should be given zero. + +The API to generate or transfer to `mixed_type` format is available on [dpdata](https://github.com/deepmodeling/dpdata) for a more convenient experience. diff --git a/doc/development/create-a-model-pt.md b/doc/development/create-a-model-pt.md index 35d81b364a..c6d372b195 100644 --- a/doc/development/create-a-model-pt.md +++ b/doc/development/create-a-model-pt.md @@ -20,6 +20,7 @@ The framework-independent backend is implemented in pure NumPy, serving as a ref When creating a new descriptor, it is essential to inherit from both the {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor` class and the {py:class}`torch.nn.Module` class. Abstract methods, including {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor.forward`, must be implemented, while others remain optional. It is crucial to adhere to the original method arguments without any modifications. Once the implementation is complete, the next step involves registering the component with a designated key: ```py +from deepmd.utils.data_system import DeepmdDataSystem from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -63,7 +64,12 @@ class SomeDescript(BaseDescriptor, torch.nn.Module): def deserialize(cls, data: dict) -> "SomeDescript": pass - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ): pass ``` diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index 0d17dfd475..3dd97df6ef 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -7,3 +7,7 @@ The DPA-2 model implementation. See https://arxiv.org/abs/2312.15492 for more details. Training example: `examples/water/dpa2/input_torch.json`. + +## Data format + +DPA-2 supports both the [standard data format](../data/system.md) and the [mixed type data format](../data/system.md#mixed-type). diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 4d55383891..5bebdf580d 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -1,7 +1,7 @@ -# Descriptor `"se_atten"` {{ tensorflow_icon }} {{ pytorch_icon }} +# Descriptor `"se_atten"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} :::{note} -**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} ::: ## DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation @@ -83,10 +83,6 @@ With the training input script, data are also provided in the example directory. An example of the DPA-1 descriptor is provided as follows -::::{tab-set} - -:::{tab-item} TensorFlow {{ tensorflow_icon }} - ```json "descriptor" :{ "type": "se_atten", @@ -116,41 +112,6 @@ An example of the DPA-1 descriptor is provided as follows - {ref}`attn_mask ` determines whether to mask the diagonal in the attention weights and False is recommended. - {ref}`attn_dotr ` determines whether to dot the relative coordinates on the attention weights as a gated scheme, True is recommended. -::: - -:::{tab-item} PyTorch {{ pytorch_icon }} - -```json - "descriptor" :{ - "type": "dpa1", - "rcut_smth": 0.50, - "rcut": 6.00, - "sel": 120, - "neuron": [25, 50, 100], - "tebd_dim": 8, - "axis_neuron": 16, - "attn": 128, - "attn_layer": 2, - "attn_mask": false, - "attn_dotr": true, - } -``` - -- The {ref}`type ` of the descriptor is set to `"dpa1"`, which will use DPA-1 structures. -- {ref}`rcut ` is the cut-off radius for neighbor searching, and the {ref}`rcut_smth ` gives where the smoothing starts. -- **{ref}`sel `** gives the maximum possible number of neighbors in the cut-off radius. It is an int. Note that this number highly affects the efficiency of training, which we usually use less than 200. (We use 120 for training 56 elements in [OC2M dataset](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md)) -- The {ref}`neuron ` specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layer from the input end to the output end, respectively. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is built between them. -- The {ref}`tebd_dim ` specifies the dimension of the type embedding. -- The {ref}`axis_neuron ` specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper](https://arxiv.org/abs/1805.09003) -- {ref}`attn ` sets the length of a hidden vector during scale-dot attention computation. -- {ref}`attn_layer ` sets the number of layers in attention mechanism. -- {ref}`attn_mask ` determines whether to mask the diagonal in the attention weights and False is recommended. -- {ref}`attn_dotr ` determines whether to dot the relative coordinates on the attention weights as a gated scheme, True is recommended. - -::: - -:::: - ### Descriptor `"se_atten_v2"` We highly recommend using the version 2.0 of the attention-based descriptor `"se_atten_v2"`, which is inherited from `"se_atten"` but with the following parameter modifications: @@ -199,37 +160,7 @@ which should include all the elements in the dataset you want to train on. ## Data format -DPA-1 supports the standard data format, which is detailed in [data-conv.md](../data/data-conv.md) and [system.md](../data/system.md). -Note that in this format, only those frames with the same fingerprint (i.e. the number of atoms of different elements) can be put together as a unified system. -This may lead to sparse frame numbers in those rare systems. - -An ideal way is to put systems with the same total number of atoms together, which is the way we trained DPA-1 on [OC2M](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md). -This system format, which is called `mixed_type`, is proper to put frame-sparse systems together and is slightly different from the standard one. -Take an example, a `mixed_type` may contain the following files: - -``` -type.raw -type_map.raw -set.*/box.npy -set.*/coord.npy -set.*/energy.npy -set.*/force.npy -set.*/real_atom_types.npy -``` - -This system contains `Nframes` frames with the same atom number `Natoms`, the total number of element types contained in all frames is `Ntypes`. Most files are the same as those in [standard formats](../data/system.md), here we only list the distinct ones: - -| ID | Property | File | Required/Optional | Shape | Description | -| -------- | -------------------------------- | ------------------- | ----------------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------ | -| / | Atom type indexes (place holder) | type.raw | Required | Natoms | All zeros to fake the type input | -| type_map | Atom type names | type_map.raw | Required | Ntypes | Atom names that map to atom type contained in all the frames, which is unnecessart to be contained in the periodic table | -| type | Atom type indexes of each frame | real_atom_types.npy | Required | Nframes \* Natoms | Integers that describe atom types in each frame, corresponding to indexes in type_map. `-1` means virtual atoms. | - -With these edited files, one can put together frames with the same `Natoms`, instead of the same formula (like `H2O`). Note that this `mixed_type` format only supports `se_atten` descriptor. - -To put frames with different `Natoms` into the same system, one can pad systems by adding virtual atoms whose type is `-1`. Virtual atoms do not contribute to any fitting property, so the atomic property of virtual atoms (e.g. forces) should be given zero. - -The API to generate or transfer to `mixed_type` format is available on [dpdata](https://github.com/deepmodeling/dpdata) for a more convenient experience. +DPA-1 supports both the [standard data format](../data/system.md) and the [mixed type data format](../data/system.md#mixed-type). ## Training example diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index 3a0c1a9547..ab3a405f57 100644 --- a/doc/model/train-se-e3.md +++ b/doc/model/train-se-e3.md @@ -1,7 +1,7 @@ -# Descriptor `"se_e3"` {{ tensorflow_icon }} +# Descriptor `"se_e3"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} :::{note} -**Supported backends**: TensorFlow {{ tensorflow_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} ::: The notation of `se_e3` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The embedding takes bond angles between a central atom and its two neighboring atoms as input (denoted by `e3`). diff --git a/source/tests/common/dpmodel/test_descriptor_se_t.py b/source/tests/common/dpmodel/test_descriptor_se_t.py new file mode 100644 index 0000000000..805ef8627b --- /dev/null +++ b/source/tests/common/dpmodel/test_descriptor_se_t.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeT, +) + +from .case_single_frame_with_nlist import ( + TestCaseSingleFrameWithNlist, +) + + +class TestDescrptSeT(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_self_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + em0 = DescrptSeT(self.rcut, self.rcut_smth, self.sel) + em0.davg = davg + em0.dstd = dstd + em1 = DescrptSeT.deserialize(em0.serialize()) + mm0 = em0.call(self.coord_ext, self.atype_ext, self.nlist) + mm1 = em1.call(self.coord_ext, self.atype_ext, self.nlist) + for ii in [0, 4]: + np.testing.assert_allclose(mm0[ii], mm1[ii]) diff --git a/source/tests/common/dpmodel/test_update_sel.py b/source/tests/common/dpmodel/test_update_sel.py new file mode 100644 index 0000000000..a2656da7de --- /dev/null +++ b/source/tests/common/dpmodel/test_update_sel.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import random +import unittest +from unittest.mock import ( + patch, +) + +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) + + +def update_sel(jdata): + type_map = jdata["model"].get("type_map") + train_data = None + jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"]) + return jdata + + +class TestTrain(unittest.TestCase): + def setUp(self) -> None: + self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() + return super().setUp() + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") + # self.assertEqual(descriptor['sel'], [11,22]) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") + # self.assertEqual(descriptor['sel'], [15,30]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": "auto"}, + {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, + ], + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}, + {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, + ], + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } + expected_out = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": "auto", + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 28, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_wrap_up_4(self): + self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) + self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4) diff --git a/source/tests/consistent/descriptor/test_se_t.py b/source/tests/consistent/descriptor/test_se_t.py new file mode 100644 index 0000000000..7579344012 --- /dev/null +++ b/source/tests/consistent/descriptor/test_se_t.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, + Tuple, +) + +import numpy as np + +from deepmd.dpmodel.descriptor.se_t import DescrptSeT as DescrptSeTDP +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, + CommonTest, + parameterized, +) +from .common import ( + DescriptorTest, +) + +if INSTALLED_PT: + from deepmd.pt.model.descriptor.se_t import DescrptSeT as DescrptSeTPT +else: + DescrptSeTPT = None +if INSTALLED_TF: + from deepmd.tf.descriptor.se_t import DescrptSeT as DescrptSeTTF +else: + DescrptSeTTF = None +from deepmd.utils.argcheck import ( + descrpt_se_t_args, +) + + +@parameterized( + (True, False), # resnet_dt + ([], [[0, 1]]), # excluded_types + ("float32", "float64"), # precision + (0.0, 1e-8, 1e-2), # env_protection +) +class TestSeT(CommonTest, DescriptorTest, unittest.TestCase): + @property + def data(self) -> dict: + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + return { + "sel": [9, 10], + "rcut_smth": 5.80, + "rcut": 6.00, + "neuron": [6, 12, 24], + "resnet_dt": resnet_dt, + "exclude_types": excluded_types, + "env_protection": env_protection, + "precision": precision, + "seed": 1145141919810, + } + + @property + def skip_pt(self) -> bool: + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + return CommonTest.skip_pt + + @property + def skip_dp(self) -> bool: + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + return CommonTest.skip_dp + + @property + def skip_tf(self) -> bool: + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + return env_protection != 0.0 or excluded_types + + tf_class = DescrptSeTTF + dp_class = DescrptSeTDP + pt_class = DescrptSeTPT + args = descrpt_se_t_args() + + def setUp(self): + CommonTest.setUp(self) + + self.ntypes = 2 + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + # TF se_e2_a type_one_side=False requires atype sorted + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + idx = np.argsort(self.atype) + self.atype = self.atype[idx] + self.coords = self.coords.reshape(-1, 3)[idx].ravel() + + def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: + return self.build_tf_descriptor( + obj, + self.natoms, + self.coords, + self.atype, + self.box, + suffix, + ) + + def eval_dp(self, dp_obj: Any) -> Any: + return self.eval_dp_descriptor( + dp_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def eval_pt(self, pt_obj: Any) -> Any: + return self.eval_pt_descriptor( + pt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def extract_ret(self, ret: Any, backend) -> Tuple[np.ndarray, ...]: + return (ret[0],) + + @property + def rtol(self) -> float: + """Relative tolerance for comparing the return value.""" + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") + + @property + def atol(self) -> float: + """Absolute tolerance for comparing the return value.""" + ( + resnet_dt, + excluded_types, + precision, + env_protection, + ) = self.param + if precision == "float64": + return 1e-10 + elif precision == "float32": + return 1e-4 + else: + raise ValueError(f"Unknown precision: {precision}") diff --git a/source/tests/pt/model/test_se_e2_a.py b/source/tests/pt/model/test_se_e2_a.py index 214fdeb00f..47ad9789e8 100644 --- a/source/tests/pt/model/test_se_e2_a.py +++ b/source/tests/pt/model/test_se_e2_a.py @@ -56,7 +56,7 @@ def test_consistency( precision=prec, resnet_dt=idt, old_impl=False, - exclude_mask=em, + exclude_types=em, ).to(env.DEVICE) dd0.sea.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) dd0.sea.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) @@ -102,7 +102,7 @@ def test_consistency( err_msg=err_msg, ) # old impl - if idt is False and prec == "float64": + if idt is False and prec == "float64" and em == []: dd3 = DescrptSeA( self.rcut, self.rcut_smth, diff --git a/source/tests/pt/model/test_se_t.py b/source/tests/pt/model/test_se_t.py new file mode 100644 index 0000000000..bc90fea8b3 --- /dev/null +++ b/source/tests/pt/model/test_se_t.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.descriptor import DescrptSeT as DPDescrptSeT +from deepmd.pt.model.descriptor.se_t import ( + DescrptSeT, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + PRECISION_DICT, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) +from .test_mlp import ( + get_tols, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +# to be merged with the tf test case +class TestDescrptSeT(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for idt, prec, em in itertools.product( + [False, True], + ["float64", "float32"], + [ + [], + [[0, 1]], + [[1, 1]], + ], + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + err_msg = f"idt={idt} prec={prec}" + # pt impl + dd0 = DescrptSeT( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + exclude_types=em, + ).to(env.DEVICE) + dd0.seat.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.seat.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + rd0, _, _, _, sw0 = dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + # serialization + dd1 = DescrptSeT.deserialize(dd0.serialize()) + rd1, _, _, _, sw1 = dd1( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy()[0][self.perm[: self.nloc]], + rd0.detach().cpu().numpy()[1], + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + # dp impl + dd2 = DPDescrptSeT.deserialize(dd0.serialize()) + rd2, _, _, _, sw2 = dd2.call( + self.coord_ext, + self.atype_ext, + self.nlist, + ) + for aa, bb in zip([rd1, sw1], [rd2, sw2]): + np.testing.assert_allclose( + aa.detach().cpu().numpy(), + bb, + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) + + def test_jit( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for idt, prec in itertools.product( + [False, True], + ["float64", "float32"], + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + err_msg = f"idt={idt} prec={prec}" + # pt impl + dd0 = DescrptSeT( + self.rcut, + self.rcut_smth, + self.sel, + precision=prec, + resnet_dt=idt, + ) + dd0.seat.mean = torch.tensor(davg, dtype=dtype, device=env.DEVICE) + dd0.seat.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) + dd1 = DescrptSeT.deserialize(dd0.serialize()) + model = torch.jit.script(dd0) + model = torch.jit.script(dd1) diff --git a/source/tests/pt/model/test_trans.py b/source/tests/pt/model/test_trans.py index d2d790d9e3..df346b10b1 100644 --- a/source/tests/pt/model/test_trans.py +++ b/source/tests/pt/model/test_trans.py @@ -64,7 +64,7 @@ def test( spins=spin.unsqueeze(0), ) ret1 = {key: result_1[key].squeeze(0) for key in test_keys} - prec = 1e-10 + prec = 1e-7 for key in test_keys: if key in ["energy", "force", "force_mag"]: torch.testing.assert_close(ret0[key], ret1[key], rtol=prec, atol=prec) diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py new file mode 100644 index 0000000000..e4eff0ebf3 --- /dev/null +++ b/source/tests/pt/test_update_sel.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import random +import unittest +from unittest.mock import ( + patch, +) + +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.utils.update_sel import ( + UpdateSel, +) + + +def update_sel(jdata): + type_map = jdata["model"].get("type_map") + train_data = None + jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"]) + return jdata + + +class TestTrain(unittest.TestCase): + def setUp(self) -> None: + self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() + return super().setUp() + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") + # self.assertEqual(descriptor['sel'], [11,22]) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") + # self.assertEqual(descriptor['sel'], [15,30]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": "auto"}, + {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, + ], + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}, + {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, + ], + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } + expected_out = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": "auto", + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 28, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_skip_frozen(self): + jdata = { + "model": { + "type": "frozen", + }, + "training": {"training_data": {}}, + } + expected_out = jdata.copy() + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_wrap_up_4(self): + self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) + self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4) diff --git a/source/tests/tf/test_get_potential.py b/source/tests/tf/test_get_potential.py index 47462a20a3..fb39d41d2e 100644 --- a/source/tests/tf/test_get_potential.py +++ b/source/tests/tf/test_get_potential.py @@ -1,8 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Test if `DeepPotential` facto function returns the right type of potential.""" +import tempfile import unittest +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, +) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) from deepmd.tf.infer import ( DeepDipole, DeepPolar, @@ -35,16 +42,19 @@ def setUp(self): str(self.work_dir / "deeppolar.pbtxt"), str(self.work_dir / "deep_polar.pb") ) - # TODO add model files for globalpolar and WFC - # convert_pbtxt_to_pb( - # str(self.work_dir / "deepglobalpolar.pbtxt"), - # str(self.work_dir / "deep_globalpolar.pb") - # ) + with open(self.work_dir / "deeppolar.pbtxt") as f: + deeppolar_pbtxt = f.read() - # convert_pbtxt_to_pb( - # str(self.work_dir / "deepwfc.pbtxt"), - # str(self.work_dir / "deep_wfc.pb") - # ) + # not an actual globalpolar and wfc model, but still good enough for testing factory + with tempfile.NamedTemporaryFile(mode="w") as f: + f.write(deeppolar_pbtxt.replace("polar", "global_polar")) + f.flush() + convert_pbtxt_to_pb(f.name, str(self.work_dir / "deep_globalpolar.pb")) + + with tempfile.NamedTemporaryFile(mode="w") as f: + f.write(deeppolar_pbtxt.replace("polar", "wfc")) + f.flush() + convert_pbtxt_to_pb(f.name, str(self.work_dir / "deep_wfc.pb")) def tearDown(self): for f in self.work_dir.glob("*.pb"): @@ -62,11 +72,10 @@ def test_factory(self): dp = DeepPotential(self.work_dir / "deep_pot.pb") self.assertIsInstance(dp, DeepPot, msg.format(DeepPot, type(dp))) - # TODO add model files for globalpolar and WFC - # dp = DeepPotential(self.work_dir / "deep_globalpolar.pb") - # self.assertIsInstance( - # dp, DeepGlobalPolar, msg.format(DeepGlobalPolar, type(dp)) - # ) + dp = DeepPotential(self.work_dir / "deep_globalpolar.pb") + self.assertIsInstance( + dp, DeepGlobalPolar, msg.format(DeepGlobalPolar, type(dp)) + ) - # dp = DeepPotential(self.work_dir / "deep_wfc.pb") - # self.assertIsInstance(dp, DeepWFC, msg.format(DeepWFC, type(dp))) + dp = DeepPotential(self.work_dir / "deep_wfc.pb") + self.assertIsInstance(dp, DeepWFC, msg.format(DeepWFC, type(dp))) diff --git a/source/tests/tf/test_train.py b/source/tests/tf/test_train.py index 3e22dc57bc..f5a57a948d 100644 --- a/source/tests/tf/test_train.py +++ b/source/tests/tf/test_train.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import random import unittest from unittest.mock import ( patch, @@ -15,6 +16,7 @@ class TestTrain(unittest.TestCase): def setUp(self) -> None: self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() return super().setUp() def test_train_parse_auto_sel(self): @@ -34,22 +36,25 @@ def test_train_parse_auto_sel_ratio(self): with self.assertRaises(RuntimeError): self.update_sel.parse_auto_sel_ratio([1, 2, 3]) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_one_sel(self, sel_mock): - sel_mock.return_value = [10, 20] - jdata = {} - descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto"} - descriptor = self.update_sel.update_one_sel(jdata, descriptor) + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") # self.assertEqual(descriptor['sel'], [11,22]) - self.assertEqual(descriptor["sel"], [12, 24]) - descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"} - descriptor = self.update_sel.update_one_sel(jdata, descriptor) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") # self.assertEqual(descriptor['sel'], [15,30]) - self.assertEqual(descriptor["sel"], [16, 32]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_hybrid(self, sel_mock): - sel_mock.return_value = [10, 20] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -59,7 +64,8 @@ def test_update_sel_hybrid(self, sel_mock): {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, ], } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -70,24 +76,33 @@ def test_update_sel_hybrid(self, sel_mock): {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, ], } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel(self, sel_mock): - sel_mock.return_value = [10, 20] - jdata = {"model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}} + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } expected_out = { - "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}} + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_auto(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -95,7 +110,8 @@ def test_update_sel_atten_auto(self, sel_mock): "sel": "auto", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -104,14 +120,17 @@ def test_update_sel_atten_auto(self, sel_mock): "sel": 28, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_int(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -119,7 +138,8 @@ def test_update_sel_atten_int(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -128,14 +148,17 @@ def test_update_sel_atten_int(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_list(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -143,7 +166,8 @@ def test_update_sel_atten_list(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -152,19 +176,23 @@ def test_update_sel_atten_list(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_loc_frame(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_loc_frame(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "descriptor": { "type": "loc_frame", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -172,22 +200,28 @@ def test_skip_loc_frame(self): "type": "loc_frame", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_frozen(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_frozen(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "type": "frozen", - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_linear_frozen(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_linear_frozen(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "type": "linear_ener", @@ -197,15 +231,18 @@ def test_skip_linear_frozen(self): {"type": "frozen"}, {"type": "frozen"}, ], - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) + @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_min_nbor_dist") - def test_pairwise_dprc(self, sel_mock): - sel_mock.return_value = 0.5 + def test_pairwise_dprc(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist + get_data_mock.return_value = None jdata = { "model": { "type": "pairwise_dprc", @@ -215,7 +252,8 @@ def test_pairwise_dprc(self, sel_mock): {"type": "frozen"}, {"type": "frozen"}, ], - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata)