checking and retrieving of AECCAR

Added positive definite checking of AECCAR and retrieve_task support for chgcar

Author: jmmshn

atomate/vasp/database.py

@@ -71,7 +71,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
@@ -83,6 +85,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:
@@ -99,15 +102,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
@@ -136,7 +144,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}})
@@ -169,6 +177,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):
@@ -244,11 +256,12 @@ 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
         """
@@ -262,6 +275,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):

atomate/vasp/workflows/tests/test_vasp_workflows.py

@@ -298,6 +298,14 @@ def test_chgcar_db_read_write(self):
         dcc = mmdb.get_aeccar(task_id=t_id)
         self.assertAlmostEqual(dcc['aeccar0'].data['total'].sum()/cc.ngridpts, 23.253588293583313, 4)
         self.assertAlmostEqual(dcc['aeccar2'].data['total'].sum()/cc.ngridpts, 8.01314480789829, 4)
+        # check the retrieve_task function for the same fake calculation
+        ret_task = mmdb.retrieve_task(t_id)
+        ret_chgcar = ret_task['calcs_reversed'][0]['chgcar']
+        ret_aeccar0 = ret_task['calcs_reversed'][0]['aeccar0']
+        ret_aeccar2 = ret_task['calcs_reversed'][0]['aeccar2']
+        ret_aeccar = ret_aeccar0 + ret_aeccar2
+        self.assertAlmostEqual(ret_chgcar.data['total'].sum()/ret_chgcar.ngridpts, 8.0, 4)
+        self.assertAlmostEqual(ret_aeccar.data['total'].sum()/ret_aeccar.ngridpts, 31.2667331015, 4)
 
     def test_chgcar_db_read(self):
         # add the workflow