fix ut

source/tests/pt/model/test_force_grad.py

@@ -19,33 +19,34 @@
 from deepmd.pt.utils import (
     env,
 )
-from deepmd.pt.utils.dataset import (
-    DeepmdDataSystem,
+from deepmd.utils.data import (
+    DeepmdData,
 )
 
 
-class CheckSymmetry(DeepmdDataSystem):
+class CheckSymmetry(DeepmdData):
     def __init__(
         self,
         sys_path: str,
-        rcut,
-        sec,
         type_map: Optional[List[str]] = None,
-        type_split=True,
     ):
-        super().__init__(sys_path, rcut, sec, type_map, type_split)
+        super().__init__(sys_path=sys_path, type_map=type_map)
+        self.add("energy", 1, atomic=False, must=False, high_prec=True)
+        self.add("force", 3, atomic=True, must=False, high_prec=False)
+        self.add("virial", 9, atomic=False, must=False, high_prec=False)
 
     def get_disturb(self, index, atom_index, axis_index, delta):
         for i in range(
-            0, len(self._dirs) + 1
+            0, len(self.dirs) + 1
         ):  # note: if different sets can be merged, prefix sum is unused to calculate
             if index < self.prefix_sum[i]:
                 break
-        frames = self._load_set(self._dirs[i - 1])
+        frames = self._load_set(self.dirs[i - 1])
         tmp = copy.deepcopy(frames["coord"].reshape(self.nframes, -1, 3))
         tmp[:, atom_index, axis_index] += delta
         frames["coord"] = tmp
-        frame = self.single_preprocess(frames, index - self.prefix_sum[i - 1])
+        frame = self._get_subdata(frames, index - self.prefix_sum[i - 1])
+        frame = self.preprocess(frame)
         return frame
 
 
@@ -78,16 +79,16 @@ def get_dataset(self, system_index=0, batch_index=0):
         sec = torch.cumsum(torch.tensor(sel), dim=0)
         type_map = self.config["model"]["type_map"]
         self.dpdatasystem = CheckSymmetry(
-            sys_path=systems[system_index], rcut=rcut, sec=sec, type_map=type_map
+            sys_path=systems[system_index], type_map=type_map
         )
-        self.origin_batch = self.dpdatasystem._get_item(batch_index)
+        self.origin_batch = self.dpdatasystem.get_item(batch_index)
 
     @unittest.skip("it can be replaced by autodiff")
     def test_force_grad(self, threshold=1e-2, delta0=1e-6, seed=20):
         result0 = self.model(**get_data(self.origin_batch))
         np.random.default_rng(seed)
-        errors = np.zeros((self.dpdatasystem._natoms, 3))
-        for atom_index in range(self.dpdatasystem._natoms):
+        errors = np.zeros((self.dpdatasystem.natoms, 3))
+        for atom_index in range(self.dpdatasystem.natoms):
             for axis_index in range(3):
                 delta = np.random.random() * delta0
                 disturb_batch = self.dpdatasystem.get_disturb(

source/tests/pt/model/test_rotation.py

@@ -21,29 +21,29 @@
 from deepmd.pt.utils import (
     env,
 )
-from deepmd.pt.utils.dataset import (
-    DeepmdDataSystem,
+from deepmd.utils.data import (
+    DeepmdData,
 )
 
 
-class CheckSymmetry(DeepmdDataSystem):
+class CheckSymmetry(DeepmdData):
     def __init__(
         self,
         sys_path: str,
-        rcut,
-        sec,
         type_map: Optional[List[str]] = None,
-        type_split=True,
     ):
-        super().__init__(sys_path, rcut, sec, type_map, type_split)
+        super().__init__(sys_path=sys_path, type_map=type_map)
+        self.add("energy", 1, atomic=False, must=False, high_prec=True)
+        self.add("force", 3, atomic=True, must=False, high_prec=False)
+        self.add("virial", 9, atomic=False, must=False, high_prec=False)
 
     def get_rotation(self, index, rotation_matrix):
         for i in range(
-            0, len(self._dirs) + 1
+            0, len(self.dirs) + 1
         ):  # note: if different sets can be merged, prefix sum is unused to calculate
             if index < self.prefix_sum[i]:
                 break
-        frames = self._load_set(self._dirs[i - 1])
+        frames = self._load_set(self.dirs[i - 1])
         frames["coord"] = np.dot(
             rotation_matrix, frames["coord"].reshape(-1, 3).T
         ).T.reshape(self.nframes, -1)
@@ -53,14 +53,16 @@ def get_rotation(self, index, rotation_matrix):
         frames["force"] = np.dot(
             rotation_matrix, frames["force"].reshape(-1, 3).T
         ).T.reshape(self.nframes, -1)
-        frame = self.single_preprocess(frames, index - self.prefix_sum[i - 1])
+        frame = self._get_subdata(frames, index - self.prefix_sum[i - 1])
+        frame = self.preprocess(frame)
         return frame
 
 
 def get_data(batch):
     inputs = {}
     for key in ["coord", "atype", "box"]:
-        inputs[key] = batch[key].unsqueeze(0).to(env.DEVICE)
+        inputs[key] = torch.as_tensor(batch[key])
+        inputs[key] = inputs[key].unsqueeze(0).to(env.DEVICE)
     return inputs
 
 
@@ -80,14 +82,11 @@ def get_model(self):
 
     def get_dataset(self, system_index=0, batch_index=0):
         systems = self.config["training"]["training_data"]["systems"]
-        rcut = self.config["model"]["descriptor"]["rcut"]
-        sel = self.config["model"]["descriptor"]["sel"]
-        sec = torch.cumsum(torch.tensor(sel), dim=0)
         type_map = self.config["model"]["type_map"]
         dpdatasystem = CheckSymmetry(
-            sys_path=systems[system_index], rcut=rcut, sec=sec, type_map=type_map
+            sys_path=systems[system_index],type_map=type_map
         )
-        self.origin_batch = dpdatasystem._get_item(batch_index)
+        self.origin_batch = dpdatasystem.get_item(batch_index)
         self.rotated_batch = dpdatasystem.get_rotation(batch_index, self.rotation)
 
     def test_rotation(self):