Fix PP upload widget (#1348)

This PR addresses several issues with the pseudopotential upload widget. 
It implements a system of warnings w.r.t invalid uploads.
If a valid, non-built-in pseudo is uploaded, the `functional`, `library`, and `family` are set to `None`.
This required support in the plugin now provided in `aiida-quantumespresso v4.12.0`.
Magnetization-to-moments conversion now uses `UpfData` nodes to extract Z values.
Lastly, the summary is updated to show selected pseudopotentials and reflect custom states.

Author: edan-bainglass

setup.cfg

@@ -26,7 +26,7 @@ packages = find:
 install_requires =
     aiida-core~=2.5,<3
     Jinja2~=3.0
-    aiida-quantumespresso~=4.10.0
+    aiida-quantumespresso~=4.12.0
     aiidalab-widgets-base[optimade] @ git+https://github.com/aiidalab/aiidalab-widgets-base@master
     aiida-pseudo~=1.4
     filelock~=3.8
@@ -37,6 +37,7 @@ install_requires =
     shakenbreak~=3.3.1
     plotly~=5.24
     kaleido~=0.2.1
+    upf_tools~=0.1.9
 
 python_requires = >=3.9
 

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

@@ -38,12 +38,12 @@ def __init__(self, model: MagnetizationConfigurationSettingsModel, **kwargs):
             "spin_type",
         )
         self._model.observe(
-            self._on_magnetization_type_change,
-            "type",
+            self._on_pseudos_dictionary_change,
+            "dictionary",
         )
         self._model.observe(
-            self._on_family_change,
-            "family",
+            self._on_magnetization_type_change,
+            "type",
         )
 
     def render(self):
@@ -123,13 +123,13 @@ def _on_electronic_type_change(self, _):
     def _on_spin_type_change(self, _):
         self.refresh(specific="spin")
 
+    def _on_pseudos_dictionary_change(self, _):
+        self.refresh(specific="dictionary")
+
     def _on_magnetization_type_change(self, _):
         self._toggle_widgets()
         self._model.update_type_help()
 
-    def _on_family_change(self, _):
-        self._model._update_default_moments()
-
     def _update(self, specific=""):
         if self.updated:
             return

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

@@ -1,13 +1,12 @@
 from copy import deepcopy
 
 import traitlets as tl
-from aiida_pseudo.groups.family import PseudoPotentialFamily
 
+from aiida import orm
 from aiida_quantumespresso.workflows.protocols.utils import (
     get_magnetization_parameters,
 )
 from aiidalab_qe.common.mixins import HasInputStructure
-from aiidalab_qe.utils import fetch_pseudo_family_by_label
 
 from ..subsettings import AdvancedCalculationSubSettingsModel
 
@@ -22,12 +21,16 @@ class MagnetizationConfigurationSettingsModel(
         "input_structure",
         "electronic_type",
         "spin_type",
-        "pseudos.family",
+        "pseudos.dictionary",
     ]
 
     electronic_type = tl.Unicode()
     spin_type = tl.Unicode()
-    family = tl.Unicode()
+    dictionary = tl.Dict(
+        key_trait=tl.Unicode(),  # kind name
+        value_trait=tl.Unicode(),  # pseudopotential node uuid
+        default_value={},
+    )
 
     type_options = tl.List(
         trait=tl.List(tl.Unicode()),
@@ -103,21 +106,25 @@ def _update_default_moments(self):
             # and should be carefully checked!
             return
 
-        family = fetch_pseudo_family_by_label(self.family)
         self._defaults["moments"] = {
-            kind.name: self._get_moment(kind.symbol, family)
-            for kind in self.input_structure.kinds
+            kind.name: self._get_moment(kind) for kind in self.input_structure.kinds
         }
 
-    def _get_moment(self, symbol: str, family: PseudoPotentialFamily) -> float:
+    def _get_moment(self, kind) -> float:
         """Convert the default magnetization to an initial magnetic moment."""
-        moment = self._DEFAULT_MOMENTS.get(symbol, {}).get("magmom", 0)
+        moment = self._DEFAULT_MOMENTS.get(kind.symbol, {}).get("magmom", 0)
         if moment != 0:
             return moment
 
+        try:
+            pseudo_uuid = self.dictionary.get(kind.name)
+            z_valence = orm.load_node(pseudo_uuid).z_valence
+        except Exception:
+            z_valence = 0
+
         # If no default moment is defined, or if it's 0, use 0.1 as default magnetization
         # and convert it to moments.
-        return round(0.1 * family.get_pseudo(symbol).z_valence, 3)
+        return round(0.1 * z_valence, 3)
 
     def _get_default_moments(self):
         return deepcopy(self._defaults.get("moments", {}))

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

@@ -183,13 +183,18 @@ def get_model_state(self):
 
     def set_model_state(self, parameters):
         pseudos: PseudosConfigurationSettingsModel = self.get_model("pseudos")  # type: ignore
-        if "pseudo_family" in parameters:
-            pseudo_family = PseudoFamily.from_string(parameters["pseudo_family"])
+        if pseudo_family_string := parameters.get("pseudo_family"):
+            pseudo_family = PseudoFamily.from_string(pseudo_family_string)
             library = pseudo_family.library
             accuracy = pseudo_family.accuracy
             pseudos.library = f"{library} {accuracy}"
             pseudos.functional = pseudo_family.functional
-            pseudos.family = parameters["pseudo_family"]
+            pseudos.family = pseudo_family_string
+        else:
+            pseudos.library = None
+            pseudos.functional = None
+            pseudos.family = None
+            pseudos.show_upload_warning = True
 
         if "pseudos" in parameters["pw"]:
             pseudos.dictionary = parameters["pw"]["pseudos"]

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

@@ -50,9 +50,13 @@ class PseudosConfigurationSettingsModel(
                 DEFAULT["advanced"]["pseudo_family"]["functional"],
                 DEFAULT["advanced"]["pseudo_family"]["accuracy"],
             ]
-        )
+        ),
+        allow_none=True,
+    )
+    functional = tl.Unicode(
+        DEFAULT["advanced"]["pseudo_family"]["functional"],
+        allow_none=True,
     )
-    functional = tl.Unicode(DEFAULT["advanced"]["pseudo_family"]["functional"])
     functional_options = tl.List(
         trait=tl.Unicode(),
         default_value=[
@@ -66,7 +70,8 @@ class PseudosConfigurationSettingsModel(
                 DEFAULT["advanced"]["pseudo_family"]["library"],
                 DEFAULT["advanced"]["pseudo_family"]["accuracy"],
             ]
-        )
+        ),
+        allow_none=True,
     )
     library_options = tl.List(
         trait=tl.Unicode(),
@@ -83,7 +88,8 @@ class PseudosConfigurationSettingsModel(
     )
     ecutwfc = tl.Float()
     ecutrho = tl.Float()
-    status_message = tl.Unicode("")
+    status_message = tl.Unicode("", allow_none=True)
+    show_upload_warning = tl.Bool(False)
 
     PSEUDO_HELP_SOC = """
         <div class="pseudo-text">
@@ -111,8 +117,6 @@ class PseudosConfigurationSettingsModel(
 
     family_help_message = tl.Unicode(PSEUDO_HELP_WO_SOC)
 
-    pseudo_filename_reset_trigger = tl.Int(0)
-
     def update(self, specific=""):  # noqa: ARG002
         with self.hold_trait_notifications():
             if not self.has_structure:
@@ -276,7 +280,6 @@ def reset(self):
             self.family = self._get_default("family")
             self.family_help_message = self._get_default("family_help_message")
             self.status_message = self._get_default("status_message")
-            self.pseudo_filename_reset_trigger += 1
 
     def _get_default(self, trait):
         if trait == "dictionary":