Added headless mode

src/tcviewer/mol_widget.py

@@ -347,11 +347,121 @@ def toggle_angle(self, a1, a2, a3):
 
 
 
+class MoleculeWidget:
+    def __new__(self, parent=None, headless=False):
+        if headless:
+            return _HeadlessMoleculeWidget()
+        else:
+            return _MoleculeWidget(parent)
+
+
+
+class _HeadlessMoleculeWidget(vtk.vtkRenderWindowInteractor):
+    def __init__(self):
+        super().__init__()
+        self.scenes = []
+
+        self.molecule_renderers = []
+        self.renWin = vtk.vtkRenderWindow()
+        self.SetRenderWindow(self.renWin)
+        # self.renWin.SetMultiSamples(4)
+        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._base_ren.DrawOff()
+        self.renWin.AddRenderer(self._base_ren)
+
+        self.Initialize()
+        # self.Start()
+
+        self._recording_mouse = False
+        self._mouse_pos = None
+
+    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]
+
+        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(-1)
+        return scene
+
+    def next_mol(self):
+        self.set_active_mol(self.active_scene_index + 1)
 
-class MoleculeWidget(QVTKRenderWindowInteractor):
+    def previous_mol(self):
+        self.set_active_mol(self.active_scene_index - 1)
+
+    def set_active_mol(self, index):
+        if len(self.scenes) == 0:
+            return
+
+        self.active_scene.save_camera()
+
+        index = index % len(self.scenes)
+        [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()
+
+    @property
+    def active_scene_index(self):
+        if len(self.scenes) == 0:
+            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 _MoleculeWidget(QVTKRenderWindowInteractor):
     def __init__(self, parent):
-        super(MoleculeWidget, self).__init__()
+        super().__init__()
         self.scenes = []
 
         self.parent = parent
@@ -367,16 +477,17 @@ def __init__(self, parent):
         self._base_ren.DrawOff()
         self.renWin.AddRenderer(self._base_ren)
         self.SetRenderWindow(self.renWin)
+
         self.installEventFilter(MoleculeWidgetKeyPressFilter(parent=self))
         self.setAcceptDrops(True)
+
         self.Initialize()
         self.Start()
 
         self.selected_actors = []
         self.selected_actor_highlights = {}
         self.picker = vtk.vtkPropPicker()
         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
@@ -579,7 +690,7 @@ def screenshots(self, paths=None, directory=None):
             raise ValueError('You should give either the paths or directory argument')
 
         if paths is None:
-            os.makedirs(directory, exist_ok=True)
+            os.makedirs(os.path.abspath(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):

src/tcviewer/screen.py

@@ -17,7 +17,40 @@
 import numpy as np
 
 
-class Screen(QtWidgets.QApplication):
+class Screen:
+    def __new__(cls, headless=False):
+        if headless:
+            return _HeadlessScreen()
+        else:
+            return _Screen()
+
+
+class _HeadlessScreen:
+    def __init__(self):
+        self.molview = mol_widget.MoleculeWidget(headless=True)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        # self.exec()
+        pass
+
+    def draw_molecule(self, *args, **kwargs):
+        self.molview.draw_molecule(*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)
+
+
+
+class _Screen(QtWidgets.QApplication):
     def __post_init__(self):
         self.window = QtWidgets.QMainWindow()
         self.window.layout = QtWidgets.QGridLayout()
@@ -74,13 +107,14 @@ def screenshots(self, *args, **kwargs):
 
 
 if __name__ == '__main__':
-    with Screen() as scr:
+    with Screen(headless=True) 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'])
+            # orbs = pyfmo.orbitals.Orbitals(res.files['adf.rkf'])
 
-            cub = orbs.mos['LUMO'].cube_file()
-            mol = cub.molecule
+            # cub = orbs.mos['LUMO'].cube_file()
+            # mol = cub.molecule
+            mol = res.molecule.input
 
             T = tcutility.geometry.MolTransform(mol)
             T.center(1)  # center on C1
@@ -90,7 +124,7 @@ def screenshots(self, *args, **kwargs):
             scene.transform = T.to_vtkTransform()
 
             actor = scene.draw_molecule(mol)
-            actor = scene.draw_isosurface(cub, -0.03, [1, .5, 0])
-            actor = scene.draw_isosurface(cub,  0.03, [0, 1, 1])
+            # actor = scene.draw_isosurface(cub, -0.03, [1, .5, 0])
+            # actor = scene.draw_isosurface(cub,  0.03, [0, 1, 1])
 
         scr.screenshots(directory='screenshots')