From 9ef8e015ac236fe2ef392cba882643e0eaceaf00 Mon Sep 17 00:00:00 2001 From: Yuman Hordijk Date: Tue, 13 Aug 2024 21:04:02 +0200 Subject: [PATCH] Current version correctly applies transforms to all actors in the scene by using vtkAssembly --- src/tcviewer/mol_widget.py | 396 +++++++++++++++++++++++++++++++++++++ src/tcviewer/screen.py | 341 ++++++++------------------------ 2 files changed, 479 insertions(+), 258 deletions(-) create mode 100644 src/tcviewer/mol_widget.py diff --git a/src/tcviewer/mol_widget.py b/src/tcviewer/mol_widget.py new file mode 100644 index 0000000..87eff23 --- /dev/null +++ b/src/tcviewer/mol_widget.py @@ -0,0 +1,396 @@ +from PySide6 import * + +import vtk +import vtkmodules.vtkRenderingOpenGL2 +from vtkmodules.util.numpy_support import numpy_to_vtk +from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor +from vtkmodules.vtkCommonColor import vtkNamedColors +from vtkmodules.vtkFiltersSources import vtkConeSource +from vtkmodules.vtkRenderingCore import vtkActor, vtkAssembly, vtkFollower, vtkPolyDataMapper, vtkRenderer, vtkRenderWindow, vtkRenderWindowInteractor, vtkLight, vtkCamera +from vtkmodules.vtkFiltersSources import vtkLineSource, vtkSphereSource, vtkRegularPolygonSource +from vtkmodules.vtkFiltersCore import vtkTubeFilter + +from tcviewer import mol_widget +import tcutility +import pyfmo +from scm import plams +import numpy as np + + +class MoleculeScene: + def __init__(self, parent): + self.parent = parent + self.renderer = vtkRenderer() + self.renderer.SetBackground([1, 1, 1]) + self.parent.renWin.AddRenderer(self.renderer) + + light = vtkLight() + light.SetPosition(-1.5, 2, 2) + light.SetLightTypeToCameraLight() + 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 + + def __exit__(self, *args): + self.post_draw() + pass + + def save_camera(self): + camera = self.renderer.GetActiveCamera() + self.cam_settings = {} + self.cam_settings['position'] = camera.GetPosition() + self.cam_settings['focal point'] = camera.GetFocalPoint() + self.cam_settings['view up'] = camera.GetViewUp() + self.cam_settings['distance'] = camera.GetDistance() + self.cam_settings['clipping range'] = camera.GetClippingRange() + self.cam_settings['orientation'] = camera.GetOrientation() + + + def load_camera(self): + self.renderer.ResetCamera() + camera = self.renderer.GetActiveCamera() + camera.SetPosition(self.cam_settings['position']) + camera.SetFocalPoint(self.cam_settings['focal point']) + camera.SetViewUp(self.cam_settings['view up']) + camera.SetDistance(self.cam_settings['distance']) + camera.SetClippingRange(self.cam_settings['clipping range']) + + def post_draw(self): + for follower in self.camera_followers: + 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)) + + self.renderer.AddActor(self.scene_assembly) + self.renderer.ResetCamera() + self.parent.renWin.Render() + self.parent.Initialize() + self.parent.Start() + + def draw_molecule(self, mol): + for atom in mol: + self.draw_atom(atom) + + mol.guess_bonds() + for bond in mol.bonds: + self.draw_single_bond(bond.atom1.coords, bond.atom2.coords) + + 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.SetNumberOfSides(50) + circleMapper = vtkPolyDataMapper() + circleMapper.SetInputConnection(circle.GetOutputPort()) + circleActor = vtkFollower() + circleActor.SetCamera(self.renderer.GetActiveCamera()) + circleActor.SetPosition(atom.coords) + circleActor.SetMapper(circleMapper) + circleActor.GetProperty().SetColor([0, 0, 0]) + circleActor.PickableOff() + + self.camera_followers.append({'actor': circleActor, 'rotatex': rotatex, 'rotatey': rotatey}) + self.renderer.AddActor(circleActor) + + sphere = vtkSphereSource() + sphere.SetPhiResolution(35) + sphere.SetThetaResolution(45) + sphere.SetRadius(tcutility.data.atom.radius(atom.symbol)/2.4) + sphereMapper = vtkPolyDataMapper() + sphereMapper.SetInputConnection(sphere.GetOutputPort()) + sphereActor = vtkActor() + sphereActor.SetMapper(sphereMapper) + sphereActor.SetPosition(atom.coords) + sphereActor.GetProperty().SetAmbient(0.65) + sphereActor.GetProperty().SetDiffuse(0.5) + sphereActor.GetProperty().SetSpecular(0.5) + 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) + + # return actors + + def draw_single_bond(self, p1, p2): + p1, p2 = np.array(p1), np.array(p2) + lineSource = vtkLineSource() + lineSource.SetPoint1(p1) + lineSource.SetPoint2(p2) + + tubeFilter = vtkTubeFilter() + tubeFilter.SetInputConnection(lineSource.GetOutputPort()) + tubeFilter.SetRadius(0.075) + tubeFilter.SetNumberOfSides(20) + + tubeMapper = vtkPolyDataMapper() + tubeMapper.SetInputConnection(tubeFilter.GetOutputPort()) + + tubeActor = vtkActor() + tubeActor.SetMapper(tubeMapper) + tubeActor.GetProperty().SetColor([0, 0, 0]) + tubeActor.type = 'bond' + + self.scene_assembly.AddPart(tubeActor) + # self.renderer.AddActor(tubeActor) + + # return tubeActor + + def draw_isosurface(self, grid, isovalue=0, color=(1, 1, 0)): + # vtkImageData is the vtk image volume type + # this is where the conversion happens + depthArray = numpy_to_vtk(grid.values.reshape(*grid.shape).ravel(order='F'), deep=True, array_type=vtk.VTK_DOUBLE) + imdata = vtk.vtkImageData() + imdata.SetDimensions(grid.shape) + imdata.SetSpacing(grid.spacing) + imdata.SetOrigin(grid.origin) + imdata.GetPointData().SetScalars(depthArray) + mcplus = vtk.vtkMarchingCubes() + mcplus.SetInputData(imdata) + mcplus.ComputeNormalsOn() + mcplus.ComputeGradientsOn() + mcplus.SetValue(0, isovalue) + mcplus.Update() + + mapper = vtkPolyDataMapper() + mapper.SetInputConnection(mcplus.GetOutputPort()) + mapper.ScalarVisibilityOff() + + actor = vtkActor() + actor.SetMapper(mapper) + actor.GetProperty().SetColor(*color) + actor.GetProperty().SetOpacity(.5) + actor.GetProperty().SetAmbient(-0) + actor.GetProperty().SetDiffuse(1) + actor.GetProperty().SetSpecular(5) + actor.GetProperty().SetSpecularPower(70) + actor.GetProperty().SetSpecularColor((1, 1, 1)) + actor.PickableOff() + + # self.renderer.AddActor(actor) + # self.post_draw() + self.scene_assembly.AddPart(actor) + + # return actor + + +class MoleculeWidget(QVTKRenderWindowInteractor): + def __init__(self, parent): + super(MoleculeWidget, self).__init__() + self.scenes = [] + + self.parent = parent + + self.molecule_renderers = [] + self.renWin = self.GetRenderWindow() + self.renWin.BordersOn() + self.interactor_style = vtk.vtkInteractorStyleTrackballCamera() + self.SetInteractorStyle(self.interactor_style) + self._base_ren = vtkRenderer() + self._base_ren.SetBackground(1, 1, 1) + self.renWin.AddRenderer(self._base_ren) + self.SetRenderWindow(self.renWin) + self.eventFilter = MoleculeWidgetKeyPressFilter(parent=self) + self.installEventFilter(self.eventFilter) + self.setAcceptDrops(True) + self.Initialize() + self.Start() + + 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) + + # def RenderCallback(self, event, interactor): + # self.Render() + # self.Render() + + # The following three methods set up dragging and dropping for the app + def dragEnterEvent(self, e): + if e.mimeData().hasUrls: + e.accept() + else: + e.ignore() + + def dragMoveEvent(self, e): + if e.mimeData().hasUrls: + e.accept() + else: + e.ignore() + + def dropEvent(self, e): + """ + Drop files directly onto the widget + File locations are stored in fname + :param e: + :return: + """ + if e.mimeData().hasUrls: + e.setDropAction(QtCore.Qt.CopyAction) + 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 add_highlight(self, actor): + ren = list(self.renWin.GetRenderers())[0] + self.selected_actors.append(actor) + if actor.type == 'atom': + actor.SetScale([0.75, 0.75, 0.75]) + highlight = vtkSphereSource() + highlight.SetPhiResolution(35) + highlight.SetThetaResolution(45) + highlight.SetCenter([0, 0, 0]) + radius = actor.GetMapper().GetInputConnection(0, 0).GetProducer().GetRadius() + highlight.SetRadius(radius) + + elif actor.type == 'bond': + source = actor.GetMapper().GetInputConnection(0, 0).GetInputConnection(0, 0).GetProducer() + line = vtkLineSource() + line.SetPoint1(*source.GetPoint1()) + line.SetPoint2(*source.GetPoint2()) + highlight = vtkTubeFilter() + highlight.SetInputConnection(line.GetOutputPort()) + highlight.SetRadius(source.GetRadius()) + highlight.SetNumberOfSides(20) + source.SetRadius(source.GetRadius() / 1.5) + + highlightMapper = vtkPolyDataMapper() + highlightMapper.SetInputConnection(highlight.GetOutputPort()) + highlightActor = vtkFollower() + highlightActor.SetCamera(ren.GetActiveCamera()) + highlightActor.SetPosition(actor.GetPosition()) + highlightActor.SetMapper(highlightMapper) + highlightActor.GetProperty().SetColor([0, 1, 1]) + highlightActor.GetProperty().SetOpacity(.5) + highlightActor.PickableOff() + 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) + highlightActor = self.selected_actor_highlights.pop(actor) + actor.SetScale([1, 1, 1]) + ren.RemoveActor(highlightActor) + + def remove_all_highlights(self): + for actor in self.selected_actors: + self.remove_highlight(actor) + + def draw_molecule(self, xyz): + if xyz.endswith('.xyz'): + mol = plams.Molecule(xyz) + orbs = False + else: + res = tcutility.results.read(xyz) + orbs = pyfmo.orbitals.Orbitals(res.files['adf.rkf']) + mol = res.molecule.output + + # 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) + + 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 + + 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) + + 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) + + + def set_active_mol(self, index): + if len(self.scenes) == 0: + return + + self.scenes[self.active_mol_index].save_camera() + + [scene.renderer.DrawOff() for scene in self.scenes] + self.scenes[index].renderer.DrawOn() + self.scenes[index].load_camera() + self.scenes[index].renderer.Render() + self.renWin.Initialize() + self.renWin.Render() + self.Initialize() + self.Render() + + self.parent.molviewslider.setMaximum(len(self.scenes) - 1) + self.parent.molviewslider.setValue(index) + # self.interactor_style.AutoAdjustCameraClippingRangeOff() + + + @property + def active_mol_index(self): + if len(self.scenes) == 0: + return + return [i for i, scene in enumerate(self.scenes) if scene.renderer.GetDraw()][0] + + +class MoleculeWidgetKeyPressFilter(QtCore.QObject): + def eventFilter(self, widget, event): + if event.type() == QtCore.QEvent.KeyPress: + if event.key() == QtCore.Qt.Key_Right: + self.parent().next_mol() + if event.key() == QtCore.Qt.Key_Left: + self.parent().previous_mol() + if event.key() == QtCore.Qt.Key_C and event.modifiers() == QtCore.Qt.ControlModifier: + 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 diff --git a/src/tcviewer/screen.py b/src/tcviewer/screen.py index 7a521c1..6d249f3 100644 --- a/src/tcviewer/screen.py +++ b/src/tcviewer/screen.py @@ -1,277 +1,102 @@ -import open3d as o3d +from PySide6 import * + +import vtk +import vtkmodules.vtkRenderingOpenGL2 +from vtkmodules.util.numpy_support import numpy_to_vtk +from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor +from vtkmodules.vtkCommonColor import vtkNamedColors +from vtkmodules.vtkFiltersSources import vtkConeSource +from vtkmodules.vtkRenderingCore import vtkActor, vtkAssembly, vtkFollower, vtkPolyDataMapper, vtkRenderer, vtkRenderWindow, vtkRenderWindowInteractor, vtkLight, vtkCamera +from vtkmodules.vtkFiltersSources import vtkLineSource, vtkSphereSource, vtkRegularPolygonSource +from vtkmodules.vtkFiltersCore import vtkTubeFilter + +from tcviewer import mol_widget +import tcutility +import pyfmo from scm import plams import numpy as np -from tcintegral import grid -import skimage -from tcviewer import materials -from tcutility import geometry, data -class Screen: - def __init__(self, **kwargs): - self.meshes = [] +class Screen(QtWidgets.QApplication): + def __post_init__(self): + self.window = QtWidgets.QMainWindow() + self.window.layout = QtWidgets.QGridLayout() + grid_widget = QtWidgets.QWidget() + grid_widget.setLayout(self.window.layout) + self.window.setCentralWidget(grid_widget) + self.window.setWindowTitle("TCViewer 2.0") - def __enter__(self): - return self - - def __exit__(self, *args, **kwargs): - o3d.visualization.draw(self.meshes, show_skybox=False, title='TCViewer', show_ui=False) - - def add_mesh(self, geometry, name=None, material=None): - self.meshes.append(dict(geometry=geometry, name=name or str(id(geometry)), material=material)) - - def draw_molecule(self, mol: str or plams.Molecule, **kwargs): - if isinstance(mol, str): - mol = plams.Molecule(mol) - mol.guess_bonds() - - for atom in mol: - sphere = o3d.geometry.TriangleMesh.create_sphere(data.atom.radius(atom.symbol)*kwargs.get('atom_scale', .5), resolution=kwargs.get('atom_resolution', 100)) - sphere.translate(atom.coords) - sphere = sphere.paint_uniform_color(np.array(data.atom.color(atom.symbol))/255) - sphere.compute_vertex_normals() - - self.add_mesh(sphere, material=kwargs.get('atom_material')) - - for bond in mol.bonds: - length = bond.length() - cylinder = o3d.geometry.TriangleMesh.create_cylinder(kwargs.get('bond_width', .07), length) - cylinder.translate((np.array(bond.atom1.coords) + np.array(bond.atom2.coords))/2) - - bond_vec = np.array(bond.atom1.coords) - np.array(bond.atom2.coords) - R = geometry.vector_align_rotmat([0, 0, 1], bond_vec) - cylinder.rotate(R) - cylinder.compute_vertex_normals() - cylinder.paint_uniform_color((0, 0, 0)) - - self.add_mesh(cylinder, material=kwargs.get('bond_material')) - - def draw_isosurface(self, gridd, isovalue=0, color=None, material=None, with_phase=True): - val = gridd.values.reshape(*gridd.shape) - if isovalue == 0: - with_phase = False - if val.min() < isovalue < val.max(): - spacing = gridd.spacing if len(gridd.spacing) == 3 else gridd.spacing * 3 - verts, faces, normals, values = skimage.measure.marching_cubes(val, isovalue, spacing=spacing, method='lorensen') - verts = o3d.cpu.pybind.utility.Vector3dVector(verts + gridd.origin) - triangles = o3d.cpu.pybind.utility.Vector3iVector(faces) - - mesh = o3d.geometry.TriangleMesh(verts, triangles) - if color is not None: - mesh.paint_uniform_color(np.atleast_2d(color)[0]) + # set up the molecule screen + self.molview = mol_widget.MoleculeWidget(self) + self.window.layout.addWidget(self.molview, 0, 0, 1, 3) - mesh.compute_vertex_normals() - self.add_mesh(mesh, material=material) + # make a scrollbar for changing the molecule + self.molviewslider = QtWidgets.QScrollBar() + self.molviewslider.setOrientation(QtCore.Qt.Horizontal) + self.molviewslider.setMinimum(0) + self.molviewslider.setMaximum(0) + self.molviewslider.valueChanged.connect(self.molview.set_active_mol) + self.window.layout.addWidget(self.molviewslider, 1, 1) - if with_phase: - self.draw_isosurface(gridd, isovalue=-isovalue, color=np.atleast_2d(color)[-1], material=material, with_phase=False) + self.molviewslider_prevbtn = QtWidgets.QPushButton('<') + self.molviewslider_prevbtn.clicked.connect(self.molview.previous_mol) + self.molviewslider_nextbtn = QtWidgets.QPushButton('>') + self.molviewslider_nextbtn.clicked.connect(self.molview.next_mol) + self.window.layout.addWidget(self.molviewslider_prevbtn, 1, 0) + self.window.layout.addWidget(self.molviewslider_nextbtn, 1, 2) - def draw_orbital(self, orb, gridd=None, isovalue=0.03, color1=[0, 0, 1], color2=[1, 0, 0], material=materials.orbital_shiny): - # define a default grid if one is not given - if gridd is None: - gridd = grid.molecule_bounding_box(orb.molecule, spacing=.1, margin=4) + self.window.layout.setColumnStretch(0, 0) + self.window.layout.setColumnStretch(1, 1) + self.window.layout.setColumnStretch(2, 0) - # evaluate the orbital on this grid - gridd.values = orb(gridd.points) - - # and draw the isosurface with phase - self.draw_isosurface(gridd, isovalue=isovalue, color=[color1, color2], material=material, with_phase=True) - - def draw_cub(self, cub: grid.Grid or str, isovalue=0.03, color1=[0, 0, 1], color2=[1, 0, 0], material='shiny'): - if isinstance(cub, str): - cub = grid.from_cub_file(cub) - if isinstance(material, str): - material = materials.orbital_material(material) - self.draw_molecule(cub.molecule) - self.draw_isosurface(cub, isovalue=isovalue, color=[color1, color2], material=material, with_phase=True) - - def draw_axes(self, center=[0, 0, 0], length=1, width=.04, **kwargs): - arrow_x = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_x.paint_uniform_color([1, 0, 0]) - arrow_x.translate(-np.array(center)) - arrow_x.rotate(geometry.vector_align_rotmat([0, 0, 1], [1, 0, 0]), -np.array(center)) - - arrow_y = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_y.paint_uniform_color([0, 1, 0]) - arrow_y.translate(-np.array(center)) - arrow_y.rotate(geometry.vector_align_rotmat([0, 0, 1], [0, 1, 0]), -np.array(center)) - - arrow_z = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_z.translate(-np.array(center)) - arrow_z.paint_uniform_color([0, 0, 1]) - - self.add_mesh(arrow_x) - self.add_mesh(arrow_y) - self.add_mesh(arrow_z) - - -class Screen2: - def __init__(self, **kwargs): - self.app = o3d.visualization.gui.Application.instance - self.app.initialize() - self.window = self.app.create_window('TCviewer', 640, 480) - self.renderer = self.window.renderer - # self.renderer.create_window() def __enter__(self): + self.__post_init__() return self - def __exit__(self, *args, **kwargs): - self.app.run() - del self.app # we have to delete this reference to fix some errors - - def add_mesh(self, geometry, name=None, material=None): - self.renderer.scene.add_geometry(dict(geometry=geometry, name=name or str(id(geometry)), material=material)) - - def draw_molecule(self, mol: str or plams.Molecule, **kwargs): - if isinstance(mol, str): - mol = plams.Molecule(mol) - mol.guess_bonds() - - for atom in mol: - sphere = o3d.geometry.TriangleMesh.create_sphere(data.atom.radius(atom.symbol)*kwargs.get('atom_scale', .5), resolution=kwargs.get('atom_resolution', 100)) - sphere.translate(atom.coords) - sphere = sphere.paint_uniform_color(np.array(data.atom.color(atom.symbol))/255) - sphere.compute_vertex_normals() + def __exit__(self, *args): + self.window.show() + self.exec() - self.add_mesh(sphere, material=kwargs.get('atom_material')) + def draw_molecule(self, *args, **kwargs): + self.molview.draw_molecule(*args, **kwargs) - for bond in mol.bonds: - length = bond.length() - cylinder = o3d.geometry.TriangleMesh.create_cylinder(kwargs.get('bond_width', .07), length) - cylinder.translate((np.array(bond.atom1.coords) + np.array(bond.atom2.coords))/2) - - bond_vec = np.array(bond.atom1.coords) - np.array(bond.atom2.coords) - R = geometry.vector_align_rotmat([0, 0, 1], bond_vec) - cylinder.rotate(R) - cylinder.compute_vertex_normals() - cylinder.paint_uniform_color((0, 0, 0)) - - self.add_mesh(cylinder, material=kwargs.get('bond_material')) - - def draw_isosurface(self, gridd, isovalue=0, color=None, material=None, with_phase=True): - val = gridd.values.reshape(*gridd.shape) - if isovalue == 0: - with_phase = False - if val.min() < isovalue < val.max(): - spacing = gridd.spacing if len(gridd.spacing) == 3 else gridd.spacing * 3 - verts, faces, normals, values = skimage.measure.marching_cubes(val, isovalue, spacing=spacing) - - edges = [] - for face in faces: - edges.extend(list(itertools.combinations(face, 2))) - g = nx.from_edgelist(edges) - - # compute connected components and print results - components = list(nx.algorithms.components.connected_components(g)) - # separate faces by component - for im, component in enumerate(components): - triangles = o3d.cpu.pybind.utility.Vector3iVector([face for face in faces if set(face) <= component]) # <= operator tests for subset relation - verts_ = o3d.cpu.pybind.utility.Vector3dVector(verts + gridd.origin) - mesh = o3d.geometry.TriangleMesh(verts_, triangles) - if color is not None: - mesh.paint_uniform_color(np.atleast_2d(color)[0]) - - mesh.compute_vertex_normals() - self.add_mesh(mesh, material=material) - - if with_phase: - self.draw_isosurface(gridd, isovalue=-isovalue, color=np.atleast_2d(color)[-1], material=material, with_phase=False) - - def draw_orbital(self, orb, gridd=None, isovalue=0.03, color1=[0, 0, 1], color2=[1, 0, 0], material=materials.orbital_shiny): - # define a default grid if one is not given - if gridd is None: - gridd = grid.molecule_bounding_box(orb.molecule, spacing=.1, margin=4) - - # evaluate the orbital on this grid - gridd.values = orb(gridd.points) - - # and draw the isosurface with phase - self.draw_isosurface(gridd, isovalue=isovalue, color=[color1, color2], material=material, with_phase=True) - - def draw_cub(self, cub: grid.Grid or str, isovalue=0.03, color1=[0, 0, 1], color2=[1, 0, 0], material=materials.orbital_shiny): - if isinstance(cub, str): - cub = grid.from_cub_file(cub) - self.draw_molecule(cub.molecule) - self.draw_isosurface(cub, isovalue=isovalue, color=[color1, color2], material=material, with_phase=True) - - def draw_axes(self, center=[0, 0, 0], length=1, width=.04, **kwargs): - arrow_x = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_x.paint_uniform_color([1, 0, 0]) - arrow_x.translate(-np.array(center)) - arrow_x.rotate(geometry.vector_align_rotmat([0, 0, 1], [1, 0, 0]), -np.array(center)) - - arrow_y = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_y.paint_uniform_color([0, 1, 0]) - arrow_y.translate(-np.array(center)) - arrow_y.rotate(geometry.vector_align_rotmat([0, 0, 1], [0, 1, 0]), -np.array(center)) - - arrow_z = o3d.geometry.TriangleMesh.create_arrow(width, width*2, cylinder_height=length, cone_height=length*.2, **kwargs) - arrow_z.translate(-np.array(center)) - arrow_z.paint_uniform_color([0, 0, 1]) - - self.add_mesh(arrow_x) - self.add_mesh(arrow_y) - self.add_mesh(arrow_z) + def add_molscene(self): + return self.molview.new_scene() if __name__ == '__main__': - from tcintegral import basis_set, MolecularOrbital - - # get a basis set from basis_sets - bs = basis_set.STO2G - - # read a molecule to use - mol = plams.Molecule('/Users/yumanhordijk/Desktop/acrolein.xyz') - mol.guess_bonds() - - # get atomic orbitals for this molecule - orb_to_get = ['1S', '2S', '3S', '1P:x', '1P:y', '1P:z'] # list of possible AO - orb_to_get = ['1P:z'] # list of possible AOs - aos = [] - for atom in mol: - for orb_name in orb_to_get: - try: # try to get the AO for this atom, it will fail if it does not exist - aos.append(bs.get(atom.symbol, orb_name, atom.coords)) - except IndexError: - pass - - # build the Hamiltonian - ionization_potentials = { # in Hartrees, obtained at the OLYP/QZ4P level in ADF - 'H(1S)': -0.238455, - 'C(1S)': -10.063764, - 'C(2S)': -0.506704, - 'C(1P:x)': -0.188651, - 'C(1P:y)': -0.188651, - 'C(1P:z)': -0.188651, - 'O(1S)': -18.9254, - 'O(2S)': -0.886887, - 'O(1P:x)': -0.329157, - 'O(1P:y)': -0.329157, - 'O(1P:z)': -0.329157, - } - - K = 1.75 - H = np.zeros((len(aos), len(aos))) - # build diagonal elements - for i, ao in enumerate(aos): - H[i, i] = ionization_potentials.get(ao.name) - - # build off-diagonal elements - for i, ao1 in enumerate(aos): - for j, ao2 in enumerate(aos[i+1:], start=i+1): - H[i, j] = H[j, i] = K * ao1.overlap(ao2) * (H[i, i] + H[j, j]) / 2 - - # get MO energies and coeffs - energies, coefficients = np.linalg.eigh(H) - - # we are now going to draw each MO - for mo_index in range(len(energies)): - # construct the MO using our AOs and coefficients - mo = MolecularOrbital(aos, coefficients[:, mo_index], mol) - # create a new screen - with Screen() as scr: - scr.draw_molecule(mol) - scr.draw_orbital(mo, material=materials.orbital_shiny, isovalue=.03) + with Screen() as scr: + with scr.add_molscene() as scene: + res = tcutility.results.read('/Users/yumanhordijk/PhD/Projects/RadicalAdditionASMEDA/data/DFT/TS_C_O/PyFrag_OLYP_TZ2P/frag_Substrate') + orbs = pyfmo.orbitals.Orbitals(res.files['adf.rkf']) + + cub = orbs.mos['LUMO'].cube_file() + mol = cub.molecule + v_cx = mol.as_array()[1] - mol.as_array()[0] + + T = vtk.vtkTransform() + T.PostMultiply() + T.Translate(*(-np.mean(mol.as_array(), axis=0)).tolist()) + + R_x = tcutility.geometry.vector_align_rotmat(v_cx, [1, 0, 0]) + angles_x = tcutility.geometry.rotmat_to_angles(R_x) + T.RotateX(angles_x[0] * 180 / np.pi) + T.RotateY(angles_x[1] * 180 / np.pi) + T.RotateZ(angles_x[2] * 180 / np.pi) + + mol_ = tcutility.geometry.apply_rotmat(mol.as_array(), R_x) + n_cxh = np.cross(mol_[1] - mol_[0], mol_[-1] - mol_[0]) + R_y = tcutility.geometry.vector_align_rotmat(n_cxh, [0, 1, 0]) + angles_y = tcutility.geometry.rotmat_to_angles(R_y) + T.RotateX(angles_y[0] * 180 / np.pi) + T.RotateY(angles_y[1] * 180 / np.pi) + T.RotateZ(angles_y[2] * 180 / np.pi) + + T.RotateX(7) + 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.draw_axes()