Source code for moldesign.widgets.trajectory

# 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 time

import ipywidgets as ipy

import moldesign as mdt

from moldesign.uibase.components import AtomInspector
from moldesign.uibase import selector

[docs]class TrajectoryViewer(selector.SelectionGroup): """ 3D representation, with animation controls, for a trajectory. Users will typically instantiate this using ``trajectory.draw()`` Args: display (bool): immediately display this to the notebook (default: False) **kwargs (dict): keyword arguments for :class:`ipywidgets.Box` """ def __init__(self, trajectory, display=False, **kwargs): from IPython.display import display as displaynow self.default_fps = 10 self.traj = trajectory self.pane = ipy.VBox() trajectory.apply_frame(trajectory.frames[0]) self.viewer, self.view_container = self.make_viewer() for frame in self.traj.frames[1:]: self.viewer.append_frame(positions=frame.positions, wfn=frame.get('wfn', None), render=False) self.make_controls() self.pane.children = [self.view_container, self.controls] super(TrajectoryViewer, self).__init__([self.pane, AtomInspector()], **kwargs) self.update_selections('initialization', {'framenum': 0}) if display: displaynow(self)
[docs] def make_viewer(self): viewer = self.traj._tempmol.draw3d(style='licorice') viewer.show_unbonded() return viewer, viewer
[docs] def make_controls(self): controls = [] self.play_button = ipy.Button(description=u"\u25B6") self.text_view = FrameInspector(self.traj) controls.append(self.text_view) self.play_button.on_click(self._play) self.slider = selector.create_value_selector(ipy.IntSlider, value_selects='framenum', value=0, description='Frame:', min=0, max=len(self.traj)-1) self.playbox = ipy.HBox([self.play_button, self.slider]) controls.append(self.playbox) self.controls = ipy.VBox(controls)
def _play(self, *args, **kwargs): self.animate()
[docs] def animate(self, fps=None): fps = mdt.utils.if_not_none(fps, self.default_fps) self.slider.value = 0 spf = 1.0 / fps for i in xrange(self.viewer.num_frames): t0 = time.time() self.slider.value = i dt = time.time() - t0 if dt < spf: time.sleep(spf-dt)
def __getattr__(self, item): """Users can run viz commands directly on the trajectory, e.g., trajectory.cpk(atoms=mol.chains['B'])""" # TODO: modify __dir__ to match return getattr(self.viewer, item)
[docs]class TrajectoryOrbViewer(TrajectoryViewer):
[docs] def make_viewer(self): viewframe = self.traj._tempmol.draw_orbitals() return viewframe.viewer, viewframe
[docs]class FrameInspector(ipy.HTML, selector.Selector): def __init__(self, traj, **kwargs): self.traj = traj super(FrameInspector, self).__init__(**kwargs)
[docs] def handle_selection_event(self, selection): if 'framenum' not in selection: return else: framenum = selection['framenum'] if hasattr(self.traj[framenum], 'time') and self.traj[framenum].time is not None: result = 'Time: %s; ' % self.traj[framenum].time.defunits() else: result = '' try: result += self.traj[framenum].annotation except (KeyError, AttributeError): pass self.value = result