cli/workchains: add more options, parse referenced structure file

Author: dev-zero

aiida_cp2k/cli/utils/cp2k_options.py

@@ -15,3 +15,17 @@
                               "--structure",
                               type=types.DataParamType(sub_classes=("aiida.data:structure",)),
                               help="StructureData node.")
+
+MAX_NUM_MACHINES = OverridableOption('-m',
+                                     '--max-num-machines',
+                                     type=int,
+                                     default=1,
+                                     show_default=True,
+                                     help='The maximum number of machines (nodes) to use for the calculations.')
+
+MAX_WALLCLOCK_SECONDS = OverridableOption('-w',
+                                          '--max-wallclock-seconds',
+                                          type=int,
+                                          default=1800,
+                                          show_default=True,
+                                          help='the maximum wallclock time in seconds to set for the calculations.')

aiida_cp2k/cli/workflows/base.py

@@ -19,12 +19,17 @@
 @click.option('-m', '--description', type=str)
 @cp2k_options.STRUCTURE()
 @cp2k_options.DAEMON()
+@cp2k_options.MAX_NUM_MACHINES()
+@cp2k_options.MAX_WALLCLOCK_SECONDS()
 @decorators.with_dbenv()
-def cmd_launch_workflow(code, daemon, cp2k_input_file, label, description, structure):
+def cmd_launch_workflow(code, daemon, cp2k_input_file, label, description, structure, max_num_machines,
+                        max_wallclock_seconds):
     """Run a CP2K calculation with a given input file through AiiDA"""
+    # pylint: disable=too-many-arguments,too-many-branches,too-many-statements
 
     from aiida.orm import Dict
     from aiida.plugins import WorkflowFactory
+    from aiida.orm.nodes.data.structure import StructureData
     from cp2k_input_tools.parser import CP2KInputParserSimplified
 
     parser = CP2KInputParserSimplified(key_trafo=str.upper,
@@ -33,33 +38,62 @@ def cmd_launch_workflow(code, daemon, cp2k_input_file, label, description, struc
                                        level_reduction_blacklist=["KIND"])
     tree = parser.parse(cp2k_input_file)
 
-    builder = WorkflowFactory('cp2k.base').get_builder()
+    try:
+        tree["FORCE_EVAL"]["SUBSYS"].pop("COORD")
+        structure_in_config = True
+    except KeyError:
+        structure_in_config = False
+
+    try:
+        coord_file_from_config = tree["FORCE_EVAL"]["SUBSYS"]["TOPOLOGY"].pop("COORD_FILE_NAME")
+    except KeyError:
+        coord_file_from_config = None
 
+    builder = WorkflowFactory('cp2k.base').get_builder()
     builder.cp2k.code = code
 
     if structure:
         builder.cp2k.structure = structure
 
-        try:
-            tree["FORCE_EVAL"]["SUBSYS"]["TOPOLOGY"].pop("COORD_FILE_NAME")
-            echo.echo_info("The explicitly given structure overrides the structure referenced in the input file")
-        except KeyError:
-            pass
+        if coord_file_from_config:
+            echo.echo_info(
+                "The explicitly given structure will override the structure file referenced in the input file")
 
-    else:
+        if structure_in_config:
+            echo.echo_info("The explicitly given structure will override the structure in the input file")
+
+    elif structure_in_config:
         try:
-            with open(cp2k_input_file.name, "r") as fhandle:
-                builder.cp2k.structure = structure_from_cp2k_inp(fhandle)
+            cp2k_input_file.seek(0)
+            builder.cp2k.structure = structure_from_cp2k_inp(cp2k_input_file)
             builder.cp2k.structure.store()
             echo.echo("Created StructureData<{}> from file {}]\n".format(builder.cp2k.structure.pk,
                                                                          cp2k_input_file.name))
         except ValueError as err:
             echo.echo_critical(str(err))
 
+    elif coord_file_from_config:
+        try:
+            import ase.io
+        except ImportError:
+            echo.echo_critical('You have not installed the package ase. \nYou can install it with: pip install ase')
+
+        try:
+            asecell = ase.io.read(coord_file_from_config)
+            builder.cp2k.structure = StructureData(ase=asecell)
+            builder.cp2k.structure.store()
+            echo.echo("Created StructureData<{}> from file {}]\n".format(builder.cp2k.structure.pk,
+                                                                         coord_file_from_config))
+        except ValueError as err:
+            echo.echo_critical(str(err))
+    else:
+        echo.echo_critical("No structure found/referenced in the input file and not set explicitly")
+
     builder.cp2k.metadata.options = {
         "resources": {
-            "num_machines": 1,
+            "num_machines": max_num_machines,
         },
+        'max_wallclock_seconds': int(max_wallclock_seconds),
     }
     builder.cp2k.parameters = Dict(dict=tree)