Fix PT DeepPot and replace ASE calculator

deepmd/infer/deep_pot.py

@@ -56,7 +56,7 @@ def __new__(cls, model_file: str, *args, **kwargs):
 
                 return super().__new__(DeepPotTF)
             elif backend == DPBackend.PyTorch:
-                from deepmd_pt.infer.deep_eval import DeepPot as DeepPotPT
+                from deepmd.pt.infer.deep_eval import DeepPot as DeepPotPT
 
                 return super().__new__(DeepPotPT)
             else:

deepmd/pt/infer/deep_eval.py

@@ -195,13 +195,27 @@ def _eval_model(
         )
         if isinstance(batch_output, tuple):
             batch_output = batch_output[0]
-        energy_out = batch_output["energy"].detach().cpu().numpy()
+        energy_out = batch_output["energy"].reshape(nframes, 1).detach().cpu().numpy()
         if "atom_energy" in batch_output:
-            atomic_energy_out = batch_output["atom_energy"].detach().cpu().numpy()
-        force_out = batch_output["force"].detach().cpu().numpy()
-        virial_out = batch_output["virial"].detach().cpu().numpy()
+            atomic_energy_out = (
+                batch_output["atom_energy"]
+                .reshape(nframes, natoms, 1)
+                .detach()
+                .cpu()
+                .numpy()
+            )
+        force_out = (
+            batch_output["force"].reshape(nframes, natoms, 3).detach().cpu().numpy()
+        )
+        virial_out = batch_output["virial"].reshape(nframes, 9).detach().cpu().numpy()
         if "atomic_virial" in batch_output:
-            atomic_virial_out = batch_output["atomic_virial"].detach().cpu().numpy()
+            atomic_virial_out = (
+                batch_output["atomic_virial"]
+                .reshape(nframes, natoms, 9)
+                .detach()
+                .cpu()
+                .numpy()
+            )
 
         if not atomic:
             return energy_out, force_out, virial_out

deepmd/pt/utils/ase_calc.py

@@ -1,65 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-from typing import (
-    ClassVar,
-)
+from deepmd.calculator import DP as DPCalculator
 
-import dpdata
-import numpy as np
-from ase import (
-    Atoms,
-)
-from ase.calculators.calculator import (
-    Calculator,
-    PropertyNotImplementedError,
-)
-
-from deepmd.pt.infer.deep_eval import (
-    DeepPot,
-)
-
-
-class DPCalculator(Calculator):
-    implemented_properties: ClassVar[list] = [
-        "energy",
-        "free_energy",
-        "forces",
-        "virial",
-        "stress",
-    ]
-
-    def __init__(self, model):
-        Calculator.__init__(self)
-        self.dp = DeepPot(model)
-        self.type_map = self.dp.type_map
-
-    def calculate(self, atoms: Atoms, properties, system_changes) -> None:
-        Calculator.calculate(self, atoms, properties, system_changes)
-        system = dpdata.System(atoms, fmt="ase/structure")
-        type_trans = np.array(
-            [self.type_map.index(i) for i in system.data["atom_names"]]
-        )
-        input_coords = system.data["coords"]
-        input_cells = system.data["cells"]
-        input_types = list(type_trans[system.data["atom_types"]])
-        model_predict = self.dp.eval(input_coords, input_cells, input_types)
-        self.results = {
-            "energy": model_predict[0].item(),
-            "free_energy": model_predict[0].item(),
-            "forces": model_predict[1].reshape(-1, 3),
-            "virial": model_predict[2].reshape(3, 3),
-        }
-
-        # convert virial into stress for lattice relaxation
-        if "stress" in properties:
-            if sum(atoms.get_pbc()) > 0 or (atoms.cell is not None):
-                # the usual convention (tensile stress is positive)
-                # stress = -virial / volume
-                stress = (
-                    -0.5
-                    * (self.results["virial"].copy() + self.results["virial"].copy().T)
-                    / atoms.get_volume()
-                )
-                # Voigt notation
-                self.results["stress"] = stress.flat[[0, 4, 8, 5, 2, 1]]
-            else:
-                raise PropertyNotImplementedError
+__all__ = [
+    "DPCalculator",
+]

source/tests/pt/test_calculator.py

@@ -66,6 +66,7 @@ def test_calculator(self):
             # positions=[tuple(item) for item in coordinate],
             cell=cell,
             calculator=self.calculator,
+            pbc=True,
         )
         e0, f0 = ase_atoms0.get_potential_energy(), ase_atoms0.get_forces()
         s0, v0 = (
@@ -79,6 +80,7 @@ def test_calculator(self):
             # positions=[tuple(item) for item in coordinate],
             cell=cell,
             calculator=self.calculator,
+            pbc=True,
         )
         e1, f1 = ase_atoms1.get_potential_energy(), ase_atoms1.get_forces()
         s1, v1 = (

source/tests/pt/test_deeppot.py

@@ -10,6 +10,7 @@
 
 import numpy as np
 
+from deepmd.infer.deep_pot import DeepPot as DeepPotUni
 from deepmd.pt.entrypoints.main import (
     get_trainer,
 )
@@ -79,3 +80,18 @@ def test_dp_test(self):
         atype = np.array([0, 0, 0, 1, 1]).reshape(1, -1)
 
         e, f, v, ae, av = dp.eval(coord, cell, atype, atomic=True)
+        self.assertEqual(e.shape, (1, 1))
+        self.assertEqual(f.shape, (1, 5, 3))
+        self.assertEqual(v.shape, (1, 9))
+        self.assertEqual(ae.shape, (1, 5, 1))
+        self.assertEqual(av.shape, (1, 5, 9))
+
+        self.assertEqual(dp.get_type_map(), ["O", "H"])
+        self.assertEqual(dp.get_ntypes(), 2)
+        self.assertEqual(dp.get_dim_fparam(), 0)
+        self.assertEqual(dp.get_dim_aparam(), 0)
+
+    def test_uni(self):
+        dp = DeepPotUni("model.pt")
+        self.assertIsInstance(dp, DeepPot)
+        # its methods has been tested in test_dp_test