Implement viewer representations (#373)

The main change of this PR is introducing the viewer's representations. Now, users can
specify what parts of the structure and how to display them.

These representations are now efficiently managed as arrays of `ase.Atoms` objects.

The PR also add a significant amount of tests for the viewer making it more reliable.

Co-authored-by: Aliaksandr Yakutovich <yakutovicha@gmail.com>

aiidalab_widgets_base/computational_resources.py

@@ -159,7 +159,7 @@ def refresh(self, _=None):
         with self.hold_trait_notifications():
             self.code_select_dropdown.options = self._get_codes()
             if not self.code_select_dropdown.options:
-                self.output.value = f"No codes found for default calcjob plugin '{self.default_calc_job_plugin}'."
+                self.output.value = f"No codes found for default calcjob plugin {self.default_calc_job_plugin!r}."
                 self.code_select_dropdown.disabled = True
             else:
                 self.code_select_dropdown.disabled = False

aiidalab_widgets_base/elns.py

@@ -23,7 +23,7 @@ def connect_to_eln(eln_instance=None, **kwargs):
     except (FileNotFoundError, json.JSONDecodeError, KeyError):
         return (
             None,
-            f"Can't open '{ELN_CONFIG}' (ELN configuration file). Instance: {eln_instance}",
+            f"Can't open {ELN_CONFIG!r} (ELN configuration file). Instance: {eln_instance}",
         )
 
     # If no ELN instance was specified, trying the default one.
@@ -35,7 +35,7 @@ def connect_to_eln(eln_instance=None, **kwargs):
             eln_config = config[eln_instance]
             eln_type = eln_config.pop("eln_type", None)
         else:  # The selected instance is not present in the config.
-            return None, f"Didn't find configuration for the '{eln_instance}' instance."
+            return None, f"Didn't find configuration for the {eln_instance!r} instance."
 
         # If the ELN type cannot be identified - aborting.
         if not eln_type:
@@ -73,7 +73,7 @@ def __init__(self, path_to_root="../", **kwargs):
 
         if eln is None:
             url = f"{path_to_root}aiidalab-widgets-base/notebooks/eln_configure.ipynb"
-            error_message.value = f"""Warning! The access to ELN is not configured. Please follow <a href="{url}" target="_blank">the link</a> to configure it.</br> More details: {msg}"""
+            error_message.value = f"""Warning! The access to ELN is not configured. Please follow <a href={url!r} target="_blank">the link</a> to configure it.</br> More details: {msg!r}"""
             return
 
         tl.dlink((eln, "node"), (self, "node"))

aiidalab_widgets_base/nodes.py

@@ -308,7 +308,7 @@ def to_html_string(self):
         return f"""
             <table style="border-collapse:separate;border-spacing:15px;">
             <tr>
-                <td style="width:200px"> <a href="{self.link}" target="_blank">  <img src="{self.logo}"> </a></td>
+                <td style="width:200px"> <a href={self.link!r} target="_blank">  <img src={self.logo!r}> </a></td>
                 <td style="width:800px"> <p style="font-size:16px;">{self.description} </p></td>
             </tr>
             </table>

aiidalab_widgets_base/structures.py

@@ -1,5 +1,4 @@
 """Module to provide functionality to import structures."""
-
 import datetime
 import functools
 import io
@@ -14,7 +13,7 @@
 from aiida import engine, orm, plugins
 
 # Local imports
-from .data import LigandSelectorWidget
+from .data import FunctionalGroupSelectorWidget
 from .utils import StatusHTML, exceptions, get_ase_from_file, get_formula
 from .viewers import StructureDataViewer
 
@@ -39,9 +38,7 @@ class StructureManagerWidget(ipw.VBox):
     input_structure = tl.Union(
         [tl.Instance(ase.Atoms), tl.Instance(orm.Data)], allow_none=True
     )
-    structure = tl.Union(
-        [tl.Instance(ase.Atoms), tl.Instance(orm.Data)], allow_none=True
-    )
+    structure = tl.Instance(ase.Atoms, allow_none=True)
     structure_node = tl.Instance(orm.Data, allow_none=True, read_only=True)
     node_class = tl.Unicode()
 
@@ -84,8 +81,8 @@ def __init__(
         if viewer:
             self.viewer = viewer
         else:
-            self.viewer = StructureDataViewer(**kwargs)
-        tl.dlink((self, "structure_node"), (self.viewer, "structure"))
+            self.viewer = StructureDataViewer()
+        tl.dlink((self, "structure"), (self.viewer, "structure"))
 
         # Store button.
         self.btn_store = ipw.Button(description="Store in AiiDA", disabled=True)
@@ -184,6 +181,10 @@ def _structure_editors(self, editors):
             for i, editor in enumerate(editors):
                 editors_tab.set_title(i, editor.title)
                 tl.link((editor, "structure"), (self, "structure"))
+                if editor.has_trait("input_selection"):
+                    tl.dlink(
+                        (editor, "input_selection"), (self.viewer, "input_selection")
+                    )
                 if editor.has_trait("selection"):
                     tl.link((editor, "selection"), (self.viewer, "selection"))
                 if editor.has_trait("camera_orientation"):
@@ -336,6 +337,7 @@ def _structure_changed(self, change=None):
         This function enables/disables `btn_store` widget if structure is provided/set to None.
         Also, the function sets `structure_node` trait to the selected node type.
         """
+
         if not self.structure_set_by_undo:
             self.history.append(change["new"])
 
@@ -1051,7 +1053,7 @@ def disable_element(_=None):
                 self.element.disabled = True
 
         # Ligand selection.
-        self.ligand = LigandSelectorWidget()
+        self.ligand = FunctionalGroupSelectorWidget()
         self.ligand.observe(disable_element, names="value")
 
         # Add atom.
@@ -1262,6 +1264,8 @@ def translate_dr(self, _=None, atoms=None, selection=None):
             self.action_vector * self.displacement.value
         )
 
+        self.input_selection = None  # Clear selection.
+
         self.structure, self.input_selection = atoms, selection
 
     @_register_structure
@@ -1271,7 +1275,7 @@ def translate_dxdydz(self, _=None, atoms=None, selection=None):
 
         # The action.
         atoms.positions[self.selection] += np.array(self.str2vec(self.dxyz.value))
-
+        self.input_selection = None  # Clear selection.
         self.structure, self.input_selection = atoms, selection
 
     @_register_structure
@@ -1281,7 +1285,7 @@ def translate_to_xyz(self, _=None, atoms=None, selection=None):
         # The action.
         geo_center = np.average(self.structure[self.selection].get_positions(), axis=0)
         atoms.positions[self.selection] += self.str2vec(self.dxyz.value) - geo_center
-
+        self.input_selection = None  # Clear selection.
         self.structure, self.input_selection = atoms, selection
 
     @_register_structure
@@ -1295,6 +1299,7 @@ def rotate(self, _=None, atoms=None, selection=None):
         center = self.str2vec(self.point.value)
         rotated_subset.rotate(self.phi.value, v=vec, center=center, rotate_cell=False)
         atoms.positions[self.selection] = rotated_subset.positions
+        self.input_selection = None  # Clear selection.
 
         self.structure, self.input_selection = atoms, selection
 
@@ -1329,6 +1334,8 @@ def mirror(self, _=None, norm=None, point=None, atoms=None, selection=None):
         # Mirror atoms.
         atoms.positions[selection] -= 2 * projections
 
+        self.input_selection = None  # Clear selection.
+
         self.structure, self.input_selection = atoms, selection
 
     def mirror_3p(self, _=None):
@@ -1375,6 +1382,7 @@ def mod_element(self, _=None, atoms=None, selection=None):
             initial_ligand = self.ligand.rotate(
                 align_to=self.action_vector, remove_anchor=True
             )
+
             for idx in self.selection:
                 position = self.structure.positions[idx].copy()
                 lgnd = initial_ligand.copy()
@@ -1390,22 +1398,23 @@ def mod_element(self, _=None, atoms=None, selection=None):
     @_register_selection
     def copy_sel(self, _=None, atoms=None, selection=None):
         """Copy selected atoms and shift by 1.0 A along X-axis."""
+
         last_atom = atoms.get_global_number_of_atoms()
 
         # The action
         add_atoms = atoms[self.selection].copy()
         add_atoms.translate([1.0, 0, 0])
         atoms += add_atoms
 
-        new_selection = list(range(last_atom, last_atom + len(selection)))
-        self.structure, self.input_selection = atoms, new_selection
+        self.structure, self.input_selection = atoms, list(
+            range(last_atom, last_atom + len(selection))
+        )
 
     @_register_structure
     @_register_selection
     def add(self, _=None, atoms=None, selection=None):
         """Add atoms."""
         last_atom = atoms.get_global_number_of_atoms()
-
         if self.ligand.value == 0:
             initial_ligand = ase.Atoms([ase.Atom(self.element.value, [0, 0, 0])])
             rad = SYMBOL_RADIUS[self.element.value]
@@ -1431,6 +1440,8 @@ def add(self, _=None, atoms=None, selection=None):
 
         new_selection = list(range(last_atom, last_atom + len(selection) * len(lgnd)))
 
+        # The order of the traitlets below is important -
+        # we must be sure trait atoms is set before trait selection
         self.structure, self.input_selection = atoms, new_selection
 
     @_register_structure
@@ -1439,6 +1450,7 @@ def remove(self, _=None, atoms=None, selection=None):
         """Remove selected atoms."""
         del [atoms[selection]]
 
+        # The order of the traitlets below is important -
+        # we must be sure trait atoms is set before trait selection
         self.structure = atoms
         self.input_selection = None
-        self.input_selection = []

aiidalab_widgets_base/utils/__init__.py

@@ -34,7 +34,7 @@ def predefine_settings(obj, **kwargs):
         if hasattr(obj, key):
             setattr(obj, key, value)
         else:
-            raise AttributeError(f"'{obj}' object has no attribute '{key}'")
+            raise AttributeError(f"{obj!r} object has no attribute {key!r}")
 
 
 def get_ase_from_file(fname, file_format=None):  # pylint: disable=redefined-builtin

aiidalab_widgets_base/utils/exceptions.py

@@ -3,5 +3,5 @@ class ListOrTuppleError(TypeError):
 
     def __init__(self, value):
         super().__init__(
-            f"The provided value '{value}' is not a list or a tupple, but a {type(value)}."
+            f"The provided value {value!r} is not a list or a tupple, but a {type(value)}."
         )