Merge pull request #721 from JaGeo/main

Add additional optimization to Lobster

Author: itsduowang

atomate/vasp/firetasks/lobster_tasks.py

@@ -8,12 +8,6 @@
 import shutil
 import warnings
 
-from fireworks import FiretaskBase, explicit_serialize, FWAction
-from fireworks.utilities.fw_serializers import DATETIME_HANDLER
-from monty.json import jsanitize
-from monty.os.path import zpath
-from monty.serialization import loadfn
-
 from atomate.common.firetasks.glue_tasks import get_calc_loc
 from atomate.utils.utils import env_chk, get_meta_from_structure
 from atomate.vasp.config import VASP_OUTPUT_FILES
@@ -25,6 +19,11 @@
     LobsterFilesValidator,
 )
 from custodian.lobster.jobs import LobsterJob
+from fireworks import FiretaskBase, explicit_serialize, FWAction
+from fireworks.utilities.fw_serializers import DATETIME_HANDLER
+from monty.json import jsanitize
+from monty.os.path import zpath
+from monty.serialization import loadfn
 from pymatgen.core.structure import Structure
 from pymatgen.io.lobster import Lobsterout, Lobsterin
 
@@ -163,7 +162,12 @@ def run_task(self, fw_spec):
         c.run()
 
         if os.path.exists(zpath("custodian.json")):
-            stored_custodian_data = {"custodian": loadfn(zpath("custodian.json"))}
+            if os.path.exists(zpath("FW_offline.json")):
+                import json
+                with open(zpath("custodian.json")) as f:
+                    stored_custodian_data = {"custodian": json.load(f)}
+            else:
+                stored_custodian_data = {"custodian": loadfn(zpath("custodian.json"))}
             return FWAction(stored_data=stored_custodian_data)
 
 
@@ -200,11 +204,14 @@ class LobsterRunToDb(FiretaskBase):
     std_additional_outputs = [
         "ICOHPLIST.lobster",
         "ICOOPLIST.lobster",
+        "ICOBILIST.lobster",
         "COHPCAR.lobster",
         "COOPCAR.lobster",
         "GROSSPOP.lobster",
         "CHARGE.lobster",
         "DOSCAR.lobster",
+        "MadelungEnergies.lobster",
+        "SitePotentials.lobster"
     ]
 
     def __init__(self, *args, **kwargs):
@@ -281,7 +288,7 @@ def run_task(self, fw_spec):
         db_file = env_chk(self.get("db_file"), fw_spec)
 
         # db insertion or taskdoc dump
-        if not db_file:
+        if not db_file or os.path.exists(zpath("FW_offline.json")):
             with open("task_lobster.json", "w") as f:
                 f.write(json.dumps(task_doc, default=DATETIME_HANDLER))
         else:
@@ -339,7 +346,7 @@ def _verify_inputs(self):
         # Check lobsterin
         if self.get("check_lobsterin", True):
             ref_lobsterin = Lobsterin.from_file(
-                os.path.join(self["ref_dir"], "inputs", "lobsterin")
+                os.path.join(self["ref_dir"], "inputs", "lobsterin.gz")
             )
             params_to_check = self.get("params_to_check", [])
             for p in params_to_check:

atomate/vasp/firetasks/parse_outputs.py

@@ -7,6 +7,7 @@
 import numpy as np
 
 from monty.json import MontyEncoder, jsanitize
+from monty.os.path import zpath
 from pydash.objects import has, get
 
 from atomate.vasp.config import DEFUSE_UNSUCCESSFUL
@@ -135,7 +136,7 @@ def run_task(self, fw_spec):
         db_file = env_chk(self.get("db_file"), fw_spec)
 
         # db insertion or taskdoc dump
-        if not db_file:
+        if not db_file or os.path.exists(zpath("FW_offline.json")):
             with open("task.json", "w") as f:
                 f.write(json.dumps(task_doc, default=DATETIME_HANDLER))
         else:

atomate/vasp/firetasks/run_calc.py

@@ -1,8 +1,7 @@
+from atomate.vasp.config import HALF_KPOINTS_FIRST_RELAX
 from monty.os.path import zpath
 from monty.serialization import loadfn
 
-from atomate.vasp.config import HALF_KPOINTS_FIRST_RELAX
-
 """
 This module defines tasks that support running vasp in various ways.
 """
@@ -299,7 +298,12 @@ def run_task(self, fw_spec):
         c.run()
 
         if os.path.exists(zpath("custodian.json")):
-            stored_custodian_data = {"custodian": loadfn(zpath("custodian.json"))}
+            if os.path.exists(zpath("FW_offline.json")):
+                import json
+                with open(zpath("custodian.json")) as f:
+                    stored_custodian_data = {"custodian": json.load(f)}
+            else:
+                stored_custodian_data = {"custodian": loadfn(zpath("custodian.json"))}
             return FWAction(stored_data=stored_custodian_data)
 
 
@@ -396,7 +400,7 @@ def _verify_inputs(self):
             defaults = {"ISPIN": 1, "ISMEAR": 1, "SIGMA": 0.2}
             for p in params_to_check:
                 if user_incar.get(p, defaults.get(p)) != ref_incar.get(
-                    p, defaults.get(p)
+                        p, defaults.get(p)
                 ):
                     raise ValueError(f"INCAR value of {p} is inconsistent!")
 
@@ -407,8 +411,8 @@ def _verify_inputs(self):
                 os.path.join(self["ref_dir"], "inputs", "KPOINTS")
             )
             if (
-                user_kpoints.style != ref_kpoints.style
-                or user_kpoints.num_kpts != ref_kpoints.num_kpts
+                    user_kpoints.style != ref_kpoints.style
+                    or user_kpoints.num_kpts != ref_kpoints.num_kpts
             ):
                 raise ValueError(
                     "KPOINT files are inconsistent! Paths are:\n{}\n{} with kpoints {} and {}".format(
@@ -426,8 +430,8 @@ def _verify_inputs(self):
                 os.path.join(self["ref_dir"], "inputs", "POSCAR")
             )
             if (
-                user_poscar.natoms != ref_poscar.natoms
-                or user_poscar.site_symbols != ref_poscar.site_symbols
+                    user_poscar.natoms != ref_poscar.natoms
+                    or user_poscar.site_symbols != ref_poscar.site_symbols
             ):
                 raise ValueError(
                     "POSCAR files are inconsistent! Paths are:\n{}\n{}".format(