debug for rebase

Author: YuzhiLiu-ai

deepmd/pt/model/model/density_model.py

@@ -60,6 +60,7 @@ def forward(
         fparam: Optional[torch.Tensor] = None,
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
+        charge_spin: Optional[torch.Tensor] = None,
     ) -> dict[str, torch.Tensor]:
         model_ret = self.forward_common(
             coord,
@@ -87,5 +88,6 @@ def forward_lower(
         aparam: Optional[torch.Tensor] = None,
         do_atomic_virial: bool = False,
         comm_dict: Optional[dict[str, torch.Tensor]] = None,
+        charge_spin: Optional[torch.Tensor] = None,
     ):
         raise NotImplementedError

deepmd/pt/model/model/make_density_model.py

@@ -101,6 +101,26 @@ def model_output_type(self) -> list[str]:
                     vars.append(kk)
             return vars
 
+        @torch.jit.export
+        def has_chg_spin_ebd(self) -> bool:
+            """Check if the model has charge spin embedding."""
+            return self.atomic_model.has_chg_spin_ebd()
+
+        @torch.jit.export
+        def get_dim_chg_spin(self) -> int:
+            """Get the dimension of charge_spin input."""
+            return self.atomic_model.get_dim_chg_spin()
+
+        @torch.jit.export
+        def has_default_chg_spin(self) -> bool:
+            """Check if the model has default charge_spin values."""
+            return self.atomic_model.has_default_chg_spin()
+
+        @torch.jit.export
+        def get_default_chg_spin(self) -> torch.Tensor | None:
+            """Get the default charge_spin values."""
+            return self.atomic_model.get_default_chg_spin()
+
         # cannot use the name forward. torch script does not work
         def forward_common(
             self,
@@ -111,6 +131,7 @@ def forward_common(
             fparam: Optional[torch.Tensor] = None,
             aparam: Optional[torch.Tensor] = None,
             do_atomic_virial: bool = False,
+            charge_spin: Optional[torch.Tensor] = None,
         ) -> dict[str, torch.Tensor]:
             """Return model prediction.
 
@@ -145,6 +166,7 @@ def forward_common(
                 coord, grid, box=box, fparam=fparam, aparam=aparam
             )
             del coord, grid, box, fparam, aparam
+            gg = gg.view(gg.shape[0], -1, 3)
             (
                 extended_coord,
                 extended_atype,
@@ -158,7 +180,7 @@ def forward_common(
                 mixed_types=self.mixed_types(),
                 box=bb,
             )
-            grid_type = torch.zeros(gg.shape[:-1], device=gg.device, dtype=atype.dtype)
+            grid_type = torch.zeros(gg.shape[0], gg.shape[1], device=gg.device, dtype=atype.dtype)
             grid_nlist = build_directional_neighbor_list(
                 gg,
                 grid_type,
@@ -233,6 +255,7 @@ def forward_common_lower(
             mapping: Optional[torch.Tensor] = None,
             fparam: Optional[torch.Tensor] = None,
             aparam: Optional[torch.Tensor] = None,
+            charge_spin: Optional[torch.Tensor] = None,
             do_atomic_virial: bool = False,
             comm_dict: Optional[dict[str, torch.Tensor]] = None,
             extra_nlist_sort: bool = False,
@@ -572,9 +595,14 @@ def compute_or_load_stat(
             self,
             sampled_func,
             stat_file_path: Optional[DPPath] = None,
+            preset_observed_type: list[str] | None = None,
         ):
             """Compute or load the statistics."""
-            return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path)
+            return self.atomic_model.compute_or_load_stat(
+                sampled_func,
+                stat_file_path,
+                preset_observed_type=preset_observed_type,
+            )
 
         def get_sel(self) -> list[int]:
             """Returns the number of selected atoms for each type."""

deepmd/pt/train/wrapper.py

@@ -2,6 +2,8 @@
 import logging
 from typing import (
     Any,
+    Optional,
+    Union,
 )
 
 import torch

deepmd/utils/data.py

@@ -966,6 +966,8 @@ def _load_single_data(
         data = mmap_obj[frame_idx].copy().astype(dtype, copy=False)
 
         try:
+            if key in ("grid", "density"):
+                return np.float32(1.0), data
             if vv["atomic"]:
                 # Handle type_sel logic
                 if vv["type_sel"] is not None: