feat(dpa3): decouple charge_spin from fparam

deepmd/calculator.py

@@ -137,8 +137,14 @@ def calculate(
 
         fparam = self.atoms.info.get("fparam", None)
         aparam = self.atoms.info.get("aparam", None)
+        charge_spin = self.atoms.info.get("charge_spin", None)
         e, f, v = self.dp.eval(
-            coords=coord, cells=cell, atom_types=atype, fparam=fparam, aparam=aparam
+            coords=coord,
+            cells=cell,
+            atom_types=atype,
+            fparam=fparam,
+            aparam=aparam,
+            charge_spin=charge_spin,
         )[:3]
         self.results["energy"] = e[0][0]
         # see https://gitlab.com/ase/ase/-/merge_requests/2485

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -156,6 +156,18 @@ def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return None
 
+    def has_chg_spin_ebd(self) -> bool:
+        """Check if the model has charge spin embedding."""
+        return False
+
+    def has_default_chg_spin(self) -> bool:
+        """Check if the model has default charge_spin values."""
+        return False
+
+    def get_default_chg_spin(self) -> list[float] | None:
+        """Get the default charge_spin values."""
+        return None
+
     def reinit_atom_exclude(
         self,
         exclude_types: list[int] = [],
@@ -231,6 +243,7 @@ def forward_common_atomic(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Common interface for atomic inference.
 
@@ -279,6 +292,7 @@ def forward_common_atomic(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            charge_spin=charge_spin,
         )
         ret_dict = self.apply_out_stat(ret_dict, atype)
 
@@ -307,6 +321,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         return self.forward_common_atomic(
             extended_coord,
@@ -315,6 +330,7 @@ def call(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            charge_spin=charge_spin,
         )
 
     def get_intensive(self) -> bool:
@@ -519,6 +535,7 @@ def model_forward(
             box: np.ndarray | None,
             fparam: np.ndarray | None = None,
             aparam: np.ndarray | None = None,
+            charge_spin: np.ndarray | None = None,
         ) -> dict[str, np.ndarray]:
             # Get reference array to determine the target array type and device
             # Use out_bias as reference since it's always present
@@ -538,6 +555,8 @@ def model_forward(
                 fparam = xp.asarray(fparam, device=device)
             if aparam is not None:
                 aparam = xp.asarray(aparam, device=device)
+            if charge_spin is not None:
+                charge_spin = xp.asarray(charge_spin, device=device)
 
             (
                 extended_coord,
@@ -559,6 +578,7 @@ def model_forward(
                 mapping=mapping,
                 fparam=fparam,
                 aparam=aparam,
+                charge_spin=charge_spin,
             )
             # Convert outputs back to numpy arrays
             return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()}

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -79,6 +79,28 @@ def __init__(
         )
         super().init_out_stat()
 
+    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) -> list[float] | None:
+        """Get the default charge_spin values."""
+        if self.add_chg_spin_ebd and self.descriptor.has_default_chg_spin():
+            return self.descriptor.get_default_chg_spin()
+        return None
+
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return self.fitting_net.output_def()
@@ -157,6 +179,7 @@ def forward_atomic(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Models' atomic predictions.
 
@@ -174,6 +197,8 @@ def forward_atomic(
             frame parameter. nf x ndf
         aparam
             atomic parameter. nf x nloc x nda
+        charge_spin
+            charge and spin parameter for descriptor. nf x 2
 
         Returns
         -------
@@ -184,37 +209,28 @@ def forward_atomic(
         nframes, nloc, nnei = nlist.shape
         atype = xp_take_first_n(extended_atype, 1, nloc)
 
-        # Handle default fparam if fitting net supports it
-        if (
-            hasattr(self.fitting_net, "get_dim_fparam")
-            and self.fitting_net.get_dim_fparam() > 0
-            and fparam is None
-        ):
-            # use default fparam
-            from deepmd.dpmodel.array_api import (
-                array_api_compat,
-            )
-
-            default_fparam = self.fitting_net.get_default_fparam()
-            assert default_fparam is not None
-            xp = array_api_compat.array_namespace(extended_coord)
-            default_fparam_array = xp.asarray(
-                default_fparam,
-                dtype=extended_coord.dtype,
-                device=array_api_compat.device(extended_coord),
-            )
-            fparam_input_for_des = xp.tile(
-                xp.reshape(default_fparam_array, (1, -1)), (nframes, 1)
-            )
-        else:
-            fparam_input_for_des = fparam
+        # Handle default charge_spin if descriptor supports it
+        if self.add_chg_spin_ebd and charge_spin is None:
+            default_cs = self.descriptor.get_default_chg_spin()
+            if default_cs is not None:
+                from deepmd.dpmodel.array_api import (
+                    array_api_compat,
+                )
+
+                xp = array_api_compat.array_namespace(extended_coord)
+                cs_array = xp.asarray(
+                    default_cs,
+                    dtype=extended_coord.dtype,
+                    device=array_api_compat.device(extended_coord),
+                )
+                charge_spin = xp.tile(xp.reshape(cs_array, (1, -1)), (nframes, 1))
 
         descriptor, rot_mat, g2, h2, sw = self.descriptor(
             extended_coord,
             extended_atype,
             nlist,
             mapping=mapping,
-            fparam=fparam_input_for_des if self.add_chg_spin_ebd else None,
+            charge_spin=charge_spin if self.add_chg_spin_ebd else None,
         )
         ret = self.fitting_net(
             descriptor,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -224,6 +224,7 @@ def forward_atomic(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         """Return atomic prediction.
 
@@ -280,6 +281,7 @@ def forward_atomic(
                     mapping,
                     fparam,
                     aparam,
+                    charge_spin=charge_spin,
                 )["energy"]
             )
         weights = self._compute_weight(extended_coord, extended_atype, nlists_)

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -138,6 +138,7 @@ def fwd(
             mapping: t_tensor | None = None,
             fparam: t_tensor | None = None,
             aparam: t_tensor | None = None,
+            charge_spin: t_tensor | None = None,
         ) -> dict[str, t_tensor]:
             pass
 

deepmd/dpmodel/atomic_model/pairtab_atomic_model.py

@@ -253,6 +253,7 @@ def forward_atomic(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> dict[str, Array]:
         xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist)
         nframes, nloc, nnei = nlist.shape

deepmd/dpmodel/descriptor/dpa1.py

@@ -500,6 +500,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> Array:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa2.py

@@ -831,6 +831,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/dpa3.py

@@ -359,6 +359,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__()
 
@@ -414,6 +415,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_tebd_bias = use_tebd_bias
         self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
@@ -480,6 +486,18 @@ def get_rcut(self) -> float:
         """Returns the cut-off radius."""
         return self.rcut
 
+    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 get_rcut_smth(self) -> float:
         """Returns the radius where the neighbor information starts to smoothly decay to 0."""
         return self.rcut_smth
@@ -616,6 +634,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array, Array, Array, Array]:
         """Compute the descriptor.
 
@@ -668,13 +687,13 @@ def call(
             )
 
         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
             chg_tebd = self.chg_embedding.call()
             spin_tebd = self.spin_embedding.call()
-            charge = xp.astype(fparam[:, 0], xp.int64) + 100
-            spin = xp.astype(fparam[:, 1], xp.int64)
+            charge = xp.astype(charge_spin[:, 0], xp.int64) + 100
+            spin = xp.astype(charge_spin[:, 1], xp.int64)
             chg_ebd = xp.reshape(
                 xp.take(chg_tebd, xp.reshape(charge, (-1,)), axis=0),
                 (nframes, self.tebd_dim),
@@ -718,6 +737,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.serialize(),
         }

deepmd/dpmodel/descriptor/hybrid.py

@@ -276,6 +276,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[
         Array,
         Array | None,
@@ -332,7 +333,9 @@ def call(
                 # mixed_types is True, but descrpt.mixed_types is False
                 assert nl_distinguish_types is not None
                 nl = nl_distinguish_types[:, :, nci]
-            odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping)
+            odescriptor, gr, g2, h2, sw = descrpt(
+                coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin
+            )
             out_descriptor.append(odescriptor)
             if gr is not None:
                 out_gr.append(gr)

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -96,6 +96,18 @@ def get_dim_emb(self) -> int:
             """Returns the embedding dimension of g2."""
             pass
 
+        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) -> Any:
+            """Returns the default charge_spin value, or None."""
+            return None
+
         @abstractmethod
         def mixed_types(self) -> bool:
             """Returns if the descriptor requires a neighbor list that distinguish different
@@ -187,6 +199,7 @@ def fwd(
             nlist: Array,
             mapping: Array | None = None,
             fparam: Array | None = None,
+            charge_spin: Array | None = None,
         ) -> Array:
             """Calculate descriptor."""
             pass

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -399,6 +399,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> Array:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/se_r.py

@@ -371,6 +371,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> Array:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/se_t.py

@@ -346,6 +346,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array]:
         """Compute the descriptor.
 

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -354,6 +354,7 @@ def call(
         nlist: Array,
         mapping: Array | None = None,
         fparam: Array | None = None,
+        charge_spin: Array | None = None,
     ) -> tuple[Array, Array]:
         """Compute the descriptor.