add reference charge density

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
@@ -142,28 +160,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<{running.pk}> to compute the charge-density"
+            f"Running PpCalculation<{pp_valence_running.pk}> to compute the valence charge-density"
         )
-        return ToContext(pp_calc=running)
+        self.report(
+            f"Running PpCalculation<{pp_all_running.pk}> to compute the all-electron charge-density"
+        )
+        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 +214,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