pre-commit run --all-files

janosh · janosh · commit a3efc15b4e87 · 2022-04-26T11:05:37.000+01:00
diff --git a/atomate/common/firetasks/glue_tasks.py b/atomate/common/firetasks/glue_tasks.py
@@ -99,9 +99,13 @@ class CopyFilesFromCalcLoc(FiretaskBase):
     """
 
     required_params = ["calc_loc"]
-    optional_params = ["filenames", "name_prepend",
-                       "name_append", "exclude_files",
-                       "decompress"]
+    optional_params = [
+        "filenames",
+        "name_prepend",
+        "name_append",
+        "exclude_files",
+        "decompress",
+    ]
 
     def run_task(self, fw_spec=None):
         calc_loc = get_calc_loc(self["calc_loc"], fw_spec["calc_locs"])
@@ -146,13 +150,16 @@ def run_task(self, fw_spec=None):
         for f in files_to_copy:
             prev_path_full = os.path.join(calc_dir, f)
             f, ext = os.path.splitext(f)
-            dest_fname = self.get("name_prepend", "") + f + self.get("name_append", "") + ext
+            dest_fname = (
+                self.get("name_prepend", "") + f + self.get("name_append", "") + ext
+            )
             dest_path = os.path.join(os.getcwd(), dest_fname)
 
             fileclient.copy(prev_path_full, dest_path)
-            if self.get("decompress",False):
+            if self.get("decompress", False):
                 monty.shutil.decompress_file(dest_path)
 
+
 @explicit_serialize
 class DeleteFiles(FiretaskBase):
     """
diff --git a/atomate/vasp/analysis/phonopy.py b/atomate/vasp/analysis/phonopy.py
@@ -179,6 +179,6 @@ def get_phonopy_thermal_expansion(
 
     # thermal expansion coefficient and temperature
     max_t_index = phonopy_qha._qha._len
-    alpha = phonopy_qha.get_thermal_expansion()[: max_t_index]
+    alpha = phonopy_qha.get_thermal_expansion()[:max_t_index]
     T = phonopy_qha._qha._temperatures[:max_t_index]
     return alpha, T
diff --git a/atomate/vasp/drones.py b/atomate/vasp/drones.py
@@ -420,7 +420,7 @@ def generate_doc(self, dir_name, vasprun_files, outcar_files):
                     d["output"][k] = d_calc_final["output"][k]
 
             # store optical data, overwrites the LOPTICS data
-            if d["input"]["incar"].get("ALGO") == 'CHI':
+            if d["input"]["incar"].get("ALGO") == "CHI":
                 for k in ["optical_absorption_coeff", "dielectric"]:
                     d["output"][k] = d_calc_final["output"][k]
 
@@ -588,7 +588,7 @@ def process_vasprun(self, dir_name, taskname, filename):
             d["output"]["optical_absorption_coeff"] = vrun.optical_absorption_coeff
 
         # parse output from response function
-        if vrun.incar.get("ALGO") == 'CHI':
+        if vrun.incar.get("ALGO") == "CHI":
             dielectric = vrun.dielectric
             d["output"]["dielectric"] = dict(
                 energy=dielectric[0], real=dielectric[1], imag=dielectric[2]
diff --git a/atomate/vasp/firetasks/absorption_tasks.py b/atomate/vasp/firetasks/absorption_tasks.py
@@ -1,16 +1,6 @@
-import os
-from importlib import import_module
-
-import numpy as np
-
-from monty.serialization import dumpfn
 from fireworks import FiretaskBase, explicit_serialize
-from fireworks.utilities.dict_mods import apply_mod
-from pymatgen.core.structure import Structure
-from pymatgen.io.vasp import Incar, Poscar, Potcar, PotcarSingle, Kpoints
 from pymatgen.io.vasp.sets import MPAbsorptionSet
-from pymatgen.io.vasp.outputs import Vasprun
-from atomate.utils.utils import env_chk, load_class
+
 
 @explicit_serialize
 class WriteVaspAbsorptionFromPrev(FiretaskBase):
@@ -25,6 +15,7 @@ class WriteVaspAbsorptionFromPrev(FiretaskBase):
         "potcar_spec"
 
     """
+
     optional_params = [
         "prev_calc_dir",
         "structure",
@@ -37,7 +28,7 @@ class WriteVaspAbsorptionFromPrev(FiretaskBase):
         "ncores",
         "nedos",
         "potcar_spec",
-        "other_params"
+        "other_params",
     ]
 
     def run_task(self, fw_spec):
diff --git a/atomate/vasp/firetasks/glue_tasks.py b/atomate/vasp/firetasks/glue_tasks.py
@@ -290,14 +290,14 @@ def interpolate_poscar(self, fw_spec):
         # use CopyFilesFromCalcLoc to get files from previous locations.
         CopyFilesFromCalcLoc(
             calc_loc=self["start"],
-            filenames=["CONTCAR","CONTCAR.gz"],
+            filenames=["CONTCAR", "CONTCAR.gz"],
             name_prepend=interpolate_folder + os.sep,
             name_append="_0",
             decompress=True,
         ).run_task(fw_spec=fw_spec)
         CopyFilesFromCalcLoc(
             calc_loc=self["end"],
-            filenames=["CONTCAR","CONTCAR.gz"],
+            filenames=["CONTCAR", "CONTCAR.gz"],
             name_prepend=interpolate_folder + os.sep,
             name_append="_1",
             decompress=True,
diff --git a/atomate/vasp/fireworks/absorption.py b/atomate/vasp/fireworks/absorption.py
@@ -1,46 +1,34 @@
-from atomate.vasp.config import (
-    VASP_CMD,
-    DB_FILE,
-)
 from fireworks import Firework
-from pymatgen.io.vasp.sets import MPStaticSet, MPAbsorptionSet
-from atomate.common.firetasks.glue_tasks import (
-    PassCalcLocs,
-    CopyFiles,
-    DeleteFiles,
-    GzipDir,
-    CreateFolder,
-    PassCalcLocs
-)
-from atomate.vasp.firetasks import (
-    CheckBandgap,
-    CopyVaspOutputs,
-    ModifyIncar,
-    RunVaspCustodian,
-    VaspToDb,
-)
-from atomate.vasp.firetasks.write_inputs import WriteVaspFromIOSet, WriteVaspStaticFromPrev
+from pymatgen.io.vasp.sets import MPStaticSet
+
+from atomate.common.firetasks.glue_tasks import PassCalcLocs
+from atomate.vasp.config import DB_FILE, VASP_CMD
+from atomate.vasp.firetasks import CopyVaspOutputs, RunVaspCustodian, VaspToDb
 from atomate.vasp.firetasks.absorption_tasks import WriteVaspAbsorptionFromPrev
+from atomate.vasp.firetasks.write_inputs import (
+    WriteVaspFromIOSet,
+    WriteVaspStaticFromPrev,
+)
 
 
 class AbsorptionFW(Firework):
     def __init__(
-            self,
-            structure,
-            name="frequency dependent dielectrics",
-            mode='STATIC',
-            nbands=None,
-            nbands_factor=2,
-            reciprocal_density=200,
-            nkred=None,
-            nedos=2001,
-            vasp_cmd=VASP_CMD,
-            prev_calc_dir=None,
-            db_file=DB_FILE,
-            vasptodb_kwargs=None,
-            parents=None,
-            vasp_input_set_params=None,
-            **kwargs,
+        self,
+        structure,
+        name="frequency dependent dielectrics",
+        mode="STATIC",
+        nbands=None,
+        nbands_factor=2,
+        reciprocal_density=200,
+        nkred=None,
+        nedos=2001,
+        vasp_cmd=VASP_CMD,
+        prev_calc_dir=None,
+        db_file=DB_FILE,
+        vasptodb_kwargs=None,
+        parents=None,
+        vasp_input_set_params=None,
+        **kwargs,
     ):
         """
         FW that calculates frequency dependent dielectric function within
@@ -85,7 +73,9 @@ def __init__(
         vasptodb_kwargs["additional_fields"]["task_label"] = name
 
         fw_name = "{}-{}-{}".format(
-            structure.composition.reduced_formula if structure else "unknown", name, mode
+            structure.composition.reduced_formula if structure else "unknown",
+            name,
+            mode,
         )
 
         # define what wavecars to copy from the previous run
@@ -102,10 +92,12 @@ def __init__(
         if mode == "IPA" or mode == "RPA":
             if prev_calc_dir:
                 # Copy the WAVECAR from previous calc directory
-                t.append(CopyVaspOutputs(
-                    calc_dir=prev_calc_dir,
-                    contcar_to_poscar=True,
-                    additional_files=wavecars)
+                t.append(
+                    CopyVaspOutputs(
+                        calc_dir=prev_calc_dir,
+                        contcar_to_poscar=True,
+                        additional_files=wavecars,
+                    )
                 )
 
                 t.append(
@@ -119,17 +111,15 @@ def __init__(
                         reciprocal_density=reciprocal_density,
                         nkred=nkred,
                         nedos=nedos,
-                        **vasp_input_set_params
+                        **vasp_input_set_params,
                     )
                 )
 
             elif parents:
                 # Copy the WAVECAR from previous calc location
                 t.append(
                     CopyVaspOutputs(
-                        calc_loc=True,
-                        contcar_to_poscar=True,
-                        additional_files=wavecars
+                        calc_loc=True, contcar_to_poscar=True, additional_files=wavecars
                     )
                 )
 
@@ -144,66 +134,80 @@ def __init__(
                         reciprocal_density=reciprocal_density,
                         nkred=nkred,
                         nedos=nedos,
-                        **vasp_input_set_params
+                        **vasp_input_set_params,
                     )
                 )
 
             else:
-                raise ValueError("Must specify previous calculation for {}".format(mode))
+                raise ValueError(f"Must specify previous calculation for {mode}")
 
         # when mode = "static"
         elif mode == "STATIC":
-            static_incar={"LWAVE": True,
-                          "ENCUT": 500,
-                          "ISMEAR": 0,
-                          "SIGMA": 0.01,
-                          "LREAL": False,
-                          "GGA": "PE",
-                          "LELF": False,
-                          "LAECHG": False,
-                          "LASPH": False,
-                          "LVHAR": False,
-                          "LVTOT": False,
-                          "METAGGA": "None",
-                          "LMIXTAU": False}
+            static_incar = {
+                "LWAVE": True,
+                "ENCUT": 500,
+                "ISMEAR": 0,
+                "SIGMA": 0.01,
+                "LREAL": False,
+                "GGA": "PE",
+                "LELF": False,
+                "LAECHG": False,
+                "LASPH": False,
+                "LVHAR": False,
+                "LVTOT": False,
+                "METAGGA": "None",
+                "LMIXTAU": False,
+            }
 
             if prev_calc_dir:
                 # Copy only the CONTCAR from previous calc directory (often a relaxation run)
                 t.append(
                     CopyVaspOutputs(
                         calc_dir=prev_calc_dir,
                         contcar_to_poscar=True,
-                        additional_files=wavecars
+                        additional_files=wavecars,
                     )
                 )
 
                 t.append(
-                    WriteVaspStaticFromPrev(reciprocal_density=reciprocal_density,
-                                            other_params={"user_incar_settings": static_incar})
+                    WriteVaspStaticFromPrev(
+                        reciprocal_density=reciprocal_density,
+                        other_params={"user_incar_settings": static_incar},
+                    )
                 )
 
             elif parents:
                 # Copy only the CONTCAR from previous calc
-                t.append(CopyVaspOutputs(calc_loc=True,
-                                         additional_files=wavecars,
-                                         contcar_to_poscar=True)
-                         )
+                t.append(
+                    CopyVaspOutputs(
+                        calc_loc=True, additional_files=wavecars, contcar_to_poscar=True
+                    )
+                )
 
                 t.append(
-                    WriteVaspStaticFromPrev(reciprocal_density=reciprocal_density,
-                                            other_params={"user_incar_settings": static_incar})
+                    WriteVaspStaticFromPrev(
+                        reciprocal_density=reciprocal_density,
+                        other_params={"user_incar_settings": static_incar},
+                    )
                 )
 
             elif structure:
-                static_input_set = MPStaticSet(structure=structure,
-                                               reciprocal_density=reciprocal_density,
-                                               user_incar_settings=static_incar)
+                static_input_set = MPStaticSet(
+                    structure=structure,
+                    reciprocal_density=reciprocal_density,
+                    user_incar_settings=static_incar,
+                )
 
-                t.append(WriteVaspFromIOSet(structure=structure, vasp_input_set=static_input_set)
-                         )
+                t.append(
+                    WriteVaspFromIOSet(
+                        structure=structure, vasp_input_set=static_input_set
+                    )
+                )
 
             else:
-                raise ValueError("Must specify structure or previous calculation for static calculation")
+                raise ValueError(
+                    "Must specify structure or previous calculation for static calculation"
+                )
 
         else:
             raise ValueError("Must specify a mode from 'STATIC', 'IPA', or 'RPA'")
@@ -216,15 +220,22 @@ def __init__(
             RunVaspCustodian(
                 vasp_cmd=vasp_cmd,
                 auto_npar=">>auto_npar<<",
-                handler_group=handler_group
+                handler_group=handler_group,
             )
         )
         t.append(PassCalcLocs(name=name))
         # Parse
-        t.append(VaspToDb(db_file=db_file,
-                          additional_fields={
-                              "task_label": structure.composition.reduced_formula + " " + name + " " + mode}))
+        t.append(
+            VaspToDb(
+                db_file=db_file,
+                additional_fields={
+                    "task_label": structure.composition.reduced_formula
+                    + " "
+                    + name
+                    + " "
+                    + mode
+                },
+            )
+        )
 
         super().__init__(t, parents=parents, name=fw_name, **kwargs)
-
-
