add reference charge density (#3)

aiida_bader/calculations/__init__.py

@@ -37,6 +37,19 @@ def define(cls, spec):
             required=True,
             help="Use a remote folder",
         )
+        spec.input(
+            "reference_charge_density_folder",
+            valid_type=RemoteData,
+            required=False,
+            help="reference_charge",
+        )
+        spec.input(
+            "reference_charge_density_filename",
+            valid_type=Str,
+            default=Str("aiida.fileout"),
+            required=False,
+            help="Name of the charge density file",
+        )
         spec.inputs["metadata"]["options"]["parser_name"].default = "bader"
         spec.inputs["metadata"]["options"]["resources"].default = {
             "num_machines": 1,
@@ -87,21 +100,34 @@ def prepare_for_submission(self, folder):
             charge_density_folder.get_remote_path(),
             self.inputs.charge_density_filename.value,
         )
-        copy_info = (comp_uuid, remote_path, "charge_density.cube")
-        if (
-            self.inputs.code.computer.uuid == comp_uuid
-        ):  # if running on the same computer - make a symlink
-            calcinfo.remote_symlink_list.append(copy_info)
-        else:  # if not - copy the folder
-            self.report(
-                f"Warning: Transferring cube file {charge_density_folder.get_remote_path()} from "
-                + f"computer {charge_density_folder.computer.label} to computer {self.inputs.code.computer.label}. "
-                + "This may put strain on your network."
+        copy_infos = [(comp_uuid, remote_path, "charge_density.cube")]
+        if self.inputs.reference_charge_density_folder:
+            reference_charge_density_folder = (
+                self.inputs.reference_charge_density_folder
+            )
+            comp_uuid = reference_charge_density_folder.computer.uuid
+            remote_path = os.path.join(
+                reference_charge_density_folder.get_remote_path(),
+                self.inputs.reference_charge_density_filename.value,
             )
-            calcinfo.remote_copy_list.append(copy_info)
+            copy_infos.append((comp_uuid, remote_path, "reference_charge_density.cube"))
+        for copy_info in copy_infos:
+            if (
+                self.inputs.code.computer.uuid == copy_info[0]
+            ):  # if running on the same computer - make a symlink
+                calcinfo.remote_symlink_list.append(copy_info)
+            else:  # if not - copy the folder
+                self.report(
+                    f"Warning: Transferring cube file {charge_density_folder.get_remote_path()} from "
+                    + f"computer {charge_density_folder.computer.label} to computer {self.inputs.code.computer.label}. "
+                    + "This may put strain on your network."
+                )
+                calcinfo.remote_copy_list.append(copy_info)
 
         codeinfo = CodeInfo()
         codeinfo.cmdline_params = ["charge_density.cube"]
+        if self.inputs.reference_charge_density_folder:
+            codeinfo.cmdline_params.extend(["-ref", "reference_charge_density.cube"])
         codeinfo.code_uuid = self.inputs.code.uuid
         calcinfo.codes_info = [codeinfo]
 

aiida_bader/workchains/qe_bader.py

@@ -38,21 +38,39 @@ def define(cls, spec):
                 "populate_defaults": False,
             },
         )
-        spec.expose_inputs(PpCalculation, namespace="pp", exclude=["parent_folder"])
         spec.expose_inputs(
-            BaderCalculation, namespace="bader", exclude=["charge_density_folder"]
+            PpCalculation, namespace="pp_valence", exclude=["parent_folder"]
+        )
+        spec.expose_inputs(PpCalculation, namespace="pp_all", exclude=["parent_folder"])
+        spec.expose_inputs(
+            BaderCalculation,
+            namespace="bader",
+            exclude=["charge_density_folder", "reference_charge_density_folder"],
         )
 
-        spec.outline(cls.run_pw, cls.run_pp, cls.run_bader, cls.return_results)
+        spec.outline(
+            cls.run_pw,
+            cls.run_pp,
+            cls.run_bader,
+            cls.return_results,
+        )
 
         spec.expose_outputs(PwBaseWorkChain, namespace="scf")
-        spec.expose_outputs(PpCalculation, namespace="pp")
+        spec.expose_outputs(PpCalculation, namespace="pp_valence")
+        spec.expose_outputs(PpCalculation, namespace="pp_all")
         spec.expose_outputs(BaderCalculation, namespace="bader")
 
         spec.exit_code(903, "ERROR_PARSING_PW_OUTPUT", "Error while parsing PW output")
-        spec.exit_code(904, "ERROR_PARSING_PP_OUTPUT", "Error while parsing PP output")
         spec.exit_code(
-            905, "ERROR_PARSING_BADER_OUTPUT", "Error while parsing bader output"
+            904,
+            "ERROR_PARSING_PP_VALENCE_OUTPUT",
+            "Error while parsing PP_VALENCE output",
+        )
+        spec.exit_code(
+            905, "ERROR_PARSING_PP_ALL_OUTPUT", "Error while parsing PP_ALL output"
+        )
+        spec.exit_code(
+            906, "ERROR_PARSING_BADER_OUTPUT", "Error while parsing bader output"
         )
 
     @classmethod
@@ -102,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")
@@ -142,28 +174,40 @@ def run_pp(self):
         # TODO extract number of core electrons from the pw pseudopotential
 
         try:
-            pp_inputs = AttributeDict(self.exposed_inputs(PpCalculation, "pp"))
-            pp_inputs["parent_folder"] = self.ctx.pw_calc.outputs.remote_folder
+            pp_valence_inputs = AttributeDict(
+                self.exposed_inputs(PpCalculation, "pp_valence")
+            )
+            pp_valence_inputs["parent_folder"] = self.ctx.pw_calc.outputs.remote_folder
+            pp_all_inputs = AttributeDict(self.exposed_inputs(PpCalculation, "pp_all"))
+            pp_all_inputs["parent_folder"] = self.ctx.pw_calc.outputs.remote_folder
         except Exception as exc:  # pylint: disable=broad-except
             self.report(f'Encountered exception "{str(exc)}" while parsing PW output')
             return self.exit_codes.ERROR_PARSING_PW_OUTPUT  # pylint: disable=no-member
 
-        pp_inputs["metadata"]["call_link_label"] = "call_pp_calc"
+        pp_valence_inputs["metadata"]["call_link_label"] = "call_pp_valence_calc"
+        pp_all_inputs["metadata"]["call_link_label"] = "call_pp_all_calc"
 
         # Create the calculation process and launch it
-        running = self.submit(PpCalculation, **pp_inputs)
+        pp_valence_running = self.submit(PpCalculation, **pp_valence_inputs)
+        pp_all_running = self.submit(PpCalculation, **pp_valence_inputs)
+        self.report(
+            f"Running PpCalculation<{pp_valence_running.pk}> to compute the valence charge-density"
+        )
         self.report(
-            f"Running PpCalculation<{running.pk}> to compute the charge-density"
+            f"Running PpCalculation<{pp_all_running.pk}> to compute the all-electron charge-density"
         )
-        return ToContext(pp_calc=running)
+        return ToContext(pp_valence_calc=pp_valence_running, pp_all_calc=pp_all_running)
 
     def run_bader(self):
         """Parse the PP ouputs cube file, and submit bader calculation."""
         try:
             bader_inputs = AttributeDict(self.exposed_inputs(BaderCalculation, "bader"))
             bader_inputs[
                 "charge_density_folder"
-            ] = self.ctx.pp_calc.outputs.remote_folder
+            ] = self.ctx.pp_valence_calc.outputs.remote_folder
+            bader_inputs[
+                "reference_charge_density_folder"
+            ] = self.ctx.pp_all_calc.outputs.remote_folder
         except Exception as exc:  # pylint: disable=broad-except
             self.report(f'Encountered exception "{str(exc)}" while parsing PP output')
             return self.exit_codes.ERROR_PARSING_PP_OUTPUT  # pylint: disable=no-member
@@ -184,7 +228,14 @@ def return_results(self):
                 self.exposed_outputs(self.ctx.pw_calc, PwBaseWorkChain, namespace="scf")
             )
             self.out_many(
-                self.exposed_outputs(self.ctx.pp_calc, PpCalculation, namespace="pp")
+                self.exposed_outputs(
+                    self.ctx.pp_valence_calc, PpCalculation, namespace="pp_valence"
+                )
+            )
+            self.out_many(
+                self.exposed_outputs(
+                    self.ctx.pp_all_calc, PpCalculation, namespace="pp_all"
+                )
             )
             self.out_many(
                 self.exposed_outputs(

aiida_bader/workgraph/cp2k_bader.py

@@ -9,7 +9,10 @@ def Cp2kBaderWorkGraph():
     from aiida_bader.calculations import BaderCalculation
 
     wt = WorkGraph("charge-density")
-    cp2k_node = wt.nodes.new(Cp2kBaseWorkChain, name="cp2k_base")
-    wt.nodes.new(BaderCalculation, name="bader",
-                              charge_density_folder = cp2k_node.outputs["remote_folder"])
+    cp2k_node = wt.nodes.new(Cp2kBaseWorkChain, name="scf")
+    wt.nodes.new(
+        BaderCalculation,
+        name="bader",
+        charge_density_folder=cp2k_node.outputs["remote_folder"],
+    )
     return wt

aiida_bader/workgraph/qe_bader.py

@@ -10,9 +10,17 @@ def QeBaderWorkGraph():
     from aiida_bader.calculations import BaderCalculation
 
     wg = WorkGraph("charge-density")
-    pw_node = wg.nodes.new(PwBaseWorkChain, name="pw_base")
-    pp_node = wg.nodes.new(PpCalculation, name="pp", 
-                           parent_folder=pw_node.outputs["remote_folder"])
-    wg.nodes.new(BaderCalculation, name="bader",
-                                charge_density_folder=pp_node.outputs["remote_folder"])
+    pw_node = wg.nodes.new(PwBaseWorkChain, name="scf")
+    pp_valence = wg.nodes.new(
+        PpCalculation, name="pp_valence", parent_folder=pw_node.outputs["remote_folder"]
+    )
+    pp_all = wg.nodes.new(
+        PpCalculation, name="pp_all", parent_folder=pw_node.outputs["remote_folder"]
+    )
+    wg.nodes.new(
+        BaderCalculation,
+        name="bader",
+        charge_density_folder=pp_valence.outputs["remote_folder"],
+        reference_charge_density_folder=pp_all.outputs["remote_folder"],
+    )
     return wg

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(