PR-R2: Load viewer representations from extras

aiidalab_widgets_base/structures.py

@@ -16,7 +16,10 @@
 # Local imports
 from .data import FunctionalGroupSelectorWidget
 from .utils import StatusHTML, exceptions, get_ase_from_file, get_formula
-from .viewers import StructureDataViewer
+from .viewers import (
+    StructureDataViewer,
+    restore_viewer_representations_from_extras,
+)
 
 CifData = plugins.DataFactory("core.cif")
 StructureData = plugins.DataFactory("core.structure")
@@ -336,11 +339,14 @@ def _observe_input_structure(self, change):
             change["new"], CifData
         ):  # Special treatement of the CifData object
             str_io = io.StringIO(change["new"].get_content())
-            self.structure = ase.io.read(
-                str_io, format="cif", reader="ase", store_tags=True
+            structure = ase.io.read(str_io, format="cif", reader="ase", store_tags=True)
+            self.structure = restore_viewer_representations_from_extras(
+                change["new"], structure
             )
         elif isinstance(change["new"], StructureData):
-            self.structure = change["new"].get_ase()
+            self.structure = restore_viewer_representations_from_extras(
+                change["new"], change["new"].get_ase()
+            )
 
         else:
             self.structure = None

aiidalab_widgets_base/viewers.py

@@ -169,6 +169,7 @@ def __init__(self, parameter, downloadable=True, **kwargs):
 
 REPRESENTATION_PREFIX = "_aiidalab_viewer_representation_"
 DEFAULT_REPRESENTATION = f"{REPRESENTATION_PREFIX}default"
+VIEWER_REPRESENTATIONS_EXTRA = "aiidalab_viewer_representations"
 _REPRESENTATION_TYPE_TO_TOKEN = {
     "ball+stick": "ballstick",
     "spacefill": "spacefill",
@@ -242,6 +243,23 @@ def parse_representation_style_id(style_id):
     }
 
 
+def restore_viewer_representations_from_extras(node, structure):
+    """Restore viewer representation masks from AiiDA node extras into ASE arrays."""
+    if node is None or structure is None:
+        return structure
+    representations = node.base.extras.get(VIEWER_REPRESENTATIONS_EXTRA, {})
+    if not isinstance(representations, dict):
+        return structure
+    for key, values in representations.items():
+        if not str(key).startswith(REPRESENTATION_PREFIX):
+            continue
+        values = np.asarray(values, dtype=int)
+        if len(values) != len(structure):
+            continue
+        structure.set_array(key, values)
+    return structure
+
+
 class NglViewerRepresentation(ipw.HBox):
     """This class represents the parameters for displaying a structure in NGLViewer.
 
@@ -1346,7 +1364,9 @@ def _valid_structure(self, change):
             self.pk = None
         elif isinstance(structure, (orm.StructureData, orm.CifData)):
             self.pk = structure.pk
-            structure = structure.get_ase()
+            structure = restore_viewer_representations_from_extras(
+                structure, structure.get_ase()
+            )
 
         # Add default representation array if it is not present.
         # This will make sure that the new structure is displayed at the beginning.
@@ -1389,11 +1409,13 @@ def _observe_structure(self, change=None):
                     indices=np.where(structure.arrays[style_id] >= 1)[0],
                     apply=False,
                 )
-        # Empty atoms selection for the representations that are not present in the structure.
-        # Typically this happens when a new structure is imported.
-        for i, style_id in enumerate(representation_ids):
-            if style_id not in structure_ids:
-                self._all_representations[i].selection.value = ""
+        # Drop representations that are not present in the new structure.
+        # Typically this happens when a different structure is imported/selected.
+        self._all_representations = [
+            representation
+            for representation in self._all_representations
+            if representation.style_id in structure_ids
+        ]
 
         self._observe_supercell()  # To trigger an update of the displayed structure
         self.set_trait("cell", structure.cell)

tests/test_structures.py

@@ -8,6 +8,7 @@
 
 import aiidalab_widgets_base as awb
 import aiidalab_widgets_base.structures as structures
+from aiidalab_widgets_base import viewers
 
 
 @pytest.fixture
@@ -135,6 +136,34 @@ def test_structure_manager_widget(structure_data_object):
     assert structure_manager_widget.structure is None
 
 
+@pytest.mark.usefixtures("aiida_profile_clean")
+def test_structure_manager_widget_restores_viewer_representations_from_extras():
+    style_id = viewers.encode_representation_style_id(
+        viewers.REPRESENTATION_PREFIX,
+        representation_type="spacefill",
+        size=2,
+        color="red",
+        token="stored",
+    )
+    structure = ase.Atoms(
+        symbols=["C", "H"],
+        positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1)],
+        cell=[5.0, 5.0, 5.0],
+        pbc=True,
+    )
+    node = orm.StructureData(ase=structure).store()
+    node.base.extras.set(viewers.VIEWER_REPRESENTATIONS_EXTRA, {style_id: [1, -1]})
+
+    structure_manager_widget = awb.StructureManagerWidget(
+        importers=[], input_structure=node
+    )
+
+    representation_ids = [
+        rep.style_id for rep in structure_manager_widget.viewer._all_representations
+    ]
+    assert style_id in representation_ids
+
+
 @pytest.mark.usefixtures("aiida_profile_clean")
 def test_structure_browser_widget(structure_data_object, monkeypatch):
     """Test the `StructureBrowserWidget`."""

tests/test_viewers.py

@@ -367,6 +367,33 @@ def test_structure_data_viewer_imports_encoded_representation_array():
     assert representation.color.value == "red"
 
 
+@pytest.mark.usefixtures("aiida_profile_clean")
+def test_structure_data_viewer_restores_representation_arrays_from_extras():
+    structure = ase.Atoms(
+        symbols=["C", "H"],
+        positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1)],
+    )
+    style_id = viewers.encode_representation_style_id(
+        viewers.StructureDataViewer.REPRESENTATION_PREFIX,
+        representation_type="spacefill",
+        size=2,
+        color="red",
+        token="stored",
+    )
+    node = orm.StructureData(ase=structure)
+    node.base.extras.set(viewers.VIEWER_REPRESENTATIONS_EXTRA, {style_id: [1, -1]})
+
+    viewer = viewers.StructureDataViewer(node)
+
+    representation_ids = [rep.style_id for rep in viewer._all_representations]
+    assert style_id in representation_ids
+    representation = viewer._all_representations[representation_ids.index(style_id)]
+    assert representation.selection.value == "1"
+    assert representation.type.value == "spacefill"
+    assert representation.size.value == 2
+    assert representation.color.value == "red"
+
+
 def test_structure_data_viewer_imports_multiple_encoded_representation_arrays():
     structure = ase.Atoms(
         symbols=["C", "H", "H"],
@@ -406,6 +433,34 @@ def test_structure_data_viewer_imports_multiple_encoded_representation_arrays():
     assert representation_2.color.value == "red"
 
 
+def test_structure_data_viewer_drops_stale_representations_on_structure_change():
+    structure = ase.Atoms(
+        symbols=["C", "H"],
+        positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1)],
+    )
+    style_id = viewers.encode_representation_style_id(
+        viewers.StructureDataViewer.REPRESENTATION_PREFIX,
+        representation_type="spacefill",
+        size=2,
+        color="red",
+        token="old",
+    )
+    structure.set_array(style_id, np.array([1, -1], dtype=int))
+    viewer = viewers.StructureDataViewer()
+    viewer.structure = structure
+    assert style_id in [rep.style_id for rep in viewer._all_representations]
+
+    viewer.structure = ase.Atoms(
+        symbols=["C", "H", "H"],
+        positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1), (0.0, 1.0, 0.0)],
+    )
+
+    assert [rep.style_id for rep in viewer._all_representations] == [
+        viewers.DEFAULT_REPRESENTATION
+    ]
+    assert viewer._all_representations[0].selection.value == "1..3"
+
+
 def test_structure_data_viewer_updates_encoded_representation_array_name():
     structure = ase.Atoms(
         symbols=["C", "H"],