fix comments

Author: iProzd

deepmd/dpmodel/descriptor/hybrid.py

@@ -349,6 +349,7 @@ def call(
                 atype_ext,
                 nl,
                 mapping,
+                fparam=fparam,
                 comm_dict=comm_dict,
                 charge_spin=charge_spin,
             )

deepmd/dpmodel/model/make_model.py

@@ -298,10 +298,10 @@ def call_common(
                 The keys are defined by the `ModelOutputDef`.
 
             """
-            cc, bb, fp, ap, input_prec = self._input_type_cast(
-                coord, box=box, fparam=fparam, aparam=aparam
+            cc, bb, fp, ap, cs, input_prec = self._input_type_cast(
+                coord, box=box, fparam=fparam, aparam=aparam, charge_spin=charge_spin
             )
-            del coord, box, fparam, aparam
+            del coord, box, fparam, aparam, charge_spin
             model_predict = model_call_from_call_lower(
                 call_lower=self.call_common_lower,
                 rcut=self.get_rcut(),
@@ -315,7 +315,7 @@ def call_common(
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
                 coord_corr_for_virial=coord_corr_for_virial,
-                charge_spin=charge_spin,
+                charge_spin=cs,
             )
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
@@ -377,10 +377,10 @@ def call_common_lower(
                 nlist,
                 extra_nlist_sort=self.need_sorted_nlist_for_lower(),
             )
-            cc_ext, _, fp, ap, input_prec = self._input_type_cast(
-                extended_coord, fparam=fparam, aparam=aparam
+            cc_ext, _, fp, ap, cs, input_prec = self._input_type_cast(
+                extended_coord, fparam=fparam, aparam=aparam, charge_spin=charge_spin
             )
-            del extended_coord, fparam, aparam
+            del extended_coord, fparam, aparam, charge_spin
             model_predict = self.forward_common_atomic(
                 cc_ext,
                 extended_atype,
@@ -391,7 +391,7 @@ def call_common_lower(
                 do_atomic_virial=do_atomic_virial,
                 extended_coord_corr=extended_coord_corr,
                 comm_dict=comm_dict,
-                charge_spin=charge_spin,
+                charge_spin=cs,
             )
             model_predict = self._output_type_cast(model_predict, input_prec)
             return model_predict
@@ -482,7 +482,8 @@ def _input_type_cast(
             box: Array | None = None,
             fparam: Array | None = None,
             aparam: Array | None = None,
-        ) -> tuple[Array, Array | None, Array | None, Array | None, Any]:
+            charge_spin: Array | None = None,
+        ) -> tuple[Array, Array | None, Array | None, Array | None, Array | None, Any]:
             """Cast the input data to global float type."""
             xp = array_api_compat.array_namespace(coord)
             input_dtype = coord.dtype
@@ -494,17 +495,20 @@ def _input_type_cast(
             ###
             _lst: list[Array | None] = [
                 xp.astype(vv, input_dtype) if vv is not None else None
-                for vv in [box, fparam, aparam]
+                for vv in [box, fparam, aparam, charge_spin]
             ]
-            box, fparam, aparam = _lst
+            box, fparam, aparam, charge_spin = _lst
             if input_dtype == global_dtype:
-                return coord, box, fparam, aparam, input_dtype
+                return coord, box, fparam, aparam, charge_spin, input_dtype
             else:
                 return (
                     xp.astype(coord, global_dtype),
                     xp.astype(box, global_dtype) if box is not None else None,
                     xp.astype(fparam, global_dtype) if fparam is not None else None,
                     xp.astype(aparam, global_dtype) if aparam is not None else None,
+                    xp.astype(charge_spin, global_dtype)
+                    if charge_spin is not None
+                    else None,
                     input_dtype,
                 )
 

deepmd/jax/model/base_model.py

@@ -68,6 +68,7 @@ def eval_output(
                     mapping: jnp.ndarray | None,
                     fparam: jnp.ndarray | None,
                     aparam: jnp.ndarray | None,
+                    charge_spin_: jnp.ndarray | None,
                     *,
                     _kk: str = kk,
                     _atom_axis: int = atom_axis,
@@ -79,7 +80,9 @@ def eval_output(
                         mapping=mapping[None, ...] if mapping is not None else None,
                         fparam=fparam[None, ...] if fparam is not None else None,
                         aparam=aparam[None, ...] if aparam is not None else None,
-                        charge_spin=charge_spin,
+                        charge_spin=charge_spin_[None, ...]
+                        if charge_spin_ is not None
+                        else None,
                     )
                     return jnp.sum(atomic_ret[_kk][0], axis=_atom_axis)
 
@@ -92,6 +95,7 @@ def eval_output(
                     mapping,
                     fparam,
                     aparam,
+                    charge_spin,
                 )
                 # extended_force: [nf, nall, *def, 3]
                 def_ndim = len(vdef.shape)
@@ -109,6 +113,7 @@ def eval_output(
                         mapping,
                         fparam,
                         aparam,
+                        charge_spin,
                     )
                     kk_hessian = get_hessian_name(kk)
                     model_predict[kk_hessian] = hessian
@@ -130,6 +135,7 @@ def eval_ce(
                         mapping: jnp.ndarray | None,
                         fparam: jnp.ndarray | None,
                         aparam: jnp.ndarray | None,
+                        charge_spin_: jnp.ndarray | None,
                         *,
                         _kk: str = kk,
                         _atom_axis: int = atom_axis - 1,
@@ -142,7 +148,9 @@ def eval_ce(
                             mapping=mapping[None, ...] if mapping is not None else None,
                             fparam=fparam[None, ...] if fparam is not None else None,
                             aparam=aparam[None, ...] if aparam is not None else None,
-                            charge_spin=charge_spin,
+                            charge_spin=charge_spin_[None, ...]
+                            if charge_spin_ is not None
+                            else None,
                         )
                         nloc = nlist.shape[0]
                         cc_loc = jax.lax.stop_gradient(cc_ext)[:nloc, ...]
@@ -160,6 +168,7 @@ def eval_ce(
                         mapping,
                         fparam,
                         aparam,
+                        charge_spin,
                     )
                     # move the first 3 to the last
                     # [nf, *def, nall, 3, 3]

deepmd/pt/model/descriptor/dpa3.py

@@ -204,14 +204,14 @@ def init_subclass_params(sub_data: Any, sub_class: Any) -> Any:
 
         if self.add_chg_spin_ebd:
             self.act = ActivationFn(activation_function)
-            # -100 ~ 100 is a conservative bound
+            # charge range [-100, 99] mapped to indices [0, 199]
             self.chg_embedding = TypeEmbedNet(
                 200,
                 self.tebd_dim,
                 precision=precision,
                 seed=child_seed(seed, 3),
             )
-            # 100 is a conservative upper bound
+            # spin range [0, 99] mapped to indices [0, 99]
             self.spin_embedding = TypeEmbedNet(
                 100,
                 self.tebd_dim,
@@ -588,9 +588,19 @@ def forward(
             assert charge_spin is not None
             assert self.chg_embedding is not None
             assert self.spin_embedding is not None
-            charge = charge_spin[:, 0].to(dtype=torch.int64) + 100
+            charge = charge_spin[:, 0].to(dtype=torch.int64)
             spin = charge_spin[:, 1].to(dtype=torch.int64)
-            chg_ebd = self.chg_embedding(charge)
+            # Validate charge range [-100, 99] (200 embedding entries)
+            if torch.any(charge < -100) or torch.any(charge > 99):
+                raise ValueError(
+                    f"charge must be in range [-100, 99], got min={charge.min().item()}, max={charge.max().item()}"
+                )
+            # Validate spin range [0, 99] (100 embedding entries)
+            if torch.any(spin < 0) or torch.any(spin >= 100):
+                raise ValueError(
+                    f"spin must be in range [0, 99], got min={spin.min().item()}, max={spin.max().item()}"
+                )
+            chg_ebd = self.chg_embedding(charge + 100)
             spin_ebd = self.spin_embedding(spin)
             sys_cs_embd = self.act(
                 self.mix_cs_mlp(torch.cat((chg_ebd, spin_ebd), dim=-1))

deepmd/pt/model/descriptor/hybrid.py

@@ -101,14 +101,20 @@ def __init__(
 
     def get_dim_chg_spin(self) -> int:
         """Returns the dimension of charge_spin input (0 if not supported)."""
-        return 0
+        return max(
+            (descrpt.get_dim_chg_spin() for descrpt in self.descrpt_list), default=0
+        )
 
     def has_default_chg_spin(self) -> bool:
         """Returns whether the descriptor has a default charge_spin value."""
-        return False
+        return any(descrpt.has_default_chg_spin() for descrpt in self.descrpt_list)
 
-    def get_default_chg_spin(self) -> None:
+    def get_default_chg_spin(self) -> list[float] | None:
         """Returns the default charge_spin value, or None."""
+        for descrpt in self.descrpt_list:
+            default = descrpt.get_default_chg_spin()
+            if default is not None:
+                return default
         return None
 
     def get_rcut(self) -> float:
@@ -346,7 +352,13 @@ def forward(
                     :, :, self.nlist_cut_idx[ii].to(atype_ext.device)
                 ]
             odescriptor, gr, g2, h2, sw = descrpt(
-                coord_ext, atype_ext, nl, mapping, charge_spin=charge_spin
+                coord_ext,
+                atype_ext,
+                nl,
+                mapping,
+                comm_dict=comm_dict,
+                fparam=fparam,
+                charge_spin=charge_spin,
             )
             out_descriptor.append(odescriptor)
             if gr is not None:

deepmd/pt_expt/infer/deep_eval.py

@@ -624,6 +624,12 @@ def eval(
         aparam
             The atomic parameter.
             The array should be of size nframes x natoms x dim_aparam.
+        charge_spin
+            The charge and spin values for each frame.
+            The array should be of size nframes x 2, where the first column is charge
+            and the second column is spin. If the model has add_chg_spin_ebd=True and
+            no default_chg_spin is set, this parameter is required. If default_chg_spin
+            is configured, this parameter is optional and will override the default.
         **kwargs
             Other parameters
 
@@ -1056,8 +1062,14 @@ def _prepare_inputs(
 
         # charge_spin handling: dedicated input, separate from fparam.
         if charge_spin is not None:
+            charge_spin_arr = np.asarray(charge_spin)
+            if charge_spin_arr.shape != (nframes, 2):
+                raise ValueError(
+                    f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. "
+                    f"Expected ({nframes}, 2) for {nframes} frame(s)."
+                )
             charge_spin_t = torch.tensor(
-                np.asarray(charge_spin).reshape(nframes, 2),
+                charge_spin_arr,
                 dtype=torch.float64,
                 device=DEVICE,
             )
@@ -1272,8 +1284,14 @@ def _eval_model_spin(
 
         # charge_spin handling: dedicated input, separate from fparam.
         if charge_spin is not None:
+            charge_spin_arr = np.asarray(charge_spin)
+            if charge_spin_arr.shape != (nframes, 2):
+                raise ValueError(
+                    f"charge_spin must have shape (nframes, 2), got {charge_spin_arr.shape}. "
+                    f"Expected ({nframes}, 2) for {nframes} frame(s)."
+                )
             charge_spin_t = torch.tensor(
-                np.asarray(charge_spin).reshape(nframes, 2),
+                charge_spin_arr,
                 dtype=torch.float64,
                 device=DEVICE,
             )