diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 4dd590e04d..668ef36043 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -121,6 +121,9 @@ class RepFlowArgs: optim_update : bool, optional Whether to enable the optimized update method. Uses a more efficient process when enabled. Defaults to True + smooth_edge_update : bool, optional + Whether to make edge update smooth. + If True, the edge update from angle message will not use self as padding. """ def __init__( @@ -147,6 +150,7 @@ def __init__( fix_stat_std: float = 0.3, skip_stat: bool = False, optim_update: bool = True, + smooth_edge_update: bool = False, ) -> None: self.n_dim = n_dim self.e_dim = e_dim @@ -172,6 +176,7 @@ def __init__( self.a_compress_e_rate = a_compress_e_rate self.a_compress_use_split = a_compress_use_split self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update def __getitem__(self, key): if hasattr(self, key): @@ -202,6 +207,7 @@ def serialize(self) -> dict: "update_residual_init": self.update_residual_init, "fix_stat_std": self.fix_stat_std, "optim_update": self.optim_update, + "smooth_edge_update": self.smooth_edge_update, } @classmethod @@ -297,6 +303,7 @@ def init_subclass_params(sub_data, sub_class): update_residual_init=self.repflow_args.update_residual_init, fix_stat_std=self.repflow_args.fix_stat_std, optim_update=self.repflow_args.optim_update, + smooth_edge_update=self.repflow_args.smooth_edge_update, exclude_types=exclude_types, env_protection=env_protection, precision=precision, diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 469b6c008f..e667d25c90 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -114,6 +114,9 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock): optim_update : bool, optional Whether to enable the optimized update method. Uses a more efficient process when enabled. Defaults to True + smooth_edge_update : bool, optional + Whether to make edge update smooth. + If True, the edge update from angle message will not use self as padding. ntypes : int Number of element types activation_function : str, optional @@ -161,6 +164,7 @@ def __init__( precision: str = "float64", fix_stat_std: float = 0.3, optim_update: bool = True, + smooth_edge_update: bool = False, seed: Optional[Union[int, list[int]]] = None, ) -> None: super().__init__() @@ -191,6 +195,7 @@ def __init__( self.set_stddev_constant = fix_stat_std != 0.0 self.a_compress_use_split = a_compress_use_split self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update self.n_dim = n_dim self.e_dim = e_dim @@ -243,6 +248,7 @@ def __init__( update_residual_init=self.update_residual_init, precision=precision, optim_update=self.optim_update, + smooth_edge_update=self.smooth_edge_update, seed=child_seed(child_seed(seed, 1), ii), ) ) @@ -563,6 +569,7 @@ def __init__( axis_neuron: int = 4, update_angle: bool = True, optim_update: bool = True, + smooth_edge_update: bool = False, activation_function: str = "silu", update_style: str = "res_residual", update_residual: float = 0.1, @@ -607,6 +614,7 @@ def __init__( self.seed = seed self.prec = PRECISION_DICT[precision] self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update assert update_residual_init in [ "norm", @@ -1136,19 +1144,23 @@ def call( ], axis=2, ) - full_mask = xp.concat( - [ - a_nlist_mask, - xp.zeros( - (nb, nloc, self.nnei - self.a_sel), - dtype=a_nlist_mask.dtype, - ), - ], - axis=-1, - ) - padding_edge_angle_update = xp.where( - xp.expand_dims(full_mask, axis=-1), padding_edge_angle_update, edge_ebd - ) + if not self.smooth_edge_update: + # will be deprecated in the future + full_mask = xp.concat( + [ + a_nlist_mask, + xp.zeros( + (nb, nloc, self.nnei - self.a_sel), + dtype=a_nlist_mask.dtype, + ), + ], + axis=-1, + ) + padding_edge_angle_update = xp.where( + xp.expand_dims(full_mask, axis=-1), + padding_edge_angle_update, + edge_ebd, + ) e_update_list.append( self.act(self.edge_angle_linear2(padding_edge_angle_update)) ) @@ -1235,7 +1247,7 @@ def serialize(self) -> dict: The serialized networks. """ data = { - "@class": "RepformerLayer", + "@class": "RepFlowLayer", "@version": 1, "e_rcut": self.e_rcut, "e_rcut_smth": self.e_rcut_smth, @@ -1259,6 +1271,7 @@ def serialize(self) -> dict: "update_residual_init": self.update_residual_init, "precision": self.precision, "optim_update": self.optim_update, + "smooth_edge_update": self.smooth_edge_update, "node_self_mlp": self.node_self_mlp.serialize(), "node_sym_linear": self.node_sym_linear.serialize(), "node_edge_linear": self.node_edge_linear.serialize(), diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index 8cdd3e15fe..545da962e7 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -149,6 +149,7 @@ def init_subclass_params(sub_data, sub_class): update_residual_init=self.repflow_args.update_residual_init, fix_stat_std=self.repflow_args.fix_stat_std, optim_update=self.repflow_args.optim_update, + smooth_edge_update=self.repflow_args.smooth_edge_update, exclude_types=exclude_types, env_protection=env_protection, precision=precision, diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 43cae8c746..9ec1bd4dee 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -52,6 +52,7 @@ def __init__( axis_neuron: int = 4, update_angle: bool = True, optim_update: bool = True, + smooth_edge_update: bool = False, activation_function: str = "silu", update_style: str = "res_residual", update_residual: float = 0.1, @@ -96,6 +97,7 @@ def __init__( self.seed = seed self.prec = PRECISION_DICT[precision] self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update assert update_residual_init in [ "norm", @@ -718,20 +720,22 @@ def forward( ], dim=2, ) - full_mask = torch.concat( - [ - a_nlist_mask, - torch.zeros( - [nb, nloc, self.nnei - self.a_sel], - dtype=a_nlist_mask.dtype, - device=a_nlist_mask.device, - ), - ], - dim=-1, - ) - padding_edge_angle_update = torch.where( - full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd - ) + if not self.smooth_edge_update: + # will be deprecated in the future + full_mask = torch.concat( + [ + a_nlist_mask, + torch.zeros( + [nb, nloc, self.nnei - self.a_sel], + dtype=a_nlist_mask.dtype, + device=a_nlist_mask.device, + ), + ], + dim=-1, + ) + padding_edge_angle_update = torch.where( + full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd + ) e_update_list.append( self.act(self.edge_angle_linear2(padding_edge_angle_update)) ) @@ -823,7 +827,7 @@ def serialize(self) -> dict: The serialized networks. """ data = { - "@class": "RepformerLayer", + "@class": "RepFlowLayer", "@version": 1, "e_rcut": self.e_rcut, "e_rcut_smth": self.e_rcut_smth, @@ -847,6 +851,7 @@ def serialize(self) -> dict: "update_residual_init": self.update_residual_init, "precision": self.precision, "optim_update": self.optim_update, + "smooth_edge_update": self.smooth_edge_update, "node_self_mlp": self.node_self_mlp.serialize(), "node_sym_linear": self.node_sym_linear.serialize(), "node_edge_linear": self.node_edge_linear.serialize(), diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py index 0fe17e64db..330336b1de 100644 --- a/deepmd/pt/model/descriptor/repflows.py +++ b/deepmd/pt/model/descriptor/repflows.py @@ -130,6 +130,9 @@ class DescrptBlockRepflows(DescriptorBlock): fix_stat_std : float, optional If non-zero (default is 0.3), use this constant as the normalization standard deviation instead of computing it from data statistics. + smooth_edge_update : bool, optional + Whether to make edge update smooth. + If True, the edge update from angle message will not use self as padding. optim_update : bool, optional Whether to enable the optimized update method. Uses a more efficient process when enabled. Defaults to True @@ -179,6 +182,7 @@ def __init__( env_protection: float = 0.0, precision: str = "float64", fix_stat_std: float = 0.3, + smooth_edge_update: bool = False, optim_update: bool = True, seed: Optional[Union[int, list[int]]] = None, ) -> None: @@ -210,6 +214,7 @@ def __init__( self.set_stddev_constant = fix_stat_std != 0.0 self.a_compress_use_split = a_compress_use_split self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update self.n_dim = n_dim self.e_dim = e_dim @@ -262,6 +267,7 @@ def __init__( update_residual_init=self.update_residual_init, precision=precision, optim_update=self.optim_update, + smooth_edge_update=self.smooth_edge_update, seed=child_seed(child_seed(seed, 1), ii), ) ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a1afcaf1d0..8d03a2d4d1 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1493,6 +1493,10 @@ def dpa3_repflow_args(): "Whether to enable the optimized update method. " "Uses a more efficient process when enabled. Defaults to True" ) + doc_smooth_edge_update = ( + "Whether to make edge update smooth. " + "If True, the edge update from angle message will not use self as padding." + ) return [ # repflow args @@ -1586,6 +1590,13 @@ def dpa3_repflow_args(): default=True, doc=doc_optim_update, ), + Argument( + "smooth_edge_update", + bool, + optional=True, + default=False, # For compatability. This will be True in the future + doc=doc_smooth_edge_update, + ), ] diff --git a/source/tests/pt/model/test_dpa3.py b/source/tests/pt/model/test_dpa3.py index 240ffad925..c2a53ba18d 100644 --- a/source/tests/pt/model/test_dpa3.py +++ b/source/tests/pt/model/test_dpa3.py @@ -90,6 +90,7 @@ def test_consistency( update_angle=ua, update_style=rus, update_residual_init=ruri, + smooth_edge_update=True, ) # dpa3 new impl @@ -190,6 +191,7 @@ def test_jit( update_angle=ua, update_style=rus, update_residual_init=ruri, + smooth_edge_update=True, ) # dpa3 new impl diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index 0354336e37..8bb6a72c1f 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -177,10 +177,10 @@ "a_dim": 8, "nlayers": 6, "e_rcut": 6.0, - "e_rcut_smth": 5.0, + "e_rcut_smth": 3.0, "e_sel": 20, "a_rcut": 4.0, - "a_rcut_smth": 3.5, + "a_rcut_smth": 2.0, "a_sel": 10, "axis_neuron": 4, "a_compress_rate": 1, @@ -190,8 +190,9 @@ "update_style": "res_residual", "update_residual": 0.1, "update_residual_init": "const", + "smooth_edge_update": True, }, - "activation_function": "silu", + "activation_function": "silut:10.0", "use_tebd_bias": False, "precision": "float32", "concat_output_tebd": False, @@ -200,7 +201,7 @@ "neuron": [24, 24], "resnet_dt": True, "precision": "float32", - "activation_function": "silu", + "activation_function": "silut:10.0", "seed": 1, }, } diff --git a/source/tests/pt/model/test_smooth.py b/source/tests/pt/model/test_smooth.py index 1c6303d14c..34479c5064 100644 --- a/source/tests/pt/model/test_smooth.py +++ b/source/tests/pt/model/test_smooth.py @@ -21,6 +21,7 @@ model_dos, model_dpa1, model_dpa2, + model_dpa3, model_hybrid, model_se_e2_a, model_spin, @@ -59,10 +60,16 @@ def test( 0.0, 4.0 - 0.5 * epsilon, 0.0, + 6.0 - 0.5 * epsilon, + 0.0, + 0.0, + 0.0, + 6.0 - 0.5 * epsilon, + 0.0, ], dtype=dtype, device=env.DEVICE, - ).view([-1, 3]) + ).view([-1, 3]) # to test descriptors with two rcuts, e.g. DPA2/3 coord1 = torch.rand( [natoms - coord0.shape[0], 3], dtype=dtype, @@ -77,11 +84,15 @@ def test( coord0 = torch.clone(coord) coord1 = torch.clone(coord) coord1[1][0] += epsilon + coord1[3][0] += epsilon coord2 = torch.clone(coord) coord2[2][1] += epsilon + coord2[4][1] += epsilon coord3 = torch.clone(coord) coord3[1][0] += epsilon + coord1[3][0] += epsilon coord3[2][1] += epsilon + coord2[4][1] += epsilon test_spin = getattr(self, "test_spin", False) if not test_spin: test_keys = ["energy", "force", "virial"] @@ -226,6 +237,17 @@ def setUp(self) -> None: self.epsilon, self.aprec = None, None +class TestEnergyModelDPA3(unittest.TestCase, SmoothTest): + def setUp(self) -> None: + model_params = copy.deepcopy(model_dpa3) + self.type_split = True + self.model = get_model(model_params).to(env.DEVICE) + # less degree of smoothness, + # error can be systematically removed by reducing epsilon + self.epsilon = 1e-5 + self.aprec = 1e-5 + + class TestEnergyModelHybrid(unittest.TestCase, SmoothTest): def setUp(self) -> None: model_params = copy.deepcopy(model_hybrid) diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 8fe6a131ef..5a4c64c803 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -723,8 +723,14 @@ def test_smooth(self) -> None: 0.0, self.expected_rcut - 0.5 * epsilon, 0.0, + self.expected_rcut / 2 - 0.5 * epsilon, + 0.0, + 0.0, + 0.0, + self.expected_rcut / 2 - 0.5 * epsilon, + 0.0, ] - ).reshape(-1, 3) + ).reshape(-1, 3) # to test descriptors with two rcuts, e.g. DPA2/3 coord1 = rng.random([natoms - coord0.shape[0], 3]) coord1 = np.matmul(coord1, cell) coord = np.concatenate([coord0, coord1], axis=0) @@ -732,11 +738,15 @@ def test_smooth(self) -> None: coord0 = deepcopy(coord) coord1 = deepcopy(coord) coord1[1][0] += epsilon + coord1[3][0] += epsilon coord2 = deepcopy(coord) coord2[2][1] += epsilon + coord2[4][1] += epsilon coord3 = deepcopy(coord) coord3[1][0] += epsilon + coord1[3][0] += epsilon coord3[2][1] += epsilon + coord2[4][1] += epsilon # reshape for input coord0 = coord0.reshape([nf, -1]) diff --git a/source/tests/universal/dpmodel/descriptor/test_descriptor.py b/source/tests/universal/dpmodel/descriptor/test_descriptor.py index 6af07f1e44..4fa0593419 100644 --- a/source/tests/universal/dpmodel/descriptor/test_descriptor.py +++ b/source/tests/universal/dpmodel/descriptor/test_descriptor.py @@ -481,6 +481,7 @@ def DescriptorParamDPA3( a_compress_e_rate=1, a_compress_use_split=False, optim_update=True, + smooth_edge_update=False, fix_stat_std=0.3, precision="float64", ): @@ -502,6 +503,7 @@ def DescriptorParamDPA3( "a_compress_e_rate": a_compress_e_rate, "a_compress_use_split": a_compress_use_split, "optim_update": optim_update, + "smooth_edge_update": smooth_edge_update, "fix_stat_std": fix_stat_std, "n_multi_edge_message": n_multi_edge_message, "axis_neuron": 2, @@ -537,6 +539,7 @@ def DescriptorParamDPA3( "a_compress_e_rate": (2,), "a_compress_use_split": (True, False), "optim_update": (True, False), + "smooth_edge_update": (True,), "fix_stat_std": (0.3,), "n_multi_edge_message": (1, 2), "env_protection": (0.0, 1e-8), diff --git a/source/tests/universal/dpmodel/model/test_model.py b/source/tests/universal/dpmodel/model/test_model.py index a07a29dcf5..a67faaa769 100644 --- a/source/tests/universal/dpmodel/model/test_model.py +++ b/source/tests/universal/dpmodel/model/test_model.py @@ -129,8 +129,7 @@ def setUpClass(cls) -> None: cls.epsilon_dict["test_smooth"] = 1e-6 if Descrpt in [DescrptSeT, DescrptSeTTebd]: # computational expensive - cls.expected_sel = [i // 4 for i in cls.expected_sel] - cls.expected_rcut = cls.expected_rcut / 2 + cls.expected_sel = [i // 2 for i in cls.expected_sel] cls.input_dict_ds = DescriptorParam( len(cls.expected_type_map), cls.expected_rcut, @@ -216,8 +215,7 @@ def setUpClass(cls) -> None: cls.epsilon_dict["test_smooth"] = 1e-8 if Descrpt in [DescrptSeT, DescrptSeTTebd]: # computational expensive - cls.expected_sel = [i // 4 for i in cls.expected_sel] - cls.expected_rcut = cls.expected_rcut / 2 + cls.expected_sel = [i // 2 for i in cls.expected_sel] spin = Spin( use_spin=cls.spin_dict["use_spin"], @@ -229,6 +227,7 @@ def setUpClass(cls) -> None: if Descrpt in [DescrptSeA, DescrptSeR, DescrptSeT]: spin_sel = cls.expected_sel + cls.expected_sel else: + cls.expected_sel = [i * 2 for i in cls.expected_sel] spin_sel = cls.expected_sel pair_exclude_types = spin.get_pair_exclude_types() atom_exclude_types = spin.get_atom_exclude_types() diff --git a/source/tests/universal/pt/model/test_model.py b/source/tests/universal/pt/model/test_model.py index 867fa48b87..e4fa6c5960 100644 --- a/source/tests/universal/pt/model/test_model.py +++ b/source/tests/universal/pt/model/test_model.py @@ -167,8 +167,7 @@ def setUpClass(cls) -> None: cls.epsilon_dict["test_smooth"] = 1e-8 if Descrpt in [DescrptSeT, DescrptSeTTebd]: # computational expensive - cls.expected_sel = [i // 4 for i in cls.expected_sel] - cls.expected_rcut = cls.expected_rcut / 2 + cls.expected_sel = [i // 2 for i in cls.expected_sel] cls.input_dict_ds = DescriptorParam( len(cls.expected_type_map), cls.expected_rcut, @@ -272,8 +271,7 @@ def setUpClass(cls) -> None: cls.epsilon_dict["test_smooth"] = 1e-8 if Descrpt in [DescrptSeT, DescrptSeTTebd]: # computational expensive - cls.expected_sel = [i // 4 for i in cls.expected_sel] - cls.expected_rcut = cls.expected_rcut / 2 + cls.expected_sel = [i // 2 for i in cls.expected_sel] cls.input_dict_ds = DescriptorParam( len(cls.expected_type_map), cls.expected_rcut, @@ -663,8 +661,7 @@ def setUpClass(cls) -> None: cls.epsilon_dict["test_smooth"] = 1e-8 if Descrpt in [DescrptSeT, DescrptSeTTebd]: # computational expensive - cls.expected_sel = [i // 4 for i in cls.expected_sel] - cls.expected_rcut = cls.expected_rcut / 2 + cls.expected_sel = [i // 2 for i in cls.expected_sel] spin = Spin( use_spin=cls.spin_dict["use_spin"], @@ -676,6 +673,7 @@ def setUpClass(cls) -> None: if Descrpt in [DescrptSeA, DescrptSeR, DescrptSeT]: spin_sel = cls.expected_sel + cls.expected_sel else: + cls.expected_sel = [i * 2 for i in cls.expected_sel] spin_sel = cls.expected_sel pair_exclude_types = spin.get_pair_exclude_types() atom_exclude_types = spin.get_atom_exclude_types()