feat/merge calcjobs (#36)

* merge calcjobs

* fixed entry-points

* format

Author: ethan-shanahan

pyproject.toml

@@ -26,8 +26,7 @@ dependencies = [
 # [project.entry-points."aiida.data"]
 # "fans" = "aiida_fans.data:FANSParameters"
 [project.entry-points."aiida.calculations"]
-"fans.stashed" = "aiida_fans.calculations:FansStashedCalculation"
-"fans.fragmented" = "aiida_fans.calculations:FansFragmentedCalculation"
+"fans" = "aiida_fans.calculations:FansCalculation"
 [project.entry-points."aiida.parsers"]
 "fans" = "aiida_fans.parsers:FansParser"
 # [project.entry-points."aiida.cmdline.data"]

src/aiida_fans/calculations.py

@@ -14,8 +14,8 @@
 from aiida_fans.helpers import make_input_dict
 
 
-class FansCalcBase(CalcJob):
-    """Base class of all calculations using FANS."""
+class FansCalculation(CalcJob):
+    """Calculations using FANS."""
 
     @classmethod
     def define(cls, spec: CalcJobProcessSpec) -> None:
@@ -39,6 +39,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
         # Custom Metadata
         spec.input("metadata.options.results_prefix", valid_type=str, default="")
         spec.input("metadata.options.results", valid_type=list, default=[])
+        spec.input("metadata.options.stashed_microstructure", valid_type=bool, default=True)
 
         # Input Ports
         ## Microstructure Definition
@@ -69,6 +70,40 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
 
     def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         """Prepare the calculation for submission."""
+        # Stashed Strategy:
+        if self.options.stashed_microstructure:
+            ms_filepath: Path = Path(self.inputs.code.computer.get_workdir()) / \
+                                "stash/microstructures" / \
+                                self.inputs.microstructure.file.filename
+            # if microstructure does not exist in stash, make it
+            if not ms_filepath.is_file():
+                ms_filepath.parent.mkdir(parents=True, exist_ok=True)
+                with self.inputs.microstructure.file.open(mode='rb') as source:
+                    with ms_filepath.open(mode='wb') as target:
+                        copyfileobj(source, target)
+
+            # input.json as dict
+            input_dict = make_input_dict(self)
+            input_dict["microstructure"]["filepath"] = str(ms_filepath)
+            # write input.json to working directory
+            with folder.open(self.options.input_filename, "w", "utf8") as json:
+                dump(input_dict, json, indent=4)
+        # Fragmented Strategy:
+        else:
+            datasetname : str = self.inputs.microstructure.datasetname.value
+            with folder.open("microstructure.h5","bw") as f_dest:
+                with h5File(f_dest,"w") as h5_dest:
+                    with self.inputs.microstructure.file.open(mode="rb") as f_src:
+                        with h5File(f_src,'r') as h5_src:
+                            h5_src.copy(datasetname, h5_dest, name=datasetname)
+
+            # input.json as dict
+            input_dict = make_input_dict(self)
+            input_dict["microstructure"]["filepath"] = "microstructure.h5"
+            # write input.json to working directory
+            with folder.open(self.options.input_filename, "w", "utf8") as json:
+                dump(input_dict, json, indent=4)
+
         # Specifying the code info:
         codeinfo = CodeInfo()
         codeinfo.code_uuid = self.inputs.code.uuid
@@ -87,60 +122,3 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo:
         ]
 
         return calcinfo
-
-
-class FansStashedCalculation(FansCalcBase):
-    """Calculations using FANS and the "Stashed" microstructure distribution strategy."""
-
-    @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
-        """Define inputs, outputs, and exit codes of the calculation."""
-        return super().define(spec)
-
-    def prepare_for_submission(self, folder: Folder) -> CalcInfo:
-        """Prepare the calculation for submission."""
-        ms_filepath: Path = Path(self.inputs.code.computer.get_workdir()) / \
-                            "stash/microstructures" / \
-                            self.inputs.microstructure.file.filename
-        # if microstructure does not exist in stash, make it
-        if not ms_filepath.is_file():
-            ms_filepath.parent.mkdir(parents=True, exist_ok=True)
-            with self.inputs.microstructure.file.open(mode='rb') as source:
-                with ms_filepath.open(mode='wb') as target:
-                    copyfileobj(source, target)
-
-        # input.json as dict
-        input_dict = make_input_dict(self)
-        input_dict["microstructure"]["filepath"] = str(ms_filepath)
-        # write input.json to working directory
-        with folder.open(self.options.input_filename, "w", "utf8") as json:
-            dump(input_dict, json, indent=4)
-
-        return super().prepare_for_submission(folder)
-
-class FansFragmentedCalculation(FansCalcBase):
-    """Calculations using FANS and the "Fragmented" microstructure distribution strategy."""
-
-    @classmethod
-    def define(cls, spec: CalcJobProcessSpec) -> None:
-        """Define inputs, outputs, and exit codes of the calculation."""
-        return super().define(spec)
-
-    def prepare_for_submission(self, folder: Folder) -> CalcInfo:
-        """Prepare the calculation for submission."""
-        # Write Microstructure Subset to Folder
-        datasetname : str = self.inputs.microstructure.datasetname.value
-        with folder.open("microstructure.h5","bw") as f_dest:
-            with h5File(f_dest,"w") as h5_dest:
-                with self.inputs.microstructure.file.open(mode="rb") as f_src:
-                    with h5File(f_src,'r') as h5_src:
-                        h5_src.copy(datasetname, h5_dest, name=datasetname)
-
-        # input.json as dict
-        input_dict = make_input_dict(self)
-        input_dict["microstructure"]["filepath"] = "microstructure.h5"
-        # write input.json to working directory
-        with folder.open(self.options.input_filename, "w", "utf8") as json:
-            dump(input_dict, json, indent=4)
-
-        return super().prepare_for_submission(folder)

src/aiida_fans/utils.py

@@ -24,21 +24,22 @@ def aiida_type(value: Any) -> type[Data]:
     """
     match value:
         case str():
-            return DataFactory("core.str") # Str
+            return DataFactory("core.str")  # Str
         case int():
-            return DataFactory("core.int") # Int
+            return DataFactory("core.int")  # Int
         case float():
-            return DataFactory("core.float") # Float
+            return DataFactory("core.float")  # Float
         case list():
-            return DataFactory("core.list") # List
+            return DataFactory("core.list")  # List
         case dict():
             if all(map(lambda t: isinstance(t, ndarray), value.values())):
-                return DataFactory("core.array") # ArrayData
+                return DataFactory("core.array")  # ArrayData
             else:
-                return DataFactory("core.dict") # Dict
+                return DataFactory("core.dict")  # Dict
         case _:
             raise NotImplementedError(f"Received an input of value:  {value}    with type:  {type(value)}")
 
+
 def fetch(label: str, value: Any) -> list[Node]:
     """Return a list of nodes matching the label and value provided.
 
@@ -50,26 +51,31 @@ def fetch(label: str, value: Any) -> list[Node]:
         list[Node]: the list of nodes matching the give criteria
     """
     datatype = aiida_type(value)
-    nodes = QueryBuilder(
-    ).append(cls=datatype, tag="n"
-    ).add_filter("n", {"label": label}
-    ).add_filter("n", {"attributes": {"==": datatype(value).base.attributes.all}}
-    ).all(flat=True)
+    nodes = (
+        QueryBuilder()
+        .append(cls=datatype, tag="n")
+        .add_filter("n", {"label": label})
+        .add_filter("n", {"attributes": {"==": datatype(value).base.attributes.all}})
+        .all(flat=True)
+    )
 
     if datatype != DataFactory("core.array"):
-        return nodes # type: ignore
+        return nodes  # type: ignore
     else:
         array_nodes = []
         for array_node in nodes:
             array_value = {
-                k: v for k, v in [
-                    (name, array_node.get_array(name)) for name in array_node.get_arraynames() # type: ignore
+                k: v
+                for k, v in [
+                    (name, array_node.get_array(name))
+                    for name in array_node.get_arraynames()  # type: ignore
                 ]
             }
             if arraydata_equal(value, array_value):
                 array_nodes.append(array_node)
         return array_nodes
 
+
 def generate(label: str, value: Any) -> Node:
     """Return a single node with the label and value provided.
 
@@ -93,6 +99,7 @@ def generate(label: str, value: Any) -> Node:
     else:
         raise RuntimeError
 
+
 def convert(ins: dict[str, Any], path: list[str] = []):
     """Takes a dictionary of inputs and converts the values to their respective Nodes.
 
@@ -108,7 +115,8 @@ def convert(ins: dict[str, Any], path: list[str] = []):
         else:
             ins[k] = generate(".".join([*path, k]), v)
 
-def compile_query(ins: dict[str,Any], qb: QueryBuilder) -> None:
+
+def compile_query(ins: dict[str, Any], qb: QueryBuilder) -> None:
     """Interate over the converted input dictionary and append to the QueryBuilder for each node.
 
     Args:
@@ -121,18 +129,14 @@ def compile_query(ins: dict[str,Any], qb: QueryBuilder) -> None:
         if k in ["microstructure", "error_parameters"] and isinstance(v, dict):
             compile_query(v, qb)
         else:
-            qb.append(
-                cls=type(v),
-                with_outgoing="calc",
-                filters={"pk": v.pk}
-            )
+            qb.append(cls=type(v), with_outgoing="calc", filters={"pk": v.pk})
 
 
 def execute_fans(
-        mode: Literal["Submit", "Run"],
-        inputs: dict[str, Any],
-        strategy: Literal["Fragmented", "Stashed"] = "Fragmented",
-    ):
+    mode: Literal["Submit", "Run"],
+    inputs: dict[str, Any],
+    strategy: Literal["Fragmented", "Stashed"] = "Fragmented",
+):
     """This utility function simplifies the process of executing aiida-fans jobs.
 
     The only nodes you must provide are the `code` and `microstructure` inputs.
@@ -191,17 +195,18 @@ def execute_fans(
     compile_query(inputs, qb)
     results = qb.all(flat=True)
     if (count := len(results)) != 0:
-        print(f"It seems this calculation has already been performed {count} time{"s" if count > 1 else ""}. {results}")
+        print(f"It seems this calculation has already been performed {count} time{'s' if count > 1 else ''}. {results}")
         confirmation = input("Are you sure you want to rerun it? [y/N] ").strip().lower() in ["y", "yes"]
     else:
         confirmation = True
 
     if confirmation:
         match mode:
             case "Run":
-                run(calcjob, inputs) # type: ignore
+                run(calcjob, inputs)  # type: ignore
             case "Submit":
-                submit(calcjob, inputs) # type: ignore
+                submit(calcjob, inputs)  # type: ignore
+
 
 def submit_fans(
     inputs: dict[str, Any],
@@ -210,6 +215,7 @@ def submit_fans(
     """See `execute_fans` for implementation and usage details."""
     execute_fans("Submit", inputs, strategy)
 
+
 def run_fans(
     inputs: dict[str, Any],
     strategy: Literal["Fragmented", "Stashed"] = "Fragmented",