diff --git a/src/tcviewer/mol_widget.py b/src/tcviewer/mol_widget.py index 87eff23..7bd2b8f 100644 --- a/src/tcviewer/mol_widget.py +++ b/src/tcviewer/mol_widget.py @@ -11,10 +11,12 @@ from vtkmodules.vtkFiltersCore import vtkTubeFilter from tcviewer import mol_widget +from tcviewer.settings import settings import tcutility import pyfmo from scm import plams import numpy as np +import os class MoleculeScene: @@ -30,10 +32,8 @@ def __init__(self, parent): self.renderer.AddLight(light) self.save_camera() - self.scene_assembly = vtkAssembly() self.transform = vtk.vtkTransform() self.camera_followers = [] - # self.renderer.AddActor(self.scene_assembly) def __enter__(self): return self @@ -42,6 +42,38 @@ def __exit__(self, *args): self.post_draw() pass + def screenshot(self, path: str): + self.post_draw() + img_filter = vtk.vtkWindowToImageFilter() + img_filter.SetInput(self.parent.renWin) + img_filter.SetScale(2) + # img_filter.SetMagnification(2) + img_filter.SetInputBufferTypeToRGBA() + img_filter.ReadFrontBufferOff() + img_filter.Update() + + writer = vtk.vtkPNGWriter() + writer.SetFileName(path) + writer.SetInputConnection(img_filter.GetOutputPort()) + writer.Write() + + # vtkNew windowToImageFilter; + # windowToImageFilter->SetInput(renderWindow); + # #if VTK_MAJOR_VERSION >= 8 || VTK_MAJOR_VERSION == 8 && VTK_MINOR_VERSION >= 90 + # windowToImageFilter->SetScale(2); // image quality + # #else + # windowToImageFilter->SetMagnification(2); // image quality + # #endif + # windowToImageFilter->SetInputBufferTypeToRGBA(); // also record the alpha + # // (transparency) channel + # windowToImageFilter->ReadFrontBufferOff(); // read from the back buffer + # windowToImageFilter->Update(); + + # vtkNew writer; + # writer->SetFileName("screenshot2.png"); + # writer->SetInputConnection(windowToImageFilter->GetOutputPort()); + # writer->Write(); + def save_camera(self): camera = self.renderer.GetActiveCamera() self.cam_settings = {} @@ -52,7 +84,6 @@ def save_camera(self): self.cam_settings['clipping range'] = camera.GetClippingRange() self.cam_settings['orientation'] = camera.GetOrientation() - def load_camera(self): self.renderer.ResetCamera() camera = self.renderer.GetActiveCamera() @@ -67,14 +98,20 @@ def post_draw(self): follower['actor'].RotateX(follower['rotatex']) follower['actor'].RotateY(follower['rotatey']) pos = follower['actor'].GetPosition() - T = self.scene_assembly.GetUserTransform() - follower['actor'].SetPosition(T.TransformPoint(pos)) + follower['actor'].SetPosition(self.transform.TransformPoint(pos)) + + # for actor in self.renderer.GetActors(): + # if actor in [follower['actor'] for follower in self.camera_followers]: + # continue + # actor.SetUserTransform(self.transform) - self.renderer.AddActor(self.scene_assembly) self.renderer.ResetCamera() self.parent.renWin.Render() self.parent.Initialize() self.parent.Start() + self.save_camera() + + self.post_draw = lambda: ... def draw_molecule(self, mol): for atom in mol: @@ -82,23 +119,26 @@ def draw_molecule(self, mol): mol.guess_bonds() for bond in mol.bonds: - self.draw_single_bond(bond.atom1.coords, bond.atom2.coords) + self.draw_single_bond(bond.atom1, bond.atom2) def draw_atom(self, atom): - # actors = [] def draw_disk(rotatex, rotatey): circle = vtkRegularPolygonSource() circle.SetCenter([0, 0, 0]) - circle.SetRadius(tcutility.data.atom.radius(atom.symbol)/2.4 + .02) + circle.SetRadius(tcutility.data.atom.radius(atom.symbol) * settings['atom']['size'] + settings['atom']['quadrant_width']) circle.SetNumberOfSides(50) circleMapper = vtkPolyDataMapper() circleMapper.SetInputConnection(circle.GetOutputPort()) - circleActor = vtkFollower() - circleActor.SetCamera(self.renderer.GetActiveCamera()) + if settings['atom']['quadrant_follow_camera']: + circleActor = vtkFollower() + circleActor.SetCamera(self.renderer.GetActiveCamera()) + else: + circleActor = vtkActor() circleActor.SetPosition(atom.coords) circleActor.SetMapper(circleMapper) circleActor.GetProperty().SetColor([0, 0, 0]) circleActor.PickableOff() + # circleActor.SetUserTransform(self.transform) self.camera_followers.append({'actor': circleActor, 'rotatex': rotatex, 'rotatey': rotatey}) self.renderer.AddActor(circleActor) @@ -106,7 +146,7 @@ def draw_disk(rotatex, rotatey): sphere = vtkSphereSource() sphere.SetPhiResolution(35) sphere.SetThetaResolution(45) - sphere.SetRadius(tcutility.data.atom.radius(atom.symbol)/2.4) + sphere.SetRadius(tcutility.data.atom.radius(atom.symbol) * settings['atom']['size']) sphereMapper = vtkPolyDataMapper() sphereMapper.SetInputConnection(sphere.GetOutputPort()) sphereActor = vtkActor() @@ -118,23 +158,25 @@ def draw_disk(rotatex, rotatey): sphereActor.GetProperty().SetSpecularPower(5.0) sphereActor.GetProperty().SetColor([x/255 for x in tcutility.data.atom.color(atom.symbol)]) sphereActor.type = 'atom' - self.scene_assembly.AddPart(sphereActor) - - draw_disk(0, 0) - draw_disk(-65, 0) - draw_disk(0, -65) + sphereActor.atom = atom + sphereActor.SetUserTransform(self.transform) + self.renderer.AddActor(sphereActor) - # return actors + if settings['atom']['draw_quadrants']: + draw_disk(0, 0) + draw_disk(-65, 0) + draw_disk(0, -65) - def draw_single_bond(self, p1, p2): - p1, p2 = np.array(p1), np.array(p2) + def draw_single_bond(self, a1, a2): + p1, p2 = np.array(a1.coords), np.array(a2.coords) lineSource = vtkLineSource() lineSource.SetPoint1(p1) lineSource.SetPoint2(p2) tubeFilter = vtkTubeFilter() + tubeFilter.source = lineSource tubeFilter.SetInputConnection(lineSource.GetOutputPort()) - tubeFilter.SetRadius(0.075) + tubeFilter.SetRadius(settings['bond']['radius']) tubeFilter.SetNumberOfSides(20) tubeMapper = vtkPolyDataMapper() @@ -142,13 +184,22 @@ def draw_single_bond(self, p1, p2): tubeActor = vtkActor() tubeActor.SetMapper(tubeMapper) - tubeActor.GetProperty().SetColor([0, 0, 0]) + tubeActor.GetProperty().SetColor(settings['bond']['color']) tubeActor.type = 'bond' + tubeActor.atoms = a1, a2 + tubeActor.SetUserTransform(self.transform) + self.renderer.AddActor(tubeActor) + + def remove_bond(self, a1, a2): + for actor in self.renderer.GetActors(): + if not hasattr(actor, 'type'): + continue - self.scene_assembly.AddPart(tubeActor) - # self.renderer.AddActor(tubeActor) + if actor.type != 'bond': + continue - # return tubeActor + if a1 in actor.atoms and a2 in actor.atoms: + self.renderer.RemoveActor(actor) def draw_isosurface(self, grid, isovalue=0, color=(1, 1, 0)): # vtkImageData is the vtk image volume type @@ -180,12 +231,11 @@ def draw_isosurface(self, grid, isovalue=0, color=(1, 1, 0)): actor.GetProperty().SetSpecularPower(70) actor.GetProperty().SetSpecularColor((1, 1, 1)) actor.PickableOff() + actor.type = 'surface' + actor.SetUserTransform(self.transform) - # self.renderer.AddActor(actor) - # self.post_draw() - self.scene_assembly.AddPart(actor) - - # return actor + self.renderer.AddActor(actor) + # self.Render() class MoleculeWidget(QVTKRenderWindowInteractor): @@ -202,10 +252,10 @@ def __init__(self, parent): self.SetInteractorStyle(self.interactor_style) self._base_ren = vtkRenderer() self._base_ren.SetBackground(1, 1, 1) + self._base_ren.DrawOff() self.renWin.AddRenderer(self._base_ren) self.SetRenderWindow(self.renWin) - self.eventFilter = MoleculeWidgetKeyPressFilter(parent=self) - self.installEventFilter(self.eventFilter) + self.installEventFilter(MoleculeWidgetKeyPressFilter(parent=self)) self.setAcceptDrops(True) self.Initialize() self.Start() @@ -213,13 +263,16 @@ def __init__(self, parent): self.selected_actors = [] self.selected_actor_highlights = {} self.picker = vtk.vtkPropPicker() - self.AddObserver('LeftButtonPressEvent', self.left_mousebutton_press_event) - # self.AddObserver('LeftButtonPressEvent', self.RenderCallback, -1) - self.CreateRepeatingTimer(500) + self.AddObserver('EndInteractionEvent', self.highlight_observer) + # self.GetInteractorStyle().AddObserver(vtk.vtkCommand.LeftButtonReleaseEvent, self.highlight_observer) + self.AddObserver('LeftButtonPressEvent', self.record_mouse_position) + + self._recording_mouse = False + self._mouse_pos = None - # def RenderCallback(self, event, interactor): - # self.Render() - # self.Render() + def record_mouse_position(self, interactor, event): + self._recording_mouse = True + self._mouse_pos = interactor.GetEventPosition() # The following three methods set up dragging and dropping for the app def dragEnterEvent(self, e): @@ -246,25 +299,36 @@ def dropEvent(self, e): e.accept() for url in e.mimeData().urls(): fname = str(url.toLocalFile()) - # try: self.draw_molecule(fname) - # except: - # pass else: e.ignore() # The following three methods set up dragging and dropping for the app - def left_mousebutton_press_event(self, interactor, event): - pick = self.picker.PickProp(*interactor.GetEventPosition(), list(self.renWin.GetRenderers())[0]) - if pick: - actor = self.picker.GetViewProp() - if actor in self.selected_actors: - self.remove_highlight(actor) - else: - self.add_highlight(actor) + def highlight_observer(self, interactor, event): + # check if an actor was picked by the user + pick = self.picker.PickProp(*self.GetEventPosition(), self.active_scene.renderer) + + # if we did not pick an actor we check if we clicked (not dragged) outside + if not pick: + if self._mouse_pos == self.GetEventPosition(): + self.remove_all_highlights() + return + + actor = self.picker.GetViewProp() + + # lets toggle the selected actor + self.toggle_highlight(actor) + # if the user did not hold the shift key then we deselect the other actors + if not self.GetShiftKey(): + self.remove_all_highlights(exception=actor) + + if self._recording_mouse: + self.LeftButtonReleaseEvent() + self._recording_mouse = False + def add_highlight(self, actor): - ren = list(self.renWin.GetRenderers())[0] + ren = self.active_scene.renderer self.selected_actors.append(actor) if actor.type == 'atom': actor.SetScale([0.75, 0.75, 0.75]) @@ -276,38 +340,53 @@ def add_highlight(self, actor): highlight.SetRadius(radius) elif actor.type == 'bond': - source = actor.GetMapper().GetInputConnection(0, 0).GetInputConnection(0, 0).GetProducer() + tube = actor.GetMapper().GetInputConnection(0, 0).GetProducer() + source = tube.source line = vtkLineSource() line.SetPoint1(*source.GetPoint1()) line.SetPoint2(*source.GetPoint2()) highlight = vtkTubeFilter() highlight.SetInputConnection(line.GetOutputPort()) - highlight.SetRadius(source.GetRadius()) + highlight.SetRadius(tube.GetRadius()) highlight.SetNumberOfSides(20) - source.SetRadius(source.GetRadius() / 1.5) + tube.SetRadius(tube.GetRadius() / 1.5) highlightMapper = vtkPolyDataMapper() highlightMapper.SetInputConnection(highlight.GetOutputPort()) - highlightActor = vtkFollower() - highlightActor.SetCamera(ren.GetActiveCamera()) + highlightActor = vtkActor() highlightActor.SetPosition(actor.GetPosition()) highlightActor.SetMapper(highlightMapper) highlightActor.GetProperty().SetColor([0, 1, 1]) highlightActor.GetProperty().SetOpacity(.5) highlightActor.PickableOff() + highlightActor.SetUserTransform(self.active_scene.transform) ren.AddActor(highlightActor) self.selected_actor_highlights[actor] = highlightActor def remove_highlight(self, actor): - ren = list(self.renWin.GetRenderers())[0] - self.selected_actors.remove(actor) + ren = self.active_scene.renderer + self.selected_actors = [actor_ for actor_ in self.selected_actors if actor_ != actor] highlightActor = self.selected_actor_highlights.pop(actor) actor.SetScale([1, 1, 1]) + + if actor.type == 'bond': + tube = actor.GetMapper().GetInputConnection(0, 0).GetProducer() + tube.SetRadius(tube.GetRadius() * 1.5) + ren.RemoveActor(highlightActor) - def remove_all_highlights(self): + def remove_all_highlights(self, exception=None): + exception = tcutility.ensure_list(exception) for actor in self.selected_actors: + if actor in exception: + continue + self.remove_highlight(actor) + + def toggle_highlight(self, actor): + if actor in self.selected_actors: self.remove_highlight(actor) + else: + self.add_highlight(actor) def draw_molecule(self, xyz): if xyz.endswith('.xyz'): @@ -321,45 +400,33 @@ def draw_molecule(self, xyz): # disable previous scenes [scene.renderer.DrawOff() for scene in self.scenes] - scene = MoleculeScene(self) - scene.renderer.SetActiveCamera(self._base_ren.GetActiveCamera()) - scene.draw_molecule(mol) - if orbs: - scene.draw_isosurface(tcutility.ensure_list(orbs.mos['LUMO'])[0].cube_file(), 0.03, [0, 1, 1]) - scene.draw_isosurface(tcutility.ensure_list(orbs.mos['LUMO'])[0].cube_file(), -0.03, [1, 1, 0]) - - self.scenes.append(scene) - self.set_active_mol(len(self.scenes) - 1) + with self.new_scene() as scene: + scene.renderer.SetActiveCamera(self._base_ren.GetActiveCamera()) + scene.draw_molecule(mol) + if orbs: + scene.draw_isosurface(tcutility.ensure_list(orbs.mos['LUMO'])[0].cube_file(), 0.03, [0, 1, 1]) + scene.draw_isosurface(tcutility.ensure_list(orbs.mos['LUMO'])[0].cube_file(), -0.03, [1, 1, 0]) + self.set_active_mol(-1) def new_scene(self): scene = MoleculeScene(self) scene.renderer.SetActiveCamera(self._base_ren.GetActiveCamera()) [scene.renderer.DrawOff() for scene in self.scenes] self.scenes.append(scene) - self.set_active_mol(len(self.scenes) - 1) - with scene as scene: - return scene + self.set_active_mol(-1) + return scene def next_mol(self): - if len(self.scenes) == 0: - return - - new_idx = (self.active_mol_index + 1) % len(self.scenes) - self.set_active_mol(new_idx) + self.set_active_mol(self.active_scene_index + 1) def previous_mol(self): - if len(self.scenes) == 0: - return - - new_idx = (self.active_mol_index - 1) % len(self.scenes) - self.set_active_mol(new_idx) - + self.set_active_mol(self.active_scene_index - 1) def set_active_mol(self, index): if len(self.scenes) == 0: return - self.scenes[self.active_mol_index].save_camera() + self.active_scene.save_camera() [scene.renderer.DrawOff() for scene in self.scenes] self.scenes[index].renderer.DrawOn() @@ -371,16 +438,37 @@ def set_active_mol(self, index): self.Render() self.parent.molviewslider.setMaximum(len(self.scenes) - 1) - self.parent.molviewslider.setValue(index) - # self.interactor_style.AutoAdjustCameraClippingRangeOff() - + self.parent.molviewslider.setValue(index % len(self.scenes)) @property - def active_mol_index(self): + def active_scene_index(self): if len(self.scenes) == 0: - return + return return [i for i, scene in enumerate(self.scenes) if scene.renderer.GetDraw()][0] + @property + def active_scene(self): + return self.scenes[self.active_scene_index] + + @property + def number_of_scenes(self): + return len(self.scenes) + + def screenshot(self, path): + self.active_scene.screenshot(path) + + def screenshots(self, paths=None, directory=None): + if paths is None and directory is None: + raise ValueError('You should give either the paths or directory argument') + + if paths is None: + os.makedirs(directory, exist_ok=True) + paths = [os.path.join(directory, f'scene{i}.png') for i in range(self.number_of_scenes)] + + for scene, path in zip(self.scenes, paths): + scene.screenshot(path) + + class MoleculeWidgetKeyPressFilter(QtCore.QObject): def eventFilter(self, widget, event): @@ -393,4 +481,25 @@ def eventFilter(self, widget, event): print('copy') if event.key() == QtCore.Qt.Key_V and event.modifiers() == QtCore.Qt.ControlModifier: print('paste') - return False \ No newline at end of file + + scene = self.parent().scenes[self.parent().active_scene_index] + bond_selected = len(self.parent().selected_actors) == 1 and self.parent().selected_actors[0].type == 'bond' + atoms2_selected = len(self.parent().selected_actors) == 2 and self.parent().selected_actors[0].type == 'atom' and self.parent().selected_actors[1].type == 'atom' + if atoms2_selected: + actor1, actor2 = self.parent().selected_actors + if event.key() == QtCore.Qt.Key_1: + scene.draw_single_bond(actor1.atom, actor2.atom) + self.parent().Render() + + if event.key() in [QtCore.Qt.Key_Backspace, QtCore.Qt.Key_Delete]: + scene.remove_bond(actor1.atom, actor2.atom) + self.parent().Render() + + if bond_selected: + actor = self.parent().selected_actors[0] + if event.key() in [QtCore.Qt.Key_Backspace, QtCore.Qt.Key_Delete]: + scene.remove_bond(*actor.atoms) + self.parent().remove_highlight(actor) + self.parent().Render() + + return False diff --git a/src/tcviewer/screen.py b/src/tcviewer/screen.py index 6d249f3..c7a1f4e 100644 --- a/src/tcviewer/screen.py +++ b/src/tcviewer/screen.py @@ -10,7 +10,7 @@ from vtkmodules.vtkFiltersSources import vtkLineSource, vtkSphereSource, vtkRegularPolygonSource from vtkmodules.vtkFiltersCore import vtkTubeFilter -from tcviewer import mol_widget +from tcviewer import mol_widget, settings import tcutility import pyfmo from scm import plams @@ -49,6 +49,9 @@ def __post_init__(self): self.window.layout.setColumnStretch(1, 1) self.window.layout.setColumnStretch(2, 0) + self.settings = settings.DefaultSettings() + self.window.layout.addWidget(self.settings, 0, 3, 2, 1) + def __enter__(self): self.__post_init__() @@ -64,7 +67,13 @@ def draw_molecule(self, *args, **kwargs): def add_molscene(self): return self.molview.new_scene() + def screenshot(self, *args, **kwargs): + self.molview.screenshot(*args, **kwargs) + + def screenshots(self, *args, **kwargs): + self.molview.screenshots(*args, **kwargs) + if __name__ == '__main__': with Screen() as scr: with scr.add_molscene() as scene: @@ -75,7 +84,7 @@ def add_molscene(self): mol = cub.molecule v_cx = mol.as_array()[1] - mol.as_array()[0] - T = vtk.vtkTransform() + T = scene.transform T.PostMultiply() T.Translate(*(-np.mean(mol.as_array(), axis=0)).tolist()) @@ -94,9 +103,10 @@ def add_molscene(self): T.RotateZ(angles_y[2] * 180 / np.pi) T.RotateX(7) - scene.scene_assembly.SetUserTransform(T) + # scene.scene_assembly.SetUserTransform(T) actor = scene.draw_molecule(mol) actor = scene.draw_isosurface(cub, -0.03, [1, 1, 0]) actor = scene.draw_isosurface(cub, 0.03, [0, 1, 1]) + scr.screenshots(directory='screenshots') \ No newline at end of file