From ef9d5a3da3296b20b3404fd6c4fdb6b96eefdb68 Mon Sep 17 00:00:00 2001 From: superstar54 Date: Wed, 20 Mar 2024 15:52:46 +0000 Subject: [PATCH] add aiidalab_qe plugin --- aiida_bader/aiidalab/__init__.py | 29 +++++++ aiida_bader/aiidalab/result.py | 67 +++++++++++++++ aiida_bader/aiidalab/workchain.py | 87 ++++++++++++++++++++ aiida_bader/workchains/protocols/__init__.py | 0 aiida_bader/workchains/protocols/bader.yaml | 34 ++++++++ aiida_bader/workchains/qe_bader.py | 77 ++++++++++++----- examples/bader-localhost.yaml | 13 +++ examples/qe_bader_workchain.py | 55 +++++++++++++ setup.py | 5 +- 9 files changed, 346 insertions(+), 21 deletions(-) create mode 100644 aiida_bader/aiidalab/__init__.py create mode 100644 aiida_bader/aiidalab/result.py create mode 100644 aiida_bader/aiidalab/workchain.py create mode 100644 aiida_bader/workchains/protocols/__init__.py create mode 100644 aiida_bader/workchains/protocols/bader.yaml create mode 100644 examples/bader-localhost.yaml create mode 100644 examples/qe_bader_workchain.py diff --git a/aiida_bader/aiidalab/__init__.py b/aiida_bader/aiidalab/__init__.py new file mode 100644 index 0000000..06b7f30 --- /dev/null +++ b/aiida_bader/aiidalab/__init__.py @@ -0,0 +1,29 @@ +from aiidalab_qe.common.panel import OutlinePanel +from aiidalab_widgets_base import ComputationalResourcesWidget + +from .result import Result +from .workchain import workchain_and_builder + + +class BaderOutline(OutlinePanel): + title = "Bader charge analysis" + help = """""" + + +pp_code = ComputationalResourcesWidget( + description="pp.x", + default_calc_job_plugin="quantumespresso.pp", +) + +bader_code = ComputationalResourcesWidget( + description="bader", + default_calc_job_plugin="bader", +) + + +bader = { + "outline": BaderOutline, + "code": {"pp": pp_code, "bader": bader_code}, + "result": Result, + "workchain": workchain_and_builder, +} diff --git a/aiida_bader/aiidalab/result.py b/aiida_bader/aiidalab/result.py new file mode 100644 index 0000000..fdf7934 --- /dev/null +++ b/aiida_bader/aiidalab/result.py @@ -0,0 +1,67 @@ +import ipywidgets as ipw +from aiidalab_qe.common.panel import ResultPanel + + +class Result(ResultPanel): + title = "Bader Charge" + workchain_labels = ["bader"] + + def __init__(self, node=None, **kwargs): + super().__init__(node=node, **kwargs) + self.summary_view = ipw.HTML() + + def _update_view(self): + structure = self.node.inputs.bader.structure + bader_charge = self.outputs.bader.bader.bader_charge.get_array("charge") + self._generate_table(structure, bader_charge) + self.children = [ + ipw.HBox( + children=[self.summary_view], + layout=ipw.Layout(justify_content="space-between", margin="10px"), + ), + ] + + def _generate_table(self, structure, bader_charge): + # get index and element form AiiDA StructureData + site_index = [site.kind_name for site in structure.sites] + + # Start of the HTML string for the table + html_str = """
+

Bader Charge Table

+ + + + + + + """ + + # Add rows to the table based on the bader_charge + for i in range(len(site_index)): + html_str += f""" + + + + + """ + + # Close the table and div tags + html_str += """ +
Site IndexElementBader Charge
{i}{site_index[i]}{bader_charge[i]:1.3f}
+
""" + self.summary_view = ipw.HTML(html_str) diff --git a/aiida_bader/aiidalab/workchain.py b/aiida_bader/aiidalab/workchain.py new file mode 100644 index 0000000..f7ee707 --- /dev/null +++ b/aiida_bader/aiidalab/workchain.py @@ -0,0 +1,87 @@ +from aiida.plugins import WorkflowFactory +from aiida_quantumespresso.common.types import ElectronicType, SpinType +from aiida import orm + +QeBaderWorkChain = WorkflowFactory("bader.qe") + + +def check_codes(pw_code, pp_code, bader_code): + """Check that the codes are installed on the same computer.""" + if ( + not any( + [ + pw_code is None, + pp_code is None, + bader_code is None, + ] + ) + and len( + set( + ( + pw_code.computer.pk, + pp_code.computer.pk, + bader_code.computer.pk, + ) + ) + ) + != 1 + ): + raise ValueError( + "All selected codes must be installed on the same computer. This is because the " + "Bader calculations rely on large files that are not retrieved by AiiDA." + ) + + +def get_builder(codes, structure, parameters, **kwargs): + from copy import deepcopy + + pw_code = codes.get("pw") + pp_code = codes.get("pp") + bader_code = codes.get("bader") + check_codes(pw_code, pp_code, bader_code) + protocol = parameters["workchain"]["protocol"] + + scf_overrides = deepcopy(parameters["advanced"]) + + overrides = { + "scf": scf_overrides, + "pp": { + "parameters": orm.Dict( + { + "INPUTPP": {"plot_num": 21}, + "PLOT": {"iflag": 3}, + } + ), + "metadata": { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": 1, + }, + } + }, + }, + } + if pp_code is not None and bader_code is not None: + builder = QeBaderWorkChain.get_builder_from_protocol( + pw_code=pw_code, + pp_code=pp_code, + bader_code=bader_code, + structure=structure, + protocol=protocol, + electronic_type=ElectronicType(parameters["workchain"]["electronic_type"]), + spin_type=SpinType(parameters["workchain"]["spin_type"]), + initial_magnetic_moments=parameters["advanced"]["initial_magnetic_moments"], + overrides=overrides, + **kwargs, + ) + else: + raise ValueError("The pp_code and bader_code are required.") + return builder + + +workchain_and_builder = { + "workchain": QeBaderWorkChain, + "exclude": ("clean_workdir", "structure", "relax"), + "get_builder": get_builder, +} diff --git a/aiida_bader/workchains/protocols/__init__.py b/aiida_bader/workchains/protocols/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/aiida_bader/workchains/protocols/bader.yaml b/aiida_bader/workchains/protocols/bader.yaml new file mode 100644 index 0000000..cfd6a8c --- /dev/null +++ b/aiida_bader/workchains/protocols/bader.yaml @@ -0,0 +1,34 @@ +default_inputs: + clean_workdir: False + scf: + pw: + parameters: + CONTROL: + restart_mode: from_scratch + pp: + parameters: + INPUTPP: + plot_num: 21 + PLOT: + ifloag: 3 + metadata: + options: + resources: + num_machines: 1 + max_wallclock_seconds: 43200 # Twelve hours + withmpi: True + bader: + metadata: + options: + resources: + num_machines: 1 + max_wallclock_seconds: 43200 # Twelve hours + withmpi: True +default_protocol: moderate +protocols: + moderate: + description: 'Protocol to perform a projected density of states calculation at normal precision at moderate computational cost.' + precise: + description: 'Protocol to perform a projected density of states structure calculation at high precision at higher computational cost.' + fast: + description: 'Protocol to perform a projected density of states structure calculation at low precision at minimal computational cost for testing purposes.' diff --git a/aiida_bader/workchains/qe_bader.py b/aiida_bader/workchains/qe_bader.py index 4c4780a..00b4398 100644 --- a/aiida_bader/workchains/qe_bader.py +++ b/aiida_bader/workchains/qe_bader.py @@ -8,6 +8,7 @@ from aiida.plugins import CalculationFactory, WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, RestartType, SpinType from aiida import orm +from aiida_quantumespresso.workflows.protocols.utils import ProtocolMixin PwBaseWorkChain = WorkflowFactory( "quantumespresso.pw.base" @@ -16,7 +17,7 @@ BaderCalculation = CalculationFactory("bader") # pylint: disable=invalid-name -class QeBaderWorkChain(WorkChain): +class QeBaderWorkChain(ProtocolMixin, WorkChain): """A workchain that computes bader charges using QE and Bader code.""" @classmethod @@ -24,7 +25,19 @@ def define(cls, spec): """Define workflow specification.""" super(QeBaderWorkChain, cls).define(spec) - spec.expose_inputs(PwBaseWorkChain, namespace="scf") + spec.input( + "structure", valid_type=orm.StructureData, help="The input structure." + ) + spec.expose_inputs( + PwBaseWorkChain, + namespace="scf", + exclude=("clean_workdir", "pw.structure", "pw.parent_folder"), + namespace_options={ + "help": "Inputs for the `PwBaseWorkChain` of the `scf` calculation.", + "required": False, + "populate_defaults": False, + }, + ) spec.expose_inputs(PpCalculation, namespace="pp", exclude=["parent_folder"]) spec.expose_inputs( BaderCalculation, namespace="bader", exclude=["charge_density_folder"] @@ -42,6 +55,15 @@ def define(cls, spec): 905, "ERROR_PARSING_BADER_OUTPUT", "Error while parsing bader output" ) + @classmethod + def get_protocol_filepath(cls): + """Return ``pathlib.Path`` to the ``.yaml`` file that defines the protocols.""" + from importlib_resources import files + + from . import protocols + + return files(protocols) / "bader.yaml" + @classmethod def get_builder_from_protocol( cls, @@ -52,7 +74,7 @@ def get_builder_from_protocol( protocol=None, overrides=None, options=None, - **kwargs + **kwargs, ): """Return a builder prepopulated with inputs selected according to the chosen protocol. @@ -61,37 +83,54 @@ def get_builder_from_protocol( :param protocol: protocol to use, if not specified, the default will be used. :param overrides: optional dictionary of inputs to override the defaults of the protocol. """ + from aiida_quantumespresso.workflows.protocols.utils import recursive_merge + + inputs = cls.get_protocol_inputs(protocol, overrides) if isinstance(pw_code, str): pw_code = orm.load_code(pw_code) if isinstance(bader_code, str): bader_code = orm.load_code(bader_code) - inputs = cls.get_protocol_inputs(protocol, overrides) - scf = PwBaseWorkChain.get_builder_from_protocol( - pw_code, structure, protocol, overrides=inputs.get('scf', None), - options=options, **kwargs - ) - pp = PpCalculation.get_builder_from_protocol( - pp_code, structure, protocol, overrides=inputs.get('pp', None), - options=options, **kwargs - ) - bader = BaderCalculation.get_builder_from_protocol( - bader_code, structure, protocol, overrides=inputs.get('bader', None), - options=options, **kwargs + pw_code, + structure, + protocol, + overrides=inputs.get("scf", None), + options=options, + **kwargs, ) + scf["pw"].pop("structure", None) + + metadata_pp = inputs.get("pp", {}).get("metadata", {"options": {}}) + metadata_bader = inputs.get("bader", {}).get("metadata", {"options": {}}) + + if options: + metadata_pp["options"] = recursive_merge(metadata_pp["options"], options) + metadata_bader["options"] = recursive_merge( + metadata_bader["options"], options + ) builder = cls.get_builder() + builder.structure = structure builder.scf = scf - builder.pp = pp - builder.bader = bader + builder.pp.code = pp_code # pylint: disable=no-member + builder.pp.parameters = orm.Dict( + inputs.get("pp", {}).get("parameters") + ) # pylint: disable=no-member + builder.pp.metadata = metadata_pp # pylint: disable=no-member + builder.bader.code = bader_code # pylint: disable=no-member + builder.bader.parameters = orm.Dict( + inputs.get("bader", {}).get("parameters") + ) # pylint: disable=no-member + builder.bader.metadata = metadata_bader # pylint: disable=no-member return builder def run_pw(self): """Run PW.""" scf_inputs = AttributeDict(self.exposed_inputs(PwBaseWorkChain, "scf")) + scf_inputs.pw.structure = self.inputs.structure scf_inputs["metadata"]["label"] = "pw_scf" scf_inputs["metadata"]["call_link_label"] = "call_pw_scf" running = self.submit(PwBaseWorkChain, **scf_inputs) @@ -142,9 +181,7 @@ def return_results(self): """Return exposed outputs and print the pk of the ArrayData w/bader""" try: self.out_many( - self.exposed_outputs( - self.ctx.pw_calc, PwBaseWorkChain, namespace="scf" - ) + self.exposed_outputs(self.ctx.pw_calc, PwBaseWorkChain, namespace="scf") ) self.out_many( self.exposed_outputs(self.ctx.pp_calc, PpCalculation, namespace="pp") diff --git a/examples/bader-localhost.yaml b/examples/bader-localhost.yaml new file mode 100644 index 0000000..49ca393 --- /dev/null +++ b/examples/bader-localhost.yaml @@ -0,0 +1,13 @@ +--- +label: bader +description: Bader charge analysis +default_calc_job_plugin: bader +computer: localhost +filepath_executable: /home/jovyan/.conda/envs/quantum-espresso-7.2/bin/bader +prepend_text: | + + eval "$(conda shell.posix hook)" + conda activate quantum-espresso-7.2 + + export OMP_NUM_THREADS=1 +append_text: ' ' diff --git a/examples/qe_bader_workchain.py b/examples/qe_bader_workchain.py new file mode 100644 index 0000000..40aebae --- /dev/null +++ b/examples/qe_bader_workchain.py @@ -0,0 +1,55 @@ +from aiida import load_profile +from aiida.orm import Dict, KpointsData, StructureData, load_code, load_group, load_node +from ase.build import molecule +from aiida.plugins import WorkflowFactory +from aiida.engine import submit + +QeBaderWorkChain = WorkflowFactory("bader.qe") + + +load_profile() +# =============================================================================== +# load the codes +pw_code = load_code("pw-7.2@localhost") +pp_code = load_code("pp-7.2@localhost") +bader_code = load_code("bader@localhost") + + +# create input structure node +h2o = molecule("H2O") +h2o.center(vacuum=3.0) +h2o.pbc = True +structure = StructureData(ase=h2o) + + +overrides = { + "pp": { + "parameters": Dict( + dict={ + "INPUTPP": {"plot_num": 21}, + "PLOT": {"iflag": 3}, + } + ), + "metadata": { + "options": { + "resources": { + "num_machines": 1, + "num_mpiprocs_per_machine": 2, + }, + "max_wallclock_seconds": 3600, + } + }, + }, +} + +builder = QeBaderWorkChain.get_builder_from_protocol( + pw_code, + pp_code, + bader_code, + structure, + protocol="fast", + overrides=overrides, + options=None, +) + +submit(builder) diff --git a/setup.py b/setup.py index f61a294..011a55a 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ def test_suite(): install_requires=[ "aiida-core", "aiida-worktree", - "aiida-quantumespresso", + "aiida-quantumespresso~=4.4", "aiida-cp2k", "pytest", "pytest-cov", @@ -47,6 +47,9 @@ def test_suite(): "aiida.workflows": [ "bader.qe = aiida_bader.workchains:QeBaderWorkChain", ], + "aiidalab_qe.properties": [ + "bader = aiida_bader.aiidalab:bader", + ], }, package_data={}, python_requires=">=3.8",