Fix bugs in pseudopotentials panel

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

@@ -202,14 +202,14 @@ def set_model_state(self, state: dict):
             if "smearing" in SYSTEM:
                 smearing.type = SYSTEM["smearing"]
 
-        num_kinds = len(self.input_structure.kinds) if self.has_structure else 1
-
         pseudos = t.cast(
             PseudosConfigurationSettingsModel,
             self.get_model("pseudos"),
         )
         with pseudos.hold_trait_notifications():
             if pseudo_family_string := state.get("pseudo_family"):
+                # If a pseudo family string exists, we can assume that no custom pseudos were uploaded.
+                # We therefore populate the UI from the pseudo family.
                 pseudo_family = PseudoFamily.from_string(pseudo_family_string)
                 library = f"{pseudo_family.library} {pseudo_family.accuracy}"
                 if relativistic := pseudo_family.relativistic:
@@ -218,22 +218,29 @@ def set_model_state(self, state: dict):
                 pseudos.library = library
                 pseudos.family = pseudo_family_string
             else:
+                # Otherwise, we assume a custom pseudo upload (warning enabled).
+                # Everything is set to None except the functional, which we extract from the pseudos.
+                pseudos.show_upload_warning = True
+                functionals = []
                 try:
-                    pp_uuid = next(iter(PW["pseudos"].values()))
-                    pseudo_info = get_pseudo_info(pp_uuid)
-                    pseudos.functional = pseudo_info["functional"]
-                except Exception:
+                    for pp_uuid in PW["pseudos"].values():
+                        pseudo_info = get_pseudo_info(pp_uuid)
+                        functionals.append(pseudo_info["functional"])
+                    if len(set(functionals)) == 1:
+                        pseudos.functional = functionals[0]
+                    else:
+                        pseudos.functional = None
+                except Exception as err:
+                    print(f"Error loading pseudos: {err}")
                     pseudos.functional = None
                 pseudos.library = None
                 pseudos.family = None
-                pseudos.show_upload_warning = True
-
-            pseudos.functionals = [pseudos.functional] * num_kinds
 
-            if pseudos_dictionary := PW.get("pseudos"):
-                pseudos.dictionary = pseudos_dictionary
-                pseudos.ecutwfc = SYSTEM.get("ecutwfc", 0.0)
-                pseudos.ecutrho = SYSTEM.get("ecutrho", 0.0)
+        if pseudos_dictionary := PW.get("pseudos"):
+            # We set these outside of the context manager to avoid a reset by other trait notifications.
+            pseudos.dictionary = pseudos_dictionary
+            pseudos.ecutwfc = SYSTEM.get("ecutwfc", 0.0)
+            pseudos.ecutrho = SYSTEM.get("ecutrho", 0.0)
 
         magnetization = t.cast(
             MagnetizationConfigurationSettingsModel,

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

@@ -4,6 +4,7 @@
 
 import traitlets as tl
 from aiida_pseudo.common.units import U
+from aiida_pseudo.data.pseudo import UpfData
 
 from aiida import orm
 from aiida.common import exceptions
@@ -46,16 +47,15 @@ class PseudosConfigurationSettingsModel(
         value_trait=tl.Unicode(allow_none=True),  # pseudopotential node uuid
         default_value={},
     )
-    functional = tl.Unicode(allow_none=True)
+    functional = tl.Unicode(None, allow_none=True)
     functional_options = tl.List(
         trait=tl.Unicode(),
         default_value=[
             "PBE",
             "PBEsol",
         ],
     )
-    functionals = tl.List(trait=tl.Unicode(allow_none=True))
-    library = tl.Unicode(allow_none=True)
+    library = tl.Unicode(None, allow_none=True)
     library_options = tl.List(
         trait=tl.Unicode(),
         default_value=[
@@ -67,7 +67,7 @@ class PseudosConfigurationSettingsModel(
             "PseudoDojo stringent (FR)",
         ],
     )
-    family = tl.Unicode(allow_none=True)
+    family = tl.Unicode(None, allow_none=True)
     family_header = tl.Unicode(allow_none=True)
     cutoffs = tl.List(
         trait=tl.List(tl.Float()),  # [[ecutwfc values], [ecutrho values]]
@@ -112,13 +112,12 @@ def update(self, specific=""):
         family = self.family
         self.update_family_parameters()
         # When the app starts, the family is not yet set. `update_family_parameters`
-        # will set the family, which will set the functionals and dictionary.
+        # will set the family, which will set the dictionary.
         # However, when the structure is changed, the family may already be set to
-        # the default, in which case, the functionals and dictionary will not be
+        # the default, in which case, the dictionary will not be
         # updated. Therefore, we need to force the update.
         if specific == "structure" and self.family == family:
             self.update_dictionary()
-            self.update_functionals()
 
     def update_family_parameters(self):
         if self.locked:
@@ -149,17 +148,8 @@ def update_family_parameters(self):
             self.functional = self._defaults["functional"]
             self.library = self._defaults["library"]
 
-    def update_functionals(self):
-        if self.locked or not (self.functional and self.family):
-            return
-        self.functionals = (
-            [self.functional] * len(self.input_structure.kinds)
-            if self.has_structure
-            else []
-        )
-
     def update_family(self):
-        if self.locked or not (self.library and self.functional):
+        if not (self.library and self.functional) or self.locked:
             return
 
         parts = self.library.split()
@@ -224,7 +214,7 @@ def update_dictionary(self):
                 f"Failed to fetch pseudo family using the '{self.family}' string"
             ) from err
 
-        pseudos = {}
+        pseudos: dict[str, UpfData | None] = {}
         for kind in self.input_structure.kinds:
             # If the kind is not in the family, we set it to None.
             # This will block the app and notify the user of the missing pseudo.
@@ -246,7 +236,7 @@ def update_dictionary(self):
         self.dictionary = self._get_default_dictionary()
 
     def update_cutoffs(self):
-        if self.locked or not self.dictionary:
+        if not self.dictionary or self.locked:
             return
 
         kinds = self.input_structure.kinds if self.has_structure else []
@@ -313,7 +303,7 @@ def update_cutoffs(self):
         self.cutoffs = self._get_default_cutoffs()
 
     def update_library_options(self):
-        if self.locked or not self.has_structure:
+        if not self.has_structure or self.locked:
             return
 
         relativistic_options = [
@@ -338,6 +328,36 @@ def update_library_options(self):
 
         self.update_family_parameters()
 
+    def get_cutoffs_by_index(self, index: int) -> list[float]:
+        return (
+            [self.cutoffs[0][index], self.cutoffs[1][index]]
+            if len(self.cutoffs[0]) > index
+            else [0.0, 0.0]
+        )
+
+    def update_functional(self):
+        pseudos: list[UpfData] = []
+        for kind_name, uuid in self.dictionary.items():
+            kind = self.input_structure.get_kind(kind_name)
+            try:
+                assert uuid is not None
+                pseudo = orm.load_node(uuid)
+                pseudos.append(pseudo)
+            except AssertionError:
+                print(
+                    f"The selected pseudopotential family does not contain a pseudopotential for {kind.symbol}. Consider changing the family or uploading a custom pseudopotential."
+                )
+                continue
+            except exceptions.NotExistent:
+                print(
+                    f"Pseudopotential with UUID {uuid} does not exist for {kind.symbol}."
+                )
+                continue
+        functional_set = {pp.base.extras.get("functional", None) for pp in pseudos}
+        functional = functional_set.pop() if len(functional_set) == 1 else None
+        self._defaults["functional"] = functional
+        self.functional = self._get_default("functional")
+
     def reset(self):
         with self.hold_trait_notifications():
             self.dictionary = self._get_default("dictionary")
@@ -386,10 +406,10 @@ def _get_default_cutoffs(self):
         return deepcopy(self._defaults["cutoffs"])
 
     def _check_blockers(self):
-        if not (self.dictionary and self.functionals):
+        if not (self.has_structure and self.dictionary):
             return
 
-        pseudos = []
+        pseudos: list[UpfData] = []
         for kind_name, uuid in self.dictionary.items():
             kind = self.input_structure.get_kind(kind_name)
             try:
@@ -403,7 +423,7 @@ def _check_blockers(self):
                 yield f"Pseudopotential with UUID {uuid} does not exist for {kind.symbol}."
                 return
 
-        functional_set = set(self.functionals)
+        functional_set = {pp.base.extras.get("functional", None) for pp in pseudos}
         if len(functional_set) != 1:
             yield "All pseudopotentials must have the same exchange-correlation (XC) functional."
         elif self.functional and self.functional not in functional_set:
@@ -413,3 +433,9 @@ def _check_blockers(self):
         if self.spin_orbit == "soc":
             if relativistic_set != {"full"}:
                 yield "For spin-orbit coupling (SOC) calculations, all pseudopotentials must be fully relativistic."
+
+        if self.ecutwfc == 0.0:
+            yield "The cutoff energy for wavefunctions (ecutwfc) cannot be zero."
+
+        if self.ecutrho == 0.0:
+            yield "The cutoff energy for charge density (ecutrho) cannot be zero."