From 55941f40326bb03b724d6323d257da707384cb57 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 8 Jun 2022 17:00:22 -0400 Subject: [PATCH 1/3] add ase calculator --- dpdata/ase_calculator.py | 76 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 dpdata/ase_calculator.py diff --git a/dpdata/ase_calculator.py b/dpdata/ase_calculator.py new file mode 100644 index 000000000..df89d05cb --- /dev/null +++ b/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 From 292c7007326cb84ac4001cd865a647ab30c33adf Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 9 Jun 2022 02:52:33 -0400 Subject: [PATCH 2/3] add ASEDriver --- dpdata/plugins/ase.py | 41 +++++++++++++++++++++++++++++++++++++++++ tests/test_predict.py | 25 +++++++++++++++++++++++-- 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/dpdata/plugins/ase.py b/dpdata/plugins/ase.py index 5f9dff077..f1745b472 100644 --- a/dpdata/plugins/ase.py +++ b/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 diff --git a/tests/test_predict.py b/tests/test_predict.py index c5fe5d76c..8535d02ea 100644 --- a/tests/test_predict.py +++ b/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 From 31dfbaa820188a4b41f815eefa47eca40946b420 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 9 Jun 2022 02:58:22 -0400 Subject: [PATCH 3/3] add ase_calculator property --- dpdata/driver.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/dpdata/driver.py b/dpdata/driver.py index 29ae2703f..97583a101 100644 --- a/dpdata/driver.py +++ b/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):