Fix tot magnetization (#512)

* updating Magnetization widget, and including logic for insulator
* adapting logic of bands and pdos for bandspdoswidget
* starting_magnetization as default

src/aiidalab_qe/app/configuration/advanced.py

@@ -211,6 +211,11 @@ def _update_input_structure(self, change):
             self.magnetization.input_structure = None
             self.pseudo_setter.structure = None
 
+    @tl.observe("electronic_type")
+    def _electronic_type_changed(self, change):
+        """Input electronic_type changed, update the widget values."""
+        self.magnetization.electronic_type = change["new"]
+
     @tl.observe("protocol")
     def _protocol_changed(self, _):
         """Input protocol changed, update the widget values."""
@@ -255,17 +260,16 @@ def get_panel_value(self):
         # XXX: start from parameters = {} and then bundle the settings by purposes (e.g. pw, bands, etc.)
         parameters = {
             "initial_magnetic_moments": None,
-            "pw": {
-                "parameters": {
-                    "SYSTEM": {},
-                },
-            },
+            "pw": {"parameters": {"SYSTEM": {}}},
+            "clean_workdir": self.clean_workdir.value,
+            "pseudo_family": self.pseudo_family_selector.value,
+            "kpoints_distance": self.value.get("kpoints_distance"),
         }
-        # add clean_workdir to the parameters
-        parameters["clean_workdir"] = self.clean_workdir.value
 
-        # add the pseudo_family to the parameters
-        parameters["pseudo_family"] = self.pseudo_family_selector.value
+        # Set total charge
+        parameters["pw"]["parameters"]["SYSTEM"]["tot_charge"] = self.total_charge.value
+
+        # Set the pseudos
         if self.pseudo_setter.pseudos:
             parameters["pw"]["pseudos"] = self.pseudo_setter.pseudos
             parameters["pw"]["parameters"]["SYSTEM"]["ecutwfc"] = (
@@ -303,8 +307,35 @@ def get_panel_value(self):
                 self.smearing.degauss_value
             )
 
+        # Set tot_magnetization for collinear simulations.
+        if self.spin_type == "collinear":
+            # Conditions for metallic systems. Select the magnetization type and set the value if override is True
+            if self.electronic_type == "metal" and self.override.value is True:
+                self.set_metallic_magnetization(parameters)
+            # Conditions for insulator systems. Default value is 0.0
+            elif self.electronic_type == "insulator":
+                self.set_insulator_magnetization(parameters)
+
         return parameters
 
+    def set_insulator_magnetization(self, parameters):
+        """Set the parameters for collinear insulator calculation. Total magnetization."""
+        parameters["pw"]["parameters"]["SYSTEM"]["tot_magnetization"] = (
+            self.magnetization.tot_magnetization.value
+        )
+
+    def set_metallic_magnetization(self, parameters):
+        """Set the parameters for magnetization calculation in metals"""
+        magnetization_type = self.magnetization.magnetization_type.value
+        if magnetization_type == "tot_magnetization":
+            parameters["pw"]["parameters"]["SYSTEM"]["tot_magnetization"] = (
+                self.magnetization.tot_magnetization.value
+            )
+        else:
+            parameters["initial_magnetic_moments"] = (
+                self.magnetization.get_magnetization()
+            )
+
     def set_panel_value(self, parameters):
         """Set the panel value from the given parameters."""
 
@@ -338,11 +369,18 @@ def set_panel_value(self, parameters):
                 parameters["pw"]["parameters"]["SYSTEM"].get("vdw_corr", "none"),
             )
 
+        # Logic to set the magnetization
         if parameters.get("initial_magnetic_moments"):
             self.magnetization._set_magnetization_values(
                 parameters.get("initial_magnetic_moments")
             )
 
+        if "tot_magnetization" in parameters["pw"]["parameters"]["SYSTEM"]:
+            self.magnetization.magnetization_type.value = "tot_magnetization"
+            self.magnetization._set_tot_magnetization(
+                parameters["pw"]["parameters"]["SYSTEM"]["tot_magnetization"]
+            )
+
     def reset(self):
         """Reset the widget and the traitlets"""
 
@@ -363,7 +401,11 @@ def reset(self):
             self.override.value = False
             self.smearing.reset()
             # reset the pseudo setter
-            self.pseudo_setter._reset()
+            if self.input_structure is None:
+                self.pseudo_setter.structure = None
+                self.pseudo_setter._reset()
+            else:
+                self.pseudo_setter._reset()
             # reset the magnetization
             self.magnetization.reset()
             # reset mesh grid
@@ -387,61 +429,79 @@ def _display_mesh(self, _=None):
 
 
 class MagnetizationSettings(ipw.VBox):
-    """Widget to set the initial magnetic moments for each kind names defined in the StructureData (StructureDtaa.get_kind_names())
+    """Widget to set the type of magnetization used in the calculation:
+    1) Tot_magnetization: Total majority spin charge - minority spin charge.
+    2) Starting magnetization: Starting spin polarization on atomic type 'i' in a spin polarized (LSDA or noncollinear/spin-orbit) calculation.
+
+    For Starting magnetization you can set each kind names defined in the StructureData (StructureDtaa.get_kind_names())
     Usually these are the names of the elements in the StructureData
     (For example 'C' , 'N' , 'Fe' . However the StructureData can have defined kinds like 'Fe1' and 'Fe2')
-
     The widget generate a dictionary that can be used to set initial_magnetic_moments in the builder of PwBaseWorkChain
 
     Attributes:
         input_structure(StructureData): trait that containes the input_strucgure (confirmed structure from previous step)
     """
 
     input_structure = tl.Instance(orm.StructureData, allow_none=True)
-
+    electronic_type = tl.Unicode()
     disabled = tl.Bool()
+    _DEFAULT_TOT_MAGNETIZATION = 0.0
+    _DEFAULT_DESCRIPTION = "<b>Magnetization: Input structure not confirmed</b>"
 
     def __init__(self, **kwargs):
         self.input_structure = orm.StructureData()
         self.input_structure_labels = []
-        self.description = ipw.HTML(
-            "Define magnetization: Input structure not confirmed"
+        self.tot_magnetization = ipw.BoundedIntText(
+            min=0,
+            max=100,
+            step=1,
+            value=self._DEFAULT_TOT_MAGNETIZATION,
+            disabled=True,
+            description="Total magnetization:",
+            style={"description_width": "initial"},
+        )
+        self.magnetization_type = ipw.ToggleButtons(
+            options=[
+                ("Starting Magnetization", "starting_magnetization"),
+                ("Tot. Magnetization", "tot_magnetization"),
+            ],
+            value="starting_magnetization",
+            style={"description_width": "initial"},
         )
+        self.description = ipw.HTML(self._DEFAULT_DESCRIPTION)
         self.kinds = self.create_kinds_widget()
         self.kinds_widget_out = ipw.Output()
+        self.magnetization_out = ipw.Output()
+        self.magnetization_type.observe(self._render, "value")
         super().__init__(
             children=[
-                ipw.HBox(
-                    [
-                        self.description,
-                        self.kinds_widget_out,
-                    ],
-                ),
+                self.description,
+                self.magnetization_out,
+                self.kinds_widget_out,
             ],
             layout=ipw.Layout(justify_content="space-between"),
             **kwargs,
         )
-        self.display_kinds()
 
     @tl.observe("disabled")
     def _disabled_changed(self, _):
         """Disable the widget"""
         if hasattr(self.kinds, "children") and self.kinds.children:
             for i in range(len(self.kinds.children)):
                 self.kinds.children[i].disabled = self.disabled
+        self.tot_magnetization.disabled = self.disabled
+        self.magnetization_type.disabled = self.disabled
 
     def reset(self):
         self.disabled = True
+        self.tot_magnetization.value = self._DEFAULT_TOT_MAGNETIZATION
+        #
         if self.input_structure is None:
-            self.description.value = (
-                "Define magnetization: Input structure not confirmed"
-            )
+            self.description.value = self._DEFAULT_DESCRIPTION
             self.kinds = None
-            with self.kinds_widget_out:
-                clear_output()
-
         else:
-            self.update_kinds_widget()
+            self.description.value = "<b>Magnetization</b>"
+            self.kinds = self.create_kinds_widget()
 
     def create_kinds_widget(self):
         if self.input_structure_labels:
@@ -462,11 +522,30 @@ def create_kinds_widget(self):
 
         return kinds_widget
 
+    @tl.observe("electronic_type")
+    def _electronic_type_changed(self, change):
+        with self.magnetization_out:
+            clear_output()
+            if change["new"] == "metal":
+                display(self.magnetization_type)
+                self._render({"new": self.magnetization_type.value})
+            else:
+                display(self.tot_magnetization)
+                with self.kinds_widget_out:
+                    clear_output()
+
     def update_kinds_widget(self):
         self.input_structure_labels = self.input_structure.get_kind_names()
         self.kinds = self.create_kinds_widget()
-        self.description.value = "Define magnetization: "
-        self.display_kinds()
+        self.description.value = "<b>Magnetization</b>"
+
+    def _render(self, value):
+        if value["new"] == "tot_magnetization":
+            with self.kinds_widget_out:
+                clear_output()
+                display(self.tot_magnetization)
+        else:
+            self.display_kinds()
 
     def display_kinds(self):
         if "PYTEST_CURRENT_TEST" not in os.environ and self.kinds:
@@ -477,6 +556,7 @@ def display_kinds(self):
     def _update_widget(self, change):
         self.input_structure = change["new"]
         self.update_kinds_widget()
+        self.display_kinds()
 
     def get_magnetization(self):
         """Method to generate the dictionary with the initial magnetic moments"""
@@ -497,6 +577,10 @@ def _set_magnetization_values(self, magnetic_moments):
                 else:
                     self.kinds.children[i].value = magnetic_moments
 
+    def _set_tot_magnetization(self, tot_magnetization):
+        """Set the total magnetization"""
+        self.tot_magnetization.value = tot_magnetization
+
 
 class SmearingSettings(ipw.VBox):
     # accept protocol as input and set the values

src/aiidalab_qe/app/result/summary_viewer.py

@@ -123,6 +123,10 @@ def generate_report_parameters(qeapp_wc):
     report["periodicity"] = PERIODICITY_MAPPING.get(
         qeapp_wc.inputs.structure.pbc, "xyz"
     )
+    report["tot_magnetization"] = pw_parameters["SYSTEM"].get(
+        "tot_magnetization", False
+    )
+
     # hard code bands and pdos
     if "bands" in qeapp_wc.inputs:
         report["bands_kpoints_distance"] = PwBandsWorkChain.get_protocol_inputs(

src/aiidalab_qe/app/static/workflow_summary.jinja

@@ -104,10 +104,18 @@
         <td>Van der Waals Correction</td>
         <td>{{ vdw_corr }}</td>
       </tr>
+
+      {% if tot_magnetization %}
+      <tr>
+        <td>Total magnetization</td>
+        <td>{{ tot_magnetization }} </td>
+      </tr>
+      {% else %}
       <tr>
         <td>Initial Magnetic Moments</td>
         <td>{{ initial_magnetic_moments }}</td>
       </tr>
+      {% endif %}
     </table>
   </div>
 </div>