Fix bug: Error intensive property output using virtual atoms

Author: Chengqian-Zhang

deepmd/pt/model/model/make_model.py

@@ -299,6 +299,7 @@ def forward_common_lower(
                 cc_ext,
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
+                mask=atomic_ret["mask"] if "mask" in atomic_ret else None,
             )
             model_predict = self.output_type_cast(model_predict, input_prec)
             return model_predict

deepmd/pt/model/model/transform_output.py

@@ -158,6 +158,7 @@ def fit_output_to_model_output(
     coord_ext: torch.Tensor,
     do_atomic_virial: bool = False,
     create_graph: bool = True,
+    mask: Optional[torch.Tensor] = None,
 ) -> dict[str, torch.Tensor]:
     """Transform the output of the fitting network to
     the model output.
@@ -172,7 +173,11 @@ def fit_output_to_model_output(
         if vdef.reducible:
             kk_redu = get_reduce_name(kk)
             if vdef.intensive:
-                model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis)
+                if (mask is not None) and (mask == 0.0).any(): # containing padding atoms
+                    mask = mask.to(dtype=torch.bool, device=vv.device) # [nbz, nreal+npadding]
+                    model_ret[kk_redu] = torch.stack([torch.mean(vv[ii].to(redu_prec)[mask[ii]], dim=atom_axis) for ii in range(mask.size(0))])
+                else:
+                    model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis)
             else:
                 model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis)
             if vdef.r_differentiable:

source/tests/pt/test_dp_test.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 
 from deepmd.entrypoints.test import test as dp_test
 from deepmd.pt.entrypoints.main import (
@@ -186,6 +187,23 @@ def test_dp_test_1_frame(self) -> None:
             to_numpy_array(result[model.get_var_name()])[0],
         )
 
+    def test_dp_test_padding_atoms(self) -> None:
+        trainer = get_trainer(deepcopy(self.config))
+        with torch.device("cpu"):
+            input_dict, label_dict, _ = trainer.get_data(is_train=False)
+        input_dict.pop("spin", None)
+        result = trainer.model(**input_dict)
+        padding_atoms_list = [1, 5, 10]
+        for padding_atoms in padding_atoms_list:
+            input_dict_padding = deepcopy(input_dict)
+            input_dict_padding["atype"] = F.pad(input_dict_padding["atype"], (0, padding_atoms), value=-1)
+            input_dict_padding["coord"] = F.pad(input_dict_padding["coord"], (0, 0, 0, padding_atoms, 0, 0), value=0)
+            result_padding = trainer.model(**input_dict_padding)
+            np.testing.assert_equal(
+                to_numpy_array(result[trainer.model.get_var_name()])[0],
+                to_numpy_array(result_padding[trainer.model.get_var_name()])[0],
+            )
+
     def tearDown(self) -> None:
         for f in os.listdir("."):
             if f.startswith("model") and f.endswith(".pt"):