feat(pt): add atomic_weight parameter to dipole model

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -156,6 +156,7 @@ def forward_common_atomic(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        atomic_weight: Array | None = None,
     ) -> dict[str, Array]:
         """Common interface for atomic inference.
 
@@ -177,6 +178,9 @@ def forward_common_atomic(
             frame parameters, shape: nf x dim_fparam
         aparam
             atomic parameter, shape: nf x nloc x dim_aparam
+        atomic_weight
+            atomic weights for scaling outputs, shape: nf x nloc x dim_aw
+            if provided, all output values will be multiplied by this weight.
 
         Returns
         -------
@@ -220,6 +224,19 @@ def forward_common_atomic(
             tmp_arr = ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2])
             tmp_arr = xp.where(atom_mask[:, :, None], tmp_arr, xp.zeros_like(tmp_arr))
             ret_dict[kk] = xp.reshape(tmp_arr, out_shape)
+            if atomic_weight is not None:
+                _out_shape = ret_dict[kk].shape
+                # Validate that atomic_weight.shape[2] matches the flattened trailing dimensions
+                # of the output to ensure correct broadcasting behavior
+                out_shape2 = math.prod(_out_shape[2:])
+                if atomic_weight.shape[2] != out_shape2:
+                    raise ValueError(
+                        f"atomic_weight shape {atomic_weight.shape} incompatible with output shape {_out_shape}. "
+                        f"Expected atomic_weight.shape[2] to be {out_shape2} for key '{kk}'"
+                    )
+                ret_dict[kk] = ret_dict[kk] * atomic_weight.reshape(
+                    [_out_shape[0], _out_shape[1], -1]
+                )
         ret_dict["mask"] = xp.astype(atom_mask, xp.int32)
 
         return ret_dict
@@ -232,6 +249,7 @@ def call(
         mapping: Array | None = None,
         fparam: Array | None = None,
         aparam: Array | None = None,
+        atomic_weight: Array | None = None,
     ) -> dict[str, Array]:
         return self.forward_common_atomic(
             extended_coord,
@@ -240,6 +258,7 @@ def call(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            atomic_weight=atomic_weight,
         )
 
     def serialize(self) -> dict:

deepmd/dpmodel/model/make_model.py

@@ -68,6 +68,7 @@ def model_call_from_call_lower(
     fparam: Array | None = None,
     aparam: Array | None = None,
     do_atomic_virial: bool = False,
+    atomic_weight: Array | None = None,
 ) -> dict[str, Array]:
     """Return model prediction from lower interface.
 
@@ -126,6 +127,7 @@ def model_call_from_call_lower(
         fparam=fp,
         aparam=ap,
         do_atomic_virial=do_atomic_virial,
+        atomic_weight=atomic_weight,
     )
     model_predict = communicate_extended_output(
         model_predict_lower,
@@ -229,6 +231,7 @@ def call(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction.
 
@@ -255,8 +258,12 @@ def call(
                 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, aw, input_prec = self.input_type_cast(
+                coord,
+                box=box,
+                fparam=fparam,
+                aparam=aparam,
+                atomic_weight=atomic_weight,
             )
             del coord, box, fparam, aparam
             model_predict = model_call_from_call_lower(
@@ -271,6 +278,7 @@ def call(
                 fparam=fp,
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
+                atomic_weight=aw,
             )
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
@@ -284,6 +292,7 @@ def call_lower(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -321,8 +330,11 @@ def call_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, aw, input_prec = self.input_type_cast(
+                extended_coord,
+                fparam=fparam,
+                aparam=aparam,
+                atomic_weight=atomic_weight,
             )
             del extended_coord, fparam, aparam
             model_predict = self.forward_common_atomic(
@@ -333,6 +345,7 @@ def call_lower(
                 fparam=fp,
                 aparam=ap,
                 do_atomic_virial=do_atomic_virial,
+                atomic_weight=aw,
             )
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict
@@ -346,6 +359,7 @@ def forward_common_atomic(
             fparam: Array | None = None,
             aparam: Array | None = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Array | None = None,
         ) -> dict[str, Array]:
             atomic_ret = self.atomic_model.forward_common_atomic(
                 extended_coord,
@@ -354,6 +368,7 @@ def forward_common_atomic(
                 mapping=mapping,
                 fparam=fparam,
                 aparam=aparam,
+                atomic_weight=atomic_weight,
             )
             return fit_output_to_model_output(
                 atomic_ret,
@@ -371,26 +386,30 @@ def input_type_cast(
             box: Array | None = None,
             fparam: Array | None = None,
             aparam: Array | None = None,
-        ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, str]:
+            atomic_weight: Array | None = None,
+        ) -> tuple[
+            Array, Array, np.ndarray | None, np.ndarray | None, np.ndarray | None, str
+        ]:
             """Cast the input data to global float type."""
             input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]]
             ###
             ### type checking would not pass jit, convert to coord prec anyway
             ###
             _lst: list[np.ndarray | None] = [
                 vv.astype(coord.dtype) if vv is not None else None
-                for vv in [box, fparam, aparam]
+                for vv in [box, fparam, aparam, atomic_weight]
             ]
-            box, fparam, aparam = _lst
+            box, fparam, aparam, atomic_weight = _lst
             if input_prec == RESERVED_PRECISION_DICT[self.global_np_float_precision]:
-                return coord, box, fparam, aparam, input_prec
+                return coord, box, fparam, aparam, atomic_weight, input_prec
             else:
                 pp = self.global_np_float_precision
                 return (
                     coord.astype(pp),
                     box.astype(pp) if box is not None else None,
                     fparam.astype(pp) if fparam is not None else None,
                     aparam.astype(pp) if aparam is not None else None,
+                    atomic_weight.astype(pp) if atomic_weight is not None else None,
                     input_prec,
                 )
 

deepmd/pt/infer/deep_eval.py

@@ -312,6 +312,7 @@ def eval(
         atomic: bool = False,
         fparam: np.ndarray | None = None,
         aparam: np.ndarray | None = None,
+        atomic_weight: np.ndarray | None = None,
         **kwargs: Any,
     ) -> dict[str, np.ndarray]:
         """Evaluate the energy, force and virial by using this DP.
@@ -341,6 +342,12 @@ def eval(
             - nframes x natoms x dim_aparam.
             - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam.
             - dim_aparam. Then all frames and atoms are provided with the same aparam.
+        atomic_weight
+            The atomic weight for weighted averaging of atomic contributions.
+            The array can be of size :
+            - nframes x natoms. Different weights for each atom in each frame.
+            - nframes x natoms x k. Different weights for each atom in each frame with k dimensions.
+            Note: This parameter is currently ignored when spin is provided.
         **kwargs
             Other parameters
 
@@ -361,9 +368,19 @@ def eval(
         request_defs = self._get_request_defs(atomic)
         if "spin" not in kwargs or kwargs["spin"] is None:
             out = self._eval_func(self._eval_model, numb_test, natoms)(
-                coords, cells, atom_types, fparam, aparam, request_defs
+                coords,
+                cells,
+                atom_types,
+                fparam,
+                aparam,
+                atomic_weight,
+                request_defs,
             )
         else:
+            if atomic_weight is not None:
+                raise ValueError(
+                    "atomic_weight is not supported when spin is provided."
+                )
             out = self._eval_func(self._eval_model_spin, numb_test, natoms)(
                 coords,
                 cells,
@@ -472,6 +489,7 @@ def _eval_model(
         atom_types: np.ndarray,
         fparam: np.ndarray | None,
         aparam: np.ndarray | None,
+        atomic_weight: np.ndarray | None,
         request_defs: list[OutputVariableDef],
     ) -> tuple[np.ndarray, ...]:
         model = self.dp.to(DEVICE)
@@ -514,6 +532,15 @@ def _eval_model(
             )
         else:
             aparam_input = None
+        if atomic_weight is not None:
+            atomic_weight = self._validate_and_reshape_atomic_weight(
+                atomic_weight, nframes, natoms
+            )
+            atomic_weight_input = to_torch_tensor(
+                atomic_weight.reshape(nframes, natoms, -1)
+            )
+        else:
+            atomic_weight_input = None
         do_atomic_virial = any(
             x.category == OutputVariableCategory.DERV_C for x in request_defs
         )
@@ -524,6 +551,7 @@ def _eval_model(
             do_atomic_virial=do_atomic_virial,
             fparam=fparam_input,
             aparam=aparam_input,
+            atomic_weight=atomic_weight_input,
         )
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
@@ -592,6 +620,7 @@ def _eval_model_spin(
             )
         else:
             aparam_input = None
+        atomic_weight_input = None
 
         do_atomic_virial = any(
             x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
@@ -604,6 +633,7 @@ def _eval_model_spin(
             do_atomic_virial=do_atomic_virial,
             fparam=fparam_input,
             aparam=aparam_input,
+            atomic_weight=atomic_weight_input,
         )
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
@@ -847,3 +877,60 @@ def eval_fitting_last_layer(
         fitting_net = model.eval_fitting_last_layer()
         model.set_eval_fitting_last_layer_hook(False)
         return to_numpy_array(fitting_net)
+
+    def _validate_and_reshape_atomic_weight(
+        self,
+        atomic_weight: np.ndarray,
+        nframes: int,
+        natoms: int,
+    ) -> np.ndarray:
+        """Validate and reshape atomic_weight to (nframes, natoms, k).
+
+        Parameters
+        ----------
+        atomic_weight : np.ndarray
+            Input atomic weights with various possible shapes
+        nframes : int
+            Number of frames
+        natoms : int
+            Number of atoms
+
+        Returns
+        -------
+        np.ndarray
+            Reshaped atomic_weight with shape (nframes, natoms, k)
+
+        Raises
+        ------
+        ValueError
+            If atomic_weight shape is incompatible
+        """
+        # Validate atomic_weight shape
+        if atomic_weight.ndim == 1:
+            # (natoms,) - broadcast across frames
+            if atomic_weight.shape[0] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with number of atoms {natoms}"
+                )
+            atomic_weight = np.tile(atomic_weight, (nframes, 1))
+        elif atomic_weight.ndim == 2:
+            # (nframes, natoms) or (natoms, k)
+            if atomic_weight.shape[0] == natoms and atomic_weight.shape[1] != nframes:
+                # (natoms, k) - broadcast across frames
+                atomic_weight = np.tile(atomic_weight, (nframes, 1, 1))
+            elif atomic_weight.shape[0] != nframes or atomic_weight.shape[1] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
+                )
+        elif atomic_weight.ndim == 3:
+            # (nframes, natoms, k)
+            if atomic_weight.shape[0] != nframes or atomic_weight.shape[1] != natoms:
+                raise ValueError(
+                    f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}"
+                )
+        else:
+            raise ValueError(
+                f"atomic_weight must have 1, 2, or 3 dimensions, got {atomic_weight.ndim}"
+            )
+
+        return atomic_weight

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -208,6 +208,7 @@ def forward_common_atomic(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         comm_dict: dict[str, torch.Tensor] | None = None,
+        atomic_weight: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Common interface for atomic inference.
 
@@ -231,6 +232,9 @@ def forward_common_atomic(
             atomic parameter, shape: nf x nloc x dim_aparam
         comm_dict
             The data needed for communication for parallel inference.
+        atomic_weight
+            atomic weights for scaling outputs, shape: nf x nloc x dim_aw
+            if provided, all output values will be multiplied by this weight.
 
         Returns
         -------
@@ -276,6 +280,10 @@ def forward_common_atomic(
                 ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2])
                 * atom_mask[:, :, None]
             ).view(out_shape)
+            if atomic_weight is not None:
+                ret_dict[kk] = ret_dict[kk] * atomic_weight.view(
+                    [out_shape[0], out_shape[1], -1]
+                )
         ret_dict["mask"] = atom_mask
 
         return ret_dict
@@ -289,6 +297,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         comm_dict: dict[str, torch.Tensor] | None = None,
+        atomic_weight: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         return self.forward_common_atomic(
             extended_coord,
@@ -298,6 +307,7 @@ def forward(
             fparam=fparam,
             aparam=aparam,
             comm_dict=comm_dict,
+            atomic_weight=atomic_weight,
         )
 
     def change_type_map(