Plugin API: add a update_inputs to plugin's workchain_and_builder

src/aiidalab_qe/plugins/bands/workchain.py

@@ -225,8 +225,19 @@ def get_builder(codes, structure, parameters, **kwargs):
     return bands
 
 
+def update_inputs(inputs, ctx):
+    """Update the inputs using context."""
+    inputs.structure = ctx.current_structure
+    inputs.scf.pw.parameters = inputs.scf.pw.parameters.get_dict()
+    if ctx.current_number_of_bands:
+        inputs.scf.pw.parameters.setdefault("SYSTEM", {}).setdefault(
+            "nbnd", ctx.current_number_of_bands
+        )
+
+
 workchain_and_builder = {
     "workchain": PwBandsWorkChain,
     "exclude": ("structure", "relax"),
     "get_builder": get_builder,
+    "update_inputs": update_inputs,
 }

src/aiidalab_qe/plugins/pdos/workchain.py

@@ -78,8 +78,19 @@ def get_builder(codes, structure, parameters, **kwargs):
     return pdos
 
 
+def update_inputs(inputs, ctx):
+    """Update the inputs using context."""
+    inputs.structure = ctx.current_structure
+    inputs.nscf.pw.parameters = inputs.nscf.pw.parameters.get_dict()
+    if ctx.current_number_of_bands:
+        inputs.nscf.pw.parameters.setdefault("SYSTEM", {}).setdefault(
+            "nbnd", ctx.current_number_of_bands
+        )
+
+
 workchain_and_builder = {
     "workchain": PdosWorkChain,
     "exclude": ("structure", "relax"),
     "get_builder": get_builder,
+    "update_inputs": update_inputs,
 }

src/aiidalab_qe/plugins/xas/workchain.py

@@ -104,8 +104,14 @@ def get_builder(codes, structure, parameters, **kwargs):
     return builder
 
 
+def update_inputs(inputs, ctx):
+    """Update the inputs using context."""
+    inputs.structure = ctx.current_structure
+
+
 workchain_and_builder = {
     "workchain": XspectraCrystalWorkChain,
     "exclude": ("structure", "relax"),
     "get_builder": get_builder,
+    "update_inputs": update_inputs,
 }

src/aiidalab_qe/plugins/xps/workchain.py

@@ -98,8 +98,14 @@ def get_builder(codes, structure, parameters, **kwargs):
     return builder
 
 
+def update_inputs(inputs, ctx):
+    """Update the inputs using context."""
+    inputs.structure = ctx.current_structure
+
+
 workchain_and_builder = {
     "workchain": XpsWorkChain,
     "exclude": ("structure", "relax"),
     "get_builder": get_builder,
+    "update_inputs": update_inputs,
 }

src/aiidalab_qe/workflows/__init__.py

@@ -238,7 +238,8 @@ def run_plugin(self):
                 self.exposed_inputs(plugin_workchain, namespace=name)
             )
             inputs.metadata.call_link_label = name
-            inputs.structure = self.ctx.current_structure
+            if entry_point.get("update_inputs"):
+                entry_point["update_inputs"](inputs, self.ctx)
             inputs = prepare_process_inputs(plugin_workchain, inputs)
             running = self.submit(plugin_workchain, **inputs)
             self.report(f"launching plugin {name} <{running.pk}>")