diff --git a/examples/pipeline/README.md b/examples/pipeline/README.md new file mode 100644 index 0000000..7abe5fe --- /dev/null +++ b/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 +``` diff --git a/examples/pipeline/cli_bands.py b/examples/pipeline/cli_bands.py new file mode 100755 index 0000000..452bdab --- /dev/null +++ b/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() diff --git a/examples/pipeline/cli_bands_settings.yaml b/examples/pipeline/cli_bands_settings.yaml new file mode 100644 index 0000000..e2c9870 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 6fd2cbd..e893f9b 100644 --- a/pyproject.toml +++ b/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", ] diff --git a/src/aiida_epw/controllers/__init__.py b/src/aiida_epw/controllers/__init__.py new file mode 100644 index 0000000..21704da --- /dev/null +++ b/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" +) diff --git a/src/aiida_epw/controllers/epw_bands.py b/src/aiida_epw/controllers/epw_bands.py new file mode 100644 index 0000000..79052f2 --- /dev/null +++ b/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 diff --git a/src/aiida_epw/controllers/prep.py b/src/aiida_epw/controllers/prep.py new file mode 100644 index 0000000..9678d05 --- /dev/null +++ b/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