implement Gaussian (#100)

Fix #99.

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>

.github/workflows/test.yml

@@ -19,7 +19,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: | 
-        pip install -e .
+        pip install -e .[test]
         pip install mock coverage pytest
     - name: Test
       run: SKIP_UT_WITH_DFLOW=0 DFLOW_DEBUG=1 coverage run --source=./dpgen2 -m unittest && coverage report

dpgen2/fp/__init__.py

@@ -3,11 +3,21 @@
     PrepVasp,
     RunVasp,
 )
+from .gaussian import (
+    GaussianInputs,
+    PrepGaussian,
+    RunGaussian,
+)
 
 fp_styles = {
     "vasp" :  {
         "inputs" : VaspInputs,
         "prep" : PrepVasp,
         "run" : RunVasp,
-    }
+    },
+    "gaussian" : {
+        "inputs" : GaussianInputs,
+        "prep" : PrepGaussian,
+        "run" : RunGaussian,
+    },
 }

dpgen2/fp/gaussian.py

@@ -0,0 +1,158 @@
+"""Prep and Run Gaussian tasks."""
+from dflow.python import (
+    TransientError,
+)
+from typing import (
+    Tuple,
+    List,
+    Any,
+)
+import dpdata
+from dargs import (
+    dargs, 
+    Argument,
+)
+
+from .prep_fp import PrepFp
+from .run_fp import RunFp
+from dpgen2.constants import (
+    fp_default_out_data_name,
+)
+from dpgen2.utils.run_command import run_command
+
+# global static variables
+gaussian_input_name = 'task.gjf'
+# this output name is generated by Gaussian
+gaussian_output_name = 'task.log'
+
+
+class GaussianInputs:
+    @staticmethod
+    def args() -> List[Argument]:
+        r"""The arguments of the GaussianInputs class."""
+        doc_keywords = "Gaussian keywords, e.g. force b3lyp/6-31g**. If a list, run multiple steps."
+        doc_multiplicity = (
+            "spin multiplicity state. It can be a number. If auto, multiplicity will be detected "
+            "automatically, with the following rules:\n\n"
+            "fragment_guesses=True multiplicity will +1 for each radical, and +2 for each oxygen molecule\n\n"
+            "fragment_guesses=False multiplicity will be 1 or 2, but +2 for each oxygen molecule."
+        )
+        doc_charge = "molecule charge. Only used when charge is not provided by the system"
+        doc_basis_set = "custom basis set"
+        doc_keywords_high_multiplicity = (
+            "keywords for points with multiple raicals. multiplicity should be auto. "
+            "If not set, fallback to normal keywords"
+        )
+        doc_fragment_guesses = (
+            "initial guess generated from fragment guesses. If True, multiplicity should be auto"
+        )
+        doc_nproc = "Number of CPUs to use"
+
+        return [
+            Argument('keywords', [str, list], optional=False, doc=doc_keywords),
+            Argument('multiplicity', [int, str], optional=True, default="auto", doc=doc_multiplicity),
+            Argument('charge', int, optional=True, default=0, doc=doc_charge),
+            Argument('basis_set', str, optional=True, doc=doc_basis_set),
+            Argument('keywords_high_multiplicity', str, optional=True, doc=doc_keywords_high_multiplicity),
+            Argument('fragment_guesses', bool, optional=True, default=False, doc=doc_fragment_guesses),
+            Argument('nproc', int, optional=True, default=1, doc=doc_nproc),
+        ]
+
+    def __init__(self, **kwargs: Any):
+        self.data = kwargs
+
+
+class PrepGaussian(PrepFp):
+    def prep_task(
+            self,
+            conf_frame: dpdata.System,
+            inputs: GaussianInputs,
+    ):
+        r"""Define how one Gaussian task is prepared.
+
+        Parameters
+        ----------
+        conf_frame : dpdata.System
+            One frame of configuration in the dpdata format.
+        inputs: GaussianInputs
+            The GaussianInputs object handels all other input files of the task.
+        """
+
+        conf_frame.to('gaussian/gjf', gaussian_input_name, **inputs.data)
+
+
+class RunGaussian(RunFp):
+    def input_files(self) -> List[str]:
+        r"""The mandatory input files to run a Gaussian task.
+        
+        Returns
+        -------
+        files: List[str]
+            A list of madatory input files names.
+
+        """
+        return [gaussian_input_name]
+
+    def optional_input_files(self) -> List[str]:
+        r"""The optional input files to run a Gaussian task.
+        
+        Returns
+        -------
+        files: List[str]
+            A list of optional input files names.
+
+        """
+        return []
+
+    def run_task(
+            self,
+            command : str,
+            out_name: str,
+    ) -> Tuple[str, str]:
+        r"""Defines how one FP task runs
+        
+        Parameters
+        ----------
+        command: str
+            The command of running gaussian task
+        out_name: str
+            The name of the output data file.
+
+        Returns
+        -------
+        out_name: str
+            The file name of the output data in the dpdata.LabeledSystem format.
+        log_name: str
+            The file name of the log.
+        """
+        # run gaussian
+        command = ' '.join([command, gaussian_input_name])
+        ret, out, err = run_command(command, shell=True)
+        if ret != 0:
+            raise TransientError(
+                'gaussian failed\n',
+                'out msg', out, '\n',
+                'err msg', err, '\n'
+            )                    
+        # convert the output to deepmd/npy format
+        sys = dpdata.LabeledSystem(gaussian_output_name, fmt='gaussian/log')
+        sys.to('deepmd/npy', out_name)
+        return out_name, gaussian_output_name
+
+
+    @staticmethod
+    def args() -> List[dargs.Argument]:
+        r"""The argument definition of the `run_task` method.
+
+        Returns
+        -------
+        arguments: List[dargs.Argument]
+            List of dargs.Argument defines the arguments of `run_task` method.
+        """
+
+        doc_gaussian_cmd = "The command of Gaussian"
+        doc_gaussian_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`."
+        return [
+            Argument("command", str, optional=True, default='g16', doc=doc_gaussian_cmd),
+            Argument("out", str, optional=True, default=fp_default_out_data_name, doc=doc_gaussian_out),
+        ]

pyproject.toml

@@ -44,6 +44,9 @@ docs = [
      'sphinx-argparse',
      "dargs>=0.3.1",
 ]
+test = [
+     'fakegaussian>=0.0.3',
+]
 
 [tool.setuptools.packages.find]
 include = ["dpgen2*"]

tests/test_prep_run_gaussian.py

@@ -0,0 +1,68 @@
+import numpy as np
+import unittest
+from pathlib import Path
+
+from dpgen2.fp.gaussian import (
+    GaussianInputs,
+    PrepGaussian,
+    RunGaussian,
+    dpdata,
+    gaussian_input_name,
+    gaussian_output_name,
+)
+from dargs import Argument
+
+class TestPrepGaussian(unittest.TestCase):
+    def test_prep_gaussian(self):
+        inputs = GaussianInputs(
+            keywords="force b3lyp/6-31g*",
+            multiplicity=1,
+        )
+        ta = GaussianInputs.args()
+        base = Argument("base", dict, ta)
+        data = base.normalize_value(inputs.data, trim_pattern="_*")
+        base.check_value(data, strict=True)
+        system = dpdata.LabeledSystem(data={
+            'atom_names': ['H'],
+            'atom_numbs': [1],
+            'atom_types': np.zeros(1, dtype=int),
+            'cells': np.eye(3).reshape(1, 3, 3),
+            'coords': np.zeros((1, 1, 3)),
+            'energies': np.zeros(1),
+            'forces': np.zeros((1, 1, 3)),
+            'orig': np.zeros(3),
+            'nopbc': True,
+        })
+        prep_gaussian = PrepGaussian()
+        prep_gaussian.prep_task(
+            conf_frame=system,
+            inputs=inputs,
+        )
+        assert Path(gaussian_input_name).exists()
+
+
+class TestRunGaussian(unittest.TestCase):
+    def test_run_gaussian(self):
+        dpdata.LabeledSystem(data={
+            'atom_names': ['H'],
+            'atom_numbs': [1],
+            'atom_types': np.zeros(1, dtype=int),
+            'cells': np.eye(3).reshape(1, 3, 3),
+            'coords': np.zeros((1, 1, 3)),
+            'energies': np.zeros(1),
+            'forces': np.zeros((1, 1, 3)),
+            'orig': np.zeros(3),
+            'nopbc': True,
+        }).to_gaussian_gjf(
+            gaussian_input_name,
+            keywords="force b3lyp/6-31g*",
+            multiplicity=1
+        )
+        run_gaussian = RunGaussian()
+        output = 'mock_output'
+        out_name, log_name = run_gaussian.run_task(
+            'g16',
+            output,
+        )
+        assert out_name == output
+        assert log_name == gaussian_output_name