Enable pencil decomposition for all pw.x calculations. (#1025)

This solves the problem when runninga  small system using many CPUs

calculation_history.ipynb

@@ -59,7 +59,7 @@
     "        </p>\n",
     "        <ul>\n",
     "            <li><b>Label Search Field:</b> Enter a job label to find matching jobs.</li>\n",
-    "            <li><b>Job State Dropdown:</b> Filter jobs based on their state (e.g., finished, waiting).</li>\n",
+    "            <li><b>Job State Dropdown:</b> Filter jobs based on their state (e.g., finished, running).</li>\n",
     "            <li><b>Date Range Picker:</b> Select a start and end date to view jobs created within that range.</li>\n",
     "            <li><b>Properties Filter:</b> Select specific properties associated with jobs.</li>\n",
     "        </ul>\n",

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"))

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):

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)]})

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):

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):

src/aiidalab_qe/plugins/utils.py → 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."]})

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)