✨ Examples: add basic pipeline and controllers

examples/pipeline/README.md

@@ -0,0 +1,46 @@
+# Pipeline example
+
+In this directory you can find an example of what we call a "pipeline", i.e. a collection of submission controllers that run different processes from a "parent group".
+This basic example runs two processes:
+
+1. The `EpwPrepWorkChain`, based on a group of initial structures.
+2. A `EpwCalculation` for interpolating the band structure, based on the `EpwPrepWorkChain` and the restart files it produces.
+
+>[!WARNING]
+> The `aiida-submission-controller` currently still uses the `has_key` filter, which means it can only be used with a PostgreSQL database, see:
+> 
+> https://github.com/aiidateam/aiida-submission-controller/issues/28
+
+## Steps
+
+The "pipeline" is basically a CLI wrapped around the submission controllers.
+I typically use `typer` for my CLI apps, so you have to install that in your environment:
+
+```
+pip install typer 
+```
+
+To make the pipeline CLI easier to run, make it executable:
+
+```
+chmod +x cli_bands.py
+```
+
+The `cli_bands_settings.yaml` file contains all the "pipeline settings", i.e.:
+
+- The AiiDA codes required to run the processes.
+- The groups used in the pipeline to store the structures/processes.
+- More detailed inputs for each of the processes.
+
+Adapt the inputs as needed.
+You can then create the groups used by the submission controllers with:
+
+```
+./cli_bands.py init cli_bands_settings.yaml
+```
+
+and run the pipeline via:
+
+```
+./cli_bands.py init cli_bands_settings.yaml
+```

examples/pipeline/cli_bands.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+"""A command-line interface to run a band structure interpolation pipeline."""
+
+import time
+import typer
+import yaml
+from pathlib import Path
+from rich import print as rprint
+
+from aiida import orm
+from aiida.cmdline.utils import with_dbenv
+from aiida_epw.controllers import EpwPrepWorkChainController, EpwBandsCalculationController
+
+app = typer.Typer(pretty_exceptions_show_locals=False)
+
+
+@app.command()
+@with_dbenv()
+def init(
+    pipeline_settings: Path,
+):
+    """Initialise the groups for the pipeline."""
+    with pipeline_settings.open("r") as handle:
+        settings = yaml.safe_load(handle)
+
+    for group_label in settings['groups'].values():
+        _, created = orm.Group.collection.get_or_create(group_label)
+
+        if created:
+            rprint(f"[bold yellow]Report:[/] created group with label '{group_label}'")
+
+
+@app.command()
+@with_dbenv()
+def run(
+    pipeline_settings: Path,
+):
+    """Run the pipeline."""
+    with pipeline_settings.open("r") as handle:
+        settings = yaml.safe_load(handle)
+
+    epw_prep_controller = EpwPrepWorkChainController(
+        **settings['codes'],
+        **settings['epw_prep']
+    )
+    epw_prep_controller.submit_new_batch()
+
+    epw_bands_controller = EpwBandsCalculationController(
+        epw_code=settings['codes']['epw_code'],
+        **settings['bands_int']
+    )
+    epw_bands_controller.submit_new_batch()
+
+    rprint(f'[bold blue]Info:[/] Sleeping for 60 seconds...')
+    time.sleep(60)
+
+
+if __name__ == "__main__":
+    app()

examples/pipeline/cli_bands_settings.yaml

@@ -0,0 +1,73 @@
+codes:
+  pw_code: qe-7.5-pw@localhost
+  ph_code: qe-7.5-ph@localhost
+  projwfc_code: qe-7.5-projwfc@localhost
+  pw2wannier90_code: qe-7.5-pw2wannier90@localhost
+  wannier90_code: wannier90-jq@localhost
+  epw_code: qe-7.5-epw@localhost
+
+groups:
+  structures: &structures structures
+  workchain_epw_prep: &workchain_epw_prep workchain/epw_prep
+  calculation_epw_bands: &calculation_epw_bands calculation/epw_bands
+
+# Other defaults
+max_wallclock_seconds: &default_walltime 86400
+pseudo_family: &pseudo_family "PseudoDojo/0.5/PBE/SR/standard/upf"
+resources: &default_resources
+  num_machines: 1
+  num_mpiprocs_per_machine: 1
+  num_cores_per_machine: 1
+
+epw_prep:
+  max_concurrent: 1
+  parent_group_label: *structures
+  group_label: *workchain_epw_prep
+  chk2ukk_path: &chk2ukk_path julia --project=/users/mbercx/code/chk2ukk /users/mbercx/code/chk2ukk/chk2ukk.jl
+  overrides: &overrides_epw
+    clean_workdir: True
+    qpoints_distance: 0.5
+    kpoints_distance_scf: 0.15
+    kpoints_factor_nscf: 2
+    epw:
+      options:
+        max_wallclock_seconds: *default_walltime
+        resources: *default_resources
+      parameters:
+        INPUTEPW:
+          use_ws: True
+    ph_base:
+      max_iterations: 1
+      ph:
+        metadata: &default_metadata
+          options:
+            max_wallclock_seconds: *default_walltime
+            resources: *default_resources
+    w90_bands:
+      pseudo_family: *pseudo_family
+      scf: &scf
+        pw:
+          metadata: *default_metadata
+      nscf: *scf
+      projwfc:
+        projwfc:
+          metadata:  &pp_metadata
+            options:
+              max_wallclock_seconds: 7200
+              resources: *default_resources
+      pw2wannier90:
+        pw2wannier90:
+          metadata: *pp_metadata
+      wannier90:
+        wannier90:
+          metadata: *pp_metadata
+
+bands_int:
+  max_concurrent: 1
+  parent_group_label: *workchain_epw_prep
+  group_label: *calculation_epw_bands
+  overrides:
+    metadata:
+        options:
+            resources: *default_resources
+            max_wallclock_seconds: *default_walltime

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "aiida-core~=2.5",
     "aiida-quantumespresso~=4.10",
     "aiida_wannier90_workflows",
+    "aiida_submission_controller~=0.3.0",
     "numpy",
 ]
 

src/aiida_epw/controllers/__init__.py

@@ -0,0 +1,10 @@
+"""Submission controllers for the processes in `aiida-epw`."""
+
+from .epw_bands import EpwBandsCalculationController
+from .prep import EpwPrepWorkChainController
+
+
+__all__ = (
+    "EpwBandsCalculationController",
+    "EpwPrepWorkChainController"
+)

src/aiida_epw/controllers/epw_bands.py

@@ -0,0 +1,66 @@
+"""Submission controller for a band structure `EpwCalculation`."""
+
+import copy
+
+from aiida import orm
+
+from aiida_quantumespresso.workflows.protocols.utils import recursive_merge
+
+from aiida_epw.calculations.epw import EpwCalculation
+from aiida_submission_controller import FromGroupSubmissionController
+from aiida_wannier90_workflows.workflows.bands import Wannier90BandsWorkChain
+
+class EpwBandsCalculationController(FromGroupSubmissionController):
+    """Submission controller for a band structure `EpwCalculation`."""
+
+    epw_code: str
+    overrides: dict | None = None
+
+    def get_inputs_and_processclass_from_extras(self, extras_values, _):
+        """Return inputs and process class for the submission of this specific process."""
+        parent_node = self.get_parent_node_from_extras(extras_values)
+
+        overrides = copy.deepcopy(self.overrides)
+
+        wannier = parent_node.get_outgoing(Wannier90BandsWorkChain).one().node
+
+        epw_source = parent_node.base.links.get_outgoing(link_label_filter="epw").first().node
+
+        builder = EpwCalculation.get_builder()
+        builder.code = orm.load_code(self.epw_code)
+
+        parameters = {
+            "INPUTEPW": {
+                "band_plot": True,
+                "elph": True,
+                "epbread": False,
+                "epbwrite": False,
+                "epwread": True,
+                "epwwrite": False,
+                "fsthick": 100,
+                "wannierize": False,
+                "vme": "dipole",
+            }
+        }
+        parameters["INPUTEPW"]["nbndsub"] = epw_source.inputs.parameters["INPUTEPW"]["nbndsub"]
+        parameters["INPUTEPW"]["use_ws"] = epw_source.inputs.parameters["INPUTEPW"].get("use_ws", False)
+        if "bands_skipped" in epw_source.inputs.parameters["INPUTEPW"]:
+            parameters["INPUTEPW"]["bands_skipped"] = epw_source.inputs.parameters["INPUTEPW"].get("bands_skipped")
+
+        parameters = recursive_merge(parameters, overrides.get("parameters", {}))
+
+        epw_folder = parent_node.outputs.epw_folder
+
+        builder.parameters = orm.Dict(parameters)
+
+        kpoints_path = orm.KpointsData()
+        kpoints_path.set_kpoints(wannier.outputs.seekpath_parameters.get_dict()["explicit_kpoints_rel"])
+        builder.kpoints = epw_source.inputs.kpoints
+        builder.qpoints = epw_source.inputs.qpoints
+        builder.kfpoints = kpoints_path
+        builder.qfpoints = kpoints_path
+        builder.parent_folder_epw = epw_folder
+        builder.settings = orm.Dict(overrides.get("settings", {}))
+        builder.metadata = overrides.get("metadata", {})
+
+        return builder

src/aiida_epw/controllers/prep.py

@@ -0,0 +1,73 @@
+"""Submission controller for the `EpwPrepWorkChain`."""
+from __future__ import annotations
+
+import copy
+
+from aiida import orm
+
+from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain
+from aiida_quantumespresso.common.types import ElectronicType, SpinType
+from aiida_wannier90_workflows.common.types import WannierProjectionType
+from aiida_submission_controller.from_group import FromGroupSubmissionController
+from aiida_epw.workflows.prep import EpwPrepWorkChain
+
+
+class EpwPrepWorkChainController(FromGroupSubmissionController):
+    """Submission controller for the `EpwPrepWorkChain`."""
+
+    pw_code: str
+    ph_code: str
+    projwfc_code: str
+    pw2wannier90_code: str
+    wannier90_code: str
+    epw_code: str
+    chk2ukk_path: str
+    protocol: str = "moderate"
+    overrides: dict | None = None
+    electronic_type: ElectronicType = ElectronicType.METAL
+    spin_type: SpinType = SpinType.NONE
+    wannier_projection_type: WannierProjectionType = WannierProjectionType.SCDM
+    bands_qe_group: str | None = None
+
+    def get_inputs_and_processclass_from_extras(self, extras_values, dry_run=False):
+        """Return inputs and process class for the submission of this specific process."""
+        parent_node = self.get_parent_node_from_extras(extras_values)
+
+        # Depending on the type of node in the parent class, grab the right inputs
+        if isinstance(parent_node, orm.StructureData):
+            structure = parent_node
+        elif parent_node.process_class == PwRelaxWorkChain:
+            structure = parent_node.outputs.output_structure
+        else:
+            raise TypeError(f"Node {parent_node} from parent group is of incorrect type: {type(parent_node)}.")
+
+        codes = {
+            "pw": self.pw_code,
+            "ph": self.ph_code,
+            "projwfc": self.projwfc_code,
+            "pw2wannier90": self.pw2wannier90_code,
+            "wannier90": self.wannier90_code,
+            "epw": self.epw_code,
+        }
+        codes = {key: orm.load_code(code_label) for key, code_label in codes.items()}
+
+        overrides = copy.deepcopy(self.overrides)
+
+        inputs = {
+            "codes": codes,
+            "structure": structure,
+            "protocol": self.protocol,
+            "overrides": overrides,
+            "electronic_type": self.electronic_type,
+            "spin_type": self.spin_type,
+            "wannier_projection_type": self.wannier_projection_type,
+        }
+
+        if self.wannier_projection_type == WannierProjectionType.ATOMIC_PROJECTORS_QE:
+            raise ValueError('Atomic projectors not yet supported!')
+
+        builder = EpwPrepWorkChain.get_builder_from_protocol(**inputs)
+        w90_script = orm.RemoteData(remote_path=self.chk2ukk_path, computer=codes["pw"].computer)
+        builder.w90_chk_to_ukk_script = w90_script
+
+        return builder