Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history 
for more information, see https://pre-commit.ci
  • Loading branch information
@pre-commit-ci
pre-commit-ci[bot] committed Oct 19, 2022
1 parent 3d642d1 commit 56a1532
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 162 deletions.
115 changes: 72 additions & 43 deletions aiidalab_widgets_base/structures.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@
import pathlib
import tempfile
from collections import OrderedDict
from copy import deepcopy

import ase
import ipywidgets as ipw
import numpy as np
from copy import deepcopy

# spglib for cell converting
import spglib
Expand All @@ -30,7 +30,19 @@
from ase import Atom, Atoms
from ase.data import chemical_symbols, covalent_radii
from sklearn.decomposition import PCA
from traitlets import Bool,Dict, Instance, Int, List, Unicode, Union, default, dlink, link, observe
from traitlets import (
Bool,
Dict,
Instance,
Int,
List,
Unicode,
Union,
default,
dlink,
link,
observe,
)

# Local imports
from .data import LigandSelectorWidget
Expand Down Expand Up @@ -58,7 +70,7 @@ class StructureManagerWidget(ipw.VBox):
structure_node = Instance(Data, allow_none=True, read_only=True)
node_class = Unicode()
list_of_representations = List()
#brand_new_structure = Bool()
# brand_new_structure = Bool()

SUPPORTED_DATA_FORMATS = {"CifData": "cif", "StructureData": "structure"}

Expand Down Expand Up @@ -163,12 +175,11 @@ 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"))
#if importers[0].has_trait("brand_new_structure"):
# if importers[0].has_trait("brand_new_structure"):
# link((importers[0], "brand_new_structure"), (self.viewer, "brand_new_structure"))
# link((importers[0], "brand_new_structure"), (self, "brand_new_structure"))

return importers[0]


# Otherwise making one tab per importer.
importers_tab = ipw.Tab()
Expand All @@ -177,7 +188,7 @@ def _structure_importers(self, importers):
# Labeling tabs.
importers_tab.set_title(i, importer.title)
dlink((importer, "structure"), (self, "input_structure"))
#if importer.has_trait("brand_new_structure"):
# if importer.has_trait("brand_new_structure"):
# link((importer, "brand_new_structure"), (self.viewer, "brand_new_structure"))
# link((importer, "brand_new_structure"), (self, "brand_new_structure"))
return importers_tab
Expand All @@ -189,8 +200,14 @@ def _structure_editors(self, editors):
if editors[0].has_trait("selection"):
link((editors[0], "selection"), (self.viewer, "selection"))
if editors[0].has_trait("list_of_representations"):
link((editors[0], "list_of_representations"), (self.viewer, "list_of_representations"))
link((editors[0], "list_of_representations"), (self, "list_of_representations"))
link(
(editors[0], "list_of_representations"),
(self.viewer, "list_of_representations"),
)
link(
(editors[0], "list_of_representations"),
(self, "list_of_representations"),
)
if editors[0].has_trait("camera_orientation"):
dlink(
(self.viewer._viewer, "_camera_orientation"),
Expand All @@ -208,8 +225,14 @@ def _structure_editors(self, editors):
if editor.has_trait("selection"):
link((editor, "selection"), (self.viewer, "selection"))
if editor.has_trait("list_of_representations"):
link((editors, "list_of_representations"), (self.viewer, "list_of_representations"))
link((editors, "list_of_representations"), (self, "list_of_representations"))
link(
(editors, "list_of_representations"),
(self.viewer, "list_of_representations"),
)
link(
(editors, "list_of_representations"),
(self, "list_of_representations"),
)
if editor.has_trait("camera_orientation"):
dlink(
(self.viewer._viewer, "_camera_orientation"),
Expand Down Expand Up @@ -253,9 +276,9 @@ def undo(self, _):
if self.history:
self.history = self.history[:-1]
if self.history:
#self.brand_new_structure = False
# self.brand_new_structure = False
self.structure = self.history[-1][0]
#print("in undo",self.history[-1][1])
# print("in undo",self.history[-1][1])
self.list_of_representations = self.history[-1][1]
else:
self.input_structure = None
Expand Down Expand Up @@ -357,9 +380,9 @@ def _observe_input_structure(self, change):
@observe("list_of_representations")
def _observe_list_of_representations(self, change=None):
"""Update list of representations in the history list."""
#print("observe list_rep of history")
self.history[-1]=(self.history[-1][0],deepcopy(self.list_of_representations))
#print(self.history)
# print("observe list_rep of history")
self.history[-1] = (self.history[-1][0], deepcopy(self.list_of_representations))
# print(self.history)

@observe("structure")
def _structure_changed(self, change=None):
Expand All @@ -369,8 +392,8 @@ def _structure_changed(self, change=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"],deepcopy(self.list_of_representations)))
#print("history: ",self.history)
self.history.append((change["new"], deepcopy(self.list_of_representations)))
# print("history: ",self.history)

# If structure trait was set to None, structure_node should become None as well.
if self.structure is None:
Expand Down Expand Up @@ -709,7 +732,7 @@ class SmilesWidget(ipw.VBox):
"""Convert SMILES into 3D structure."""

structure = Instance(Atoms, allow_none=True)
#brand_new_structure = Bool()
# brand_new_structure = Bool()

SPINNER = """<i class="fa fa-spinner fa-pulse" style="color:red;" ></i>"""

Expand Down Expand Up @@ -1213,8 +1236,7 @@ def disable_element(_=None):
]
)


def find_index(self,list_of_lists, element):
def find_index(self, list_of_lists, element):
for i, x in enumerate(list_of_lists):
if element in x:
return i
Expand Down Expand Up @@ -1296,7 +1318,7 @@ def translate_dr(self, _=None, atoms=None, selection=None):
self.action_vector * self.displacement.value
)

#self.brand_new_structure = False
# self.brand_new_structure = False
self.structure, self.selection = atoms, selection

@_register_structure
Expand All @@ -1307,7 +1329,7 @@ def translate_dxdydz(self, _=None, atoms=None, selection=None):
# The action.
atoms.positions[self.selection] += np.array(self.str2vec(self.dxyz.value))

#self.brand_new_structure = False
# self.brand_new_structure = False

self.structure, self.selection = atoms, selection

Expand All @@ -1319,7 +1341,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.brand_new_structure = False
# self.brand_new_structure = False
self.structure, self.selection = atoms, selection

@_register_structure
Expand All @@ -1334,7 +1356,7 @@ def rotate(self, _=None, atoms=None, selection=None):
rotated_subset.rotate(self.phi.value, v=vec, center=center, rotate_cell=False)
atoms.positions[list(self.selection)] = rotated_subset.positions

#self.brand_new_structure = False
# self.brand_new_structure = False
self.structure, self.selection = atoms, selection

@_register_structure
Expand Down Expand Up @@ -1368,7 +1390,7 @@ def mirror(self, _=None, norm=None, point=None, atoms=None, selection=None):
# Mirror atoms.
atoms.positions[selection] -= 2 * projections

#self.brand_new_structure = False
# self.brand_new_structure = False
self.structure, self.selection = atoms, selection

def mirror_3p(self, _=None):
Expand All @@ -1393,7 +1415,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.brand_new_structure = False
# self.brand_new_structure = False
self.structure, self.selection = atoms, selection

@_register_structure
Expand All @@ -1417,20 +1439,22 @@ def mod_element(self, _=None, atoms=None, selection=None):
initial_ligand = self.ligand.rotate(
align_to=self.action_vector, remove_anchor=True
)
end_atom=last_atom
end_atom = last_atom
new_list_of_representations = deepcopy(self.list_of_representations)
for idx in self.selection:
position = self.structure.positions[idx].copy()
lgnd = initial_ligand.copy()
lgnd.translate(position)
atoms += lgnd
rep_of_idx = self.find_index(new_list_of_representations, idx)
new_list_of_representations[rep_of_idx]+=[i for i in range(end_atom, end_atom + len(lgnd))]
end_atom += len(lgnd)
new_list_of_representations[rep_of_idx] += [
i for i in range(end_atom, end_atom + len(lgnd))
]
end_atom += len(lgnd)
new_selection = [
i for i in range(last_atom, last_atom + len(selection) * len(lgnd))
]

self.brand_new_structure = False

# the order of the traitlets below is important
Expand All @@ -1452,13 +1476,13 @@ def copy_sel(self, _=None, atoms=None, selection=None):
add_atoms.translate([1.0, 0, 0])
atoms += add_atoms

for i,id in enumerate(selection):
for i, id in enumerate(selection):
rep_of_idx = self.find_index(new_list_of_representations, id)
new_list_of_representations[rep_of_idx]+=[last_atom+i]
new_list_of_representations[rep_of_idx] += [last_atom + i]

new_selection = [i for i in range(last_atom, last_atom + len(selection))]

#self.brand_new_structure = False
# self.brand_new_structure = False

# the order of the traitlets below is important
self.selection = []
Expand All @@ -1471,7 +1495,7 @@ def copy_sel(self, _=None, atoms=None, selection=None):
def add(self, _=None, atoms=None, selection=None):
"""Add atoms."""
last_atom = atoms.get_global_number_of_atoms()
end_atom=last_atom
end_atom = last_atom
new_list_of_representations = deepcopy(self.list_of_representations)
if self.ligand.value == 0:
initial_ligand = Atoms([Atom(self.element.value, [0, 0, 0])])
Expand All @@ -1496,14 +1520,14 @@ def add(self, _=None, atoms=None, selection=None):

atoms += lgnd
rep_of_idx = self.find_index(new_list_of_representations, idx)
new_list_of_representations[rep_of_idx]+=[i for i in range(end_atom, end_atom + len(lgnd))]
new_list_of_representations[rep_of_idx] += [
i for i in range(end_atom, end_atom + len(lgnd))
]
end_atom += len(lgnd)

new_selection = [
i for i in range(last_atom, end_atom)
]
new_selection = [i for i in range(last_atom, end_atom)]

#self.brand_new_structure = False
# self.brand_new_structure = False

# the order of the traitlets below is important
self.selection = []
Expand All @@ -1516,15 +1540,20 @@ def add(self, _=None, atoms=None, selection=None):
def remove(self, _, atoms=None, selection=None):
"""Remove selected atoms."""
del [atoms[selection]]
new_list_of_representations = [[ele for ele in sub if ele not in selection] for sub in self.list_of_representations]
new_list_of_representations = [
[ele for ele in sub if ele not in selection]
for sub in self.list_of_representations
]
# shifts atoms ids to account for preceding atoms removed
for sel_atom in selection:
for idrep in range(len(new_list_of_representations)):
for idatom in range(len(new_list_of_representations[idrep])):
if sel_atom < new_list_of_representations[idrep][idatom]:
new_list_of_representations[idrep][idatom]=new_list_of_representations[idrep][idatom]-1

#self.brand_new_structure = False
new_list_of_representations[idrep][idatom] = (
new_list_of_representations[idrep][idatom] - 1
)

# self.brand_new_structure = False

# the order of the traitlets below is important
self.selection = []
Expand Down
Loading

0 comments on commit 56a1532

Please sign in to comment.