udpate get_builder

superstar54 · May 10, 2024 · 0c18b46 · 0c18b46
1 parent 4ca6c05
commit 0c18b46
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 9 deletions.
diff --git a/aiida_bader/workchains/qe_bader.py b/aiida_bader/workchains/qe_bader.py
@@ -120,23 +120,37 @@ def get_builder_from_protocol(
         )
         scf["pw"].pop("structure", None)
 
-        metadata_pp = inputs.get("pp", {}).get("metadata", {"options": {}})
+        metadata_pp_valence = inputs.get("pp_valence", {}).get(
+            "metadata", {"options": {}}
+        )
+        metadata_pp_all = inputs.get("pp_all", {}).get("metadata", {"options": {}})
         metadata_bader = inputs.get("bader", {}).get("metadata", {"options": {}})
 
         if options:
-            metadata_pp["options"] = recursive_merge(metadata_pp["options"], options)
+            metadata_pp_valence["options"] = recursive_merge(
+                metadata_pp_valence["options"], options
+            )
+            metadata_pp_all["options"] = recursive_merge(
+                metadata_pp_all["options"], options
+            )
             metadata_bader["options"] = recursive_merge(
                 metadata_bader["options"], options
             )
 
         builder = cls.get_builder()
         builder.structure = structure
         builder.scf = scf
-        builder.pp.code = pp_code  # pylint: disable=no-member
-        builder.pp.parameters = orm.Dict(
-            inputs.get("pp", {}).get("parameters")
+        builder.pp_valence.code = pp_code  # pylint: disable=no-member
+        builder.pp_valence.parameters = orm.Dict(
+            inputs.get("pp_valence", {}).get("parameters")
+        )  # pylint: disable=no-member
+        builder.pp_valence.metadata = metadata_pp_valence  # pylint: disable=no-member
+        #
+        builder.pp_all.code = pp_code  # pylint: disable=no-member
+        builder.pp_all.parameters = orm.Dict(
+            inputs.get("pp_all", {}).get("parameters")
         )  # pylint: disable=no-member
-        builder.pp.metadata = metadata_pp  # pylint: disable=no-member
+        builder.pp_all.metadata = metadata_pp_all  # pylint: disable=no-member
         builder.bader.code = bader_code  # pylint: disable=no-member
         builder.bader.parameters = orm.Dict(
             inputs.get("bader", {}).get("parameters")

diff --git a/examples/qe_bader_workchain.py b/examples/qe_bader_workchain.py
@@ -10,8 +10,8 @@
 load_profile()
 # ===============================================================================
 # load the codes
-pw_code = load_code("pw-7.2@localhost")
-pp_code = load_code("pp-7.2@localhost")
+pw_code = load_code("qe-7.2-pw@localhost")
+pp_code = load_code("qe-7.2-pp@localhost")
 bader_code = load_code("bader@localhost")
 
 
@@ -23,7 +23,7 @@
 
 
 overrides = {
-    "pp": {
+    "pp_valence": {
         "parameters": Dict(
             dict={
                 "INPUTPP": {"plot_num": 21},
@@ -40,6 +40,23 @@
             }
         },
     },
+    "pp_all": {
+        "parameters": Dict(
+            dict={
+                "INPUTPP": {"plot_num": 0},
+                "PLOT": {"iflag": 3},
+            }
+        ),
+        "metadata": {
+            "options": {
+                "resources": {
+                    "num_machines": 1,
+                    "num_mpiprocs_per_machine": 2,
+                },
+                "max_wallclock_seconds": 3600,
+            }
+        },
+    },
 }
 
 builder = QeBaderWorkChain.get_builder_from_protocol(