From 94741beb95433eec13e9e284db17a3ff61b6a42d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 31 Oct 2023 20:18:14 -0400 Subject: [PATCH] add psi4/inp format (#564) Signed-off-by: Jinzhe Zeng Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- dpdata/plugins/psi4.py | 51 +++++++++++++++++++++++++++++++++++++ dpdata/psi4/input.py | 57 ++++++++++++++++++++++++++++++++++++++++++ tests/comp_sys.py | 3 +++ tests/test_psi4.py | 35 +++++++++++++++++++++++++- 4 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 dpdata/psi4/input.py diff --git a/dpdata/plugins/psi4.py b/dpdata/plugins/psi4.py index 932d118c..68919b13 100644 --- a/dpdata/plugins/psi4.py +++ b/dpdata/plugins/psi4.py @@ -1,6 +1,7 @@ import numpy as np from dpdata.format import Format +from dpdata.psi4.input import write_psi4_input from dpdata.psi4.output import read_psi4_output from dpdata.unit import EnergyConversion, ForceConversion, LengthConversion @@ -50,3 +51,53 @@ def from_labeled_system(self, file_name: str, **kwargs) -> dict: "orig": np.zeros(3), "nopbc": True, } + + +@Format.register("psi4/inp") +class PSI4InputFormat(Format): + """Psi4 input file.""" + + def to_system( + self, + data: dict, + file_name: str, + method: str, + basis: str, + charge: int = 0, + multiplicity: int = 1, + frame_idx=0, + **kwargs, + ): + """Write PSI4 input. + + Parameters + ---------- + data : dict + system data + file_name : str + file name + method : str + computational method + basis : str + basis set; see https://psicode.org/psi4manual/master/basissets_tables.html + charge : int, default=0 + charge of system + multiplicity : int, default=1 + multiplicity of system + frame_idx : int, default=0 + The index of the frame to dump + **kwargs + keyword arguments + """ + types = np.array(data["atom_names"])[data["atom_types"]] + with open(file_name, "w") as fout: + fout.write( + write_psi4_input( + types=types, + coords=data["coords"][frame_idx], + method=method, + basis=basis, + charge=charge, + multiplicity=multiplicity, + ) + ) diff --git a/dpdata/psi4/input.py b/dpdata/psi4/input.py new file mode 100644 index 00000000..ad053281 --- /dev/null +++ b/dpdata/psi4/input.py @@ -0,0 +1,57 @@ +import numpy as np + +# Angston is used in Psi4 by default +template = """molecule {{ +{atoms:s} +{charge:d} {multiplicity:d} +}} +set basis {basis:s} +set gradient_write on +G, wfn = gradient("WB97M-D3BJ", return_wfn=True) +wfn.energy() +wfn.gradient().print_out() +""" + + +def write_psi4_input( + types: np.ndarray, + coords: np.ndarray, + method: str, + basis: str, + charge: int = 0, + multiplicity: int = 1, +) -> str: + """Write Psi4 input file. + + Parameters + ---------- + types : np.ndarray + atomic symbols + coords : np.ndarray + atomic coordinates + method : str + computational method + basis : str + basis set; see https://psicode.org/psi4manual/master/basissets_tables.html + charge : int, default=0 + charge of system + multiplicity : int, default=1 + multiplicity of system + + Returns + ------- + str + content of Psi4 input file + """ + return template.format( + atoms="\n".join( + [ + "{:s} {:16.9f} {:16.9f} {:16.9f}".format(*ii) + for ii in zip(types, *coords.T) + ] + ), + charge=charge, + multiplicity=multiplicity, + method=method, + basis=basis, + ) diff --git a/tests/comp_sys.py b/tests/comp_sys.py index 426df412..cfa92c49 100644 --- a/tests/comp_sys.py +++ b/tests/comp_sys.py @@ -47,6 +47,9 @@ def test_cell(self): def test_coord(self): self.assertEqual(self.system_1.get_nframes(), self.system_2.get_nframes()) # think about direct coord + if self.system_1.nopbc: + # nopbc doesn't need to test cells + return tmp_cell = self.system_1.data["cells"] tmp_cell = np.reshape(tmp_cell, [-1, 3]) tmp_cell_norm = np.reshape(np.linalg.norm(tmp_cell, axis=1), [-1, 1, 3]) diff --git a/tests/test_psi4.py b/tests/test_psi4.py index 618e0517..b9c2124e 100644 --- a/tests/test_psi4.py +++ b/tests/test_psi4.py @@ -1,3 +1,5 @@ +import tempfile +import textwrap import unittest import numpy as np @@ -5,7 +7,7 @@ from context import dpdata -class TestDeepmdLoadDumpHDF5(unittest.TestCase, CompLabeledSys, IsNoPBC): +class TestPsi4Output(unittest.TestCase, CompLabeledSys, IsNoPBC): def setUp(self): length_convert = dpdata.unit.LengthConversion("bohr", "angstrom").value() energy_convert = dpdata.unit.EnergyConversion("hartree", "eV").value() @@ -60,3 +62,34 @@ def setUp(self): self.e_places = 6 self.f_places = 6 self.v_places = 6 + + +class TestPsi4Input(unittest.TestCase): + def test_psi4_input(self): + system = dpdata.LabeledSystem("psi4/psi4.out", fmt="psi4/out") + with tempfile.NamedTemporaryFile("r") as f: + system.to_psi4_inp(f.name, method="WB97M-D3BJ", basis="def2-TZVPPD") + content = f.read() + self.assertEqual( + content, + textwrap.dedent( + """\ + molecule { + C 0.692724290 -0.280972290 0.149966626 + C -0.690715864 0.280527594 -0.157432416 + H 1.241584247 -0.707702380 -0.706342645 + H 0.492994289 -1.086482665 0.876517411 + H 1.330104482 0.478557878 0.633157279 + H -1.459385451 -0.498843014 -0.292862623 + H -0.623545813 0.873377524 -1.085142510 + H -1.005665735 0.946387574 0.663566976 + 0 1 + } + set basis def2-TZVPPD + set gradient_write on + G, wfn = gradient("WB97M-D3BJ", return_wfn=True) + wfn.energy() + wfn.gradient().print_out() + """ + ), + )