Add generic high throughput calculations (#169)

* added hts wc, need to fix pre-commits * ht workflow submit only * hts run without making a wc * trying some stuff * workgraph * some mods to workchain + some files to delete later * working workgraph+ submission * workchain works, test fail cause of submit * small fixes * remove generated things from docs to see if it works * fix docs? * Apply ruff fixes * Update for ruff and remove screening refs * Update aiida workgraph * Bump janus version * Generalise ht workgraph * Test ht workgraph * Update example ht submit * Revert unused changes * Fix workgraph tests * Add check for missing files * Add test for invalid struct dir * Add high throughput calc tutorial * Apply suggestions from code review * Update pyproject.toml * Update docstrings * Apply suggestions from code review Co-authored-by: Jacob Wilkins <[email protected]> * Allow strings and AiiDA strings for path * Print folder for FileNotFoundError * Rename workgraph builder * Allow non-recursive folder search --------- Co-authored-by: federica <[email protected]> Co-authored-by: federicazanca <federicazanca*@gmail.com> Co-authored-by: Jacob Wilkins <[email protected]>
stfc · Nov 29, 2024 · ade9f21 · ade9f21
1 parent 0ac7f65
commit ade9f21
Show file tree

Hide file tree

Showing 13 changed files with 673 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -112,19 +112,23 @@ See the [developer guide](https://stfc.github.io/aiida-mlip/developer_guide/inde
     * [`train_parser.py`](aiida_mlip/parsers/train_parser.py): `Parser` for `Train` calculation.
     * [`descriptors_parser.py`](aiida_mlip/parsers/descriptors_parser.py): `Parser` for `Descriptors` calculation.
   * [`helpers/`](aiida_mlip/helpers/): `Helpers` to run calculations.
+  * [`workflows/`](aiida_mlip/workflows/): `WorkGraphs` or `WorkChains` for common workflows with mlips.
+    * [`ht_workgraph.py`](aiida_mlip/workflows/ht_workgraph.py): A `WorkGraph` to run high-throughput optimisations.
 * [`docs/`](docs/source/): Code documentation
   * [`apidoc/`](docs/source/apidoc/): API documentation
   * [`developer_guide/`](docs/source/developer_guide/): Documentation for developers
   * [`user_guide/`](docs/source/user_guide/): Documentation for users
   * [`images/`](docs/source/images/): Logos etc used in the documentation
 * [`examples/`](examples/): Examples for submitting calculations using this plugin
-  * [`tutorials/`](examples/tutorials/): Scripts for submitting calculations
-  * [`calculations/`](examples/calculations/): Jupyter notebooks with tutorials for running calculations and other files that are used in the tutorial
+  * [`tutorials/`](examples/tutorials/): Jupyter notebooks with tutorials for running calculations and other files that are used in the tutorial
+  * [`calculations/`](examples/calculations/): Scripts for submitting calculations
     * [`submit_singlepoint.py`](examples/calculations/submit_singlepoint.py): Script for submitting a singlepoint calculation
     * [`submit_geomopt.py`](examples/calculations/submit_geomopt.py): Script for submitting a geometry optimisation calculation
     * [`submit_md.py`](examples/calculations/submit_md.py): Script for submitting a molecular dynamics calculation
     * [`submit_train.py`](examples/calculations/submit_train.py): Script for submitting a train calculation.
     * [`submit_descriptors.py`](examples/calculations/submit_descriptors.py): Script for submitting a descriptors calculation.
+  * [`workflows/`](examples/workflows/): Scripts for submitting workflows
+    * [`submit_ht_workgraph.py`](examples/workflows/submit_ht_workgraph.py): Script for submitting a high-throughput WorkGraph for singlepoint calculation.
 * [`tests/`](tests/): Basic regression tests using the [pytest](https://docs.pytest.org/en/latest/) framework (submitting a calculation, ...). Install `pip install -e .[testing]` and run `pytest`.
   * [`conftest.py`](tests/conftest.py): Configuration of fixtures for [pytest](https://docs.pytest.org/en/latest/)
   * [`calculations/`](tests/calculations): Calculations
@@ -133,9 +137,11 @@ See the [developer guide](https://stfc.github.io/aiida-mlip/developer_guide/inde
     * [`test_md.py`](tests/calculations/test_md.py): Test `MD` calculation
     * [`test_train.py`](tests/calculations/test_train.py): Test `Train` calculation
     * [`test_descriptors.py`](tests/calculations/test_descriptors.py): Test `Descriptors` calculation
-  * [`data/`](tests/data): `ModelData`
+  * [`data/`](tests/data): Data
     * [`test_model.py`](tests/data/test_model.py): Test `ModelData` type
     * [`test_config.py`](tests/data/test_config.py): Test `JanusConfigfile` type
+  * [`workflows/`](tests/workflows): Workflows
+    * [`test_ht.py`](tests/workflows/test_ht.py): Test high throughput workgraph.
 * [`.gitignore`](.gitignore): Telling git which files to ignore
 * [`.pre-commit-config.yaml`](.pre-commit-config.yaml): Configuration of [pre-commit hooks](https://pre-commit.com/) that sanitize coding style and check for syntax errors. Enable via `pip install -e .[pre-commit] && pre-commit install`
 * [`LICENSE`](LICENSE): License for the plugin

diff --git a/aiida_mlip/workflows/__init__.py b/aiida_mlip/workflows/__init__.py
@@ -0,0 +1 @@
+"""Workflows for aiida-mlip."""
diff --git a/aiida_mlip/workflows/ht_workgraph.py b/aiida_mlip/workflows/ht_workgraph.py
@@ -0,0 +1,140 @@
+"""Workgraph to run high-throughput calculations."""
+
+from pathlib import Path
+from typing import Callable, Union
+
+from aiida.engine import CalcJob, WorkChain
+from aiida.orm import Str
+from aiida_workgraph import WorkGraph, task
+from ase.io import read
+
+from aiida_mlip.helpers.help_load import load_structure
+
+
+@task.graph_builder(outputs=[{"name": "final_structures", "from": "context.structs"}])
+def build_ht_calc(
+    calc: Union[CalcJob, Callable, WorkChain, WorkGraph],
+    folder: Union[Path, str, Str],
+    calc_inputs: dict,
+    input_struct_key: str = "struct",
+    final_struct_key: str = "final_structure",
+    recursive: bool = True,
+) -> WorkGraph:
+    """
+    Build high throughput calculation WorkGraph.
+
+    The `calc` must take a structure, by default `struct`, as one of its inputs.
+    Tasks will then be created to carry out the calculation for each structure file in
+    `folder`.
+
+    Parameters
+    ----------
+    calc : Union[CalcJob, Callable, WorkChain, WorkGraph]
+        Calculation to be performed on all structures.
+    folder : Union[Path, str, Str]
+        Path to the folder containing input structure files.
+    calc_inputs : dict
+        Dictionary of inputs, shared by all the calculations. Must not contain
+        `struct_key`.
+    input_struct_key : str
+        Keyword for input structure for `calc`. Default is "struct".
+    final_struct_key : str
+        Key for final structure output from `calc`. Default is "final_structure".
+    recursive : bool
+        Whether to search `folder` recursively. Default is True.
+
+    Returns
+    -------
+    WorkGraph
+        The workgraph with calculation tasks for each structure.
+
+    Raises
+    ------
+    FileNotFoundError
+        If `folder` has no valid structure files.
+    """
+    wg = WorkGraph()
+    structure = None
+
+    if isinstance(folder, Str):
+        folder = Path(folder.value)
+    if isinstance(folder, str):
+        folder = Path(folder)
+
+    pattern = "**/*" if recursive else "*"
+    for file in filter(Path.is_file, folder.glob(pattern)):
+        try:
+            read(file)
+        except Exception:
+            continue
+        structure = load_structure(file)
+        calc_inputs[input_struct_key] = structure
+        calc_task = wg.add_task(
+            calc,
+            name=f"calc_{file.stem}",
+            **calc_inputs,
+        )
+        calc_task.set_context({final_struct_key: f"structs.{file.stem}"})
+
+    if structure is None:
+        raise FileNotFoundError(
+            f"{folder} is empty or has no readable structure files."
+        )
+
+    return wg
+
+
+def get_ht_workgraph(
+    calc: Union[CalcJob, Callable, WorkChain, WorkGraph],
+    folder: Union[Path, str, Str],
+    calc_inputs: dict,
+    input_struct_key: str = "struct",
+    final_struct_key: str = "final_structure",
+    recursive: bool = True,
+    max_number_jobs: int = 10,
+) -> WorkGraph:
+    """
+    Get WorkGraph to carry out calculation on all structures in a directory.
+
+    Parameters
+    ----------
+    calc : Union[CalcJob, Callable, WorkChain, WorkGraph]
+        Calculation to be performed on all structures.
+    folder : Union[Path, str, Str]
+        Path to the folder containing input structure files.
+    calc_inputs : dict
+        Dictionary of inputs, shared by all the calculations. Must not contain
+        `struct_key`.
+    input_struct_key : str
+        Keyword for input structure for `calc`. Default is "struct".
+    final_struct_key : str
+        Key for final structure output from `calc`. Default is "final_structure".
+    recursive : bool
+        Whether to search `folder` recursively. Default is True.
+    max_number_jobs : int
+        Max number of subprocesses running within the WorkGraph. Default is 10.
+
+    Returns
+    -------
+    WorkGraph
+        The workgraph ready to be submitted.
+    """
+    wg = WorkGraph("ht_calculation")
+
+    wg.add_task(
+        build_ht_calc,
+        name="ht_calc",
+        calc=calc,
+        folder=folder,
+        calc_inputs=calc_inputs,
+        input_struct_key=input_struct_key,
+        final_struct_key=final_struct_key,
+        recursive=recursive,
+    )
+
+    wg.group_outputs = [
+        {"name": "final_structures", "from": "ht_calc.final_structures"}
+    ]
+    wg.max_number_jobs = max_number_jobs
+
+    return wg
diff --git a/docs/source/apidoc/aiida_mlip.rst b/docs/source/apidoc/aiida_mlip.rst
@@ -11,6 +11,7 @@ Subpackages
    aiida_mlip.data
    aiida_mlip.helpers
    aiida_mlip.parsers
+   aiida_mlip.workflows
 
 Module contents
 ---------------

diff --git a/docs/source/apidoc/aiida_mlip.workflows.rst b/docs/source/apidoc/aiida_mlip.workflows.rst
@@ -0,0 +1,25 @@
+aiida\_mlip.workflows package
+=============================
+
+Submodules
+----------
+
+aiida\_mlip.workflows.ht\_workgraph module
+------------------------------------------
+
+.. automodule:: aiida_mlip.workflows.ht_workgraph
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
+
+Module contents
+---------------
+
+.. automodule:: aiida_mlip.workflows
+   :members:
+   :special-members:
+   :private-members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -196,7 +196,11 @@
 
 # Warnings to ignore when using the -n (nitpicky) option
 # We should ignore any python built-in exception, for instance
-nitpick_ignore = [("py:class", "Logger"), ("py:class", "QbFields")]
+nitpick_ignore = [
+    ("py:class", "Logger"),
+    ("py:class", "QbFields"),
+    ("py:class", "aiida_workgraph.workgraph.WorkGraph"),
+]
 
 
 def run_apidoc(_):