From 0ca74303a78060345f26c239b8e7bd9f87a8de32 Mon Sep 17 00:00:00 2001 From: Xing Wang Date: Mon, 23 Dec 2024 17:10:44 +0100 Subject: [PATCH] Enable pencil decomposition for all pw.x calculations. (#1025) This solves the problem when runninga small system using many CPUs --- calculation_history.ipynb | 2 +- docs/source/development/plugin.rst | 2 +- src/aiidalab_qe/plugins/bands/workchain.py | 9 ++++++++- src/aiidalab_qe/plugins/pdos/workchain.py | 8 +++++++- src/aiidalab_qe/plugins/xas/workchain.py | 6 +++++- src/aiidalab_qe/plugins/xps/workchain.py | 6 +++++- src/aiidalab_qe/{plugins => }/utils.py | 6 ++++++ src/aiidalab_qe/workflows/__init__.py | 2 ++ 8 files changed, 35 insertions(+), 6 deletions(-) rename src/aiidalab_qe/{plugins => }/utils.py (87%) diff --git a/calculation_history.ipynb b/calculation_history.ipynb index c917f7205..151a8dd99 100644 --- a/calculation_history.ipynb +++ b/calculation_history.ipynb @@ -59,7 +59,7 @@ "

\n", " \n", diff --git a/docs/source/development/plugin.rst b/docs/source/development/plugin.rst index 5199e70cf..7c37a53c5 100644 --- a/docs/source/development/plugin.rst +++ b/docs/source/development/plugin.rst @@ -249,7 +249,7 @@ feed the `metadata` of each Calcjob which is submitted in the workchain. For exa .. code-block:: python - from aiidalab_qe.plugins.utils import set_component_resources + from aiidalab_qe.utils import set_component_resources def update_resources(builder, codes): set_component_resources(builder.pw, codes.get("pw")) diff --git a/src/aiidalab_qe/plugins/bands/workchain.py b/src/aiidalab_qe/plugins/bands/workchain.py index 863a7dc0b..042bf4e97 100644 --- a/src/aiidalab_qe/plugins/bands/workchain.py +++ b/src/aiidalab_qe/plugins/bands/workchain.py @@ -1,6 +1,9 @@ from aiida.plugins import WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType -from aiidalab_qe.plugins.utils import set_component_resources +from aiidalab_qe.utils import ( + enable_pencil_decomposition, + set_component_resources, +) BandsWorkChain = WorkflowFactory("aiidalab_qe.bands_workchain") # from .bands_workchain import BandsWorkChain @@ -33,12 +36,16 @@ def update_resources(builder, codes): if "bands" in builder: set_component_resources(builder.bands.scf.pw, codes.get("pw")) set_component_resources(builder.bands.bands.pw, codes.get("pw")) + enable_pencil_decomposition(builder.bands.scf.pw) + enable_pencil_decomposition(builder.bands.bands.pw) elif "bands_projwfc" in builder: set_component_resources(builder.bands_projwfc.scf.pw, codes.get("pw")) set_component_resources(builder.bands_projwfc.bands.pw, codes.get("pw")) set_component_resources( builder.bands_projwfc.projwfc.projwfc, codes.get("projwfc_bands") ) + enable_pencil_decomposition(builder.bands_projwfc.scf.pw) + enable_pencil_decomposition(builder.bands_projwfc.bands.pw) def get_builder(codes, structure, parameters, **kwargs): diff --git a/src/aiidalab_qe/plugins/pdos/workchain.py b/src/aiidalab_qe/plugins/pdos/workchain.py index f63c6ea1c..91a877343 100644 --- a/src/aiidalab_qe/plugins/pdos/workchain.py +++ b/src/aiidalab_qe/plugins/pdos/workchain.py @@ -1,7 +1,10 @@ from aiida import orm from aiida.plugins import WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType -from aiidalab_qe.plugins.utils import set_component_resources +from aiidalab_qe.utils import ( + enable_pencil_decomposition, + set_component_resources, +) PdosWorkChain = WorkflowFactory("quantumespresso.pdos") PwBandsWorkChain = WorkflowFactory("quantumespresso.pw.bands") @@ -37,6 +40,9 @@ def update_resources(builder, codes): set_component_resources(builder.nscf.pw, codes.get("pw")) set_component_resources(builder.dos, codes.get("dos")) set_component_resources(builder.projwfc, codes.get("projwfc")) + enable_pencil_decomposition(builder.scf.pw) + enable_pencil_decomposition(builder.nscf.pw) + # disable the parallelization setting for projwfc # npool = codes["pw"]["parallelization"]["npool"] # builder.projwfc.settings = orm.Dict(dict={"cmdline": ["-nk", str(npool)]}) diff --git a/src/aiidalab_qe/plugins/xas/workchain.py b/src/aiidalab_qe/plugins/xas/workchain.py index 9ad99c207..bcbc13965 100644 --- a/src/aiidalab_qe/plugins/xas/workchain.py +++ b/src/aiidalab_qe/plugins/xas/workchain.py @@ -6,7 +6,10 @@ from aiida.plugins import WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType from aiidalab_qe.plugins import xas as xas_folder -from aiidalab_qe.plugins.utils import set_component_resources +from aiidalab_qe.utils import ( + enable_pencil_decomposition, + set_component_resources, +) XspectraCrystalWorkChain = WorkflowFactory("quantumespresso.xspectra.crystal") PSEUDO_TOC = yaml.safe_load(resources.read_text(xas_folder, "pseudo_toc.yaml")) @@ -18,6 +21,7 @@ def update_resources(builder, codes): """Update the resources for the builder.""" set_component_resources(builder.core.scf.pw, codes.get("pw")) set_component_resources(builder.core.xs_prod.xspectra, codes.get("xspectra")) + enable_pencil_decomposition(builder.core.scf.pw) def get_builder(codes, structure, parameters, **kwargs): diff --git a/src/aiidalab_qe/plugins/xps/workchain.py b/src/aiidalab_qe/plugins/xps/workchain.py index 06379b99c..c72d76b2c 100644 --- a/src/aiidalab_qe/plugins/xps/workchain.py +++ b/src/aiidalab_qe/plugins/xps/workchain.py @@ -1,7 +1,10 @@ from aiida.orm import Bool, Dict, Float, Group, QueryBuilder from aiida.plugins import WorkflowFactory from aiida_quantumespresso.common.types import ElectronicType, SpinType -from aiidalab_qe.plugins.utils import set_component_resources +from aiidalab_qe.utils import ( + enable_pencil_decomposition, + set_component_resources, +) XpsWorkChain = WorkflowFactory("quantumespresso.xps") @@ -16,6 +19,7 @@ def update_resources(builder, codes): """Update the resources for the builder.""" set_component_resources(builder.ch_scf.pw, codes.get("pw")) + enable_pencil_decomposition(builder.ch_scf.pw) def get_builder(codes, structure, parameters, **kwargs): diff --git a/src/aiidalab_qe/plugins/utils.py b/src/aiidalab_qe/utils.py similarity index 87% rename from src/aiidalab_qe/plugins/utils.py rename to src/aiidalab_qe/utils.py index a1d6ee841..17e8013f4 100644 --- a/src/aiidalab_qe/plugins/utils.py +++ b/src/aiidalab_qe/utils.py @@ -24,3 +24,9 @@ def set_component_resources(component, code_info): ] if "parallelization" in code_info: component.parallelization = orm.Dict(dict=code_info["parallelization"]) + + +def enable_pencil_decomposition(component): + """Enable the pencil decomposition for the given component.""" + + component.settings = orm.Dict({"CMDLINE": ["-pd", ".true."]}) diff --git a/src/aiidalab_qe/workflows/__init__.py b/src/aiidalab_qe/workflows/__init__.py index 5cb0f0658..6a9334f5f 100644 --- a/src/aiidalab_qe/workflows/__init__.py +++ b/src/aiidalab_qe/workflows/__init__.py @@ -8,6 +8,7 @@ from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData from aiida_quantumespresso.utils.mapping import prepare_process_inputs from aiida_quantumespresso.workflows.pw.relax import PwRelaxWorkChain +from aiidalab_qe.utils import enable_pencil_decomposition XyData = DataFactory("core.array.xy") StructureData = DataFactory("core.structure") @@ -187,6 +188,7 @@ def get_builder_from_protocol( overrides=relax_overrides, **kwargs, ) + enable_pencil_decomposition(relax_builder.base.pw) # pop the inputs that are excluded from the expose_inputs relax_builder.pop("structure", None) relax_builder.pop("clean_workdir", None)