add psi4/out format (#457)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz

dpdata/plugins/psi4.py

@@ -0,0 +1,52 @@
+import numpy as np
+
+from dpdata.format import Format
+from dpdata.psi4.output import read_psi4_output
+from dpdata.unit import EnergyConversion, ForceConversion, LengthConversion
+
+length_convert = LengthConversion("bohr", "angstrom").value()
+energy_convert = EnergyConversion("hartree", "eV").value()
+force_convert = ForceConversion("hartree/bohr", "eV/angstrom").value()
+
+
+@Format.register("psi4/out")
+class PSI4OutFormat(Format):
+    """Psi4 output.
+
+    Note that both the energy and the gradient should be
+    printed into the output file.
+    """
+
+    def from_labeled_system(self, file_name: str, **kwargs) -> dict:
+        """Read from Psi4 output.
+
+        Parameters
+        ----------
+        file_name : str
+            file name
+        **kwargs
+            keyword arguments
+
+        Returns
+        -------
+        dict
+            system data
+        """
+        symbols, coord, energy, forces = read_psi4_output(file_name)
+
+        atom_names, atom_types, atom_numbs = np.unique(
+            symbols, return_inverse=True, return_counts=True
+        )
+        natoms = coord.shape[0]
+
+        return {
+            "atom_types": atom_types,
+            "atom_names": list(atom_names),
+            "atom_numbs": list(atom_numbs),
+            "coords": (coord * length_convert).reshape((1, natoms, 3)),
+            "energies": np.array([energy * energy_convert]),
+            "forces": (forces * force_convert).reshape((1, natoms, 3)),
+            "cells": np.zeros((1, 3, 3)),
+            "orig": np.zeros(3),
+            "nopbc": True,
+        }

dpdata/psi4/__init__.py

@@ -0,0 +1,66 @@
+from typing import Tuple
+
+import numpy as np
+
+
+def read_psi4_output(fn: str) -> Tuple[str, np.ndarray, float, np.ndarray]:
+    """Read from Psi4 output.
+
+    Note that both the energy and the gradient should be printed.
+
+    Parameters
+    ----------
+    fn : str
+        file name
+
+    Returns
+    -------
+    str
+        atomic symbols
+    np.ndarray
+        atomic coordinates
+    float
+        total potential energy
+    np.ndarray
+        atomic forces
+    """
+    coord = None
+    symbols = None
+    forces = None
+    energy = None
+    with open(fn) as f:
+        flag = 0
+        for line in f:
+            if flag in (1, 3, 4, 5, 6):
+                flag += 1
+            elif flag == 2:
+                s = line.split()
+                if not len(s):
+                    flag = 0
+                else:
+                    symbols.append(s[0].capitalize())
+                    coord.append([float(s[1]), float(s[2]), float(s[3])])
+            elif flag == 7:
+                s = line.split()
+                if not len(s):
+                    flag = 0
+                else:
+                    forces.append([float(s[1]), float(s[2]), float(s[3])])
+            elif line.startswith(
+                "       Center              X                  Y                   Z               Mass"
+            ):
+                # coord
+                flag = 1
+                coord = []
+                symbols = []
+            elif line.startswith("  ## Total Gradient"):
+                flag = 3
+                forces = []
+            elif line.startswith("    Total Energy ="):
+                energy = float(line.split()[-1])
+    symbols = np.array(symbols)
+    forces = -np.array(forces)
+    coord = np.array(coord)
+    assert coord.shape == forces.shape
+
+    return symbols, coord, energy, forces