Merge pull request #119 from pshriwise/batch-mode

Batch mode

Author: kkiesling

README.md

@@ -30,7 +30,7 @@ pip install openmc-plotter
 From a directory containing an OpenMC model run:
 
 ```console
-$ openmc-plotter -d <path_to_openmc_model_dir>
+$ openmc-plotter <path_to_openmc_model_dir>
 ```
 
 or simply run
@@ -45,6 +45,15 @@ Once the viewer has opened, press `?` to view a variety of keyboard shortcuts.
 
 ![shortcuts](./screenshots/shortcuts.png)
 
+## Batch Mode
+
+Plot view (`.pltvw`) files generated in interactive sessions (see [Saving and Exporting](#saving-and-exporting))
+can be used to generate images from the command line without opening the plotter GUI.
+
+```console
+$ openmc-plotter -b view1.pltvw view1.pltvw view1.pltvw
+```
+
 ## Features
 
 ### Tally Visualization

openmc_plotter/__main__.py

@@ -24,6 +24,8 @@ def main():
     ap.add_argument('model_path', nargs='?', default=os.curdir,
                     help='Location of model XML file or a directory containing '
                     'XML files (default is current dir)')
+    ap.add_argument('-b', '--batch-mode', nargs='+', default=False,
+                    help='View files used to generate plots in batch mode')
 
     args = ap.parse_args()
 
@@ -44,7 +46,8 @@ def run_app(user_args):
     splash_pix = QtGui.QPixmap(path_splash)
     splash = QSplashScreen(splash_pix, QtCore.Qt.WindowStaysOnTopHint)
     splash.setMask(splash_pix.mask())
-    splash.show()
+    if not user_args.batch_mode:
+        splash.show()
     app.processEvents()
     splash.setMask(splash_pix.mask())
     splash.showMessage("Loading Model...",
@@ -68,9 +71,17 @@ def run_app(user_args):
     mainWindow = MainWindow(font_metric, screen_size, user_args.model_path)
     # connect splashscreen to main window, close when main window opens
     mainWindow.loadGui(use_settings_pkl=user_args.ignore_settings)
-    mainWindow.show()
+
     splash.close()
 
+    if user_args.batch_mode:
+        for view_file in user_args.batch_mode:
+            mainWindow.saveBatchImage(view_file)
+        mainWindow.close()
+        sys.exit()
+    else:
+        mainWindow.show()
+
     # connect interrupt signal to close call
     signal.signal(signal.SIGINT, lambda *args: mainWindow.close())
     # create timer that interrupts the Qt event loop

openmc_plotter/main_window.py

@@ -36,7 +36,7 @@ def _openmcReload(threads=None, model_path='.'):
     args = ["-c"]
     if threads is not None:
         args += ["-s", str(threads)]
-    args.append(model_path)
+    args.append(str(model_path))
     openmc.lib.init(args)
     openmc.lib.settings.verbosity = 1
 
@@ -159,7 +159,8 @@ def createMenuBar(self):
         self.saveImageAction.setShortcut("Ctrl+Shift+S")
         self.saveImageAction.setToolTip('Save plot image')
         self.saveImageAction.setStatusTip('Save plot image')
-        self.saveImageAction.triggered.connect(self.saveImage)
+        save_image_connector = partial(self.saveImage, filename=None)
+        self.saveImageAction.triggered.connect(save_image_connector)
 
         self.saveViewAction = QAction("Save &View...", self)
         self.saveViewAction.setShortcut(QtGui.QKeySequence.Save)
@@ -449,6 +450,21 @@ def updateWindowMenu(self):
         self.colorDialogAction.setChecked(self.colorDialog.isActiveWindow())
         self.mainWindowAction.setChecked(self.isActiveWindow())
 
+    def saveBatchImage(self, view_file):
+        """
+        Loads a view in the GUI and generates an image
+
+        Parameters
+        ----------
+        view_file : str or pathlib.Path
+            The path to a view file that is compatible with the loaded model.
+        """
+        # store the
+        cv = self.model.currentView
+        # load the view from file
+        self.loadViewFile(view_file)
+        self.plotIm.saveImage(view_file.replace('.pltvw', ''))
+
     # Menu and shared methods
     def loadModel(self, reload=False, use_settings_pkl=True):
         if reload:
@@ -472,15 +488,14 @@ def loadModel(self, reload=False, use_settings_pkl=True):
             self.plotIm.model = self.model
             self.applyChanges()
 
-    def saveImage(self):
-        filename, ext = QFileDialog.getSaveFileName(self,
-                                                    "Save Plot Image",
-                                                    "untitled",
-                                                    "Images (*.png)")
+    def saveImage(self, filename=None):
+        if filename is None:
+            filename, ext = QFileDialog.getSaveFileName(self,
+                                                        "Save Plot Image",
+                                                        "untitled",
+                                                        "Images (*.png)")
         if filename:
-            if "." not in filename:
-                filename += ".png"
-            self.plotIm.figure.savefig(filename, transparent=True)
+            self.plotIm.saveImage(filename)
             self.statusBar().showMessage('Plot Image Saved', 5000)
 
     def saveView(self):
@@ -497,26 +512,29 @@ def saveView(self):
             with open(filename, 'wb') as file:
                 pickle.dump(saved, file)
 
+    def loadViewFile(self, filename):
+        try:
+            with open(filename, 'rb') as file:
+                saved = pickle.load(file)
+        except Exception:
+            message = 'Error loading plot settings'
+            saved = {'version': None,
+                        'current': None}
+        if saved['version'] == self.model.version:
+            self.model.activeView = saved['current']
+            self.dock.updateDock()
+            self.colorDialog.updateDialogValues()
+            self.applyChanges()
+            message = '{} settings loaded'.format(filename)
+        else:
+            message = 'Error loading plot settings. Incompatible model.'
+        self.statusBar().showMessage(message, 5000)
+
     def openView(self):
         filename, ext = QFileDialog.getOpenFileName(self, "Open View Settings",
                                                     ".", "*.pltvw")
         if filename:
-            try:
-                with open(filename, 'rb') as file:
-                    saved = pickle.load(file)
-            except Exception:
-                message = 'Error loading plot settings'
-                saved = {'version': None,
-                         'current': None}
-            if saved['version'] == self.model.version:
-                self.model.activeView = saved['current']
-                self.dock.updateDock()
-                self.colorDialog.updateDialogValues()
-                self.applyChanges()
-                message = '{} settings loaded'.format(filename)
-            else:
-                message = 'Error loading plot settings. Incompatible model.'
-            self.statusBar().showMessage(message, 5000)
+            self.loadViewFile(filename)
 
     def openStatePoint(self):
         # check for an alread-open statepoint

openmc_plotter/plotgui.py

@@ -111,6 +111,18 @@ def _resize(self):
         self.resize(self.parent.width() * z,
                     self.parent.height() * z)
 
+    def saveImage(self, filename):
+        """Save an image of the current view
+
+        Parameters
+        ----------
+        filename : str or pathlib.Path
+            Name of the image to save
+        """
+        if "." not in str(filename):
+            filename += ".png"
+        self.figure.savefig(filename, transparent=True)
+
     def getDataIndices(self, event):
         cv = self.model.currentView
 

tests/setup_test/ref.png

@@ -1,9 +1,57 @@
+import filecmp
+import shutil
+
+import pytest
+
 from openmc_plotter.main_window import MainWindow, _openmcReload
 
-def test_window(qtbot):
-    _openmcReload()
+@pytest.fixture
+def run_in_tmpdir(tmpdir):
+    orig = tmpdir.chdir()
+    try:
+        yield
+    finally:
+        orig.chdir()
 
-    mw = MainWindow()
+def test_window(tmpdir, qtbot):
+    orig = tmpdir.chdir()
+    mw = MainWindow(model_path=orig)
+    _openmcReload(model_path=orig)
     mw.loadGui()
-    mw.plotIm.figure.savefig("test.png")
-    qtbot.addWidget(mw)
+
+    try:
+        mw.saveImage(tmpdir / 'test.png')
+
+        qtbot.addWidget(mw)
+    finally:
+        orig.chdir()
+
+    filecmp.cmp(orig / 'ref.png', tmpdir / 'test.png')
+
+    mw.close()
+
+def test_batch_image(tmpdir, qtbot):
+    orig = tmpdir.chdir()
+
+    # move view file into tmpdir
+    shutil.copy2(orig / 'test.pltvw', tmpdir)
+    shutil.copy2(orig / 'test1.pltvw', tmpdir)
+
+    _openmcReload(model_path=orig)
+
+    mw = MainWindow(model_path=orig)
+    mw.loadGui()
+
+    try:
+        mw.saveBatchImage('test.pltvw')
+        qtbot.addWidget(mw)
+
+        mw.saveBatchImage('test1.pltvw')
+        qtbot.addWidget(mw)
+    finally:
+        orig.chdir()
+
+    filecmp.cmp(orig / 'ref.png', tmpdir / 'test.png')
+    filecmp.cmp(orig / 'ref1.png', tmpdir / 'test1.png')
+
+    mw.close()