diff --git a/dpdata/plugins/amber.py b/dpdata/plugins/amber.py index 2c6fa56a5..d2a13c5bf 100644 --- a/dpdata/plugins/amber.py +++ b/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 diff --git a/tests/comp_sys.py b/tests/comp_sys.py index 2b87828c5..3996b0f42 100644 --- a/tests/comp_sys.py +++ b/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) : diff --git a/tests/test_sqm_driver.py b/tests/test_sqm_driver.py new file mode 100644 index 000000000..3755fac97 --- /dev/null +++ b/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))