Implement structure panel toggle

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

@@ -1,27 +1,92 @@
+from __future__ import annotations
+
+import traitlets as tl
+
+from aiida import orm
 from aiidalab_qe.common.panel import ResultsModel
+from aiidalab_qe.common.time import format_time, relative_time
 
 
 class StructureResultsModel(ResultsModel):
     title = "Structure"
     identifier = "structure"
 
+    structure = tl.Instance(orm.StructureData, allow_none=True)
+    selected_view = tl.Unicode("initial")
+    header = tl.Unicode()
+    source = tl.Instance(orm.utils.managers.NodeLinksManager, allow_none=True)
+    info = tl.Unicode()
+    table_data = tl.List(tl.List())
+
     _this_process_label = "PwRelaxWorkChain"
 
-    source = None
+    header_template = "<h1 style='margin: 0;'>{title}</h1>"
 
     @property
     def include(self):
         return True
 
+    @property
+    def is_relaxed(self):
+        return "relax" in self.properties
+
     def update(self):
-        is_relaxed = "relax" in self.properties
-        self.title = "Relaxed structure" if is_relaxed else "Initial structure"
-        self.source = self.outputs if is_relaxed else self.inputs
-        self.auto_render = not is_relaxed or self.has_results
+        with self.hold_trait_notifications():
+            if not self.is_relaxed or self.selected_view == "initial":
+                self.header = self.header_template.format(title="Initial")
+                self.source = self.inputs
+            else:
+                self.header = self.header_template.format(title="Relaxed")
+                self.source = self.outputs
+            self.structure = self._get_structure()
+            if self.structure:
+                self.info = self._get_structure_info()
+                self.table_data = self._get_atom_table_data()
 
-    def get_structure(self):
+    def toggle_selected_view(self):
+        self.selected_view = "relaxed" if self.selected_view == "initial" else "initial"
+
+    def _get_structure(self) -> orm.StructureData | None:
         try:
             return self.source.structure if self.source else None
         except AttributeError:
             # If source is outputs but job failed, there may not be a structure
             return None
+
+    def _get_structure_info(self):
+        structure = self.structure
+        formatted = format_time(structure.ctime)
+        relative = relative_time(structure.ctime)
+        return f"""
+            <div style='line-height: 1.4;'>
+                <strong>PK:</strong> {structure.pk}<br>
+                <strong>Label:</strong> {structure.label}<br>
+                <strong>Description:</strong> {structure.description}<br>
+                <strong>Number of atoms:</strong> {len(structure.sites)}<br>
+                <strong>Creation time:</strong> {formatted} ({relative})<br>
+            </div>
+        """
+
+    def _get_atom_table_data(self):
+        structure = self.structure.get_ase()
+        data = [
+            [
+                "Atom index",
+                "Chemical symbol",
+                "Tag",
+                "x (Å)",
+                "y (Å)",
+                "z (Å)",
+            ]
+        ]
+        positions = structure.positions
+        chemical_symbols = structure.get_chemical_symbols()
+        tags = structure.get_tags()
+
+        for index, (symbol, tag, position) in enumerate(
+            zip(chemical_symbols, tags, positions), start=1
+        ):
+            formatted_position = [f"{coord:.2f}" for coord in position]
+            data.append([index, symbol, tag, *formatted_position])
+
+        return data

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

@@ -1,43 +1,110 @@
 import ipywidgets as ipw
 
 from aiidalab_qe.common.panel import ResultsPanel
-from aiidalab_qe.common.time import format_time, relative_time
 from aiidalab_qe.common.widgets import TableWidget
 from aiidalab_widgets_base.viewers import StructureDataViewer
 
 from .model import StructureResultsModel
 
 
 class StructureResultsPanel(ResultsPanel[StructureResultsModel]):
+    def __init__(self, model: StructureResultsModel, **kwargs):
+        super().__init__(model, **kwargs)
+        self._model.observe(
+            self._on_selected_view_change,
+            "selected_view",
+        )
+
     def _render(self):
         if hasattr(self, "widget"):
             # HACK to resize the NGL viewer in cases where it auto-rendered when its
-            # container was not displayed, which leads to a null width. This hack restores
-            # the original dimensions.
+            # container was not displayed, which leads to a null width. This hack
+            # restores the original dimensions.
             ngl = self.widget._viewer
             ngl._set_size("100%", "300px")
             ngl.control.zoom(0.0)
             return
 
-        structure = self._model.get_structure()
-        self.widget = StructureDataViewer(structure=structure)
+        self.header = ipw.HTML()
+        ipw.dlink(
+            (self._model, "header"),
+            (self.header, "value"),
+        )
+
+        self.view_toggle_button = ipw.Button(
+            icon="eye",
+            layout=ipw.Layout(
+                display="block" if self._model.is_relaxed else "none",
+                width="125px",
+            ),
+        )
+        ipw.dlink(
+            (self._model, "selected_view"),
+            (self.view_toggle_button, "description"),
+            lambda view: f"View {'initial' if view == 'relaxed' else 'relaxed'}",
+        )
+        ipw.dlink(
+            (self._model, "monitor_counter"),
+            (self.view_toggle_button, "disabled"),
+            lambda _: not self._model.has_results,
+        )
+        ipw.dlink(
+            (self.view_toggle_button, "disabled"),
+            (self.view_toggle_button, "tooltip"),
+            lambda disabled: "Waiting for results"
+            if disabled
+            else "Toggle between the initial and relaxed structures",
+        )
+
+        self.view_toggle_button.on_click(self._toggle_view)
+
+        self.structure_info = ipw.HTML(layout=ipw.Layout(margin="0"))
+        ipw.dlink(
+            (self._model, "info"),
+            (self.structure_info, "value"),
+        )
+
+        self.header_box = ipw.HBox(
+            children=[
+                ipw.VBox(
+                    children=[
+                        self.header,
+                        self.view_toggle_button,
+                    ],
+                    layout=ipw.Layout(justify_content="space-between"),
+                ),
+                ipw.VBox(
+                    children=[
+                        self.structure_info,
+                    ],
+                    layout=ipw.Layout(justify_content="flex-end"),
+                ),
+            ],
+            layout=ipw.Layout(grid_gap="1em"),
+        )
+
+        self.widget = StructureDataViewer()
+        ipw.dlink(
+            (self._model, "structure"),
+            (self.widget, "structure"),
+        )
 
         self.widget.configuration_box.selected_index = 2  # select the Cell tab
 
         self.atom_coordinates_table = TableWidget()
         self.atom_coordinates_table.add_class("atom-coordinates-table")
-
-        self._generate_table(structure.get_ase())
-
-        structure_info = self._get_structure_info(structure)
+        ipw.dlink(
+            (self._model, "table_data"),
+            (self.atom_coordinates_table, "data"),
+        )
 
         ipw.link(
             (self.widget, "displayed_selection"),
             (self.atom_coordinates_table, "selected_rows"),
         )
 
         self.results_container.children = [
-            structure_info,
+            self.header_box,
             self.widget,
             ipw.HTML("""
                 <h4 style='margin: 10px 0;'>
@@ -52,40 +119,8 @@ def _render(self):
             self.atom_coordinates_table,
         ]
 
-    def _get_structure_info(self, structure):
-        formatted = format_time(structure.ctime)
-        relative = relative_time(structure.ctime)
-        return ipw.HTML(
-            f"""
-            <div style='line-height: 1.4;'>
-                <strong>PK:</strong> {structure.pk}<br>
-                <strong>Label:</strong> {structure.label}<br>
-                <strong>Description:</strong> {structure.description}<br>
-                <strong>Number of atoms:</strong> {len(structure.sites)}<br>
-                <strong>Creation time:</strong> {formatted} ({relative})<br>
-            </div>
-            """
-        )
-
-    def _generate_table(self, structure):
-        data = [
-            [
-                "Atom index",
-                "Chemical symbol",
-                "Tag",
-                "x (Å)",
-                "y (Å)",
-                "z (Å)",
-            ]
-        ]
-        positions = structure.positions
-        chemical_symbols = structure.get_chemical_symbols()
-        tags = structure.get_tags()
-
-        for index, (symbol, tag, position) in enumerate(
-            zip(chemical_symbols, tags, positions), start=1
-        ):
-            formatted_position = [f"{coord:.2f}" for coord in position]
-            data.append([index, symbol, tag, *formatted_position])
+    def _on_selected_view_change(self, _):
+        self._model.update()
 
-        self.atom_coordinates_table.data = data
+    def _toggle_view(self, _):
+        self._model.toggle_selected_view()

src/aiidalab_qe/common/panel.py

@@ -627,6 +627,8 @@ def render(self):
         else:
             self._render_controls()
             self.children += (self.results_container,)
+            if self._model.identifier == "structure":
+                self._load_results()
 
     def _on_process_change(self, _):
         self._model.update()
@@ -666,7 +668,7 @@ def _render_controls(self):
 
         self.load_controls = ipw.HBox(
             children=[]
-            if self._model.auto_render
+            if self._model.auto_render or self._model.identifier == "structure"
             else [
                 self.load_results_button,
                 ipw.HTML("""