Correctly handle supercell selection (#371)

When dealing with the displayed structure, a user often interacts with the supercell, but not the unit cell.
As the result, clicking on the "virtual" atoms leads to an error, since the selected index is out of bounds
w.r.t. to the real structure. This PR fixes the issue and introduces the correct handling of the selection. 

Additionally, we also add automatic tests to make sure the selection works as expected.

Co-authored-by: Aliaksandr Yakutovich <[email protected]>

aiidalab_widgets_base/structures.py

@@ -170,13 +170,19 @@ def _structure_editors(self, editors):
         """Preparing structure editors."""
         if editors and len(editors) == 1:
             link((editors[0], "structure"), (self, "structure"))
+
+            if editors[0].has_trait("input_selection"):
+                dlink((editors[0], "input_selection"), (self.viewer, "input_selection"))
+
             if editors[0].has_trait("selection"):
-                link((editors[0], "selection"), (self.viewer, "selection"))
+                dlink((self.viewer, "selection"), (editors[0], "selection"))
+
             if editors[0].has_trait("camera_orientation"):
                 dlink(
                     (self.viewer._viewer, "_camera_orientation"),
                     (editors[0], "camera_orientation"),
                 )  # pylint: disable=protected-access
+
             return editors[0]
 
         # If more than one editor was defined.
@@ -942,6 +948,7 @@ class BasicStructureEditor(ipw.VBox):  # pylint: disable=too-many-instance-attri
     position of periodic structure in cell) editing."""
 
     structure = Instance(Atoms, allow_none=True)
+    input_selection = List(Int, allow_none=True)
     selection = List(Int)
     camera_orientation = List()
 
@@ -1206,13 +1213,15 @@ def def_point(self, _=None):
         """Define the action point."""
         self.point.value = self.vec2str(self.sel2com())
         if self.autoclear_selection.value:
-            self.selection = []
+            self.input_selection = None
+            self.input_selection = []
 
     def def_axis_p1(self, _=None):
         """Define the first point of axis."""
         self.axis_p1.value = self.vec2str(self.sel2com())
         if self.autoclear_selection.value:
-            self.selection = []
+            self.input_selection = None
+            self.input_selection = []
 
     def def_axis_p2(self, _=None):
         """Define the second point of axis."""
@@ -1230,7 +1239,8 @@ def def_axis_p2(self, _=None):
             )
             self.axis_p2.value = self.vec2str(com)
             if self.autoclear_selection.value:
-                self.selection = []
+                self.input_selection = None
+                self.input_selection = []
 
     def def_perpendicular_to_screen(self, _=None):
         """Define a normalized vector perpendicular to the screen."""
@@ -1251,7 +1261,7 @@ def translate_dr(self, _=None, atoms=None, selection=None):
             self.action_vector * self.displacement.value
         )
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     @_register_structure
     @_register_selection
@@ -1261,7 +1271,7 @@ def translate_dxdydz(self, _=None, atoms=None, selection=None):
         # The action.
         atoms.positions[self.selection] += np.array(self.str2vec(self.dxyz.value))
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     @_register_structure
     @_register_selection
@@ -1271,7 +1281,7 @@ def translate_to_xyz(self, _=None, atoms=None, selection=None):
         geo_center = np.average(self.structure[self.selection].get_positions(), axis=0)
         atoms.positions[self.selection] += self.str2vec(self.dxyz.value) - geo_center
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     @_register_structure
     @_register_selection
@@ -1283,9 +1293,9 @@ def rotate(self, _=None, atoms=None, selection=None):
         vec = self.str2vec(self.vec2str(self.action_vector))
         center = self.str2vec(self.point.value)
         rotated_subset.rotate(self.phi.value, v=vec, center=center, rotate_cell=False)
-        atoms.positions[list(self.selection)] = rotated_subset.positions
+        atoms.positions[self.selection] = rotated_subset.positions
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     @_register_structure
     @_register_selection
@@ -1318,7 +1328,7 @@ def mirror(self, _=None, norm=None, point=None, atoms=None, selection=None):
         # Mirror atoms.
         atoms.positions[selection] -= 2 * projections
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     def mirror_3p(self, _=None):
         """Mirror atoms on the plane containing action vector and action point."""
@@ -1342,7 +1352,7 @@ def align(self, _=None, atoms=None, selection=None):
         subset.rotate(self.action_vector, self.str2vec(self.dxyz.value), center=center)
         atoms.positions[selection] = subset.positions
 
-        self.structure, self.selection = atoms, selection
+        self.structure, self.input_selection = atoms, selection
 
     @_register_structure
     @_register_selection
@@ -1373,7 +1383,7 @@ def mod_element(self, _=None, atoms=None, selection=None):
                 range(last_atom, last_atom + len(selection) * len(lgnd))
             )
 
-        self.structure, self.selection = atoms, new_selection
+        self.structure, self.input_selection = atoms, new_selection
 
     @_register_structure
     @_register_selection
@@ -1387,8 +1397,7 @@ def copy_sel(self, _=None, atoms=None, selection=None):
         atoms += add_atoms
 
         new_selection = list(range(last_atom, last_atom + len(selection)))
-
-        self.structure, self.selection = atoms, new_selection
+        self.structure, self.input_selection = atoms, new_selection
 
     @_register_structure
     @_register_selection
@@ -1421,12 +1430,14 @@ def add(self, _=None, atoms=None, selection=None):
 
         new_selection = list(range(last_atom, last_atom + len(selection) * len(lgnd)))
 
-        self.structure, self.selection = atoms, new_selection
+        self.structure, self.input_selection = atoms, new_selection
 
     @_register_structure
     @_register_selection
     def remove(self, _, atoms=None, selection=None):
         """Remove selected atoms."""
         del [atoms[selection]]
 
-        self.structure, self.selection = atoms, []
+        self.structure = atoms
+        self.input_selection = None
+        self.input_selection = []