(np) add atomic_weight to atomic model and dipole model; add UT for pt/np

Author: ChiahsinChu

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -149,6 +149,7 @@ def forward_common_atomic(
         mapping: Optional[np.ndarray] = None,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
+        atomic_weight: Optional[np.ndarray] = None,
     ) -> dict[str, np.ndarray]:
         """Common interface for atomic inference.
 
@@ -213,6 +214,11 @@ 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
+                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
@@ -225,6 +231,7 @@ def call(
         mapping: Optional[np.ndarray] = None,
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
+        atomic_weight: Optional[np.ndarray] = None,
     ) -> dict[str, np.ndarray]:
         return self.forward_common_atomic(
             extended_coord,
@@ -233,6 +240,7 @@ def call(
             mapping=mapping,
             fparam=fparam,
             aparam=aparam,
+            atomic_weight=atomic_weight,
         )
 
     def serialize(self) -> dict:

deepmd/dpmodel/model/make_model.py

@@ -63,6 +63,7 @@ def model_call_from_call_lower(
     fparam: Optional[np.ndarray] = None,
     aparam: Optional[np.ndarray] = None,
     do_atomic_virial: bool = False,
+    atomic_weight: Optional[np.ndarray] = None,
 ):
     """Return model prediction from lower interface.
 
@@ -121,6 +122,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,
@@ -224,6 +226,7 @@ def call(
             fparam: Optional[np.ndarray] = None,
             aparam: Optional[np.ndarray] = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Optional[np.ndarray] = None,
         ) -> dict[str, np.ndarray]:
             """Return model prediction.
 
@@ -250,8 +253,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(
@@ -266,6 +273,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
@@ -279,6 +287,7 @@ def call_lower(
             fparam: Optional[np.ndarray] = None,
             aparam: Optional[np.ndarray] = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Optional[np.ndarray] = None,
         ):
             """Return model prediction. Lower interface that takes
             extended atomic coordinates and types, nlist, and mapping
@@ -316,8 +325,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(
@@ -328,6 +340,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
@@ -341,6 +354,7 @@ def forward_common_atomic(
             fparam: Optional[np.ndarray] = None,
             aparam: Optional[np.ndarray] = None,
             do_atomic_virial: bool = False,
+            atomic_weight: Optional[np.ndarray] = None,
         ):
             atomic_ret = self.atomic_model.forward_common_atomic(
                 extended_coord,
@@ -349,6 +363,7 @@ def forward_common_atomic(
                 mapping=mapping,
                 fparam=fparam,
                 aparam=aparam,
+                atomic_weight=atomic_weight,
             )
             return fit_output_to_model_output(
                 atomic_ret,
@@ -365,11 +380,13 @@ def input_type_cast(
             box: Optional[np.ndarray] = None,
             fparam: Optional[np.ndarray] = None,
             aparam: Optional[np.ndarray] = None,
+            atomic_weight: Optional[np.ndarray] = None,
         ) -> tuple[
             np.ndarray,
             Optional[np.ndarray],
             Optional[np.ndarray],
             Optional[np.ndarray],
+            Optional[np.ndarray],
             str,
         ]:
             """Cast the input data to global float type."""
@@ -379,18 +396,19 @@ def input_type_cast(
             ###
             _lst: list[Optional[np.ndarray]] = [
                 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,
                 )
 

source/tests/pt/model/test_dp_atomic_model.py

@@ -73,6 +73,13 @@ def test_self_consistency(self) -> None:
                 to_numpy_array(ret0["energy"]),
                 to_numpy_array(ret1["energy"]),
             )
+            # add test for atomic_weight
+            aw = torch.rand([nf, nloc, 1], dtype=dtype, device=env.DEVICE)
+            ret2 = md0.forward_common_atomic(*args, atomic_weight=aw)
+            np.testing.assert_allclose(
+                to_numpy_array(ret0["energy"] * aw.reshape(nf, nloc, -1)),
+                to_numpy_array(ret2["energy"]),
+            )
 
     def test_dp_consistency(self) -> None:
         nf, nloc, nnei = self.nlist.shape
@@ -101,6 +108,14 @@ def test_dp_consistency(self) -> None:
             ret0["energy"],
             to_numpy_array(ret1["energy"]),
         )
+        # add test for atomic_weight
+        aw = torch.rand([nf, nloc, 1], dtype=dtype, device=env.DEVICE)
+        ret2 = md0.forward_common_atomic(*args0, atomic_weight=to_numpy_array(aw))
+        ret3 = md1.forward_common_atomic(*args1, atomic_weight=aw)
+        np.testing.assert_allclose(
+            ret2["energy"],
+            to_numpy_array(ret3["energy"]),
+        )
 
     def test_jit(self) -> None:
         nf, nloc, nnei = self.nlist.shape