From fe9f6262de11c7664b0070249f57b4cdcb610184 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 27 May 2026 03:30:18 +0000 Subject: [PATCH 1/4] Sync Paddle DPA3 with dpmodel: add default_chg_spin, rename fparam to charge_spin Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com> --- deepmd/pd/model/descriptor/dpa3.py | 27 +++++++++++++++++++++++---- 1 file changed, 23 insertions(+), 4 deletions(-) diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index bed2d45022..095ba0fe5b 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -122,6 +122,7 @@ def __init__( use_loc_mapping: bool = True, type_map: list[str] | None = None, add_chg_spin_ebd: bool = False, + default_chg_spin: list[float] | None = None, ) -> None: super().__init__() @@ -177,6 +178,11 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: self.use_econf_tebd = use_econf_tebd self.add_chg_spin_ebd = add_chg_spin_ebd + if default_chg_spin is not None and len(default_chg_spin) != 2: + raise ValueError( + "default_chg_spin must have exactly 2 values [charge, spin]" + ) + self.default_chg_spin = default_chg_spin self.use_loc_mapping = use_loc_mapping self.use_tebd_bias = use_tebd_bias self.type_map = type_map @@ -447,6 +453,18 @@ def get_stat_mean_and_stddev( stddev_list = [self.repflows.stddev] return mean_list, stddev_list + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input.""" + return 2 if self.add_chg_spin_ebd else 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether default charge_spin values are set.""" + return self.default_chg_spin is not None + + def get_default_chg_spin(self) -> list[float] | None: + """Returns the default charge_spin values.""" + return self.default_chg_spin + def serialize(self) -> dict: repflows = self.repflows data = { @@ -465,6 +483,7 @@ def serialize(self) -> dict: "use_tebd_bias": self.use_tebd_bias, "use_loc_mapping": self.use_loc_mapping, "add_chg_spin_ebd": self.add_chg_spin_ebd, + "default_chg_spin": self.default_chg_spin, "type_map": self.type_map, "type_embedding": self.type_embedding.embedding.serialize(), } @@ -541,7 +560,7 @@ def forward( nlist: paddle.Tensor, mapping: paddle.Tensor | None = None, comm_dict: list[paddle.Tensor] | None = None, - fparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, paddle.Tensor | None, @@ -593,11 +612,11 @@ def forward( node_ebd_ext = self.type_embedding(extended_atype) if self.add_chg_spin_ebd: - assert fparam is not None + assert charge_spin is not None assert self.chg_embedding is not None assert self.spin_embedding is not None - charge = fparam[:, 0].to(dtype=paddle.int64) + 100 - spin = fparam[:, 1].to(dtype=paddle.int64) + charge = charge_spin[:, 0].to(dtype=paddle.int64) + 100 + spin = charge_spin[:, 1].to(dtype=paddle.int64) chg_ebd = self.chg_embedding(charge) spin_ebd = self.spin_embedding(spin) sys_cs_embd = self.act( From 0e792aca3c3ff082aa1e5c2b588bfcb75719dd59 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 27 May 2026 03:50:50 +0000 Subject: [PATCH 2/4] fix(pd): add charge_spin stub methods and param to dpa1/dpa2/se_a/se_t_tebd descriptors Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com> --- deepmd/pd/model/descriptor/dpa1.py | 13 +++++++++++++ deepmd/pd/model/descriptor/dpa2.py | 16 +++++++++++++--- deepmd/pd/model/descriptor/se_a.py | 13 +++++++++++++ deepmd/pd/model/descriptor/se_t_tebd.py | 13 +++++++++++++ 4 files changed, 52 insertions(+), 3 deletions(-) diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index d6fb1e3d33..ccd8944cc8 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -366,6 +366,18 @@ def get_buffer_type_map(self) -> paddle.Tensor: """ return self.buffer_type_map + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_dim_out(self) -> int: """Returns the output dimension.""" ret = self.se_atten.get_dim_out() @@ -627,6 +639,7 @@ def forward( mapping: paddle.Tensor | None = None, comm_dict: list[paddle.Tensor] | None = None, fparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, paddle.Tensor | None, diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index cf5be0c41d..27fa3158a2 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -333,6 +333,18 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: param.stop_gradient = not trainable self.compress = False + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.rcut @@ -734,10 +746,8 @@ def forward( mapping: paddle.Tensor | None = None, comm_dict: list[paddle.Tensor] | None = None, fparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> tuple[ - paddle.Tensor, - paddle.Tensor | None, - paddle.Tensor | None, paddle.Tensor | None, paddle.Tensor | None, ]: diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index a4722210a4..0098ca2186 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -119,6 +119,18 @@ def __init__( seed=seed, ) + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.sea.get_rcut() @@ -289,6 +301,7 @@ def forward( mapping: paddle.Tensor | None = None, comm_dict: list[paddle.Tensor] | None = None, fparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, paddle.Tensor | None, diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index 8f2c79f795..0c16eb0ef0 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -190,6 +190,18 @@ def __init__( for param in self.parameters(): param.stop_gradient = not trainable + def get_dim_chg_spin(self) -> int: + """Returns the dimension of charge_spin input (0 if not supported).""" + return 0 + + def has_default_chg_spin(self) -> bool: + """Returns whether the descriptor has a default charge_spin value.""" + return False + + def get_default_chg_spin(self) -> None: + """Returns the default charge_spin value, or None.""" + return None + def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_ttebd.get_rcut() @@ -438,6 +450,7 @@ def forward( mapping: paddle.Tensor | None = None, comm_dict: list[paddle.Tensor] | None = None, fparam: paddle.Tensor | None = None, + charge_spin: paddle.Tensor | None = None, ) -> paddle.Tensor: """Compute the descriptor. From db26f5e6b2a22525ec2c00ad9c8ada7016960671 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 27 May 2026 05:51:44 +0000 Subject: [PATCH 3/4] fix(pd): update dp_atomic_model to call descriptor forward with charge_spin keyword Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com> --- deepmd/pd/model/atomic_model/dp_atomic_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 0c7ba9b56c..0b545ef7b6 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -332,7 +332,7 @@ def forward_atomic( nlist, mapping=mapping, comm_dict=comm_dict, - fparam=fparam if self.add_chg_spin_ebd else None, + charge_spin=fparam if self.add_chg_spin_ebd else None, ) assert descriptor is not None if self.enable_eval_descriptor_hook: From 6ef79f2b8349b01a8e0a27ed58c80ebc2688f39b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 1 Jun 2026 06:19:35 +0000 Subject: [PATCH 4/4] test(pd): sync unit tests with charge_spin interface changes Co-authored-by: HydrogenSulfate <23737287+HydrogenSulfate@users.noreply.github.com> --- source/tests/consistent/descriptor/common.py | 14 +++++++++++-- .../tests/consistent/descriptor/test_dpa3.py | 21 +------------------ source/tests/pd/model/test_dpa3.py | 20 +++++++++--------- 3 files changed, 23 insertions(+), 32 deletions(-) diff --git a/source/tests/consistent/descriptor/common.py b/source/tests/consistent/descriptor/common.py index 078db4829f..4a09db526b 100644 --- a/source/tests/consistent/descriptor/common.py +++ b/source/tests/consistent/descriptor/common.py @@ -269,13 +269,23 @@ def eval_pd_descriptor( pd_obj.get_sel(), distinguish_types=(not mixed_types), ) + fparam_pd = ( + paddle.to_tensor(fparam).to(PD_DEVICE) if fparam is not None else None + ) + charge_spin_pd = ( + paddle.to_tensor(charge_spin).to(PD_DEVICE) + if charge_spin is not None + else None + ) + kwargs = {"nlist": nlist, "mapping": mapping} + if charge_spin_pd is not None: + kwargs["charge_spin"] = charge_spin_pd return [ x.detach().cpu().numpy() if paddle.is_tensor(x) else x for x in pd_obj( ext_coords, ext_atype, - nlist=nlist, - mapping=mapping, + **kwargs, ) ] diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index 3f30d59435..467afba57f 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -287,26 +287,7 @@ def skip_pt(self) -> bool: @property def skip_pd(self) -> bool: - ( - _update_residual_init, - _exclude_types, - _update_angle, - _a_compress_rate, - _a_compress_e_rate, - _a_compress_use_split, - _optim_update, - _edge_init_use_dist, - _use_exp_switch, - _use_dynamic_sel, - _use_loc_mapping, - _fix_stat_std, - _n_multi_edge_message, - _precision, - add_chg_spin_ebd, - _default_chg_spin, - _sequential_update, - ) = self.param - return True if add_chg_spin_ebd else CommonTest.skip_pd + return CommonTest.skip_pd @property def skip_dp(self) -> bool: diff --git a/source/tests/pd/model/test_dpa3.py b/source/tests/pd/model/test_dpa3.py index a582181085..7b3c8ba18f 100644 --- a/source/tests/pd/model/test_dpa3.py +++ b/source/tests/pd/model/test_dpa3.py @@ -66,7 +66,7 @@ def test_consistency( [1, 2], # n_multi_edge_message ["float64"], # precision [False], # use_econf_tebd - [False], # add_chg_spin_ebd (PD backend does not support charge_spin) + [True, False], # add_chg_spin_ebd ): dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) @@ -111,21 +111,21 @@ def test_consistency( dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype, place=env.DEVICE) dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype, place=env.DEVICE) - # Prepare fparam if needed - fparam = None - fparam_np = None + # Prepare charge_spin if needed + charge_spin = None + charge_spin_np = None if add_chg_spin: - fparam = paddle.to_tensor( + charge_spin = paddle.to_tensor( [[5, 1]], dtype=dtype, place=env.DEVICE - ).expand(nf, -1) - fparam_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) + ).expand([nf, -1]) + charge_spin_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0) rd0, _, _, _, _ = dd0( paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE), paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), - fparam=fparam, + charge_spin=charge_spin, ) # serialization dd1 = DescrptDPA3.deserialize(dd0.serialize()) @@ -134,7 +134,7 @@ def test_consistency( paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), - fparam=fparam, + charge_spin=charge_spin, ) np.testing.assert_allclose( rd0.numpy(), @@ -149,7 +149,7 @@ def test_consistency( self.atype_ext, self.nlist, self.mapping, - fparam=fparam_np, + charge_spin=charge_spin_np, ) np.testing.assert_allclose( rd0.numpy(),