dpdata driver <--> ase calculator (#302)

* add ase calculator

* add ASEDriver

* add ase_calculator property

Author: njzjz

dpdata/ase_calculator.py

@@ -0,0 +1,76 @@
+from typing import List, Optional, TYPE_CHECKING
+
+from ase.calculators.calculator import (
+    Calculator, all_changes, PropertyNotImplementedError
+)
+
+import dpdata
+from .driver import Driver
+
+if TYPE_CHECKING:
+    from ase import Atoms
+
+
+class DPDataCalculator(Calculator):
+    """Implementation of ASE deepmd calculator based on a driver. 
+
+    Parameters
+    ----------
+    driver : Driver
+        dpdata driver
+    """
+
+    name = "dpdata"
+    implemented_properties = [
+        "energy", "free_energy", "forces", "virial", "stress"]
+
+    def __init__(
+        self,
+        driver: Driver,
+        **kwargs
+    ) -> None:
+        Calculator.__init__(self, label=Driver.__name__, **kwargs)
+        self.driver = driver
+
+    def calculate(
+        self,
+        atoms: Optional["Atoms"] = None,
+        properties: List[str] = ["energy", "forces"],
+        system_changes: List[str] = all_changes,
+    ):
+        """Run calculation with a driver.
+
+        Parameters
+        ----------
+        atoms : Optional[Atoms], optional
+            atoms object to run the calculation on, by default None
+        properties : List[str], optional
+            unused, only for function signature compatibility,
+            by default ["energy", "forces"]
+        system_changes : List[str], optional
+            unused, only for function signature compatibility, by default all_changes
+        """
+        if atoms is not None:
+            self.atoms = atoms.copy()
+
+        system = dpdata.System(self.atoms, fmt="ase/structure")
+        data = system.predict(driver=self.driver).data
+
+        self.results['energy'] = data['energies'][0]
+        # see https://gitlab.com/ase/ase/-/merge_requests/2485
+        self.results['free_energy'] = data['energies'][0]
+        self.results['forces'] = data['forces'][0]
+        if 'virials' in data:
+            self.results['virial'] = data['virials'][0].reshape(3, 3)
+
+        # convert virial into stress for lattice relaxation
+        if "stress" in properties:
+            if sum(atoms.get_pbc()) > 0:
+                # the usual convention (tensile stress is positive)
+                # stress = -virial / volume
+                stress = -0.5 * (data['virials'][0].copy() + data['virials'][0].copy().T) / \
+                    atoms.get_volume()
+                # Voigt notation
+                self.results['stress'] = stress.flat[[0, 4, 8, 5, 2, 1]]
+            else:
+                raise PropertyNotImplementedError

dpdata/driver.py

@@ -1,8 +1,10 @@
 """Driver plugin system."""
-from typing import Callable, List, Union
+from typing import Callable, List, Union, TYPE_CHECKING
 from .plugin import Plugin
 from abc import ABC, abstractmethod
 
+if TYPE_CHECKING:
+    import ase
 
 class Driver(ABC):
     """The base class for a driver plugin. A driver can
@@ -79,6 +81,12 @@ def label(self, data: dict) -> dict:
         """
         return NotImplemented
 
+    @property
+    def ase_calculator(self) -> "ase.calculators.calculator.Calculator":
+        """Returns an ase calculator based on this driver."""
+        from .ase_calculator import DPDataCalculator
+        return DPDataCalculator(self)
+
 
 @Driver.register("hybrid")
 class HybridDriver(Driver):

dpdata/plugins/ase.py

@@ -1,5 +1,7 @@
+from dpdata.driver import Driver
 from dpdata.format import Format
 import numpy as np
+import dpdata
 try:
     import ase.io
     from ase.calculators.calculator import PropertyNotImplementedError
@@ -162,3 +164,42 @@ def to_labeled_system(self, data, *args, **kwargs):
             structures.append(structure)
 
         return structures
+
+
+@Driver.register("ase")
+class ASEDriver(Driver):
+    """ASE Driver.
+    
+    Parameters
+    ----------
+    calculator : ase.calculators.calculator.Calculato
+        ASE calculator
+    """
+
+    def __init__(self, calculator: "ase.calculators.calculator.Calculator") -> None:
+        """Setup the driver."""
+        self.calculator = calculator
+
+    def label(self, data: dict) -> dict:
+        """Label a system data. Returns new data with energy, forces, and virials.
+        
+        Parameters
+        ----------
+        data : dict
+            data with coordinates and atom types
+        
+        Returns
+        -------
+        dict
+            labeled data with energies and forces
+        """
+        # convert data to ase data
+        system = dpdata.System(data=data)
+        # list[Atoms]
+        structures = system.to_ase_structure()
+        labeled_system = dpdata.LabeledSystem()
+        for atoms in structures:
+            atoms.calc = self.calculator
+            ls = dpdata.LabeledSystem(atoms, fmt="ase/structure", type_map=data['atom_names'])
+            labeled_system.append(ls)
+        return labeled_system.data

tests/test_predict.py

@@ -1,8 +1,14 @@
 import unittest
 import numpy as np
 
-from comp_sys import CompLabeledSys
+from comp_sys import CompLabeledSys, IsPBC
 from context import dpdata
+try:
+    import ase
+except ModuleNotFoundError:
+    skip_ase = True
+else:
+    skip_ase = False
 
 
 @dpdata.driver.Driver.register("zero")
@@ -17,7 +23,7 @@ def label(self, data):
 
 
 @dpdata.driver.Driver.register("one")
-class ZeroDriver(dpdata.driver.Driver):
+class OneDriver(dpdata.driver.Driver):
     def label(self, data):
         nframes = data['coords'].shape[0]
         natoms = data['coords'].shape[1]
@@ -69,3 +75,18 @@ def setUp(self) :
         self.e_places = 6
         self.f_places = 6
         self.v_places = 6
+
+
+@unittest.skipIf(skip_ase,"skip ase related test. install ase to fix")
+class TestASEtraj1(unittest.TestCase, CompLabeledSys, IsPBC):
+    def setUp (self) :
+        ori_sys = dpdata.LabeledSystem('poscars/deepmd.h2o.md', 
+                                        fmt = 'deepmd/raw', 
+                                        type_map = ['O', 'H'])
+        one_driver = OneDriver()
+        self.system_1 = ori_sys.predict(driver=one_driver)
+        self.system_2 = ori_sys.predict(one_driver.ase_calculator, driver="ase")
+        self.places = 6
+        self.e_places = 6
+        self.f_places = 6
+        self.v_places = 4