Implement structure view toggle

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

@@ -1,27 +1,93 @@
+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
+        self.auto_render = 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,93 +1,133 @@
 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 not hasattr(self, "widget"):
-            structure = self._model.get_structure()
-            self.widget = StructureDataViewer(structure=structure)
-            # Select the Cell tab by default
-            self.widget.configuration_box.selected_index = 2
-            self.table_description = ipw.HTML("""
+        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.
+            ngl = self.widget._viewer
+            ngl._set_size("100%", "300px")
+            ngl.control.zoom(0.0)
+            return
+
+        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")
+        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 = [
+            self.header_box,
+            self.widget,
+            ipw.HTML("""
                 <h4 style='margin: 10px 0;'>
-                    Structure table information: Atom coordinates in Å
+                    Structure information: Atom coordinates in Å
                 </h4>
                 <p style='margin: 5px 0; color: #555;'>
-                    You can click on a row to select an atom. Multiple atoms
-                    can be selected by clicking on additional rows. To unselect
-                    an atom, click on the selected row again.
+                    You can click on a row to select an atom. Multiple atoms can be
+                    selected by clicking on additional rows. To unselect an atom, click
+                    on the selected row again.
                 </p>
-            """)
-            self.atom_coordinates_table = TableWidget()
-            self._generate_table(structure.get_ase())
-
-            structure_info = self._get_structure_info(structure)
-
-            self.results_container.children = [
-                structure_info,
-                self.widget,
-                self.table_description,
-                self.atom_coordinates_table,
-            ]
-
-            self.atom_coordinates_table.observe(self._change_selection, "selected_rows")
-            # Listen for changes in self.widget.displayed_selection and update the table
-            self.widget.observe(self._update_table_selection, "displayed_selection")
-
-        # 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.
-        ngl = self.widget._viewer
-        ngl._set_size("100%", "300px")
-        ngl.control.zoom(0.0)
-
-    def _get_structure_info(self, structure):
-        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> {format_time(structure.ctime)} ({relative_time(structure.ctime)})<br>
-            </div>
-            """
-        )
-
-    def _generate_table(self, structure):
-        data = [
-            [
-                "Atom index",
-                "Chemical symbol",
-                "Tag",
-                "x (Å)",
-                "y (Å)",
-                "z (Å)",
-            ]
+            """),
+            self.atom_coordinates_table,
         ]
-        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
-        ):
-            # Format position values to two decimal places
-            formatted_position = [f"{coord:.2f}" for coord in position]
-            data.append([index, symbol, tag, *formatted_position])
-        self.atom_coordinates_table.data = data
-
-    def _change_selection(self, _):
-        selected_indices = self.atom_coordinates_table.selected_rows
-        self.widget.displayed_selection = selected_indices
-
-    def _update_table_selection(self, change):
-        selected_indices = change.new
-        self.atom_coordinates_table.selected_rows = selected_indices
+
+    def _on_process_change(self, _):
+        super()._on_process_change(_)
+        if self.rendered:
+            self.view_toggle_button.layout.display = (
+                "block" if self._model.is_relaxed else "none"
+            )
+
+    def _on_selected_view_change(self, _):
+        self._model.update()
+
+    def _toggle_view(self, _):
+        self._model.toggle_selected_view()

src/aiidalab_qe/app/static/styles/results.css

@@ -2,10 +2,12 @@
   display: flex;
   margin: 10px auto;
 }
+
 .results-step-toggles div {
   display: flex;
   column-gap: 8px;
 }
+
 .results-step-toggles i {
   margin-left: 4px;
 }

src/aiidalab_qe/app/static/styles/structure.css

@@ -0,0 +1,35 @@
+/* Structure manager */
+
+.structure-viewer .view-box {
+  justify-content: center !important;
+}
+
+/* Structure coordinates table */
+
+.atom-coordinates-table > div {
+  width: fit-content;
+  max-height: 300px;
+  overflow: auto;
+}
+
+.atom-coordinates-table table {
+  width: 100% !important;
+}
+
+.atom-coordinates-table table tr:first-child {
+  background-color: #ccc !important;
+}
+
+.atom-coordinates-table table tr:not(:first-child):hover {
+  background-color: #ddd !important;
+}
+
+.atom-coordinates-table table tr:not(:first-child):active {
+  background-color: #ccc !important;
+}
+
+.atom-coordinates-table table th,
+.atom-coordinates-table table td {
+  text-align: center !important;
+  padding: 0.5em 1em !important;
+}

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("""