From a0518c30a433614b32a5d648b20dd8329625c54c Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Fri, 28 Jan 2022 15:33:38 +0000
Subject: [PATCH 1/2] add optional_param jobs: list[VaspJob] to
 RunVaspCustodian

---
 atomate/vasp/firetasks/run_calc.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/atomate/vasp/firetasks/run_calc.py b/atomate/vasp/firetasks/run_calc.py
index b15f85a89..46caf1489 100644
--- a/atomate/vasp/firetasks/run_calc.py
+++ b/atomate/vasp/firetasks/run_calc.py
@@ -79,21 +79,22 @@ class RunVaspCustodian(FiretaskBase):
         vasp_cmd (str): the name of the full executable for running VASP. Supports env_chk.
 
     Optional params:
-        job_type: (str) - choose from "normal" (default), "double_relaxation_run" (two consecutive
+        jobs (list[VaspJob]): VaspJobs to run.
+        job_type (str): choose from "normal" (default), "double_relaxation_run" (two consecutive
             jobs), "full_opt_run" (multiple optimizations), and "neb"
         handler_group: (str | list[ErrorHandler]) - group of handlers to use. See handler_groups dict in the code or
             the groups and complete list of handlers in each group. Alternatively, you can
             specify a list of ErrorHandler objects.
         max_force_threshold: (float) - if >0, adds MaxForceErrorHandler. Not recommended for
             nscf runs.
-        scratch_dir: (str) - if specified, uses this directory as the root scratch dir.
+        scratch_dir (str): if specified, uses this directory as the root scratch dir.
             Supports env_chk.
         gzip_output: (bool) - gzip output (default=T)
         max_errors: (int) - maximum # of errors to fix before giving up (default=5)
         ediffg: (float) shortcut for setting EDIFFG in special custodian jobs
         auto_npar: (bool) - use auto_npar (default=F). Recommended set to T
             for single-node jobs only. Supports env_chk.
-        gamma_vasp_cmd: (str) - cmd for Gamma-optimized VASP compilation.
+        gamma_vasp_cmd (str): cmd for Gamma-optimized VASP compilation.
             Supports env_chk.
         wall_time (int): Total wall time in seconds. Activates WalltimeHandler if set.
         half_kpts_first_relax (bool): Use half the k-points for the first relaxation
@@ -101,6 +102,7 @@ class RunVaspCustodian(FiretaskBase):
 
     required_params = ["vasp_cmd"]
     optional_params = [
+        "jobs",
         "job_type",
         "handler_group",
         "max_force_threshold",
@@ -167,7 +169,8 @@ def run_task(self, fw_spec):
             vasp_cmd = shlex.split(vasp_cmd)
 
         # initialize variables
-        job_type = self.get("job_type", "normal")
+        jobs = self.get("jobs", [])
+        job_type = self.get("job_type", "normal" if jobs == [] else None)
         scratch_dir = env_chk(self.get("scratch_dir"), fw_spec)
         gzip_output = self.get("gzip_output", True)
         max_errors = self.get("max_errors", CUSTODIAN_MAX_ERRORS)

From f7a9ecd8cebaa61cf26db048fc5f9cead41aea65 Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Fri, 28 Jan 2022 15:33:49 +0000
Subject: [PATCH 2/2] add type hints to DFPTFW

---
 atomate/vasp/fireworks/core.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/atomate/vasp/fireworks/core.py b/atomate/vasp/fireworks/core.py
index b9bb79d36..a8bcf9e67 100644
--- a/atomate/vasp/fireworks/core.py
+++ b/atomate/vasp/fireworks/core.py
@@ -3,6 +3,7 @@
 sequences of VASP calculations.
 """
 import warnings
+from typing import Any, Dict, List, Union
 
 from fireworks import Firework
 from pymatgen.core import Structure
@@ -591,20 +592,20 @@ def __init__(
 class DFPTFW(Firework):
     def __init__(
         self,
-        structure=None,
-        prev_calc_dir=None,
-        name="static dielectric",
-        vasp_cmd=VASP_CMD,
-        copy_vasp_outputs=True,
-        lepsilon=True,
-        db_file=DB_FILE,
-        parents=None,
-        user_incar_settings=None,
-        pass_nm_results=False,
+        structure: Structure = None,
+        prev_calc_dir: str = None,
+        name: str = "static dielectric",
+        vasp_cmd: str = VASP_CMD,
+        copy_vasp_outputs: bool = True,
+        lepsilon: bool = True,
+        db_file: str = DB_FILE,
+        parents: Union[Firework, List[Firework]] = None,
+        user_incar_settings: Dict[str, Any] = None,
+        pass_nm_results: bool = False,
         **kwargs,
     ):
         """
-         Static DFPT calculation Firework
+        Static DFPT calculation Firework
 
         Args:
             structure (Structure): Input structure. If copy_vasp_outputs, used only to set the