Avoid storing structure nodes - use UUID

edan-bainglass · edan-bainglass · commit adb111e1b603 · 2025-09-03T09:56:42.000Z
diff --git a/src/aiidalab_qe/app/configuration/__init__.py b/src/aiidalab_qe/app/configuration/__init__.py
@@ -63,7 +63,7 @@ def __init__(self, model: ConfigurationStepModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
 
         self.settings = {
@@ -254,7 +254,7 @@ def _fetch_plugin_calculation_settings(self):
 
             if identifier == "bands":
                 ipw.dlink(
-                    (self._model, "input_structure"),
+                    (self._model, "structure_uuid"),
                     (outline.include, "disabled"),
                     lambda _: not self._model.has_pbc,
                 )
diff --git a/src/aiidalab_qe/app/configuration/advanced/advanced.py b/src/aiidalab_qe/app/configuration/advanced/advanced.py
@@ -79,7 +79,7 @@ def __init__(self, model: AdvancedConfigurationSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_spin_type_change,
@@ -150,7 +150,7 @@ def _update_tabs(self):
         self.advanced_tabs.selected_index = 0
 
     def _on_advanced_tab_change(self, change):
-        tab: ConfigurationSettingsPanel = self.advanced_tabs.children[change["new"]]
+        tab: ConfigurationSettingsPanel = self.advanced_tabs.children[change["new"]]  # type: ignore
         tab.render()
 
     def _on_reset_to_defaults_button_click(self, _):
diff --git a/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py b/src/aiidalab_qe/app/configuration/advanced/convergence/convergence.py
@@ -13,8 +13,8 @@ def __init__(self, model: ConvergenceConfigurationSettingsModel, **kwargs):
         super().__init__(model, **kwargs)
 
         self._model.observe(
-            self._on_structure_change,
-            "input_structure",
+            self._on_input_structure_change,
+            "structure_uuid",
         )
         self._model.observe(
             self._on_protocol_change,
@@ -107,7 +107,7 @@ def render(self):
             (self.kpoints_distance, "value"),
         )
         ipw.dlink(
-            (self._model, "input_structure"),
+            (self._model, "structure_uuid"),
             (self.kpoints_distance, "disabled"),
             lambda _: not self._model.has_pbc,
         )
@@ -239,7 +239,7 @@ def render(self):
 
         self.rendered = True
 
-    def _on_structure_change(self, _):
+    def _on_input_structure_change(self, _):
         self.refresh(specific="structure")
 
     def _on_protocol_change(self, _):
diff --git a/src/aiidalab_qe/app/configuration/advanced/convergence/model.py b/src/aiidalab_qe/app/configuration/advanced/convergence/model.py
@@ -20,7 +20,7 @@ class ConvergenceConfigurationSettingsModel(
     identifier = "convergence"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "protocol",
     ]
 
diff --git a/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py b/src/aiidalab_qe/app/configuration/advanced/hubbard/hubbard.py
@@ -14,7 +14,7 @@ def __init__(self, model: HubbardConfigurationSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_hubbard_activation,
diff --git a/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py b/src/aiidalab_qe/app/configuration/advanced/hubbard/model.py
@@ -19,7 +19,7 @@ class HubbardConfigurationSettingsModel(
     identifier = "hubbard"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
     ]
 
     is_active = tl.Bool(False)
@@ -95,7 +95,7 @@ def set_active_eigenvalues(self, eigenvalues: list):
         self.has_eigenvalues = True
 
     def get_parameters_from_hubbard_structure(self):
-        hubbard_parameters = self.input_structure.hubbard.dict()["parameters"]
+        hubbard_parameters = self.input_structure.hubbard.model_dump()["parameters"]
         sites = self.input_structure.sites
         return {
             f"{sites[hp['atom_index']].kind_name} - {hp['atom_manifold']}": hp["value"]
diff --git a/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py b/src/aiidalab_qe/app/configuration/advanced/magnetization/magnetization.py
@@ -27,7 +27,7 @@ def __init__(self, model: MagnetizationConfigurationSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_electronic_type_change,
@@ -151,7 +151,7 @@ def _build_kinds_widget(self):
 
         kind_names = (
             self._model.input_structure.get_kind_names()
-            if self._model.input_structure
+            if self._model.has_structure
             else []
         )
 
diff --git a/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py b/src/aiidalab_qe/app/configuration/advanced/magnetization/model.py
@@ -18,7 +18,7 @@ class MagnetizationConfigurationSettingsModel(
     identifier = "magnetization"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "electronic_type",
         "spin_type",
         "pseudos.dictionary",
diff --git a/src/aiidalab_qe/app/configuration/advanced/model.py b/src/aiidalab_qe/app/configuration/advanced/model.py
@@ -26,7 +26,7 @@ class AdvancedConfigurationSettingsModel(
     identifier = "advanced"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "workchain.protocol",
         "workchain.spin_type",
         "workchain.electronic_type",
@@ -233,7 +233,7 @@ def _set_pw_parameters(self, pw_parameters):
         control_params: dict = pw_parameters.get("CONTROL", {})
         electron_params: dict = pw_parameters.get("ELECTRONS", {})
 
-        num_atoms = len(self.input_structure.sites) if self.input_structure else 1
+        num_atoms = len(self.input_structure.sites) if self.has_structure else 1
 
         general = t.cast(
             GeneralConfigurationSettingsModel,
diff --git a/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py b/src/aiidalab_qe/app/configuration/advanced/pseudos/model.py
@@ -31,7 +31,7 @@ class PseudosConfigurationSettingsModel(
     identifier = "pseudos"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "protocol",
         "spin_orbit",
     ]
@@ -249,7 +249,7 @@ def update_cutoffs(self):
         if self.locked or not self.dictionary:
             return
 
-        kinds = self.input_structure.kinds if self.input_structure else []
+        kinds = self.input_structure.kinds if self.has_structure else []
         self.status_message = ""
 
         if self.family:
diff --git a/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py b/src/aiidalab_qe/app/configuration/advanced/pseudos/pseudos.py
@@ -22,7 +22,7 @@ def __init__(self, model: PseudosConfigurationSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_protocol_change,
diff --git a/src/aiidalab_qe/app/configuration/basic/basic.py b/src/aiidalab_qe/app/configuration/basic/basic.py
@@ -17,7 +17,7 @@ def __init__(self, model: BasicConfigurationSettingsModel, **kwargs):
         super().__init__(model, **kwargs)
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
 
     def render(self):
diff --git a/src/aiidalab_qe/app/configuration/basic/model.py b/src/aiidalab_qe/app/configuration/basic/model.py
@@ -20,7 +20,7 @@ class BasicConfigurationSettingsModel(
     identifier = "workchain"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
     ]
 
     protocol_options = tl.List(
diff --git a/src/aiidalab_qe/app/structure/__init__.py b/src/aiidalab_qe/app/structure/__init__.py
@@ -17,6 +17,7 @@
     ShakeNBreakEditor,
 )
 from aiidalab_qe.common.infobox import InAppGuide
+from aiidalab_qe.common.setup_pseudos import PseudosInstallWidget
 from aiidalab_qe.common.widgets import CategorizedStructureExamplesWidget
 from aiidalab_qe.common.wizard import QeConfirmableWizardStep
 from aiidalab_widgets_base import (
@@ -66,7 +67,7 @@ def __init__(self, model: StructureStepModel, auto_setup=True, **kwargs):
         )
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._install_sssp(auto_setup)
 
@@ -133,7 +134,8 @@ def _render(self):
 
         ipw.dlink(
             (self.manager, "structure_node"),
-            (self._model, "input_structure"),
+            (self._model, "structure_uuid"),
+            lambda node: node.uuid if node else None,
         )
         ipw.link(
             (self._model, "manager_output"),
@@ -198,8 +200,6 @@ def _on_input_structure_change(self, _):
         self._update_state()
 
     def _install_sssp(self, auto_setup):
-        from aiidalab_qe.common.setup_pseudos import PseudosInstallWidget
-
         self.sssp_installation = PseudosInstallWidget(auto_start=False)
         ipw.dlink(
             (self.sssp_installation, "busy"),
@@ -224,7 +224,7 @@ def _toggle_sssp_installation_widget(self):
     def _update_state(self):
         if self._model.confirmed:
             self.state = self.State.SUCCESS
-        elif self._model.input_structure is None:
+        elif not self._model.has_structure:
             self.state = self.State.READY
         else:
             self.state = self.State.CONFIGURED
diff --git a/src/aiidalab_qe/app/structure/model.py b/src/aiidalab_qe/app/structure/model.py
@@ -31,7 +31,7 @@ def update_widget_text(self):
             self.structure_name = str(self.input_structure.get_formula())
 
     def reset(self):
-        self.input_structure = None
+        self.structure_uuid = None
         self.structure_name = ""
         self.manager_output = ""
 
diff --git a/src/aiidalab_qe/app/submission/global_settings/model.py b/src/aiidalab_qe/app/submission/global_settings/model.py
@@ -20,7 +20,7 @@ class GlobalResourceSettingsModel(
     identifier = "global"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "input_parameters",
     ]
 
diff --git a/src/aiidalab_qe/app/submission/global_settings/setting.py b/src/aiidalab_qe/app/submission/global_settings/setting.py
@@ -20,7 +20,7 @@ def __init__(self, model: GlobalResourceSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_input_parameters_change,
diff --git a/src/aiidalab_qe/app/submission/model.py b/src/aiidalab_qe/app/submission/model.py
@@ -68,7 +68,7 @@ def update(self):
             model.update()
 
     def update_process_label(self):
-        if not self.input_structure:
+        if not self.has_structure:
             self.process_label = ""
             return
         structure_label = (
@@ -170,7 +170,6 @@ def set_model_state(self, state: dict):
 
     def reset(self):
         with self.hold_trait_notifications():
-            self.input_structure = None
             self.input_parameters = {}
             self.process_node = None
             for identifier, model in self.get_models():
diff --git a/src/aiidalab_qe/app/wizard_app.py b/src/aiidalab_qe/app/wizard_app.py
@@ -145,16 +145,16 @@ def _render_step(self, step_index):
 
     def _update_configuration_step(self):
         if self.structure_model.confirmed:
-            self.configure_model.input_structure = self.structure_model.input_structure
+            self.configure_model.structure_uuid = self.structure_model.structure_uuid
         else:
-            self.configure_model.input_structure = None
+            self.configure_model.structure_uuid = None
 
     def _update_submission_step(self):
         if self.configure_model.confirmed:
-            self.submit_model.input_structure = self.structure_model.input_structure
+            self.submit_model.structure_uuid = self.structure_model.structure_uuid
             self.submit_model.input_parameters = self.configure_model.get_model_state()
         else:
-            self.submit_model.input_structure = None
+            self.submit_model.structure_uuid = None
             self.submit_model.input_parameters = {}
 
     def _update_results_step(self):
@@ -179,7 +179,7 @@ def _update_from_process(self, pk):
         else:
             self._show_process_loading_message()
             process_node = load_node(pk)
-            self.structure_model.input_structure = process_node.inputs.structure
+            self.structure_model.structure_uuid = process_node.inputs.structure.uuid
             self.structure_model.confirm()
             parameters = process_node.base.extras.get("ui_parameters", {})
             if parameters and isinstance(parameters, str):
diff --git a/src/aiidalab_qe/common/mixins.py b/src/aiidalab_qe/common/mixins.py
@@ -7,17 +7,24 @@
 from aiida import orm
 from aiida.common.exceptions import NotExistent
 from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
+from aiidalab_qe.common.decorators import cache_per_thread
 from aiidalab_qe.common.mvc import Model
 
+StructureType = t.Union[orm.StructureData, HubbardStructureData]
+
 
 class HasInputStructure(tl.HasTraits):
-    input_structure = tl.Union(
-        [
-            tl.Instance(orm.StructureData),
-            tl.Instance(HubbardStructureData),
-        ],
-        allow_none=True,
-    )
+    structure_uuid = tl.Unicode(None, allow_none=True)
+
+    @property
+    @cache_per_thread(invalidator="structure_uuid")
+    def input_structure(self) -> StructureType | None:
+        if not self.structure_uuid:
+            return None
+        try:
+            return t.cast(StructureType, orm.load_node(self.structure_uuid))
+        except NotExistent:
+            return None
 
     @property
     def has_structure(self):
@@ -29,7 +36,7 @@ def has_pbc(self):
 
     @property
     def has_tags(self):
-        return any(
+        return self.has_structure and any(
             not kind_name.isalpha()
             for kind_name in self.input_structure.get_kind_names()
         )
diff --git a/src/aiidalab_qe/plugins/pdos/model.py b/src/aiidalab_qe/plugins/pdos/model.py
@@ -14,7 +14,7 @@ class PdosConfigurationSettingsModel(PanelModel, HasInputStructure):
     identifier = "pdos"
 
     dependencies = [
-        "input_structure",
+        "structure_uuid",
         "workchain.protocol",
     ]
 
diff --git a/src/aiidalab_qe/plugins/pdos/setting.py b/src/aiidalab_qe/plugins/pdos/setting.py
@@ -18,7 +18,7 @@ def __init__(self, model: PdosConfigurationSettingsModel, **kwargs):
 
         self._model.observe(
             self._on_input_structure_change,
-            "input_structure",
+            "structure_uuid",
         )
         self._model.observe(
             self._on_protocol_change,
@@ -44,7 +44,7 @@ def render(self):
             (self.nscf_kpoints_distance, "value"),
         )
         ipw.dlink(
-            (self._model, "input_structure"),
+            (self._model, "structure_uuid"),
             (self.nscf_kpoints_distance, "disabled"),
             lambda _: not self._model.has_pbc,
         )
@@ -65,7 +65,7 @@ def render(self):
             (self.use_pdos_degauss, "value"),
         )
         ipw.dlink(
-            (self._model, "input_structure"),
+            (self._model, "structure_uuid"),
             (self.use_pdos_degauss, "disabled"),
             lambda _: not self._model.has_pbc,
         )
diff --git a/tests/configuration/test_advanced.py b/tests/configuration/test_advanced.py
diff --git a/tests/conftest.py b/tests/conftest.py
diff --git a/tests/test_app.py b/tests/test_app.py
diff --git a/tests/test_pseudo.py b/tests/test_pseudo.py
diff --git a/tests/test_submit_qe_workchain.py b/tests/test_submit_qe_workchain.py

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,7 @@ def __init__(self, model: ConfigurationStepModel, **kwargs):`
`63`	`63`
`64`	`64`	`self._model.observe(`
`65`	`65`	`self._on_input_structure_change,`
`66`		`- "input_structure",`
	`66`	`+ "structure_uuid",`
`67`	`67`	`)`
`68`	`68`
`69`	`69`	`self.settings = {`
`@@ -254,7 +254,7 @@ def _fetch_plugin_calculation_settings(self):`
`254`	`254`
`255`	`255`	`if identifier == "bands":`
`256`	`256`	`ipw.dlink(`
`257`		`- (self._model, "input_structure"),`
	`257`	`+ (self._model, "structure_uuid"),`
`258`	`258`	`(outline.include, "disabled"),`
`259`	`259`	`lambda _: not self._model.has_pbc,`
`260`	`260`	`)`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ class ConvergenceConfigurationSettingsModel(`
`20`	`20`	`identifier = "convergence"`
`21`	`21`
`22`	`22`	`dependencies = [`
`23`		`- "input_structure",`
	`23`	`+ "structure_uuid",`
`24`	`24`	`"protocol",`
`25`	`25`	`]`
`26`	`26`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ def __init__(self, model: HubbardConfigurationSettingsModel, **kwargs):`
`14`	`14`
`15`	`15`	`self._model.observe(`
`16`	`16`	`self._on_input_structure_change,`
`17`		`- "input_structure",`
	`17`	`+ "structure_uuid",`
`18`	`18`	`)`
`19`	`19`	`self._model.observe(`
`20`	`20`	`self._on_hubbard_activation,`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ def __init__(self, model: MagnetizationConfigurationSettingsModel, **kwargs):`
`27`	`27`
`28`	`28`	`self._model.observe(`
`29`	`29`	`self._on_input_structure_change,`
`30`		`- "input_structure",`
	`30`	`+ "structure_uuid",`
`31`	`31`	`)`
`32`	`32`	`self._model.observe(`
`33`	`33`	`self._on_electronic_type_change,`
`@@ -151,7 +151,7 @@ def _build_kinds_widget(self):`
`151`	`151`
`152`	`152`	`kind_names = (`
`153`	`153`	`self._model.input_structure.get_kind_names()`
`154`		`- if self._model.input_structure`
	`154`	`+ if self._model.has_structure`
`155`	`155`	`else []`
`156`	`156`	`)`
`157`	`157`