add psi4/inp format (#564)

Signed-off-by: Jinzhe Zeng <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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,
+                )
+            )

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,
+    )

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])

tests/test_psi4.py

@@ -1,11 +1,13 @@
+import tempfile
+import textwrap
 import unittest
 
 import numpy as np
 from comp_sys import CompLabeledSys, IsNoPBC
 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()
+            """
+            ),
+        )