fix(pd): add default chg spin parameter

deepmd/pd/model/atomic_model/base_atomic_model.py

@@ -181,6 +181,22 @@ def reinit_pair_exclude(
         else:
             self.pair_excl = PairExcludeMask(self.get_ntypes(), self.pair_exclude_types)
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        return False
+
+    def get_dim_chg_spin(self) -> int:
+        """Get the dimension of charge_spin input."""
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        return False
+
+    def get_default_chg_spin(self) -> paddle.Tensor | None:
+        """Get the default charge_spin values."""
+        return None
+
     # to make jit happy...
     def make_atom_mask(
         self,
@@ -229,6 +245,7 @@ def forward_common_atomic(
         fparam: paddle.Tensor | None = None,
         aparam: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> dict[str, paddle.Tensor]:
         """Common interface for atomic inference.
 
@@ -252,6 +269,8 @@ def forward_common_atomic(
             atomic parameter, shape: nf x nloc x dim_aparam
         comm_dict
             The data needed for communication for parallel inference.
+        charge_spin
+            charge and spin parameters, shape: nf x 2
 
         Returns
         -------
@@ -282,6 +301,7 @@ def forward_common_atomic(
             fparam=fparam,
             aparam=aparam,
             comm_dict=comm_dict,
+            charge_spin=charge_spin,
         )
         ret_dict = self.apply_out_stat(ret_dict, atype)
         # nf x nloc
@@ -311,6 +331,7 @@ def forward(
         fparam: paddle.Tensor | None = None,
         aparam: paddle.Tensor | None = None,
         comm_dict: list[paddle.Tensor] | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> dict[str, paddle.Tensor]:
         return self.forward_common_atomic(
             extended_coord,
@@ -320,6 +341,7 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             comm_dict=comm_dict,
+            charge_spin=charge_spin,
         )
 
     def change_type_map(

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -298,6 +298,7 @@ def forward_atomic(
         fparam: paddle.Tensor | None = None,
         aparam: paddle.Tensor | None = None,
         comm_dict: dict[str, paddle.Tensor] | None = None,
+        charge_spin: paddle.Tensor | None = None,
     ) -> dict[str, paddle.Tensor]:
         """Return atomic prediction.
 
@@ -315,6 +316,8 @@ def forward_atomic(
             frame parameter. nf x ndf
         aparam
             atomic parameter. nf x nloc x nda
+        charge_spin
+            charge and spin parameters. nf x 2
 
         Returns
         -------
@@ -326,13 +329,21 @@ def forward_atomic(
         atype = extended_atype[:, :nloc]
         if self.do_grad_r() or self.do_grad_c():
             extended_coord.stop_gradient = False
+
+        # Handle default chg_spin if descriptor supports it
+        if self.add_chg_spin_ebd and charge_spin is None:
+            default_cs_tensor = self.descriptor.get_default_chg_spin()
+            if default_cs_tensor is not None:
+                default_cs_tensor = default_cs_tensor.to(extended_coord.place)
+                charge_spin = paddle.tile(default_cs_tensor.unsqueeze(0), [nframes, 1])
+
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
             comm_dict=comm_dict,
-            fparam=fparam if self.add_chg_spin_ebd else None,
+            charge_spin=charge_spin if self.add_chg_spin_ebd else None,
         )
         assert descriptor is not None
         if self.enable_eval_descriptor_hook:
@@ -466,3 +477,25 @@ def is_aparam_nall(self) -> bool:
         If False, the shape is (nframes, nloc, ndim).
         """
         return False
+
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        return self.add_chg_spin_ebd
+
+    def get_dim_chg_spin(self) -> int:
+        """Get the dimension of charge_spin input."""
+        if self.add_chg_spin_ebd:
+            return self.descriptor.get_dim_chg_spin()
+        return 0
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        if self.add_chg_spin_ebd:
+            return self.descriptor.has_default_chg_spin()
+        return False
+
+    def get_default_chg_spin(self) -> paddle.Tensor | None:
+        """Get the default charge_spin values as a tensor."""
+        if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin():
+            return self.descriptor.get_default_chg_spin()
+        return None

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,

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,6 +746,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,

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__()
 
@@ -167,6 +168,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
             use_exp_switch=self.repflow_args.use_exp_switch,
             use_dynamic_sel=self.repflow_args.use_dynamic_sel,
             sel_reduce_factor=self.repflow_args.sel_reduce_factor,
+            sequential_update=self.repflow_args.sequential_update,
             use_loc_mapping=use_loc_mapping,
             exclude_types=exclude_types,
             env_protection=env_protection,
@@ -177,6 +179,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 +454,23 @@ 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) -> paddle.Tensor | None:
+        """Get the default charge_spin values as a tensor."""
+        if self.default_chg_spin is None:
+            return None
+        return paddle.to_tensor(
+            self.default_chg_spin,
+            dtype=self.prec,
+        )
+
     def serialize(self) -> dict:
         repflows = self.repflows
         data = {
@@ -465,6 +489,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 +566,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 +618,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(

deepmd/pd/model/descriptor/repflow_layer.py

@@ -53,6 +53,7 @@ def __init__(
         optim_update: bool = True,
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
+        sequential_update: bool = False,
         smooth_edge_update: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
@@ -102,8 +103,15 @@ def __init__(
         self.smooth_edge_update = smooth_edge_update
         self.use_dynamic_sel = use_dynamic_sel
         self.sel_reduce_factor = sel_reduce_factor
+        self.sequential_update = sequential_update
         self.dynamic_e_sel = self.nnei / self.sel_reduce_factor
         self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
+        if self.sequential_update and self.update_style != "res_residual":
+            raise NotImplementedError(
+                "sequential_update only supports update_style='res_residual'!"
+            )
+        if self.sequential_update and not self.update_angle:
+            raise NotImplementedError("sequential_update requires update_angle=True!")
 
         assert update_residual_init in [
             "norm",
@@ -1229,6 +1237,7 @@ def serialize(self) -> dict:
             "smooth_edge_update": self.smooth_edge_update,
             "use_dynamic_sel": self.use_dynamic_sel,
             "sel_reduce_factor": self.sel_reduce_factor,
+            "sequential_update": self.sequential_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(),

deepmd/pd/model/descriptor/repflows.py

@@ -224,6 +224,7 @@ def __init__(
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
         use_loc_mapping: bool = True,
+        sequential_update: bool = False,
         optim_update: bool = True,
         seed: int | list[int] | None = None,
         trainable: bool = True,
@@ -267,6 +268,7 @@ def __init__(
         self.use_exp_switch = use_exp_switch
         self.use_dynamic_sel = use_dynamic_sel
         self.sel_reduce_factor = sel_reduce_factor
+        self.sequential_update = sequential_update
         if self.use_dynamic_sel and not self.smooth_edge_update:
             raise NotImplementedError(
                 "smooth_edge_update must be True when use_dynamic_sel is True!"
@@ -338,6 +340,7 @@ def __init__(
                     optim_update=self.optim_update,
                     use_dynamic_sel=self.use_dynamic_sel,
                     sel_reduce_factor=self.sel_reduce_factor,
+                    sequential_update=self.sequential_update,
                     smooth_edge_update=self.smooth_edge_update,
                     seed=child_seed(child_seed(seed, 1), ii),
                     trainable=trainable,

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,

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.