diff --git a/pyproject.toml b/pyproject.toml index 81e1f08..9c5c771 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/aiida_fans/calculations.py b/src/aiida_fans/calculations.py index a53965a..65178d1 100644 --- a/src/aiida_fans/calculations.py +++ b/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) diff --git a/src/aiida_fans/utils.py b/src/aiida_fans/utils.py index 6aa478a..f496a96 100644 --- a/src/aiida_fans/utils.py +++ b/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,7 +195,7 @@ 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 @@ -199,9 +203,10 @@ def execute_fans( 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",