Fix/traitlet structure fixes #381 blocked by #377

aiidalab_widgets_base/dicts.py

@@ -221,3 +221,21 @@
     "Hs": 1.75,
     "Mt": 1.75,
 }
+
+RGB_colors = {
+    "red": [1, 0, 0],
+    "green": [0, 1, 0],
+    "blue": [0, 0, 1],
+    "yellow": [1, 1, 0],
+    "cyan": [0, 1, 1],
+    "magenta": [1, 0, 1],
+    "white": [1, 1, 1],
+    "black": [0, 0, 0],
+    "grey": [0.5, 0.5, 0.5],
+    "purple": [0.5, 0, 0.5],
+    "orange": [1, 0.5, 0],
+    "brown": [0.5, 0.25, 0],
+    "pink": [1, 0.75, 0.8],
+    "lime": [0.75, 1, 0],
+    "olive": [0.5, 0.5, 0],
+}

aiidalab_widgets_base/structures.py

@@ -97,7 +97,7 @@ def __init__(
             self.viewer = viewer
         else:
             self.viewer = StructureDataViewer(**kwargs)
-        dlink((self, "structure_node"), (self.viewer, "structure"))
+        dlink((self, "structure"), (self.viewer, "structure"))
 
         # Store button.
         self.btn_store = ipw.Button(description="Store in AiiDA", disabled=True)
@@ -159,6 +159,7 @@ def _structure_importers(self, importers):
         if len(importers) == 1:
             # Assigning a function which will be called when importer provides a structure.
             dlink((importers[0], "structure"), (self, "input_structure"))
+
             return importers[0]
 
         # Otherwise making one tab per importer.
@@ -315,6 +316,7 @@ def _observe_input_structure(self, change):
         # If the `input_structure` trait is set to Atoms object, then the `structure` trait should be set to it as well.
         self.history = []
 
+        self.structure = None
         if isinstance(change["new"], Atoms):
             self.structure = change["new"]
 
@@ -330,16 +332,14 @@ def _observe_input_structure(self, change):
         elif isinstance(change["new"], StructureData):
             self.structure = change["new"].get_ase()
 
-        else:
-            self.structure = None
-
     @observe("structure")
     def _structure_changed(self, change=None):
         """Perform some operations that depend on the value of `structure` trait.
 
         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"])
 
@@ -408,6 +408,7 @@ def _validate_and_fix_ase_cell(self, ase_structure, vacuum_ang=10.0):
 
     def _on_file_upload(self, change=None):
         """When file upload button is pressed."""
+        self.structure = None
         for fname, item in change["new"].items():
             self.structure = self._read_structure(fname, item["content"])
             self.file_upload.value.clear()
@@ -801,6 +802,7 @@ def _on_button_pressed(self, change):  # pylint: disable=unused-argument
         """Convert SMILES to ase structure when button is pressed."""
         self.output.value = ""
 
+        self.structure = None
         if not self.smiles.value:
             return
         spinner = f"Screening possible conformers {self.SPINNER}"  # font-size:20em;
@@ -1182,6 +1184,12 @@ def disable_element(_=None):
             ]
         )
 
+    def find_index(self, list_of_lists, element):
+        for i, x in enumerate(list_of_lists):
+            if element in x:
+                return i
+        return -1
+
     def str2vec(self, string):
         return np.array(list(map(float, string.split())))
 
@@ -1371,6 +1379,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()
@@ -1380,7 +1389,10 @@ def mod_element(self, _=None, atoms=None, selection=None):
                 i for i in range(last_atom, last_atom + len(selection) * len(lgnd))
             ]
 
-        self.structure, self.selection = atoms, new_selection
+        # the order of the traitlets below is important
+        # we must be sure trait atoms is set BEFORE trait selection
+        self.structure = atoms
+        self.selection = new_selection
 
     @_register_structure
     @_register_selection
@@ -1395,14 +1407,16 @@ def copy_sel(self, _=None, atoms=None, selection=None):
 
         new_selection = [i for i in range(last_atom, last_atom + len(selection))]
 
-        self.structure, self.selection = atoms, new_selection
+        # the order of the traitlets below is important
+        # we must be sure trait atoms is set BEFORE trait selection
+        self.structure = atoms
+        self.selection = new_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 = Atoms([Atom(self.element.value, [0, 0, 0])])
             rad = SYMBOL_RADIUS[self.element.value]
@@ -1430,12 +1444,18 @@ def add(self, _=None, atoms=None, selection=None):
             i for i in range(last_atom, last_atom + len(selection) * len(lgnd))
         ]
 
-        self.structure, self.selection = atoms, new_selection
+        # the order of the traitlets below is important
+        # we must be sure trait atoms is set BEFORE trait selection
+        self.structure = atoms
+        self.selection = new_selection
 
     @_register_structure
     @_register_selection
     def remove(self, _, atoms=None, selection=None):
         """Remove selected atoms."""
         del [atoms[selection]]
 
-        self.structure, self.selection = atoms, list()
+        # the order of the traitlets below is important
+        # we must be sure trait atoms is set BEFORE trait selection
+        self.selection = []
+        self.structure = atoms