Source code for moldesign.widgets.geombuilder

# Copyright 2016 Autodesk Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
import ipywidgets as ipy

import moldesign as mdt
from moldesign import utils
from moldesign.viewer import BondClicker
from moldesign import units as u

from moldesign.uibase import ViewerToolBase, ReadoutFloatSlider

def exports(o):
    return o
__all__ = []

[docs]class GeometryBuilder(ViewerToolBase): VIEWERTYPE = BondClicker MAXDIST = 20.0 # TODO: we need to set this dynamically NBR2HIGHLIGHT = '#C5AED8' NBR1HIGHLIGHT = '#AFC6A8' HIGHLIGHTOPACITY = 0.6 POSFMT = u'{:.3f} \u212B' DEGFMT = u'{:.1f}\u00B0' def __init__(self, mol): super(GeometryBuilder, self).__init__(mol) # All numbers here are assumed angstroms and radians for now ... self._selection = utils.DotDict(blank=True, type=None) self._highlighted_bonds = [] self._highlighted_atoms = [] self.original_position = self.mol.positions.copy() self.clear_button = ipy.Button(description='Clear selection') self.clear_button.on_click(self.clear_selection) self.label_box = ipy.Checkbox(description='Label atoms', value=False) self.label_box.observe(self.label_atoms, 'value') # Viewer self.viewer.atom_callbacks.append(self.atom_click) self.viewer.bond_callbacks.append(self.bond_click) self.selection_description = ipy.HTML() self.subtools.children = (ipy.HBox([self.clear_button, self.label_box]), self.selection_description) # Atom manipulation tools self.x_slider = ReadoutFloatSlider(min=-self.MAXDIST, max=self.MAXDIST, description='<b>x</b>', format=self.POSFMT) self.x_slider.observe(self.set_atom_x, 'value') self.y_slider = ReadoutFloatSlider(min=-self.MAXDIST, max=self.MAXDIST, description='<b>y</b>', format=self.POSFMT) self.y_slider.observe(self.set_atom_y, 'value') self.z_slider = ReadoutFloatSlider(min=-self.MAXDIST, max=self.MAXDIST, description='<b>z</b>', format=self.POSFMT) self.z_slider.observe(self.set_atom_z, 'value') # Bond manipulation tools self.adjust_button = ipy.Checkbox(description='Adjust entire molecule', align='end', value=True) self.length_slider = ReadoutFloatSlider(min=0.1, max=self.MAXDIST, format=self.POSFMT) self.length_slider.observe(self.set_distance, 'value') self.angle_slider = ReadoutFloatSlider(min=1.0, max=179.0, step=2.0, format=self.DEGFMT) self.angle_slider.observe(self.set_angle, 'value') self.dihedral_slider = ReadoutFloatSlider(min=-90.0, max=360.0, step=4.0, format=self.DEGFMT) self.dihedral_slider.observe(self.set_dihedral, 'value') self.bond_tools = ipy.VBox((self.adjust_button, self.length_slider, self.angle_slider, self.dihedral_slider)) self.atom_tools = ipy.VBox((self.adjust_button, self.x_slider, self.y_slider, self.z_slider)) self.reset_button = ipy.Button(description='Reset geometry') self.reset_button.on_click(self.reset_geometry) self.tool_holder = ipy.VBox() self.toolpane.children = (self.tool_holder, self.reset_button)
[docs] def set_distance(self, *args): sel = self._selection assert sel.type == 'bond' dist_in_angstrom = self.length_slider.value mdt.set_distance(sel.a1, sel.a2, dist_in_angstrom*u.angstrom, adjustmol=self.adjust_button.value) self.viewer.set_positions()
[docs] def set_angle(self, *args): sel = self._selection assert sel.type == 'bond' angle = self.angle_slider.value mdt.set_angle(sel.a1, sel.a2, sel.nbr_a2, angle*u.pi/180.0, adjustmol=self.adjust_button.value) self.viewer.set_positions()
[docs] def set_dihedral(self, *args): sel = self._selection assert sel.type == 'bond' angle = self.dihedral_slider.value mdt.set_dihedral(sel.nbr_a1, sel.a1, sel.a2, sel.nbr_a2, angle*u.pi/180.0, adjustmol=self.adjust_button.value) self.viewer.set_positions()
[docs] def set_atom_x(self, *args): pass
[docs] def set_atom_y(self, *args): pass
[docs] def set_atom_z(self, *args): pass
[docs] def label_atoms(self, *args): if self.label_box.value: self.viewer.label_atoms() else: self.viewer.remove_all_labels()
[docs] def atom_click(self, atom): sel = self._selection if sel.blank: # select this atom sel.blank = False sel.type = 'atom' sel.atom = atom elif sel.type == 'atom': # We've selected 2 atoms - i.e. a bond if atom is sel.atom: # clicked twice -> deselect the thing return self.clear_selection() elif atom in sel.atom.bond_graph: # select the bond return self.bond_click(mdt.Bond(sel.atom, atom)) # turn this into a bond selection else: # select a new atom self.clear_selection(render=False) sel = self._selection sel.blank = False sel.type = 'atom' sel.atom = atom elif sel.type == 'bond': if atom in sel.a1_neighbors: # change the neighboring selection sel.nbr_a1 = atom elif atom in sel.a2_neighbors: sel.nbr_a2 = atom else: # select a new atom self.clear_selection(render=False) return self.atom_click(atom) self._redraw_selection()
def _set_tool_state(self): # start with everything disabled for tool in self.atom_tools.children + self.bond_tools.children: tool.disabled = True if self._selection.blank: self.tool_holder.children = (ipy.HTML('Please click on a bond or atom'),) elif self._selection.type == 'atom': self.adjust_button.disabled = False self.tool_holder.children = (self.atom_tools,) x, y, z = self._selection.atom.position.value_in(u.angstrom) self.x_slider.value = x self.x_slider.disabled = False # for now self.y_slider.value = y self.y_slider.disabled = False # for now self.z_slider.value = z self.z_slider.disabled = False # for now for tool in self.atom_tools.children: tool.disabled = True elif self._selection.type == 'bond': sel = self._selection self.adjust_button.disabled = False # bond length self.length_slider.value = sel.a1.distance(sel.a2).value_in(u.angstrom) self.length_slider.disabled = False self.length_slider.description = '<b>Bond distance</b> <span style="color:{c1}">{}' \ ' - {}</span>'.format( a1=sel.a1, a2=sel.a2, c1=self.viewer.HIGHLIGHT_COLOR) # Bond angle if sel.nbr_a2: self.angle_slider.value = mdt.angle(sel.a1, sel.a2, sel.nbr_a2).value_in(u.degrees) # self.angle_slider.observe(self.set_angle, 'value') self.angle_slider.disabled = False self.angle_slider.description = '<b>Bond angle</b> <span style="color:{c1}">{}' \ ' - {}</span> ' \ '- <span style="color:{c2}">{}</span>'.format( a1=sel.a1, a2=sel.a2, a3=sel.nbr_a2, c1=self.viewer.HIGHLIGHT_COLOR, c2=self.NBR2HIGHLIGHT) else: self.angle_slider.description = 'no angle associated with this bond' # self.angle_slider.unobserve(self.set_angle) self.angle_slider.disabled = True # Dihedral twist if sel.nbr_a2 and sel.nbr_a1: self.dihedral_slider.value = mdt.dihedral(sel.nbr_a1, sel.a1, sel.a2, sel.nbr_a2).value_in(u.degrees) # self.dihedral_slider.observe(self.set_dihedral, 'value') self.dihedral_slider.disabled = False self.dihedral_slider.description = '<b>Dihedral angle</b> <span style="color:{c0}">{}</span>' \ ' - <span style="color:{c1}">{}' \ ' - {}</span> ' \ '- <span style="color:{c2}">{}</span>'.format( a4=sel.nbr_a1, a1=sel.a1, a2=sel.a2, a3=sel.nbr_a2, c0=self.NBR1HIGHLIGHT, c1=self.viewer.HIGHLIGHT_COLOR, c2=self.NBR2HIGHLIGHT) else: self.dihedral_slider.description = 'not a torsion bond' # self.dihedral_slider.unobserve(self.set_dihedral) self.dihedral_slider.disabled = True self.tool_holder.children = [self.bond_tools] else: raise ValueError('Unknown selection type %s' % self._selection.type)
[docs] def bond_click(self, bond): sel = self._selection if sel.type == 'bond': # check if this bond is already selected a1, a2 = bond.a1, bond.a2 if (a1 is sel.a1 and a2 is sel.a2) or (a1 is sel.a2 and a2 is sel.a1): return self.clear_selection() self.clear_selection(render=False) sel = self._selection sel.blank = False sel.type = 'bond' = bond sel.a1 = bond.a1 sel.a2 = bond.a2 sel.a1_neighbors = set([a for a in bond.a1.bond_graph if a is not bond.a2]) sel.a2_neighbors = set([a for a in bond.a2.bond_graph if a is not bond.a1]) sel.nbr_a1 = sel.nbr_a2 = None if sel.a1_neighbors: sel.nbr_a1 = max(sel.a1_neighbors, key=lambda x: x.mass) if sel.a2_neighbors: sel.nbr_a2 = max(sel.a2_neighbors, key=lambda x: x.mass) self._redraw_selection()
def _highlight_atoms(self, atoms, color=None, render=True): color = utils.if_not_none(color, self.viewer.HIGHLIGHT_COLOR) self._highlighted_atoms += atoms self.viewer.add_style('vdw', atoms=atoms, radius=self.viewer.ATOMRADIUS * 1.1, color=color, opacity=self.HIGHLIGHTOPACITY, render=render) def _unhighlight_atoms(self, atoms, render=True): self.viewer.set_style('vdw', atoms=atoms, radius=self.viewer.ATOMRADIUS, render=render) def _redraw_selection(self): # unhighlight any previous selections if self._highlighted_atoms: self._unhighlight_atoms(self._highlighted_atoms, render=False) self._highlighted_atoms = [] for bond in self._highlighted_bonds: self.viewer.unset_bond_color(bond, render=False) self._highlighted_bonds = [] # Set the selection view sel = self._selection if sel.type == 'atom': self._highlight_atoms([sel.atom], render=False) self.selection_description.value = \ u"<b>Atom</b> {} at coordinates " \ u"x:{p[0]:.3f}, y:{p[1]:.3f}, z:{p[2]:.3f} \u212B".format( atom=sel.atom, p=sel.atom.position.value_in(u.angstrom)) elif sel.type == 'bond': self.selection_description.value = "<b>Bond:</b> %s - %s" % (, self._highlighted_bonds = [] self.viewer.set_bond_color(self.viewer.HIGHLIGHT_COLOR,, render=False) self._highlight_atoms([sel.a1, sel.a2], render=False) if sel.nbr_a1 is not None: nmdtond = mdt.Bond(sel.a1, sel.nbr_a1) self._highlight_atoms([sel.nbr_a1], color=self.NBR1HIGHLIGHT, render=False) self.viewer.set_bond_color(self.NBR1HIGHLIGHT, nmdtond, render=False) self._highlighted_bonds.append(nmdtond) if sel.nbr_a2 is not None: nmdtond = mdt.Bond(sel.a2, sel.nbr_a2) self._highlight_atoms([sel.nbr_a2], color=self.NBR2HIGHLIGHT, render=False) self.viewer.set_bond_color(self.NBR2HIGHLIGHT, nmdtond, render=False) self._highlighted_bonds.append(nmdtond) elif sel.type is not None: raise ValueError('Unknown selection type %s' % self._selection.type) self.viewer.render() self._set_tool_state()
[docs] def clear_selection(self, render=True, *args): self._selection = utils.DotDict(blank=True, type=None) self.selection_description.value = "" if render: self._redraw_selection()
[docs] def reset_geometry(self, *args): self.clear_selection(render=False) self.mol.positions = self.original_position self.viewer.set_positions() self._redraw_selection()