Start implementing the CubeHandlerCalculation plugin

aiida_nanotech_empa/plugins/__init__.py

@@ -1,6 +1,13 @@
 from .afm import AfmCalculation
+from .cubehandler import CubeHandlerCalculation
 from .hrstm import HrstmCalculation
 from .overlap import OverlapCalculation
 from .stm import StmCalculation
 
-__all__ = ("AfmCalculation", "HrstmCalculation", "OverlapCalculation", "StmCalculation")
+__all__ = (
+    "AfmCalculation",
+    "HrstmCalculation",
+    "OverlapCalculation",
+    "StmCalculation",
+    "CubeHandlerCalculation",
+)

aiida_nanotech_empa/plugins/cubehandler.py

@@ -0,0 +1,52 @@
+from aiida import common, engine, orm
+
+
+class CubeHandlerCalculation(engine.CalcJob):
+
+    @classmethod
+    def define(cls, spec):
+        super().define(spec)
+        spec.input(
+            "parameters", valid_type=orm.Dict, help="CubeHandler input parameters."
+        )
+        spec.input(
+            "parent_calc_folder",
+            valid_type=orm.RemoteData,
+            help="Parent folder containing original cube files.",
+        )
+
+        spec.input("metadata.options.withmpi", valid_type=bool, default=False)
+
+    def prepare_for_submission(self, folder):
+        """Create the input files from the input nodes passed to this instance of the `CalcJob`."""
+
+        # Create code info.
+        codeinfo = common.CodeInfo()
+        codeinfo.code_uuid = self.inputs.code.uuid
+        self.inputs.parameters.get_dict()
+        cmdline = ["shrink", "parent_calc_folder/test.cube", "shrinked.cube"]
+        codeinfo.cmdline_params = cmdline
+
+        # Create calc info.
+        calcinfo = common.CalcInfo()
+        calcinfo.uuid = self.uuid
+        calcinfo.cmdline_params = codeinfo.cmdline_params
+        calcinfo.codes_info = [codeinfo]
+
+        # File lists.
+        calcinfo.remote_symlink_list = []
+        calcinfo.local_copy_list = []
+        calcinfo.remote_copy_list = []
+        calcinfo.retrieve_list = ["*.cube"]
+
+        # Symlinks.
+        if "parent_calc_folder" in self.inputs:
+            comp_uuid = self.inputs.parent_calc_folder.computer.uuid
+            remote_path = self.inputs.parent_calc_folder.get_remote_path()
+            copy_info = (comp_uuid, remote_path, "parent_calc_folder/")
+            if self.inputs.code.computer.uuid == comp_uuid:
+                calcinfo.remote_symlink_list.append(copy_info)
+            else:
+                calcinfo.remote_copy_list.append(copy_info)
+
+        return calcinfo

examples/plugins/example_cubehandler.py

@@ -0,0 +1,36 @@
+import click
+from aiida import engine, orm
+from aiida.common.datastructures import StashMode
+
+
+@click.command("cli")
+@click.argument("cubehandler_code", default="cubehandler@localhost")
+def example_cubehander(cubehandler_code):
+    remote_folder = orm.load_node(2470)
+    builder = orm.load_code(cubehandler_code).get_builder()
+
+    builder.parameters = orm.Dict(dict={"some": "parameters"})
+    builder.parent_calc_folder = remote_folder
+
+    builder.metadata.options = {
+        "resources": {
+            "num_machines": 1,
+            "num_mpiprocs_per_machine": 1,
+        },
+        "max_wallclock_seconds": 600,
+        "prepend_text": "for file in parent_calc_folder/*.cube; do",
+        "append_text": "done",
+        "stash": {
+            "source_list": ["parent_calc_folder/*.cube"],
+            "target_base": "/project/s1267/yaa/aiida_stash/",
+            "stash_mode": StashMode.COPY.value,
+        },
+    }
+
+    _, calc_node = engine.run_get_node(builder)
+
+    assert calc_node.is_finished_ok
+
+
+if __name__ == "__main__":
+    example_cubehander()

pyproject.toml

@@ -58,6 +58,7 @@ dev = [
 "nanotech_empa.overlap" = "aiida_nanotech_empa.plugins:OverlapCalculation"
 "nanotech_empa.afm" = "aiida_nanotech_empa.plugins:AfmCalculation"
 "nanotech_empa.hrstm" = "aiida_nanotech_empa.plugins:HrstmCalculation"
+"nanotech_empa.cubehandler" = "aiida_nanotech_empa.plugins:CubeHandlerCalculation"
 
 [project.entry-points."aiida.workflows"]
 "nanotech_empa.nanoribbon" = "aiida_nanotech_empa.workflows.qe:NanoribbonWorkChain"