diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 8abc8c2777..9f10aa15e5 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -804,7 +804,10 @@ def call( nf, nloc, nnei, _ = dmatrix.shape exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) # nfnl x nnei + exclude_mask = exclude_mask.reshape(nf * nloc, nnei) + # nfnl x nnei nlist = nlist.reshape(nf * nloc, nnei) + nlist = np.where(exclude_mask, nlist, -1) # nfnl x nnei x 4 dmatrix = dmatrix.reshape(nf * nloc, nnei, 4) # nfnl x nnei x 1 @@ -824,8 +827,6 @@ def call( nf * nloc, nnei, self.tebd_dim ) ng = self.neuron[-1] - # nfnl x nnei - exclude_mask = exclude_mask.reshape(nf * nloc, nnei) # nfnl x nnei x 4 rr = dmatrix.reshape(nf * nloc, nnei, 4) rr = rr * exclude_mask[:, :, None] diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index c9ac48efec..3f7f57478b 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -338,7 +338,7 @@ def call( mapping: Optional[np.ndarray] = None, ): exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) - nlist = nlist * exclude_mask + nlist = np.where(exclude_mask, nlist, -1) # nf x nloc x nnei x 4 dmatrix, diff, sw = self.env_mat.call( coord_ext, atype_ext, nlist, self.mean, self.stddev diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 7a8899b4e5..d26cba2aff 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -117,6 +117,9 @@ class DescrptSeA(NativeOP, BaseDescriptor): The precision of the embedding net parameters. Supported options are |PRECISION| spin The deepspin object. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. Limitations ----------- @@ -150,9 +153,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, spin: Optional[Any] = None, + ntypes: Optional[int] = None, # to be compat with input # consistent with argcheck, not used though seed: Optional[int] = None, ) -> None: + del ntypes ## seed, uniform_seed, not included. if spin is not None: raise NotImplementedError("spin is not implemented") diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index dbb6e104fb..9197859294 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -75,6 +75,9 @@ class DescrptSeR(NativeOP, BaseDescriptor): The precision of the embedding net parameters. Supported options are |PRECISION| spin The deepspin object. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. Limitations ----------- @@ -107,9 +110,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, spin: Optional[Any] = None, + ntypes: Optional[int] = None, # to be compat with input # consistent with argcheck, not used though seed: Optional[int] = None, ) -> None: + del ntypes ## seed, uniform_seed, not included. if not type_one_side: raise NotImplementedError("type_one_side == False not implemented") diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index cdbeb701ce..2ed951f440 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -78,6 +78,9 @@ class DescrptSeT(NativeOP, BaseDescriptor): If the weights of embedding net are trainable. seed : int, Optional Random seed for initializing the network parameters. + ntypes : int + Number of element types. + Not used in this descriptor, only to be compat with input. """ def __init__( @@ -94,7 +97,9 @@ def __init__( precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: Optional[int] = None, + ntypes: Optional[int] = None, # to be compat with input ) -> None: + del ntypes self.rcut = rcut self.rcut_smth = rcut_smth self.sel = sel diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index a66693653e..2e091d6bb8 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -410,7 +410,7 @@ def forward( atype = extended_atype[:, :nloc] # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) - nlist = nlist * exclude_mask + nlist = torch.where(exclude_mask != 0, nlist, -1) # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1 dmatrix, diff, sw = prod_env_mat( extended_coord, diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index a59eaca409..181277b637 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -478,11 +478,12 @@ def forward( self.rcut_smth, protection=self.env_protection, ) + # nb x nloc x nnei + exclude_mask = self.emask(nlist, extended_atype) + nlist = torch.where(exclude_mask != 0, nlist, -1) nlist_mask = nlist != -1 nlist = torch.where(nlist == -1, 0, nlist) sw = torch.squeeze(sw, -1) - # beyond the cutoff sw should be 0.0 - sw = sw.masked_fill(~nlist_mask, 0.0) # nf x nloc x nt -> nf x nloc x nnei x nt atype_tebd = extended_atype_embd[:, :nloc, :] atype_tebd_nnei = atype_tebd.unsqueeze(2).expand(-1, -1, self.nnei, -1) @@ -495,8 +496,10 @@ def forward( atype_tebd_nlist = torch.gather(atype_tebd_ext, dim=1, index=index) # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) + # beyond the cutoff sw should be 0.0 + sw = sw.masked_fill(~nlist_mask, 0.0) # (nb x nloc) x nnei - exclude_mask = self.emask(nlist, extended_atype).view(nb * nloc, nnei) + exclude_mask = exclude_mask.view(nb * nloc, nnei) if self.old_impl: assert self.filter_layers_old is not None dmatrix = dmatrix.view( diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 6d3cfeaa6e..b240f00647 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -718,6 +718,12 @@ def _pass_filter( tf.shape(inputs_i)[0], self.nei_type_vec, # extra input for atten ) + # (nframes * nloc * nnei, 1) + nei_exclude_mask = tf.slice( + tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), + [0, 0], + [-1, 1], + ) if self.smooth: inputs_i = tf.where( tf.cast(mask, tf.bool), @@ -727,15 +733,18 @@ def _pass_filter( tf.reshape(self.avg_looked_up, [-1, 1]), [1, self.ndescrpt] ), ) + # (nframes, nloc, nnei) self.recovered_switch *= tf.reshape( - tf.slice( - tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]), - [0, 0], - [-1, 1], - ), + nei_exclude_mask, [-1, natoms[0], self.sel_all_a[0]], ) else: + # (nframes * nloc, 1, nnei) + self.nmask *= tf.reshape( + nei_exclude_mask, + [-1, 1, self.sel_all_a[0]], + ) + self.negative_mask = -(2 << 32) * (1.0 - self.nmask) inputs_i *= mask if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor: inputs_i = descrpt2r4(inputs_i, atype) diff --git a/source/tests/universal/common/cases/cases.py b/source/tests/universal/common/cases/cases.py new file mode 100644 index 0000000000..a8c9a7cd71 --- /dev/null +++ b/source/tests/universal/common/cases/cases.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + + +# originally copied from source/tests/pt/model/test_env_mat.py +class TestCaseSingleFrameWithNlist: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 3 + self.nall = 4 + self.nf, self.nt = 2, 2 + self.coord_ext = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, -2, 0], + ], + dtype=np.float64, + ).reshape([1, self.nall, 3]) + self.atype_ext = np.array([0, 0, 1, 0], dtype=int).reshape([1, self.nall]) + self.mapping = np.array([0, 1, 2, 0], dtype=int).reshape([1, self.nall]) + # sel = [5, 2] + self.sel = [5, 2] + self.sel_mix = [7] + self.natoms = [3, 3, 2, 1] + self.nlist = np.array( + [ + [1, 3, -1, -1, -1, 2, -1], + [0, -1, -1, -1, -1, 2, -1], + [0, 1, -1, -1, -1, -1, -1], + ], + dtype=int, + ).reshape([1, self.nloc, sum(self.sel)]) + self.rcut = 2.2 + self.rcut_smth = 0.4 + # permutations + self.perm = np.array([2, 0, 1, 3], dtype=np.int32) + inv_perm = np.array([1, 2, 0, 3], dtype=np.int32) + # permute the coord and atype + self.coord_ext = np.concatenate( + [self.coord_ext, self.coord_ext[:, self.perm, :]], axis=0 + ).reshape(self.nf, self.nall * 3) + self.atype_ext = np.concatenate( + [self.atype_ext, self.atype_ext[:, self.perm]], axis=0 + ) + self.mapping = np.concatenate( + [self.mapping, self.mapping[:, self.perm]], axis=0 + ) + + # permute the nlist + nlist1 = self.nlist[:, self.perm[: self.nloc], :] + mask = nlist1 == -1 + nlist1 = inv_perm[nlist1] + nlist1 = np.where(mask, -1, nlist1) + self.nlist = np.concatenate([self.nlist, nlist1], axis=0) + self.atol = 1e-12 + + +class TestCaseSingleFrameWithNlistWithVirtual: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 4 + self.nall = 5 + self.nf, self.nt = 2, 2 + self.coord_ext = np.array( + [ + [0, 0, 0], + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + [0, -2, 0], + ], + dtype=np.float64, + ).reshape([1, self.nall, 3]) + self.atype_ext = np.array([0, -1, 0, 1, 0], dtype=int).reshape([1, self.nall]) + # sel = [5, 2] + self.sel = [5, 2] + self.sel_mix = [7] + self.natoms = [3, 3, 2, 1] + self.nlist = np.array( + [ + [2, 4, -1, -1, -1, 3, -1], + [-1, -1, -1, -1, -1, -1, -1], + [0, -1, -1, -1, -1, 3, -1], + [0, 2, -1, -1, -1, -1, -1], + ], + dtype=int, + ).reshape([1, self.nloc, sum(self.sel)]) + self.rcut = 2.2 + self.rcut_smth = 0.4 + # permutations + self.perm = np.array([3, 0, 1, 2, 4], dtype=np.int32) + inv_perm = np.argsort(self.perm) + # permute the coord and atype + self.coord_ext = np.concatenate( + [self.coord_ext, self.coord_ext[:, self.perm, :]], axis=0 + ).reshape(self.nf, self.nall * 3) + self.atype_ext = np.concatenate( + [self.atype_ext, self.atype_ext[:, self.perm]], axis=0 + ) + # permute the nlist + nlist1 = self.nlist[:, self.perm[: self.nloc], :] + mask = nlist1 == -1 + nlist1 = inv_perm[nlist1] + nlist1 = np.where(mask, -1, nlist1) + self.nlist = np.concatenate([self.nlist, nlist1], axis=0) + self.get_real_mapping = np.array([[0, 2, 3], [0, 1, 3]], dtype=np.int32) + self.atol = 1e-12 + + +class TestCaseSingleFrameWithoutNlist: + def setUp(self): + # nloc == 3, nall == 4 + self.nloc = 3 + self.nf, self.nt = 1, 2 + self.coord = np.array( + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + dtype=np.float64, + ).reshape([1, self.nloc * 3]) + self.atype = np.array([0, 0, 1], dtype=int).reshape([1, self.nloc]) + self.cell = 2.0 * np.eye(3).reshape([1, 9]) + # sel = [5, 2] + self.sel = [16, 8] + self.sel_mix = [24] + self.natoms = [3, 3, 2, 1] + self.rcut = 2.2 + self.rcut_smth = 0.4 + self.atol = 1e-12 diff --git a/source/tests/universal/common/cases/descriptor/__init__.py b/source/tests/universal/common/cases/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/common/cases/descriptor/descriptor.py b/source/tests/universal/common/cases/descriptor/descriptor.py new file mode 100644 index 0000000000..2f3bdbb1ee --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/descriptor.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + + +from .utils import ( + DescriptorTestCase, +) + + +class DescriptorTest(DescriptorTestCase): + def setUp(self) -> None: + DescriptorTestCase.setUp(self) diff --git a/source/tests/universal/common/cases/descriptor/utils.py b/source/tests/universal/common/cases/descriptor/utils.py new file mode 100644 index 0000000000..aa1a8c21d4 --- /dev/null +++ b/source/tests/universal/common/cases/descriptor/utils.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) + +import numpy as np + +from deepmd.dpmodel.utils import ( + PairExcludeMask, +) + +from ..cases import ( + TestCaseSingleFrameWithNlist, +) + + +class DescriptorTestCase(TestCaseSingleFrameWithNlist): + """Common test case for descriptor.""" + + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + + def test_forward_consistency(self): + ret = [] + for module in self.modules_to_test: + module = self.forward_wrapper(module) + ret.append( + module( + self.coord_ext, + self.atype_ext, + self.nlist, + mapping=self.mapping, + ) + ) + for kk, vv in enumerate(ret[0]): + subret = [] + for rr in ret: + if rr is not None: + subret.append(rr[kk]) + if len(subret): + for ii, rr in enumerate(subret[1:]): + if subret[0] is None: + assert rr is None + else: + np.testing.assert_allclose( + subret[0], + rr, + err_msg=f"compare {kk} output between 0 and {ii}", + ) + + def test_exclude_types( + self, + ): + coord_ext_device = self.coord_ext + atype_ext_device = self.atype_ext + nlist_device = self.nlist + mapping_device = self.mapping + dd = self.forward_wrapper(self.module) + # only equal when set_davg_zero is True + serialize_dict = self.module.serialize() + + for em in [[[0, 1]], [[1, 1]]]: + ex_pair = PairExcludeMask(self.nt, em) + pair_mask = ex_pair.build_type_exclude_mask(nlist_device, atype_ext_device) + # exclude neighbors in the nlist + nlist_exclude = np.where(pair_mask == 1, nlist_device, -1) + rd_ex, _, _, _, sw_ex = dd( + coord_ext_device, + atype_ext_device, + nlist_exclude, + mapping=mapping_device, + ) + + # normal nlist but use exclude_types params + serialize_dict_em = deepcopy(serialize_dict) + if "list" not in serialize_dict_em: + serialize_dict_em.update({"exclude_types": em}) + else: + # for hybrid + for sd in serialize_dict_em["list"]: + sd.update({"exclude_types": em}) + dd0 = self.forward_wrapper(self.module.deserialize(serialize_dict_em)) + rd0, _, _, _, sw0 = dd0( + coord_ext_device, + atype_ext_device, + nlist_device, + mapping=mapping_device, + ) + np.testing.assert_allclose(rd0, rd_ex) diff --git a/source/tests/universal/dpmodel/descriptor/__init__.py b/source/tests/universal/dpmodel/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/dpmodel/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py new file mode 100644 index 0000000000..9d0253c54c --- /dev/null +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.dpmodel.descriptor import ( + DescrptDPA1, + DescrptDPA2, + DescrptHybrid, + DescrptSeA, + DescrptSeR, + DescrptSeT, +) + +from ...common.cases.descriptor.descriptor import ( + DescriptorTest, +) +from ..backend import ( + DPTestCase, +) + + +class TestDescriptorSeADP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeA(**self.input_dict) + + +class TestDescriptorSeRDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeR(**self.input_dict) + + +class TestDescriptorSeTDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeT(**self.input_dict) + + +class TestDescriptorDPA1DP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptDPA1(**self.input_dict) + + +class TestDescriptorDPA2DP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "repinit": { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix, + }, + "repformer": { + "rcut": self.rcut / 2, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix[0] // 2, + }, + } + self.module = DescrptDPA2(**self.input_dict) + + +class TestDescriptorHybridDP(unittest.TestCase, DescriptorTest, DPTestCase): + def setUp(self): + DescriptorTest.setUp(self) + ddsub0 = { + "type": "se_e2_a", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + ddsub1 = { + "type": "dpa1", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel_mix, + } + self.input_dict = { + "list": [ddsub0, ddsub1], + } + self.module = DescrptHybrid(**self.input_dict) diff --git a/source/tests/universal/pt/descriptor/__init__.py b/source/tests/universal/pt/descriptor/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/universal/pt/descriptor/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/universal/pt/descriptor/test_descriptor.py b/source/tests/universal/pt/descriptor/test_descriptor.py new file mode 100644 index 0000000000..87107a2f90 --- /dev/null +++ b/source/tests/universal/pt/descriptor/test_descriptor.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.pt.model.descriptor import ( + DescrptDPA1, + DescrptDPA2, + DescrptHybrid, + DescrptSeA, + DescrptSeR, + DescrptSeT, +) + +from ...common.cases.descriptor.descriptor import ( + DescriptorTest, +) +from ..backend import ( + PTTestCase, +) + + +class TestDescriptorSeAPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeA(**self.input_dict) + + +class TestDescriptorSeRPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeR(**self.input_dict) + + +class TestDescriptorSeTPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptSeT(**self.input_dict) + + +class TestDescriptorDPA1PT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.module = DescrptDPA1(**self.input_dict) + + +class TestDescriptorDPA2PT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + self.input_dict = { + "ntypes": self.nt, + "repinit": { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix, + }, + "repformer": { + "rcut": self.rcut / 2, + "rcut_smth": self.rcut_smth, + "nsel": self.sel_mix[0] // 2, + }, + } + self.module = DescrptDPA2(**self.input_dict) + + +class TestDescriptorHybridPT(unittest.TestCase, DescriptorTest, PTTestCase): + def setUp(self): + DescriptorTest.setUp(self) + ddsub0 = { + "type": "se_e2_a", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + } + ddsub1 = { + "type": "dpa1", + "ntypes": self.nt, + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel_mix, + } + self.input_dict = { + "list": [ddsub0, ddsub1], + } + self.module = DescrptHybrid(**self.input_dict)