Fix imported viewer representation arrays

aiidalab_widgets_base/viewers.py

@@ -693,7 +693,7 @@ def _apply_representations(self, change=None):
             representation_ids.append(representation.style_id)
 
         # Remove missing representations from the structure.
-        for array in self.structure.arrays:
+        for array in list(self.structure.arrays):
             if (
                 array.startswith(self.REPRESENTATION_PREFIX)
                 and array not in representation_ids
@@ -1278,7 +1278,7 @@ def _observe_structure(self, change=None):
             except ValueError:
                 self._add_representation(
                     style_id=style_id,
-                    indices=np.where(structure.arrays[self.style_id] >= 1)[0],
+                    indices=np.where(structure.arrays[style_id] >= 1)[0],
                 )
         # Empty atoms selection for the representations that are not present in the structure.
         # Typically this happens when a new structure is imported.

tests/test_viewers.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 import ase
+import numpy as np
 import pytest
 import traitlets as tl
 from aiida import orm
@@ -320,6 +321,23 @@ def test_structure_data_viewer_representation(structure_data_object):
         v.structure = orm.Int(1)
 
 
+def test_structure_data_viewer_imports_unknown_representation_array():
+    structure = ase.Atoms(
+        symbols=["C", "H"],
+        positions=[(0.0, 0.0, 0.0), (0.0, 0.0, 1.1)],
+    )
+    style_id = f"{viewers.StructureDataViewer.REPRESENTATION_PREFIX}custom"
+    structure.set_array(style_id, np.array([1, -1], dtype=int))
+
+    viewer = viewers.StructureDataViewer()
+    viewer.structure = structure
+
+    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"
+
+
 def test_structure_data_viewer_clears_bond_shape_components():
     water = ase.Atoms(
         symbols=["O", "H", "H"],