Save and load result panel state

src/aiidalab_qe/app/result/components/viewer/structure/model.py

@@ -97,3 +97,11 @@ def _get_atom_table_data(self):
             data.append([index, symbol, tag, *formatted_position])
 
         return data
+
+    def get_model_state(self):
+        return {
+            "selected_view": self.selected_view,
+        }
+
+    def set_model_state(self, parameters):
+        self.selected_view = parameters.get("selected_view", "initial")

src/aiidalab_qe/common/panel.py

@@ -507,7 +507,7 @@ def _link_override_to_widget_disable(self, code_model, code_widget):
             )
 
 
-class ResultsModel(PanelModel, HasProcess):
+class ResultsModel(SettingsModel, HasProcess):
     process_status_notification = tl.Unicode("")
 
     _this_process_label = ""
@@ -592,6 +592,20 @@ def _get_child_outputs(self, which="this"):
             return getattr(outputs, child) if child in outputs else AttributeDict({})
         return AttributeDict({key: getattr(node.outputs, key) for key in node.outputs})
 
+    def save_state(self):
+        """Saves the current state of the model to the AiiDA database."""
+        node = self.fetch_process_node()
+        results = node.base.extras.get("results", {})
+        results[self.identifier] = self.get_model_state()
+        node.base.extras.set("results", results)
+
+    def load_state(self):
+        """Loads the state of the model from the AiiDA database."""
+        node = self.fetch_process_node()
+        results = node.base.extras.get("results", {})
+        if self.identifier in results:
+            self.set_model_state(results[self.identifier])
+
 
 RM = t.TypeVar("RM", bound=ResultsModel)
 
@@ -630,6 +644,23 @@ def render(self):
             identifier=f"{self._model.identifier}-results",
             classes=["results-panel-guide"],
         )
+        self.save_state_button = ipw.Button(
+            description="Save state",
+            tooltip="Save the current visualization settings",
+            button_style="primary",
+            icon="save",
+        )
+        self.save_state_button.on_click(self._save_state)
+        self.load_state_button = ipw.Button(
+            description="Load state",
+            tooltip="Load previously saved visualization settings",
+            button_style="primary",
+            icon="download",
+        )
+        self.load_state_button.on_click(self._load_state)
+        self.state_buttons = ipw.HBox(
+            children=[self.save_state_button, self.load_state_button],
+        )
 
         self.results_container = ipw.VBox()
 
@@ -714,3 +745,11 @@ def _render(self):
 
     def _post_render(self):
         pass
+
+    def _save_state(self, _=None):
+        """Save the current state of the results panel."""
+        self._model.save_state()
+
+    def _load_state(self, _=None):
+        """Load a previously saved state of the results panel."""
+        self._model.load_state()

src/aiidalab_qe/plugins/electronic_structure/result/model.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import traitlets as tl
+
 from aiidalab_qe.common.panel import ResultsModel
 
 
@@ -20,6 +22,30 @@ class ElectronicStructureResultsModel(ResultsModel):
         "pdos": "PDOS",
     }
 
+    dos_atoms_group_options = tl.List(
+        trait=tl.List(tl.Unicode()),
+        default_value=[
+            ("Group by element (atomic species)", "kinds"),
+            ("No grouping (each site separately)", "atoms"),
+        ],
+    )
+    dos_atoms_group = tl.Unicode("kinds")
+    dos_plot_group_options = tl.List(
+        trait=tl.List(tl.Unicode()),
+        default_value=[
+            ("Group all orbitals per atom", "total"),
+            ("Group by angular momentum ", "angular_momentum"),
+            ("No grouping (each orbital separately)", "orbital"),
+        ],
+    )
+    dos_plot_group = tl.Unicode("angular_momentum")
+    selected_atoms = tl.Unicode("")
+    project_bands_box = tl.Bool(False)
+    proj_bands_width = tl.Float(0.5)
+
+    horizontal_width_percentage = tl.Int(100)
+    bands_width_percentage = tl.Int(70)
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._completed_processes = set()
@@ -74,3 +100,21 @@ def _has_bands(self):
     def _has_pdos(self):
         outputs = self._get_child_outputs("pdos")
         return all(output in outputs for output in ("dos", "projwfc"))
+
+    def get_model_state(self):
+        return {
+            "dos_atoms_group": self.dos_atoms_group,
+            "dos_plot_group": self.dos_plot_group,
+            "selected_atoms": self.selected_atoms,
+            "horizontal_width_percentage": self.horizontal_width_percentage,
+            "bands_width_percentage": self.bands_width_percentage,
+        }
+
+    def set_model_state(self, parameters):
+        self.dos_atoms_group = parameters.get("dos_atoms_group", "kinds")
+        self.dos_plot_group = parameters.get("dos_plot_group", "angular_momentum")
+        self.selected_atoms = parameters.get("selected_atoms", "")
+        self.horizontal_width_percentage = parameters.get(
+            "horizontal_width_percentage", 100
+        )
+        self.bands_width_percentage = parameters.get("bands_width_percentage", 70)

src/aiidalab_qe/plugins/electronic_structure/result/result.py

@@ -14,7 +14,11 @@ class ElectronicStructureResultsPanel(ResultsPanel[ElectronicStructureResultsMod
 
     def _render(self):
         self.bands_pdos_container = ipw.VBox()
-        children = []
+        # If the model implements `get_model_state` and `set_model_state`
+        # we can add the `state_buttons` to the results container.
+        # If we force all models to implement these methods, we can
+        # move this logic to the base class.
+        children = [self.state_buttons]
         if self._model.needs_property_selector:
             children.append(self._render_property_selector())
             self.has_property_selector = True
@@ -102,3 +106,24 @@ def _render_bands_pdos_widget(self, node_identifiers):
         widget = BandsPdosWidget(model=model)
         widget.render()
         self.bands_pdos_container.children = [widget]
+
+        ipw.link(
+            (self._model, "dos_atoms_group"),
+            (model, "dos_atoms_group"),
+        )
+        ipw.link(
+            (self._model, "dos_plot_group"),
+            (model, "dos_plot_group"),
+        )
+        ipw.link(
+            (self._model, "selected_atoms"),
+            (model, "selected_atoms"),
+        )
+        ipw.link(
+            (self._model, "horizontal_width_percentage"),
+            (model, "horizontal_width_percentage"),
+        )
+        ipw.link(
+            (self._model, "bands_width_percentage"),
+            (model, "bands_width_percentage"),
+        )

tests/test_plugins_bands.py

@@ -18,7 +18,8 @@ def test_result(generate_qeapp_workchain):
 
     panel.render()
 
-    assert len(panel.results_container.children) == 1  # only bands, so no controls
+    # only state buttons container and bands, so no controls
+    assert len(panel.results_container.children) == 2
 
     widget = panel.bands_pdos_container.children[0]  # type: ignore
     model = widget._model

tests/test_plugins_electronic_structure.py

@@ -17,7 +17,8 @@ def test_electronic_structure(generate_qeapp_workchain):
 
     panel.render()
 
-    assert len(panel.results_container.children) == 2  # has controls
+    # only state buttons container and bands, so no controls
+    assert len(panel.results_container.children) == 3
 
     widget = panel.bands_pdos_container.children[0]  # type: ignore
     model = widget._model

tests/test_plugins_pdos.py

@@ -17,7 +17,8 @@ def test_result(generate_qeapp_workchain):
 
     panel.render()
 
-    assert len(panel.results_container.children) == 1  # only pdos, so no controls
+    # only state buttons container and pdos, so no controls
+    assert len(panel.results_container.children) == 2
 
     widget = panel.bands_pdos_container.children[0]  # type: ignore
     model = widget._model