From d042aefccc2dcdd44b5f6c0c095d04fd1d1d919c Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 23 May 2024 17:42:27 +0800 Subject: [PATCH 01/10] feat(dp/pt): refactor `se_e3` descriptor --- deepmd/dpmodel/descriptor/__init__.py | 4 + deepmd/dpmodel/descriptor/se_t.py | 359 ++++++++++ deepmd/pt/model/descriptor/__init__.py | 4 + deepmd/pt/model/descriptor/se_t.py | 662 ++++++++++++++++++ deepmd/tf/descriptor/se_t.py | 271 ++++++- deepmd/tf/env.py | 1 + deepmd/utils/argcheck.py | 16 + doc/model/train-se-e3.md | 2 +- examples/water/se_e3/input_torch.json | 80 +++ .../common/dpmodel/test_descriptor_se_t.py | 35 + source/tests/common/test_examples.py | 1 + .../tests/consistent/descriptor/test_se_t.py | 205 ++++++ source/tests/pt/model/test_se_e2_a.py | 4 +- source/tests/pt/model/test_se_t.py | 134 ++++ 14 files changed, 1770 insertions(+), 8 deletions(-) create mode 100644 deepmd/dpmodel/descriptor/se_t.py create mode 100644 deepmd/pt/model/descriptor/se_t.py create mode 100644 examples/water/se_e3/input_torch.json create mode 100644 source/tests/common/dpmodel/test_descriptor_se_t.py create mode 100644 source/tests/consistent/descriptor/test_se_t.py create mode 100644 source/tests/pt/model/test_se_t.py diff --git a/deepmd/dpmodel/descriptor/__init__.py b/deepmd/dpmodel/descriptor/__init__.py index 563fb9d149..1a7b376a36 100644 --- a/deepmd/dpmodel/descriptor/__init__.py +++ b/deepmd/dpmodel/descriptor/__init__.py @@ -17,10 +17,14 @@ from .se_r import ( DescrptSeR, ) +from .se_t import ( + DescrptSeT, +) __all__ = [ "DescrptSeA", "DescrptSeR", + "DescrptSeT", "DescrptDPA1", "DescrptDPA2", "DescrptHybrid", diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py new file mode 100644 index 0000000000..062d865ec9 --- /dev/null +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -0,0 +1,359 @@ +# 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.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 tj < ti: + # avoid repeat calculation + continue + # nfnl x nt_i x 3 + rr_i = rr[:, sec[ti] : sec[ti + 1], 1:] + # nfnl x nt_j x 3 + rr_j = rr[:, sec[tj] : sec[tj + 1], 1:] + # 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, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) diff --git a/deepmd/pt/model/descriptor/__init__.py b/deepmd/pt/model/descriptor/__init__.py index 5fb988bb18..e5298ba3ef 100644 --- a/deepmd/pt/model/descriptor/__init__.py +++ b/deepmd/pt/model/descriptor/__init__.py @@ -32,6 +32,9 @@ from .se_r import ( DescrptSeR, ) +from .se_t import ( + DescrptSeT, +) __all__ = [ "BaseDescriptor", @@ -41,6 +44,7 @@ "DescrptBlockSeAtten", "DescrptSeA", "DescrptSeR", + "DescrptSeT", "DescrptDPA1", "DescrptDPA2", "DescrptHybrid", diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py new file mode 100644 index 0000000000..22c905034b --- /dev/null +++ b/deepmd/pt/model/descriptor/se_t.py @@ -0,0 +1,662 @@ +# 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.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 + ): + del ntypes + 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, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + local_jdata_cpy = local_jdata.copy() + return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + + +@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 tj < ti: + # avoid repeat calculation + continue + # nfnl x nt_i x 3 + rr_i = dmatrix[:, self.sec[ti] : self.sec[ti + 1], 1:] + # nfnl x nt_j x 3 + rr_j = dmatrix[:, self.sec[tj] : self.sec[tj + 1], 1:] + # 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/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index cd0a9c0a19..2ff5eb7b2d 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("env_protection != [] 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 has 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 has 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/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/utils/argcheck.py b/deepmd/utils/argcheck.py index 0f88a33773..7b6e13be3a 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/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index 3a0c1a9547..466bdae642 100644 --- a/doc/model/train-se-e3.md +++ b/doc/model/train-se-e3.md @@ -1,4 +1,4 @@ -# Descriptor `"se_e3"` {{ tensorflow_icon }} +# Descriptor `"se_e3"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }} :::{note} **Supported backends**: TensorFlow {{ tensorflow_icon }} diff --git a/examples/water/se_e3/input_torch.json b/examples/water/se_e3/input_torch.json new file mode 100644 index 0000000000..eb631f87b6 --- /dev/null +++ b/examples/water/se_e3/input_torch.json @@ -0,0 +1,80 @@ +{ + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_e3", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.50, + "rcut": 5.80, + "neuron": [ + 2, + 4, + 8 + ], + "resnet_dt": false, + "precision": "float64", + "seed": 1, + "_comment": " that's all" + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" + }, + "data_stat_nbatch": 20, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "_comment": " that's all" + }, + "training": { + "stat_file": "./se_e3.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" + }, + "numb_steps": 100000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 10000, + "_comment": "that's all" + }, + "_comment": "that's all" +} 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/test_examples.py b/source/tests/common/test_examples.py index f7f1593f6f..4b44b1ee99 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -47,6 +47,7 @@ p_examples / "dprc" / "pairwise" / "input.json", p_examples / "dprc" / "generalized_force" / "input.json", p_examples / "water" / "se_e2_a" / "input_torch.json", + p_examples / "water" / "se_e3" / "input_torch.json", p_examples / "water" / "se_atten" / "input_torch.json", p_examples / "water" / "dpa2" / "input_torch.json", ) 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..5b90077f17 --- /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 + ([],), # 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 + + 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..078ce41289 --- /dev/null +++ b/source/tests/pt/model/test_se_t.py @@ -0,0 +1,134 @@ +# 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"], + [ + [], + ], + ): + 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) From b367d4d5cb1b9c91c50073855dce2f1e78799ba4 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 23 May 2024 17:43:55 +0800 Subject: [PATCH 02/10] Update train-se-e3.md --- doc/model/train-se-e3.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index 466bdae642..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 }} {{ 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`). From cb7e757b55f63011fd4e7f7d6e44a84f68a4ba5d Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 23 May 2024 17:59:49 +0800 Subject: [PATCH 03/10] Update se_t.py --- deepmd/tf/descriptor/se_t.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 2ff5eb7b2d..610539686b 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -755,7 +755,7 @@ def serialize_network( dict The converted network data """ - assert ndim == 2, "Embeddings in descriptor 'se_e3' must has two dimensions." + assert ndim == 2, "Embeddings in descriptor 'se_e3' must have two dimensions." embeddings = NetworkCollection( ntypes=ntypes, ndim=ndim, @@ -831,7 +831,7 @@ def deserialize_network(cls, data: dict, suffix: str = "") -> dict: embeddings = NetworkCollection.deserialize(data) assert ( embeddings.ndim == 2 - ), "Embeddings in descriptor 'se_e3' must has two dimensions." + ), "Embeddings in descriptor 'se_e3' must have two dimensions." for ii in range(embeddings.ntypes**embeddings.ndim): net_idx = [] rest_ii = ii From 5f22783539148df0be79ff6ed26460f52edb91ba Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 23 May 2024 22:01:00 +0800 Subject: [PATCH 04/10] fix jit --- deepmd/dpmodel/descriptor/se_t.py | 31 +++++++++++++++--------------- deepmd/pt/model/descriptor/se_t.py | 31 +++++++++++++++--------------- 2 files changed, 30 insertions(+), 32 deletions(-) diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index 062d865ec9..7d6548e842 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -275,23 +275,22 @@ def call( ti, tj = embedding_idx nei_type_i = self.sel[ti] nei_type_j = self.sel[tj] - if tj < ti: + if ti <= tj: # avoid repeat calculation - continue - # nfnl x nt_i x 3 - rr_i = rr[:, sec[ti] : sec[ti + 1], 1:] - # nfnl x nt_j x 3 - rr_j = rr[:, sec[tj] : sec[tj + 1], 1:] - # 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 + # nfnl x nt_i x 3 + rr_i = rr[:, sec[ti] : sec[ti + 1], 1:] + # nfnl x nt_j x 3 + rr_j = rr[:, sec[tj] : sec[tj + 1], 1:] + # 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 diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 22c905034b..831afd311d 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -634,23 +634,22 @@ def forward( nei_type_j = self.sel[ti] tj = embedding_idx // self.ntypes nei_type_i = self.sel[tj] - if tj < ti: + if ti <= tj: # avoid repeat calculation - continue - # nfnl x nt_i x 3 - rr_i = dmatrix[:, self.sec[ti] : self.sec[ti + 1], 1:] - # nfnl x nt_j x 3 - rr_j = dmatrix[:, self.sec[tj] : self.sec[tj + 1], 1:] - # 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 + # nfnl x nt_i x 3 + rr_i = dmatrix[:, self.sec[ti] : self.sec[ti + 1], 1:] + # nfnl x nt_j x 3 + rr_j = dmatrix[:, self.sec[tj] : self.sec[tj + 1], 1:] + # 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 ( From 9fd8a11ea207e81e96b2020709167a097ae58fad Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 27 May 2024 14:06:56 +0800 Subject: [PATCH 05/10] Update deepmd/tf/descriptor/se_t.py Co-authored-by: Jinzhe Zeng Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com> --- deepmd/tf/descriptor/se_t.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index 610539686b..b1a278703a 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -115,7 +115,7 @@ def __init__( f"rcut_smth ({rcut_smth:f}) should be no more than rcut ({rcut:f})!" ) if exclude_types: - raise NotImplementedError("env_protection != [] is not supported.") + 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 From 5b1fe5ecb40f3e6dcb2f19bd5fad8ff1155ff00d Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 27 May 2024 15:20:17 +0800 Subject: [PATCH 06/10] rm pt example --- examples/water/se_e3/input_torch.json | 80 --------------------------- source/tests/common/test_examples.py | 1 - 2 files changed, 81 deletions(-) delete mode 100644 examples/water/se_e3/input_torch.json diff --git a/examples/water/se_e3/input_torch.json b/examples/water/se_e3/input_torch.json deleted file mode 100644 index eb631f87b6..0000000000 --- a/examples/water/se_e3/input_torch.json +++ /dev/null @@ -1,80 +0,0 @@ -{ - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_e3", - "sel": [ - 46, - 92 - ], - "rcut_smth": 0.50, - "rcut": 5.80, - "neuron": [ - 2, - 4, - 8 - ], - "resnet_dt": false, - "precision": "float64", - "seed": 1, - "_comment": " that's all" - }, - "fitting_net": { - "neuron": [ - 240, - 240, - 240 - ], - "resnet_dt": true, - "seed": 1, - "_comment": " that's all" - }, - "data_stat_nbatch": 20, - "_comment": " that's all" - }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-8, - "_comment": "that's all" - }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "_comment": " that's all" - }, - "training": { - "stat_file": "./se_e3.hdf5", - "training_data": { - "systems": [ - "../data/data_0", - "../data/data_1", - "../data/data_2" - ], - "batch_size": 1, - "_comment": "that's all" - }, - "validation_data": { - "systems": [ - "../data/data_3" - ], - "batch_size": 1, - "numb_btch": 3, - "_comment": "that's all" - }, - "numb_steps": 100000, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 100, - "save_freq": 10000, - "_comment": "that's all" - }, - "_comment": "that's all" -} diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 4b44b1ee99..f7f1593f6f 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -47,7 +47,6 @@ p_examples / "dprc" / "pairwise" / "input.json", p_examples / "dprc" / "generalized_force" / "input.json", p_examples / "water" / "se_e2_a" / "input_torch.json", - p_examples / "water" / "se_e3" / "input_torch.json", p_examples / "water" / "se_atten" / "input_torch.json", p_examples / "water" / "dpa2" / "input_torch.json", ) From 4108d10e9316141e3b74639fbd9ca196459e385e Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 27 May 2024 15:24:48 +0800 Subject: [PATCH 07/10] Update se_a.py --- deepmd/pt/model/descriptor/se_a.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 0035eddba6..e0ab656d8d 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -82,9 +82,11 @@ 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 ): + del ntypes super().__init__() self.sea = DescrptBlockSeA( rcut, @@ -100,8 +102,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: From ce9c07ba7860a6cc7321b37c222bdc6f02ea854a Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 27 May 2024 15:40:35 +0800 Subject: [PATCH 08/10] Update se_a.py --- deepmd/pt/model/descriptor/se_a.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index e0ab656d8d..598d13452a 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -85,8 +85,11 @@ def __init__( trainable: bool = True, seed: Optional[int] = None, ntypes: Optional[int] = None, # to be compat with input + spin=None, ): del ntypes + if spin is not None: + raise NotImplementedError("old implementation of spin is not supported.") super().__init__() self.sea = DescrptBlockSeA( rcut, From ae96f6ce9d53b3cee1bc816fb86961a03a98dc8c Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 28 May 2024 21:18:48 +0800 Subject: [PATCH 09/10] add exclude_types --- deepmd/dpmodel/descriptor/se_t.py | 4 ++++ deepmd/pt/model/descriptor/se_a.py | 1 + deepmd/pt/model/descriptor/se_t.py | 8 ++++++++ source/tests/pt/model/test_se_t.py | 2 ++ 4 files changed, 15 insertions(+) diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index 7d6548e842..d9b2741b3e 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -279,8 +279,12 @@ def call( # 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 diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 598d13452a..0054360593 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -85,6 +85,7 @@ def __init__( trainable: bool = True, seed: Optional[int] = None, ntypes: Optional[int] = None, # to be compat with input + # not implemented spin=None, ): del ntypes diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 831afd311d..db6244000d 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -115,8 +115,12 @@ def __init__( 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, @@ -638,8 +642,12 @@ def forward( # 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 diff --git a/source/tests/pt/model/test_se_t.py b/source/tests/pt/model/test_se_t.py index 078ce41289..bc90fea8b3 100644 --- a/source/tests/pt/model/test_se_t.py +++ b/source/tests/pt/model/test_se_t.py @@ -45,6 +45,8 @@ def test_consistency( ["float64", "float32"], [ [], + [[0, 1]], + [[1, 1]], ], ): dtype = PRECISION_DICT[prec] From e0d14988800aa54d7513ae606db932a43aa88827 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 28 May 2024 21:33:52 +0800 Subject: [PATCH 10/10] Update test_se_t.py --- source/tests/consistent/descriptor/test_se_t.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/consistent/descriptor/test_se_t.py b/source/tests/consistent/descriptor/test_se_t.py index 5b90077f17..7579344012 100644 --- a/source/tests/consistent/descriptor/test_se_t.py +++ b/source/tests/consistent/descriptor/test_se_t.py @@ -37,7 +37,7 @@ @parameterized( (True, False), # resnet_dt - ([],), # excluded_types + ([], [[0, 1]]), # excluded_types ("float32", "float64"), # precision (0.0, 1e-8, 1e-2), # env_protection ) @@ -90,7 +90,7 @@ def skip_tf(self) -> bool: precision, env_protection, ) = self.param - return env_protection != 0.0 + return env_protection != 0.0 or excluded_types tf_class = DescrptSeTTF dp_class = DescrptSeTDP