store volumetric data @jmmshn @rkingsbury

materialsproject · Oct 30, 2020 · 3d069e0 · 3d069e0
1 parent 60d9c14
commit 3d069e0
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/emmet-cli/emmet/cli/tasks.py b/emmet-cli/emmet/cli/tasks.py
@@ -36,6 +36,7 @@
     for f in FILE_FILTERS
     for d in ["", "relax1", "relax2"]
 ]
+STORE_VOLUMETRIC_DATA = []
 
 
 @click.group()
@@ -339,9 +340,14 @@ def restore(inputfile, file_filter):
     show_default=True,
     help="Number of processes for parallel parsing.",
 )
-@click.option("--parse-chgcar", is_flag=True, help="Parse CHGCAR.")
-@click.option("--parse-aeccar", is_flag=True, help="Parse AECCAR.")
-def parse(task_ids, nproc, parse_chgcar, parse_aeccar):
+@click.option(
+    "-s",
+    "--store-volumetric-data",
+    multiple=True,
+    default=STORE_VOLUMETRIC_DATA,
+    help="Store any of CHGCAR, LOCPOT, AECCAR0, AECCAR1, AECCAR2, ELFCAR.",
+)
+def parse(task_ids, nproc, store_volumetric_data):
     """Parse VASP launchers into tasks"""
     ctx = click.get_current_context()
     if "CLIENT" not in ctx.obj:

diff --git a/emmet-cli/emmet/cli/utils.py b/emmet-cli/emmet/cli/utils.py
@@ -339,13 +339,10 @@ def parse_vasp_dirs(vaspdirs, tag, task_ids):
     run = ctx.parent.parent.params["run"]
     projection = {"tags": 1, "task_id": 1}
     count = 0
-    kwargs = {}
-
-    for p in ["chgcar", "aeccar"]:
-      par = f"parse_{p}"
-      kwargs[par] = ctx.params[par]
-
-    drone = VaspDrone(additional_fields={"tags": tags}, **kwargs)
+    drone = VaspDrone(
+        additional_fields={"tags": tags},
+        store_volumetric_data=ctx.params['store_volumetric_data']
+    )
 
     for vaspdir in vaspdirs:
         logger.info(f"{name} VaspDir: {vaspdir}")