From 19c70c202f7dbfd8f3d24fd4b1ae334af021ed11 Mon Sep 17 00:00:00 2001 From: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Date: Sat, 11 May 2024 18:30:41 +0800 Subject: [PATCH] feat: descriptor provide get_rcut_smth and get_env_protection (#3761) ## Summary by CodeRabbit - **New Features** - Introduced methods to enhance the calculation of smooth decay radius and environment matrix protection across various descriptor modules. This will improve the precision and reliability of environmental interactions in models. --------- Signed-off-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> Co-authored-by: Han Wang Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- deepmd/dpmodel/descriptor/dpa1.py | 16 +++++++ deepmd/dpmodel/descriptor/dpa2.py | 8 ++++ deepmd/dpmodel/descriptor/hybrid.py | 17 +++++++ .../descriptor/make_base_descriptor.py | 10 +++++ deepmd/dpmodel/descriptor/se_e2_a.py | 8 ++++ deepmd/dpmodel/descriptor/se_r.py | 8 ++++ deepmd/pt/model/descriptor/descriptor.py | 10 +++++ deepmd/pt/model/descriptor/dpa1.py | 8 ++++ deepmd/pt/model/descriptor/dpa2.py | 10 +++++ deepmd/pt/model/descriptor/hybrid.py | 17 +++++++ deepmd/pt/model/descriptor/repformers.py | 8 ++++ deepmd/pt/model/descriptor/se_a.py | 16 +++++++ deepmd/pt/model/descriptor/se_atten.py | 8 ++++ deepmd/pt/model/descriptor/se_r.py | 8 ++++ deepmd/pt/utils/env_mat_stat.py | 4 +- .../common/dpmodel/test_descriptor_hybrid.py | 45 ++++++++++++++++++- .../tests/pt/model/test_descriptor_hybrid.py | 34 ++++++++++++++ 17 files changed, 232 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index e5e79a8984..14f4851023 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -302,6 +302,10 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_atten.get_rcut() + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.se_atten.get_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return self.se_atten.get_nsel() @@ -336,6 +340,10 @@ def mixed_types(self) -> bool: """ return self.se_atten.mixed_types() + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.se_atten.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. @@ -635,6 +643,10 @@ 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) @@ -687,6 +699,10 @@ def mixed_types(self) -> bool: """ return True + 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.""" diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index b76530cf2f..e9c3033945 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -369,6 +369,10 @@ 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) @@ -404,6 +408,10 @@ def mixed_types(self) -> bool: """ return True + 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. diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 96640d75c8..15825ecc10 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import math from typing import ( Any, Dict, @@ -96,6 +97,12 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return np.max([descrpt.get_rcut() for descrpt in self.descrpt_list]).item() + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + # may not be a good idea... + # Note: Using the minimum rcut_smth might not be appropriate in all scenarios. Consider using a different approach or provide detailed documentation on why the minimum value is chosen. + return np.min([descrpt.get_rcut_smth() for descrpt in self.descrpt_list]).item() + def get_sel(self) -> List[int]: """Returns the number of selected atoms for each type.""" if self.mixed_types(): @@ -127,6 +134,16 @@ def mixed_types(self): """ return any(descrpt.mixed_types() for descrpt in self.descrpt_list) + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix. All descriptors should be the same.""" + all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list] + same_as_0 = [math.isclose(ii, all_protection[0]) for ii in all_protection] + if not all(same_as_0): + raise ValueError( + "Hybrid descriptor requires the same environment matrix protection for all descriptors. Found differing values." + ) + return all_protection[0] + 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. diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index a984c96f33..cba9eebe4b 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -51,6 +51,11 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" pass + @abstractmethod + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + pass + @abstractmethod def get_sel(self) -> List[int]: """Returns the number of selected neighboring atoms for each type.""" @@ -86,6 +91,11 @@ def mixed_types(self) -> bool: """ pass + @abstractmethod + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + pass + @abstractmethod def share_params(self, base_class, shared_level, resume=False): """ diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index c50fdb4cb3..adc1913e96 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -239,6 +239,10 @@ 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 @@ -249,6 +253,10 @@ def mixed_types(self): """ 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. diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 6b50c3ba68..ad802d5b25 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -195,6 +195,10 @@ 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 @@ -205,6 +209,10 @@ def mixed_types(self): """ 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. diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index d586fc988f..5e0cdac72b 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -58,6 +58,11 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" pass + @abstractmethod + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + pass + @abstractmethod def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" @@ -88,6 +93,11 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension.""" pass + @abstractmethod + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + pass + def compute_input_stats( self, merged: Union[Callable[[], List[dict]], List[dict]], diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index e955fab048..6a3a947cb1 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -282,6 +282,10 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_atten.get_rcut() + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.se_atten.get_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return self.se_atten.get_nsel() @@ -316,6 +320,10 @@ def mixed_types(self) -> bool: """ return self.se_atten.mixed_types() + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.se_atten.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. diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 28ff1b6848..baa1d10c41 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -363,6 +363,7 @@ def __init__( self.tebd_dim = repinit_tebd_dim self.rcut = self.repinit.get_rcut() + self.rcut_smth = self.repinit.get_rcut_smth() self.ntypes = ntypes self.sel = self.repinit.sel # set trainable @@ -373,6 +374,10 @@ 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) @@ -408,6 +413,11 @@ def mixed_types(self) -> bool: """ return True + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + # the env_protection of repinit is the same as that of the repformer + return self.repinit.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. diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index a5353b67d2..bce52fd8ca 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import math from typing import ( Any, Dict, @@ -101,6 +102,12 @@ def get_rcut(self) -> float: # do not use numpy here - jit is not happy return max([descrpt.get_rcut() for descrpt in self.descrpt_list]) + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + # may not be a good idea... + # Note: Using the minimum rcut_smth might not be appropriate in all scenarios. Consider using a different approach or provide detailed documentation on why the minimum value is chosen. + return min([descrpt.get_rcut_smth() for descrpt in self.descrpt_list]) + def get_sel(self) -> List[int]: """Returns the number of selected atoms for each type.""" if self.mixed_types(): @@ -132,6 +139,16 @@ def mixed_types(self): """ return any(descrpt.mixed_types() for descrpt in self.descrpt_list) + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix. All descriptors should be the same.""" + all_protection = [descrpt.get_env_protection() for descrpt in self.descrpt_list] + same_as_0 = [math.isclose(ii, all_protection[0]) for ii in all_protection] + if not all(same_as_0): + raise ValueError( + "Hybrid descriptor requires the same environment matrix protection for all descriptors. Found differing values." + ) + return all_protection[0] + 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. diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 2d6a61d264..b8adc0d71e 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -306,6 +306,10 @@ 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) @@ -358,6 +362,10 @@ def mixed_types(self) -> bool: """ return True + 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.""" diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 3316ed5de7..50393e8a03 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -106,6 +106,10 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.sea.get_rcut() + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.sea.get_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return self.sea.get_nsel() @@ -132,6 +136,10 @@ def mixed_types(self): """ return self.sea.mixed_types() + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.sea.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. @@ -400,6 +408,10 @@ 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) @@ -436,6 +448,10 @@ def mixed_types(self) -> bool: """ 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.""" diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 5de0aeffab..d87ad76e3c 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -286,6 +286,10 @@ 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) @@ -338,6 +342,10 @@ def mixed_types(self) -> bool: """ return True + 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.""" diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index ff922e0649..4ac510e2a7 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -124,6 +124,10 @@ 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) @@ -160,6 +164,10 @@ def mixed_types(self) -> bool: """ 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. diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index a853b1722a..a3279f1727 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -142,9 +142,9 @@ def iter( one_stddev, self.descriptor.get_rcut(), # TODO: export rcut_smth from DescriptorBlock - self.descriptor.rcut_smth, + self.descriptor.get_rcut_smth(), radial_only, - protection=self.descriptor.env_protection, + protection=self.descriptor.get_env_protection(), ) # apply excluded_types exclude_mask = self.descriptor.emask(nlist, extended_atype) diff --git a/source/tests/common/dpmodel/test_descriptor_hybrid.py b/source/tests/common/dpmodel/test_descriptor_hybrid.py index da29fe12a1..668c87b1c6 100644 --- a/source/tests/common/dpmodel/test_descriptor_hybrid.py +++ b/source/tests/common/dpmodel/test_descriptor_hybrid.py @@ -26,6 +26,49 @@ def setUp(self): unittest.TestCase.setUp(self) TestCaseSingleFrameWithNlist.setUp(self) + def test_get_parameters( + 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) + ddsub0 = DescrptSeA( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + sel=self.sel, + ) + ddsub0.davg = davg + ddsub0.dstd = dstd + ddsub1 = DescrptDPA1( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + sel=np.sum(self.sel).item() - 1, + ntypes=len(self.sel), + ) + ddsub1.davg = davg[:, :6] + ddsub1.dstd = dstd[:, :6] + ddsub2 = DescrptSeR( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth - 0.1, + sel=[3, 1], + ) + ddsub2.davg = davg[:, :4, :1] + ddsub2.dstd = dstd[:, :4, :1] + em0 = DescrptHybrid(list=[ddsub0, ddsub1, ddsub2]) + self.assertAlmostEqual(em0.get_env_protection(), 0.0) + self.assertAlmostEqual(em0.get_rcut_smth(), self.rcut_smth - 0.1) + ddsub3 = DescrptSeR( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth - 0.1, + sel=[3, 1], + env_protection=0.1, + ) + em0 = DescrptHybrid(list=[ddsub0, ddsub1, ddsub3]) + with self.assertRaises(ValueError): + self.assertAlmostEqual(em0.get_env_protection(), 0.0) + def test_self_consistency( self, ): @@ -52,7 +95,7 @@ def test_self_consistency( ddsub1.dstd = dstd[:, :6] ddsub2 = DescrptSeR( rcut=self.rcut / 2, - rcut_smth=self.rcut_smth, + rcut_smth=self.rcut_smth / 2, sel=[3, 1], ) ddsub2.davg = davg[:, :4, :1] diff --git a/source/tests/pt/model/test_descriptor_hybrid.py b/source/tests/pt/model/test_descriptor_hybrid.py index 6742388bd9..f8a2b8d99f 100644 --- a/source/tests/pt/model/test_descriptor_hybrid.py +++ b/source/tests/pt/model/test_descriptor_hybrid.py @@ -53,6 +53,40 @@ def test_jit( dd0 = torch.jit.script(dd0) dd1 = torch.jit.script(dd1) + def test_get_parameters( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + ddsub0 = DescrptSeA( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + sel=self.sel, + ) + ddsub1 = DescrptDPA1( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + sel=np.sum(self.sel).item() - 1, + ntypes=len(self.sel), + ) + ddsub2 = DescrptSeR( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth - 0.1, + sel=[3, 1], + ) + em0 = DescrptHybrid(list=[ddsub0, ddsub1, ddsub2]) + self.assertAlmostEqual(em0.get_env_protection(), 0.0) + self.assertAlmostEqual(em0.get_rcut_smth(), self.rcut_smth - 0.1) + ddsub3 = DescrptSeR( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth - 0.1, + sel=[3, 1], + env_protection=0.1, + ) + em0 = DescrptHybrid(list=[ddsub0, ddsub1, ddsub3]) + with self.assertRaises(ValueError): + self.assertAlmostEqual(em0.get_env_protection(), 0.0) + def test_hybrid_mixed_and_no_mixed(self): coord_ext = to_torch_tensor(self.coord_ext) atype_ext = to_torch_tensor(self.atype_ext)