Add formula to structure results view

edan-bainglass · Jan 13, 2025 · 037e47b · 037e47b
1 parent 97bf5ca
commit 037e47b
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 7 deletions.
diff --git a/src/aiidalab_qe/app/result/components/viewer/structure/model.py b/src/aiidalab_qe/app/result/components/viewer/structure/model.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import traitlets as tl
+from ase.formula import Formula
 
 from aiida import orm
 from aiidalab_qe.common.panel import ResultsModel
@@ -14,13 +15,15 @@ class StructureResultsModel(ResultsModel):
     structure = tl.Instance(orm.StructureData, allow_none=True)
     selected_view = tl.Unicode("initial")
     header = tl.Unicode()
+    sub_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"
 
-    header_template = "<h1 style='margin: 0;'>{title}</h1>"
+    HEADER_TEMPLATE = "<h2 style='margin: 0;'>{content}</h2>"
+    _SUB_HEADER_TEMPLATE = "<h4 style='margin: 0; font-size: 18px'>{content}</h4>"
 
     @property
     def include(self):
@@ -34,15 +37,18 @@ def update(self):
         super().update()
         with self.hold_trait_notifications():
             if not self.is_relaxed or self.selected_view == "initial":
-                self.header = self.header_template.format(title="Initial")
+                self.sub_header = self._SUB_HEADER_TEMPLATE.format(content="Initial")
                 self.source = self.inputs
             else:
-                self.header = self.header_template.format(title="Relaxed")
+                self.sub_header = self._SUB_HEADER_TEMPLATE.format(content="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()
+                if not self.header:
+                    formatted = Formula(self.structure.get_formula()).format("html")
+                    self.header = self.HEADER_TEMPLATE.format(content=formatted)
 
     def toggle_selected_view(self):
         self.selected_view = "relaxed" if self.selected_view == "initial" else "initial"

diff --git a/src/aiidalab_qe/app/result/components/viewer/structure/structure.py b/src/aiidalab_qe/app/result/components/viewer/structure/structure.py
@@ -31,6 +31,12 @@ def _render(self):
             (self.header, "value"),
         )
 
+        self.sub_header = ipw.HTML()
+        ipw.dlink(
+            (self._model, "sub_header"),
+            (self.sub_header, "value"),
+        )
+
         self.view_toggle_button = ipw.Button(
             icon="eye",
             layout=ipw.Layout(
@@ -68,7 +74,12 @@ def _render(self):
             children=[
                 ipw.VBox(
                     children=[
-                        self.header,
+                        ipw.VBox(
+                            children=[
+                                self.header,
+                                self.sub_header,
+                            ],
+                        ),
                         self.view_toggle_button,
                     ],
                     layout=ipw.Layout(justify_content="space-between"),

diff --git a/tests/test_result.py b/tests/test_result.py
@@ -151,7 +151,8 @@ def test_table_data(model):
     wc = generate_qeapp_workchain(relax_type="none")
     model.process_uuid = wc.node.uuid
     node = model.fetch_process_node()
-    assert "Initial" in model.header
+    assert "Si<sub>2</sub>" in model.header
+    assert "Initial" in model.sub_header
     assert "properties" in model.source  # inputs
     assert model.structure.pk == node.inputs.structure.pk
     assert str(node.inputs.structure.pk) in model.info
@@ -163,12 +164,12 @@ def test_table_data(model):
     wc = generate_qeapp_workchain(relax_type="positions_cell")
     model.process_uuid = wc.node.uuid
     node = model.fetch_process_node()
-    assert "Initial" in model.header
+    assert "Initial" in model.sub_header
     assert panel.view_toggle_button.layout.display == "block"
     assert panel.view_toggle_button.description == "View relaxed"
     panel.view_toggle_button.click()
     assert panel.view_toggle_button.description == "View initial"
-    assert "Relaxed" in model.header
+    assert "Relaxed" in model.sub_header
     assert "properties" not in model.source  # outputs
     assert model.structure.pk == node.outputs.structure.pk
     assert str(node.outputs.structure.pk) in model.info