Use plugin utility to set default starting magnetization (#1041)

This PR uses the QE plugin `get_starting_magenetization` utility function to set default initial magnetic moments for the structure when magnetism is turned on. The defaults are set to the default `magmom` provided by the utility scaled by the Z valence value for the element defined in the selected pseudopotential. Note that for those elements of a zero `magmom` default, a general default magnetization of 0.1 is returned.

If multiple tags are provided for the same species, a warning is displayed providing some instructions for antiferromagnetic configurations.

src/aiidalab_qe/app/configuration/advanced/advanced.py

@@ -50,6 +50,11 @@ def __init__(self, model: AdvancedConfigurationSettingsModel, **kwargs):
             "kpoints_distance",
         )
 
+        # NOTE connect pseudos first, as some settings depend on it
+        pseudos_model = PseudosConfigurationSettingsModel()
+        self.pseudos = PseudosConfigurationSettingsPanel(model=pseudos_model)
+        model.add_model("pseudos", pseudos_model)
+
         smearing_model = SmearingConfigurationSettingsModel()
         self.smearing = SmearingConfigurationSettingsPanel(model=smearing_model)
         model.add_model("smearing", smearing_model)
@@ -64,10 +69,6 @@ def __init__(self, model: AdvancedConfigurationSettingsModel, **kwargs):
         self.hubbard = HubbardConfigurationSettingsPanel(model=hubbard_model)
         model.add_model("hubbard", hubbard_model)
 
-        pseudos_model = PseudosConfigurationSettingsModel()
-        self.pseudos = PseudosConfigurationSettingsPanel(model=pseudos_model)
-        model.add_model("pseudos", pseudos_model)
-
     def render(self):
         if self.rendered:
             return

src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py

@@ -8,10 +8,10 @@ class MagnetizationConfigurationSettingsPanel(
     AdvancedConfigurationSubSettingsPanel[MagnetizationConfigurationSettingsModel],
 ):
     """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.
+    1) Total magnetization: Total majority spin charge - minority spin charge.
+    2) Magnetic moments: 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 (StructureData.get_kind_names())
+    For Magnetic moments you can set each kind names defined in the StructureData (StructureData.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
@@ -39,18 +39,34 @@ def __init__(self, model: MagnetizationConfigurationSettingsModel, **kwargs):
             self._on_magnetization_type_change,
             "type",
         )
+        self._model.observe(
+            self._on_family_change,
+            "family",
+        )
 
     def render(self):
         if self.rendered:
             return
 
-        self.description = ipw.HTML("<b>Magnetization:</b>")
+        self.header = ipw.HTML("<b>Magnetization:</b>")
+
+        self.unit = ipw.HTML(
+            value="µ<sub>B</sub>",
+            layout=ipw.Layout(margin="2px 2px 5px"),
+        )
+
+        self.magnetization_type_help = ipw.HTML()
+        ipw.dlink(
+            (self._model, "type_help"),
+            (self.magnetization_type_help, "value"),
+        )
 
         self.magnetization_type = ipw.ToggleButtons(
             style={
                 "description_width": "initial",
                 "button_width": "initial",
             },
+            layout=ipw.Layout(margin="0 0 10px 0"),
         )
         ipw.dlink(
             (self._model, "type_options"),
@@ -73,11 +89,19 @@ def render(self):
             (self.tot_magnetization, "value"),
         )
 
+        self.tot_magnetization_with_unit = ipw.HBox(
+            children=[
+                self.tot_magnetization,
+                self.unit,
+            ],
+            layout=ipw.Layout(align_items="center"),
+        )
+
         self.kind_moment_widgets = ipw.VBox()
 
         self.container = ipw.VBox(
             children=[
-                self.tot_magnetization,
+                self.tot_magnetization_with_unit,
             ]
         )
 
@@ -92,12 +116,17 @@ def _on_input_structure_change(self, _):
 
     def _on_electronic_type_change(self, _):
         self._switch_widgets()
+        self._model.update_type_help()
 
     def _on_spin_type_change(self, _):
         self.refresh(specific="spin")
 
     def _on_magnetization_type_change(self, _):
         self._toggle_widgets()
+        self._model.update_type_help()
+
+    def _on_family_change(self, _):
+        self._model.update_default_starting_magnetization()
 
     def _update(self, specific=""):
         if self.updated:
@@ -129,8 +158,8 @@ def _build_kinds_widget(self):
         for kind_name in kind_names:
             kind_moment_widget = ipw.BoundedFloatText(
                 description=kind_name,
-                min=-4,
-                max=4,
+                min=-7,
+                max=7,
                 step=0.1,
             )
             link = ipw.link(
@@ -145,7 +174,15 @@ def _build_kinds_widget(self):
                 ],
             )
             self.links.append(link)
-            children.append(kind_moment_widget)
+            children.append(
+                ipw.HBox(
+                    children=[
+                        kind_moment_widget,
+                        self.unit,
+                    ],
+                    layout=ipw.Layout(align_items="center"),
+                )
+            )
 
         self.kind_moment_widgets.children = children
 
@@ -155,23 +192,29 @@ def _switch_widgets(self):
         if self._model.spin_type == "none":
             children = []
         else:
-            children = [self.description]
+            children = [self.header]
             if self._model.electronic_type == "metal":
                 children.extend(
                     [
                         self.magnetization_type,
+                        self.magnetization_type_help,
                         self.container,
                     ]
                 )
             else:
-                children.append(self.tot_magnetization)
+                children.extend(
+                    [
+                        self.magnetization_type_help,
+                        self.tot_magnetization_with_unit,
+                    ],
+                )
         self.children = children
 
     def _toggle_widgets(self):
         if self._model.spin_type == "none" or not self.rendered:
             return
         self.container.children = [
-            self.tot_magnetization
+            self.tot_magnetization_with_unit
             if self._model.type == "tot_magnetization"
             else self.kind_moment_widgets
         ]

src/aiidalab_qe/app/configuration/advanced/magnetization/model.py

@@ -1,8 +1,14 @@
 from copy import deepcopy
 
 import traitlets as tl
+from aiida_pseudo.groups.family import PseudoPotentialFamily
 
+from aiida_quantumespresso.workflows.protocols.utils import (
+    get_magnetization_parameters,
+    get_starting_magnetization,
+)
 from aiidalab_qe.common.mixins import HasInputStructure
+from aiidalab_qe.utils import fetch_pseudo_family_by_label
 
 from ..subsettings import AdvancedCalculationSubSettingsModel
 
@@ -17,41 +23,109 @@ class MagnetizationConfigurationSettingsModel(
         "input_structure",
         "electronic_type",
         "spin_type",
+        "pseudos.family",
     ]
 
     electronic_type = tl.Unicode()
     spin_type = tl.Unicode()
+    family = tl.Unicode()
 
     type_options = tl.List(
         trait=tl.List(tl.Unicode()),
         default_value=[
-            ["Starting magnetization", "starting_magnetization"],
+            ["Initial magnetic moments", "starting_magnetization"],
             ["Total magnetization", "tot_magnetization"],
         ],
     )
     type = tl.Unicode("starting_magnetization")
+    type_help = tl.Unicode("")
     total = tl.Float(0.0)
     moments = tl.Dict(
         key_trait=tl.Unicode(),  # kind name
         value_trait=tl.Float(),  # magnetic moment
         default_value={},
     )
 
+    _TYPE_HELP_TEMPLATE = """
+        <div style="line-height: 1.4; margin-bottom: 5px">
+            {content}
+        </div>
+    """
+
+    _DEFAULT_MOMENTS = get_magnetization_parameters()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._default_starting_magnetization = {}
+
     def update(self, specific=""):  # noqa: ARG002
         if self.spin_type == "none" or not self.has_structure:
             self._defaults["moments"] = {}
         else:
-            self._defaults["moments"] = {
-                kind_name: 0.0 for kind_name in self.input_structure.get_kind_names()
-            }
+            self.update_type_help()
+            self.update_default_starting_magnetization()
+            self._update_default_moments()
         with self.hold_trait_notifications():
             self.moments = self._get_default_moments()
 
+    def update_type_help(self):
+        """Update the type field help text w.r.t the current model state."""
+        if self.electronic_type == "insulator" or self.type == "tot_magnetization":
+            self.type_help = self._TYPE_HELP_TEMPLATE.format(
+                content="""
+                    Constrain the desired total electronic magnetization (difference
+                    between majority and minority spin charge).
+                """
+            )
+        else:
+            self.type_help = self._TYPE_HELP_TEMPLATE.format(
+                content="""
+                    If a nonzero ground-state magnetization is expected, you
+                    <strong>must</strong> assign a nonzero value to at least one atomic
+                    type (note that the app already provide tentative initial values).
+                    To simulate an antiferromagnetic state, first, if you have not
+                    done so already, please use the atom tag editor <b>(Select
+                    structure -> Edit structure -> Edit atom tags)</b> to mark atoms of
+                    the species of interest as distinct by assigning each a different
+                    integer tag. Once tagged, assign each an initial magnetic moments
+                    of opposite sign.
+                """
+            )
+
+    def update_default_starting_magnetization(self):
+        """Update the default starting magnetization based on the structure and
+        pseudopotential family.
+        """
+        if not self.has_structure:
+            # TODO this guard shouldn't be here! It IS here only because in the present
+            # implementation, an update is called on app start. This breaks lazy loading
+            # and should be carefully checked!
+            return
+        family = fetch_pseudo_family_by_label(self.family)
+        initial_guess = get_starting_magnetization(self.input_structure, family)
+        self._default_starting_magnetization = initial_guess
+
     def reset(self):
         with self.hold_trait_notifications():
             self.type = self.traits()["type"].default_value
             self.total = self.traits()["total"].default_value
             self.moments = self._get_default_moments()
 
+    def _update_default_moments(self):
+        family = fetch_pseudo_family_by_label(self.family)
+        self._defaults["moments"] = {
+            kind.name: self._to_moment(symbol=kind.symbol, family=family)
+            for kind in self.input_structure.kinds
+        }
+
+    def _to_moment(self, symbol: str, family: PseudoPotentialFamily) -> float:
+        """Convert the default magnetization to an initial magnetic moment."""
+        magnetization = (
+            self._default_starting_magnetization.get(symbol, 0.1)
+            if self._DEFAULT_MOMENTS.get(symbol, {}).get("magmom")
+            else 0.1
+        )
+        return round(magnetization * family.get_pseudo(symbol).z_valence, 3)
+
     def _get_default_moments(self):
         return deepcopy(self._defaults.get("moments", {}))

src/aiidalab_qe/app/configuration/advanced/model.py

@@ -255,11 +255,7 @@ def _link_model(self, model: AdvancedCalculationSubSettingsModel):
             self._on_any_change,
             tl.All,
         )
-        for trait in model.dependencies:
-            ipw.dlink(
-                (self, trait),
-                (model, trait),
-            )
+        super()._link_model(model)
 
     def _update_kpoints_mesh(self, _=None):
         if not self.has_structure:

src/aiidalab_qe/app/configuration/advanced/pseudos/model.py

@@ -5,13 +5,13 @@
 import traitlets as tl
 from aiida_pseudo.common.units import U
 
-from aiida import orm
 from aiida.common import exceptions
 from aiida.plugins import GroupFactory
 from aiida_quantumespresso.workflows.pw.base import PwBaseWorkChain
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
 from aiidalab_qe.common.mixins import HasInputStructure
 from aiidalab_qe.setup.pseudos import PSEUDODOJO_VERSION, SSSP_VERSION, PseudoFamily
+from aiidalab_qe.utils import fetch_pseudo_family_by_label
 
 from ..subsettings import AdvancedCalculationSubSettingsModel
 
@@ -132,7 +132,7 @@ def update_default_pseudos(self):
         self.status_message = ""
 
         try:
-            pseudo_family = self._get_pseudo_family_from_database()
+            pseudo_family = fetch_pseudo_family_by_label(self.family)
             pseudos = pseudo_family.get_pseudos(structure=self.input_structure)
         except ValueError as exception:
             self.status_message = f"""
@@ -156,7 +156,7 @@ def update_default_cutoffs(self):
         self.status_message = ""
 
         try:
-            pseudo_family = self._get_pseudo_family_from_database()
+            pseudo_family = fetch_pseudo_family_by_label(self.family)
             current_unit = pseudo_family.get_cutoffs_unit()
             cutoff_dict = pseudo_family.get_cutoffs()
         except exceptions.NotExistent:
@@ -301,21 +301,6 @@ def _get_default(self, trait):
             )
         return self._defaults.get(trait, self.traits()[trait].default_value)
 
-    def _get_pseudo_family_from_database(self):
-        """Get the pseudo family from the database."""
-        return (
-            orm.QueryBuilder()
-            .append(
-                (
-                    PseudoDojoFamily,
-                    SsspFamily,
-                    CutoffsPseudoPotentialFamily,
-                ),
-                filters={"label": self.family},
-            )
-            .one()[0]
-        )
-
     def _get_default_dictionary(self):
         return deepcopy(self._defaults["dictionary"])