Use plugin utility to set default starting magnetization

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
@@ -44,13 +44,20 @@ def render(self):
         if self.rendered:
             return
 
-        self.description = ipw.HTML("<b>Magnetization:</b>")
+        self.description = ipw.HTML("""
+            <div style="margin-bottom: 5px;">
+                <b>Magnetization:</b>
+                <br>
+                Default magnetic moments correspond to theoretical values.
+            </div>
+        """)
 
         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"),

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

@@ -2,7 +2,10 @@
 
 import traitlets as tl
 
+from aiida.common.exceptions import NotExistent
+from aiida_quantumespresso.workflows.protocols.utils import 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,15 +20,17 @@ 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"],
+            ["Magnetic moments", "starting_magnetization"],
             ["Total magnetization", "tot_magnetization"],
         ],
     )
@@ -41,9 +46,22 @@ 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()
-            }
+            try:
+                family = fetch_pseudo_family_by_label(self.family)
+                initial_guess = get_starting_magnetization(self.input_structure, family)
+                self._defaults["moments"] = {
+                    kind.name: round(
+                        initial_guess[kind.name]
+                        * family.get_pseudo(kind.symbol).z_valence,
+                        3,
+                    )
+                    for kind in self.input_structure.kinds
+                }
+            except NotExistent:
+                self._defaults["moments"] = {
+                    kind_name: 0.0
+                    for kind_name in self.input_structure.get_kind_names()
+                }
         with self.hold_trait_notifications():
             self.moments = self._get_default_moments()
 

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

@@ -250,11 +250,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"])
 

src/aiidalab_qe/app/configuration/basic/basic.py

@@ -45,17 +45,18 @@ def render(self):
             (self._model, "spin_type"),
             (self.spin_type, "value"),
         )
+        self.spin_type.observe(
+            self._on_spin_type_change,
+            "value",
+        )
+
         self.magnetization_info = ipw.HTML(
             value="""
                 <div style="margin-left: 10px;">
                     Set the desired magnetic configuration in <b>advanced</b> settings
                 </div>
             """,
-            layout=ipw.Layout(visibility="hidden"),
-        )
-        self.spin_type.observe(
-            self._on_spin_type_change,
-            "value",
+            layout=ipw.Layout(display="none"),
         )
 
         # Spin-Orbit calculation
@@ -80,6 +81,29 @@ def render(self):
             (self.protocol, "value"),
         )
 
+        self.warning = ipw.HTML(
+            value="""
+                <div
+                    class="alert alert-warning"
+                    style="line-height: 140%; margin: 10px 0 0"
+                >
+                    <p>
+                        <b>Warning:</b> detected multiples atoms with different tags.
+                        You may be interested in an antiferromagnetic system. Note that
+                        default starting magnetic moments do not distinguish tagged
+                        atoms and are set to the same value.
+                    </p>
+                    <p>
+                        Please go to <b>Advanced settings</b> and override the default
+                        values, specifying appropriate magnetic moments for each
+                        species (e.g. with different signs for an antiferromagnetic
+                        configuration).
+                    </p>
+                </div>
+            """,
+            layout=ipw.Layout(display="none"),
+        )
+
         self.children = [
             InAppGuide(identifier="basic-settings"),
             ipw.HTML("""
@@ -153,6 +177,7 @@ def render(self):
                     (at the price of longer/costlier calculations).
                 </div>
             """),
+            self.warning,
         ]
 
         self.rendered = True
@@ -161,7 +186,10 @@ def _on_input_structure_change(self, _):
         self.refresh(specific="structure")
 
     def _on_spin_type_change(self, _):
-        if self.spin_type.value == "none":
-            self.magnetization_info.layout.visibility = "hidden"
+        if self._model.spin_type == "collinear":
+            self.magnetization_info.layout.display = "block"
+            if self._model.has_tags:
+                self.warning.layout.display = "flex"
         else:
-            self.magnetization_info.layout.visibility = "visible"
+            self.magnetization_info.layout.display = "none"
+            self.warning.layout.display = "none"

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

@@ -1,22 +1,23 @@
 import traitlets as tl
 
-from aiida import orm
 from aiidalab_qe.app.parameters import DEFAULT_PARAMETERS
+from aiidalab_qe.common.mixins import HasInputStructure
 from aiidalab_qe.common.panel import ConfigurationSettingsModel
 
 DEFAULT: dict = DEFAULT_PARAMETERS  # type: ignore
 
 
-class BasicConfigurationSettingsModel(ConfigurationSettingsModel):
+class BasicConfigurationSettingsModel(
+    ConfigurationSettingsModel,
+    HasInputStructure,
+):
     title = "Basic settings"
     identifier = "workchain"
 
     dependencies = [
         "input_structure",
     ]
 
-    input_structure = tl.Union([tl.Instance(orm.StructureData)], allow_none=True)
-
     protocol_options = tl.List(
         trait=tl.Tuple(tl.Unicode(), tl.Unicode()),
         default_value=[

src/aiidalab_qe/app/configuration/model.py

@@ -129,18 +129,7 @@ def _link_model(self, model: ConfigurationSettingsModel):
             (self, "confirmed"),
             (model, "confirmed"),
         )
-        for dependency in model.dependencies:
-            dependency_parts = dependency.split(".")
-            if len(dependency_parts) == 1:  # from parent, e.g. input_structure
-                target_model = self
-                trait = dependency
-            else:  # from sibling, e.g. workchain.protocol
-                sibling, trait = dependency_parts
-                target_model = self.get_model(sibling)
-            ipw.dlink(
-                (target_model, trait),
-                (model, trait),
-            )
+        super()._link_model(model)
 
     def _get_properties(self):
         properties = []

src/aiidalab_qe/common/mixins.py

@@ -26,6 +26,13 @@ def has_structure(self):
     def has_pbc(self):
         return not self.has_structure or any(self.input_structure.pbc)
 
+    @property
+    def has_tags(self):
+        return any(
+            not kind_name.isalpha()
+            for kind_name in self.input_structure.get_kind_names()
+        )
+
 
 class HasModels(t.Generic[T]):
     def __init__(self):
@@ -51,7 +58,20 @@ def get_models(self) -> t.Iterable[tuple[str, T]]:
         return self._models.items()
 
     def _link_model(self, model: T):
-        pass
+        if not hasattr(model, "dependencies"):
+            return
+        for dependency in model.dependencies:
+            dependency_parts = dependency.split(".")
+            if len(dependency_parts) == 1:  # from parent
+                target_model = self
+                trait = dependency
+            else:  # from sibling
+                sibling, trait = dependency_parts
+                target_model = self.get_model(sibling)
+            tl.dlink(
+                (target_model, trait),
+                (model, trait),
+            )
 
 
 class HasProcess(tl.HasTraits):

src/aiidalab_qe/utils.py

@@ -1,3 +1,5 @@
+from aiida_pseudo.groups.family import PseudoPotentialFamily
+
 from aiida import orm
 
 
@@ -30,3 +32,8 @@ def enable_pencil_decomposition(component):
     """Enable the pencil decomposition for the given component."""
 
     component.settings = orm.Dict({"CMDLINE": ["-pd", ".true."]})
+
+
+def fetch_pseudo_family_by_label(label) -> PseudoPotentialFamily:
+    """Fetch the pseudo family by label."""
+    return orm.Group.collection.get(label=label)  # type: ignore