Merge remote-tracking branch 'mp/master'

Author: dbroberg

atomate/__init__.py

@@ -1 +1 @@
-__version__ = "0.8.5"
+__version__ = "0.8.8"

atomate/vasp/config.py

@@ -14,4 +14,5 @@
 HALF_KPOINTS_FIRST_RELAX = False  # whether to use only half the kpoint density in the initial relaxation of a structure optimization for faster performance
 RELAX_MAX_FORCE = 0.25  # maximum force allowed on atom for successful structure optimization
 REMOVE_WAVECAR = False # Remove Wavecar after the calculation is finished. Only used for SCAN structure optimizations right now.
-DEFUSE_UNSUCCESSFUL = "fizzle"  # this is a three-way toggle on what to do if your job looks OK, but is actually unconverged (either electronic or ionic). True -> mark job as COMPLETED, but defuse children. False --> do nothing, continue with workflow as normal. "fizzle" --> throw an error (mark this job as FIZZLED)
+DEFUSE_UNSUCCESSFUL = "fizzle"  # this is a three-way toggle on what to do if your job looks OK, but is actually unconverged (either electronic or ionic). True -> mark job as COMPLETED, but defuse children. False --> do nothing, continue with workflow as normal. "fizzle" --> throw an error (mark this job as FIZZLED)
+CUSTODIAN_MAX_ERRORS = 5  # maximum number of errors to correct before custodian gives up

atomate/vasp/database.py

@@ -70,7 +70,9 @@ def build_indexes(self, indexes=None, background=True):
     def insert_task(self, task_doc, use_gridfs=False):
         """
         Inserts a task document (e.g., as returned by Drone.assimilate()) into the database.
-        Handles putting DOS and band structure into GridFS as needed.
+        Handles putting DOS, band structure and charge density into GridFS as needed.
+        During testing, a percentage of runs on some clusters had corrupted AECCAR files when even if everything else about the calculation looked OK.
+        So we do a quick check here and only record the AECCARs if they are valid
 
         Args:
             task_doc: (dict) the task document
@@ -82,6 +84,7 @@ def insert_task(self, task_doc, use_gridfs=False):
         bs = None
         chgcar = None
         aeccar0 = None
+        write_aeccar = False
 
         # move dos BS and CHGCAR from doc to gridfs
         if use_gridfs and "calcs_reversed" in task_doc:
@@ -98,15 +101,20 @@ def insert_task(self, task_doc, use_gridfs=False):
                 chgcar = json.dumps(task_doc["calcs_reversed"][0]["chgcar"], cls=MontyEncoder)
                 del task_doc["calcs_reversed"][0]["chgcar"]
 
-            if "aeccar0" in task_doc["calcs_reversed"][0]:  # only store idx=0 DOS
+            if "aeccar0" in task_doc["calcs_reversed"][0]:
                 aeccar0 = task_doc["calcs_reversed"][0]["aeccar0"]
-                aeccar0 = json.dumps(task_doc["calcs_reversed"][0]["aeccar0"], cls=MontyEncoder)
-                del task_doc["calcs_reversed"][0]["aeccar0"]
-                try:
-                    # aeccar2 should also be in the task_doc
+                aeccar2 = task_doc["calcs_reversed"][0]["aeccar2"]
+                # check if the aeccar is valid before insertion
+                if (aeccar0.data['total'] + aeccar2.data['total']).min() < 0:
+                    logger.warning(f"The AECCAR seems to be corrupted for task_in directory {task_doc['dir_name']}\nSkipping storage of AECCARs")
+                    write_aeccar = False
+                else:
+                    # overwrite the aeccar variable with their string representations to be inserted in GridFS
+                    aeccar0 = json.dumps(task_doc["calcs_reversed"][0]["aeccar0"], cls=MontyEncoder)
                     aeccar2 = json.dumps(task_doc["calcs_reversed"][0]["aeccar2"], cls=MontyEncoder)
-                except:
-                    raise KeyError('aeccar2 data is missing from task_doc')
+                    write_aeccar = True
+
+                del task_doc["calcs_reversed"][0]["aeccar0"]
                 del task_doc["calcs_reversed"][0]["aeccar2"]
 
         # insert the task document
@@ -135,7 +143,7 @@ def insert_task(self, task_doc, use_gridfs=False):
             self.collection.update_one({"task_id": t_id}, {"$set": {"calcs_reversed.0.chgcar_fs_id": chgcar_gfs_id}})
 
         # insert the AECCARs file into gridfs and update the task documents
-        if aeccar0:
+        if write_aeccar:
             aeccar0_gfs_id, compression_type = self.insert_gridfs(aeccar0, "aeccar0_fs", task_id=t_id)
             self.collection.update_one(
                 {"task_id": t_id}, {"$set": {"calcs_reversed.0.aeccar0_compression": compression_type}})
@@ -168,6 +176,10 @@ def retrieve_task(self, task_id):
         if 'chgcar_fs_id' in calc:
             chgcar = self.get_chgcar(task_id)
             calc["chgcar"] = chgcar
+        if 'aeccar0_fs_id' in calc:
+            aeccar = self.get_aeccar(task_id)
+            calc["aeccar0"] = aeccar['aeccar0']
+            calc["aeccar2"] = aeccar['aeccar2']
         return task_doc
 
     def insert_gridfs(self, d, collection="fs", compress=True, oid=None, task_id=None):
@@ -243,13 +255,14 @@ def get_chgcar(self, task_id):
         chgcar= json.loads(chgcar_json, cls=MontyDecoder)
         return chgcar
 
-    def get_aeccar(self, task_id):
+    def get_aeccar(self, task_id, check_valid = True):
         """
         Read the AECCAR0 + AECCAR2 grid_fs data into a Chgcar object
         Args:
             task_id(int or str): the task_id containing the gridfs metadata
+            check_valid (bool): make sure that the aeccar is positive definite
         Returns:
-            (aeccar0, aeccar2): Chgcar objects
+            {"aeccar0" : Chgcar, "aeccar2" : Chgcar}: dict of Chgcar objects
         """
         m_task = self.collection.find_one({"task_id": task_id}, {"calcs_reversed": 1})
         fs_id = m_task['calcs_reversed'][0]['aeccar0_fs_id']
@@ -261,6 +274,9 @@ def get_aeccar(self, task_id):
         aeccar_json = zlib.decompress(fs.get(fs_id).read())
         aeccar2 = json.loads(aeccar_json, cls=MontyDecoder)
 
+        if check_valid and (aeccar0.data['total'] + aeccar2.data['total']).min() < 0:
+            ValueError(f"The AECCAR seems to be corrupted for task_id = {task_id}")
+
         return {'aeccar0': aeccar0, 'aeccar2': aeccar2}
 
     def reset(self):