diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index b863a7ddc2..546c0f3c7e 100644 --- a/deepmd/infer/deep_pot.py +++ b/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: diff --git a/deepmd/pt/utils/ase_calc.py b/deepmd/pt/utils/ase_calc.py index 8d5fe8bce9..6bcb9cdc5e 100644 --- a/deepmd/pt/utils/ase_calc.py +++ b/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", +] diff --git a/source/tests/pt/test_calculator.py b/source/tests/pt/test_calculator.py index e8382b22b8..a35538250b 100644 --- a/source/tests/pt/test_calculator.py +++ b/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 = (