fix code for xps, xas, and update the code resources (#711)

This PR updates the codes and resources of the plugins
- Fix the code for XAS and XPS
- Update the resources for the builder

src/aiidalab_qe/plugins/bands/workchain.py

@@ -1,7 +1,7 @@
 import numpy as np
-from aiida import orm
 from aiida.plugins import DataFactory, WorkflowFactory
 from aiida_quantumespresso.common.types import ElectronicType, SpinType
+from aiidalab_qe.plugins.utils import set_component_resources
 
 GAMMA = "\u0393"
 
@@ -173,18 +173,8 @@ def generate_kpath_2d(structure, kpoints_distance, kpath_2d):
 
 
 def update_resources(builder, codes):
-    builder.scf.pw.metadata.options.resources = {
-        "num_machines": codes.get("pw")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"],
-    }
-    builder.scf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"])
-    builder.bands.pw.metadata.options.resources = {
-        "num_machines": codes.get("pw")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"],
-    }
-    builder.bands.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"])
+    set_component_resources(builder.scf.pw, codes.get("pw"))
+    set_component_resources(builder.bands.pw, codes.get("pw"))
 
 
 def get_builder(codes, structure, parameters, **kwargs):

src/aiidalab_qe/plugins/pdos/workchain.py

@@ -1,6 +1,6 @@
-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
 
 PdosWorkChain = WorkflowFactory("quantumespresso.pdos")
 
@@ -33,28 +33,10 @@ def check_codes(pw_code, dos_code, projwfc_code):
 
 
 def update_resources(builder, codes):
-    builder.scf.pw.metadata.options.resources = {
-        "num_machines": codes.get("pw")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"],
-    }
-    builder.scf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"])
-    builder.nscf.pw.metadata.options.resources = {
-        "num_machines": codes.get("pw")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("pw")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("pw")["cpus_per_task"],
-    }
-    builder.nscf.pw.parallelization = orm.Dict(dict=codes["pw"]["parallelization"])
-    builder.dos.metadata.options.resources = {
-        "num_machines": codes.get("dos")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("dos")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("dos")["cpus_per_task"],
-    }
-    builder.projwfc.metadata.options.resources = {
-        "num_machines": codes.get("projwfc")["nodes"],
-        "num_mpiprocs_per_machine": codes.get("projwfc")["ntasks_per_node"],
-        "num_cores_per_mpiproc": codes.get("projwfc")["cpus_per_task"],
-    }
+    set_component_resources(builder.scf.pw, codes.get("pw"))
+    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"))
     # 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/utils.py

@@ -0,0 +1,13 @@
+from aiida import orm
+
+
+def set_component_resources(component, code_info):
+    """Set the resources for a given component based on the code info."""
+    if code_info:  # Ensure code_info is not None or empty
+        component.metadata.options.resources = {
+            "num_machines": code_info["nodes"],
+            "num_mpiprocs_per_machine": code_info["ntasks_per_node"],
+            "num_cores_per_mpiproc": code_info["cpus_per_task"],
+        }
+        if "parallelization" in code_info:
+            component.parallelization = orm.Dict(dict=code_info["parallelization"])

src/aiidalab_qe/plugins/xas/workchain.py

@@ -4,6 +4,7 @@
 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.plugins import xas as xas_folder
 
@@ -13,6 +14,12 @@
 xch_elements = PSEUDO_TOC["xas_xch_elements"]
 
 
+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"))
+
+
 def get_builder(codes, structure, parameters, **kwargs):
     from copy import deepcopy
 
@@ -50,8 +57,8 @@ def get_builder(codes, structure, parameters, **kwargs):
     }
     spglib_settings = orm.Dict({"symprec": 1.0e-3})
 
-    pw_code = codes["pw"]
-    xs_code = codes["xspectra"]
+    pw_code = codes["pw"]["code"]
+    xs_code = codes["xspectra"]["code"]
     overrides = {
         "core": {
             "scf": deepcopy(parameters["advanced"]),
@@ -100,6 +107,8 @@ def get_builder(codes, structure, parameters, **kwargs):
     builder.pop("clean_workdir", None)
     builder.spglib_settings = spglib_settings
     builder.structure_preparation_settings = structure_preparation_settings
+    # update resources
+    update_resources(builder, codes)
 
     return builder
 

src/aiidalab_qe/plugins/xps/workchain.py

@@ -1,6 +1,7 @@
 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
 
 XpsWorkChain = WorkflowFactory("quantumespresso.xps")
 
@@ -12,6 +13,11 @@
 }
 
 
+def update_resources(builder, codes):
+    """Update the resources for the builder."""
+    set_component_resources(builder.ch_scf.pw, codes.get("pw"))
+
+
 def get_builder(codes, structure, parameters, **kwargs):
     from copy import deepcopy
 
@@ -62,7 +68,7 @@ def get_builder(codes, structure, parameters, **kwargs):
         "supercell_min_parameter": Float(supercell_min_parameter_map[protocol]),
         "is_molecule_input": Bool(is_molecule_input),
     }
-    pw_code = codes.get("pw", None)
+    pw_code = codes["pw"]["code"]
     overrides_ch_scf = deepcopy(parameters["advanced"])
     if is_molecule_input:
         overrides_ch_scf["pw"]["parameters"]["SYSTEM"]["assume_isolated"] = "mt"
@@ -91,6 +97,8 @@ def get_builder(codes, structure, parameters, **kwargs):
     )
     builder.pop("relax")
     builder.pop("clean_workdir", None)
+    # update resources
+    update_resources(builder, codes)
     if is_molecule_input:
         # set a large kpoints_distance value to set the kpoints to 1x1x1
         builder.ch_scf.kpoints_distance = Float(5)

tests/test_submit_qe_workchain.py

@@ -20,6 +20,9 @@ def test_create_builder_default(
     # regression test for the parameters generated by the app
     # this parameters are passed to the workchain
     data_regression.check(ui_parameters)
+    # test if create builder successfully
+    submit_step._create_builder()
+    # In the future, we will check the builder parameters using regresion test
 
 
 @pytest.mark.usefixtures("sssp")