Merge pull request #260 from materialsproject/lobster_tidying

Tidy lobster workflow
materialsproject · Mar 12, 2023 · 69f4743 · 69f4743
2 parents 4668fc3 + bf5dbf1
commit 69f4743
Show file tree

Hide file tree

Showing 23 changed files with 456 additions and 561 deletions.
diff --git a/src/atomate2/lobster/jobs.py b/src/atomate2/lobster/jobs.py
@@ -25,32 +25,29 @@
 class LobsterMaker(Maker):
     """
     LOBSTER job maker.
-    The maker copies the DFT output files
-    necessary for the LOBSTER run.
-    It will create all lobsterin files, run LOBSTER,
-    zip the outputs and parse the LOBSTER outputs.
+
+    The maker copies DFT output files necessary for the LOBSTER run. It will create all
+    lobsterin files, run LOBSTER, zip the outputs and parse the LOBSTER outputs.
 
     Parameters
     ----------
     name : str
         Name of jobs produced by this maker.
-    resubmit : bool
-        Maybe useful.
     task_document_kwargs : dict
         Keyword arguments passed to :obj:`.LobsterTaskDocument.from_directory`.
-    user_lobsterin_settings: dict
+    user_lobsterin_settings : dict
         Dict including additional information on the Lobster settings.
-    additional_outputs: list[str]
-        A list including additional output files.
-    calculation_type: str
-        Type of calculation for the Lobster run.
+    run_lobster_kwargs : dict
+        Keyword arguments that will get passed to :obj:`.run_lobster`.
+    calculation_type : str
+        Type of calculation for the Lobster run that will get passed to
+        :obj:`.Lobsterin.standard_calculations_from_vasp_files`.
     """
 
     name: str = "lobster"
-    resubmit: bool = False
     task_document_kwargs: dict = field(default_factory=dict)
     user_lobsterin_settings: dict | None = None
-    additional_outputs: list | None = None
+    run_lobster_kwargs: dict = field(default_factory=dict)
     calculation_type: str = "standard"
 
     @job(output_schema=LobsterTaskDocument, data=[CompleteCohp, LobsterCompleteDos])
@@ -84,17 +81,16 @@ def make(
                     lobsterin[key] = parameter
 
         lobsterin.write_lobsterin("lobsterin")
+
         # run lobster
         logger.info("Running LOBSTER")
-        run_lobster()
+        run_lobster(**self.run_lobster_kwargs)
 
         # gzip folder
         gzip_dir(".")
 
         # parse lobster outputs
-        task_doc = LobsterTaskDocument.from_directory(
+        return LobsterTaskDocument.from_directory(
             Path.cwd(),
             **self.task_document_kwargs,
-            additional_fields=self.additional_outputs,
         )
-        return task_doc
diff --git a/src/atomate2/lobster/run.py b/src/atomate2/lobster/run.py
@@ -19,10 +19,7 @@
 
 from atomate2 import SETTINGS
 
-__all__ = [
-    "JobType",
-    "run_lobster",
-]
+__all__ = ["JobType", "run_lobster"]
 
 _DEFAULT_VALIDATORS = (LobsterFilesValidator(), EnoughBandsValidator())
 _DEFAULT_HANDLERS = ()
@@ -57,17 +54,17 @@ def run_lobster(
 
     Parameters
     ----------
-    job_type: str or .JobType
+    job_type : str or .JobType
          The job type.
-    lobster_cmd: str
+    lobster_cmd : str
         Command to run lobster.
-    max_errors:
+    max_errors : int
         Maximum number of errors.
-    scratch_dir: str or Path
+    scratch_dir : str or Path
         Scratch directory.
-    validators: list of .Validator
+    validators : list of .Validator
         The validators handlers used by custodian.
-    lobster_job_kwargs: dict
+    lobster_job_kwargs : dict
         Keyword arguments that are passed to :obj:`.LosterJob`.
     custodian_kwargs : dict
          Keyword arguments that are passed to :obj:`.Custodian`.