Merge pull request #121 from njzjz/ase

add an ASE calculator interface

Author: amcadmus

README.md

@@ -32,6 +32,7 @@
 	    - [Include deepmd in the pair style](#include-deepmd-in-the-pair-style)
 	    - [Long-range interaction](#long-range-interaction)
 	- [Run path-integral MD with i-PI](#run-path-integral-md-with-i-pi)
+	- [Use deep potential with ASE](#use-deep-potential-with-ase)
 - [Troubleshooting](#troubleshooting)
 
 # About DeePMD-kit
@@ -553,6 +554,30 @@ The option **`graph_file`** provides the file name of the frozen model.
 
 The `dp_ipi` gets the atom names from an [XYZ file](https://en.wikipedia.org/wiki/XYZ_file_format) provided by **`coord_file`** (meanwhile ignores all coordinates in it), and translates the names to atom types by rules provided by **`atom_type`**.
 
+## Use deep potential with ASE
+
+Deep potential can be set up as a calculator with ASE to obtain potential energies and forces.
+```python
+from ase import Atoms
+from deepmd.calculator import DP
+
+water = Atoms('H2O',
+              positions=[(0.7601, 1.9270, 1),
+                         (1.9575, 1, 1),
+                         (1., 1., 1.)],
+              cell=[100, 100, 100],
+              calculator=DP(model="frozen_model.pb"))
+print(water.get_potential_energy())
+print(water.get_forces())
+```
+
+Optimization is also available:
+```python
+from ase.optimize import BFGS
+dyn = BFGS(water)
+dyn.run(fmax=1e-6)
+print(water.get_positions())
+```
 
 # Troubleshooting
 In consequence of various differences of computers or systems, problems may occur. Some common circumstances are listed as follows. 

source/train/CMakeLists.txt

@@ -2,7 +2,7 @@
 
 configure_file("RunOptions.py.in" "${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py" @ONLY)
 
-file(GLOB LIB_PY main.py common.py env.py compat.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py)
+file(GLOB LIB_PY main.py common.py env.py compat.py calculator.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py)
 
 file(GLOB CLS_PY  Local.py Slurm.py)
 

source/train/calculator.py

@@ -0,0 +1,49 @@
+"""An ASE calculator interface.
+
+Example:
+```python
+from ase import Atoms
+from deepmd.calculator import DP
+
+water = Atoms('H2O',
+              positions=[(0.7601, 1.9270, 1),
+                         (1.9575, 1, 1),
+                         (1., 1., 1.)],
+              cell=[100, 100, 100],
+              calculator=DP(model="frozen_model.pb"))
+print(water.get_potential_energy())
+print(water.get_forces())
+```
+
+Optimization is also available:
+```python
+from ase.optimize import BFGS
+dyn = BFGS(water)
+dyn.run(fmax=1e-6)
+print(water.get_positions())
+```
+"""
+
+from ase.calculators.calculator import Calculator, all_changes
+import deepmd.DeepPot as DeepPot
+
+
+class DP(Calculator):
+    name = "DP"
+    implemented_properties = ["energy", "forces", "stress"]
+
+    def __init__(self, model, label="DP", **kwargs):
+        Calculator.__init__(self, label=label, **kwargs)
+        self.dp = DeepPot(model)
+        self.type_dict = dict(zip(self.dp.get_type_map(), range(self.dp.get_ntypes())))
+
+    def calculate(self, atoms=None, properties=["energy", "forces", "stress"], system_changes=all_changes):
+        coord = atoms.get_positions().reshape([1, -1])
+        cell = atoms.get_cell().reshape([1, -1])
+        symbols = atoms.get_chemical_symbols()
+        atype = [self.type_dict[k] for k in symbols]
+        e, f, v = self.dp.eval(coord, cell, atype)
+        self.results['energy'] = e[0]
+        self.results['forces'] = f[0]
+        self.results['stress'] = v[0]
+