Add a default entry_id for get_computed_entry. Helps with #2181

materialsproject · Jun 24, 2021 · ee143e5 · ee143e5
1 parent 56fa39f
commit ee143e5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/pymatgen/io/vasp/outputs.py b/pymatgen/io/vasp/outputs.py
@@ -19,6 +19,7 @@
 from collections import defaultdict
 from typing import Optional, Tuple, List, Union, DefaultDict
 from pathlib import Path
+import datetime
 
 import numpy as np
 from monty.dev import deprecated
@@ -758,9 +759,11 @@ def is_spin(self):
         """
         return self.parameters.get("ISPIN", 1) == 2
 
-    def get_computed_entry(self, inc_structure=True, parameters=None, data=None):
+    def get_computed_entry(
+        self, inc_structure=True, parameters=None, data=None, entry_id=f"vasprun-{datetime.datetime.now()}"
+    ):
         """
-        Returns a ComputedStructureEntry from the vasprun.
+        Returns a ComputedEntry or ComputedStructureEntry from the vasprun.
 
         Args:
             inc_structure (bool): Set to True if you want
@@ -772,6 +775,8 @@ def get_computed_entry(self, inc_structure=True, parameters=None, data=None):
                 necessary for typical post-processing will be set.
             data (list): Output data to include. Has to be one of the properties
                 supported by the Vasprun object.
+            entry_id (object): Specify an entry id for the ComputedEntry. Defaults to
+                "vasprun-{current datetime}"
 
         Returns:
             ComputedStructureEntry/ComputedEntry
@@ -789,12 +794,11 @@ def get_computed_entry(self, inc_structure=True, parameters=None, data=None):
         data = {p: getattr(self, p) for p in data} if data is not None else {}
 
         if inc_structure:
-            return ComputedStructureEntry(self.final_structure, self.final_energy, parameters=params, data=data)
+            return ComputedStructureEntry(
+                self.final_structure, self.final_energy, parameters=params, data=data, entry_id=entry_id
+            )
         return ComputedEntry(
-            self.final_structure.composition,
-            self.final_energy,
-            parameters=params,
-            data=data,
+            self.final_structure.composition, self.final_energy, parameters=params, data=data, entry_id=entry_id
         )
 
     def get_band_structure(

diff --git a/pymatgen/io/vasp/tests/test_outputs.py b/pymatgen/io/vasp/tests/test_outputs.py
@@ -232,6 +232,10 @@ def test_standard(self):
         self.assertEqual(d["elements"], ["Fe", "Li", "O", "P"])
         self.assertEqual(d["nelements"], 4)
 
+        entry = vasprun.get_computed_entry(inc_structure=True)
+        self.assertTrue(entry.entry_id.startswith("vasprun"))
+        self.assertEqual(entry.parameters["run_type"], "PBEO or other Hybrid Functional")
+
     def test_unconverged(self):
         filepath = self.TEST_FILES_DIR / "vasprun.xml.unconverged"
         with warnings.catch_warnings(record=True) as w: