add sqm driver (#286)

Author: njzjz

dpdata/plugins/amber.py

@@ -1,6 +1,11 @@
+import tempfile
+import os
+import subprocess as sp
+
 import dpdata.amber.md
 import dpdata.amber.sqm
 from dpdata.format import Format
+from dpdata.driver import Driver
 
 
 @Format.register("amber/md")
@@ -49,5 +54,70 @@ class SQMINFormat(Format):
     def to_system(self, data, fname=None, frame_idx=0, **kwargs):
         """
         Generate input files for semi-emperical calculation in sqm software
+
+        Parameters
+        ----------
+        data : dict
+            system data
+        fname : str
+            output file name
+        frame_idx : int, default=0
+            index of frame to write
+
+        Other Parameters
+        ----------------
+        **kwargs : dict
+            valid parameters are:
+                theory : str, default=dftb3
+                    level of theory. Options includes AM1, RM1, MNDO, PM3-PDDG, MNDO-PDDG,
+                    PM3-CARB1, MNDO/d, AM1/d, PM6, DFTB2, DFTB3
+                charge : int, default=0
+                    total charge in electron units
+                maxcyc : int, default=0
+                    maximum number of minimization cycles to allow. 0 represents a
+                    single-point calculation
+                mult : int, default=1
+                    multiplicity. Only 1 is allowed.
         """
         return dpdata.amber.sqm.make_sqm_in(data, fname, frame_idx, **kwargs)
+
+
+@Driver.register("sqm")
+class SQMDriver(Driver):
+    """AMBER sqm program driver.
+    
+    Parameters
+    ----------
+    sqm_exec : str, default=sqm
+        path to sqm program
+    **kwargs : dict
+        other arguments to make input files. See :class:`SQMINFormat`
+    
+    Examples
+    --------
+    Use DFTB3 method to calculate potential energy:
+    >>> labeled_system = system.predict(theory="DFTB3", driver="sqm")
+    >>> labeled_system['energies'][0]
+    -15.41111246
+    """
+    def __init__(self, sqm_exec: str="sqm", **kwargs: dict) -> None:
+        self.sqm_exec = sqm_exec
+        self.kwargs = kwargs
+
+    def label(self, data: dict) -> dict:
+        ori_system = dpdata.System(data=data)
+        labeled_system = dpdata.LabeledSystem()
+        with tempfile.TemporaryDirectory() as d:
+            for ii, ss in enumerate(ori_system):
+                inp_fn = os.path.join(d, "%d.in" % ii)
+                out_fn = os.path.join(d, "%d.out" % ii)
+                ss.to("sqm/in", inp_fn, **self.kwargs)
+                try:
+                    sp.check_output([*self.sqm_exec.split(), "-O", "-i", inp_fn, "-o", out_fn])
+                except sp.CalledProcessError as e:
+                    with open(out_fn) as f:
+                        raise RuntimeError(
+                            "Run sqm failed! Output:\n" + f.read()
+                            ) from e
+                labeled_system.append(dpdata.LabeledSystem(out_fn, fmt="sqm/out"))
+        return labeled_system.data

tests/comp_sys.py

@@ -18,10 +18,8 @@ def test_atom_names(self):
                          self.system_2.data['atom_names'])
 
     def test_atom_types(self):
-        self.assertEqual(self.system_1.data['atom_types'][0], 
-                         self.system_2.data['atom_types'][0])
-        self.assertEqual(self.system_1.data['atom_types'][1],
-                         self.system_2.data['atom_types'][1])
+        np.testing.assert_array_equal(self.system_1.data['atom_types'],
+                                      self.system_2.data['atom_types'])
 
     def test_orig(self):
         for d0 in range(3) :

tests/test_sqm_driver.py

@@ -0,0 +1,31 @@
+import unittest
+import shutil
+
+import numpy as np
+from context import dpdata
+from comp_sys import CompSys, IsNoPBC
+
+
+@unittest.skipIf(shutil.which("sqm") is None, "sqm is not installed")
+class TestSQMdriver(unittest.TestCase, CompSys, IsNoPBC):
+    """Test sqm with a hydrogen ion."""
+    @classmethod
+    def setUpClass(cls):
+        cls.system_1 = dpdata.System(data={
+            "atom_names": ["H"],
+            "atom_numbs": [1],
+            "atom_types": np.zeros((1,), dtype=int),
+            "coords": np.zeros((1, 1, 3), dtype=np.float32),
+            "cells": np.zeros((1, 3, 3), dtype=np.float32),
+            "orig": np.zeros(3, dtype=np.float32),
+            "nopbc": True,
+        })
+        cls.system_2 = cls.system_1.predict(theory="DFTB3", charge=1, driver="sqm")
+        cls.places = 6
+    
+    def test_energy(self):
+        self.assertAlmostEqual(self.system_2['energies'].ravel()[0], 6.549447)
+    
+    def test_forces(self):
+        forces = self.system_2['forces']
+        np.testing.assert_allclose(forces, np.zeros_like(forces))