Add support for arbitrarily named model.xml file

Author: paulromano

openmc_plotter/__main__.py

@@ -1,5 +1,3 @@
-#!/usr/bin/env python
-
 from argparse import ArgumentParser
 from pathlib import Path
 from threading import Thread
@@ -13,28 +11,28 @@
 from . import __version__
 from .main_window import MainWindow, _openmcReload
 
+
 def main():
     ap = ArgumentParser(description='OpenMC Plotter GUI')
     version_str = f'OpenMC Plotter Version: {__version__}'
     ap.add_argument('-v', '--version', action='version', version=version_str,
                     help='Display version info.')
-    ap.add_argument('-d', '--model-directory', default=os.curdir,
-                    help='Location of model dir (default is current dir)')
     ap.add_argument('-e','--ignore-settings', action='store_false',
                     help='Ignore plot_settings.pkl file if present.')
     ap.add_argument('-s', '--threads', type=int, default=None,
                     help='If present, number of threads used to generate plots.')
+    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)')
 
     args = ap.parse_args()
 
-    os.chdir(args.model_directory)
-
     run_app(args)
 
 
 def run_app(user_args):
-    path_icon = str(Path(__file__).parent / 'assets/openmc_logo.png')
-    path_splash = str(Path(__file__).parent / 'assets/splash.png')
+    path_icon = str(Path(__file__).parent / 'assets' / 'openmc_logo.png')
+    path_splash = str(Path(__file__).parent / 'assets' / 'splash.png')
 
     app = QApplication(sys.argv)
     app.setOrganizationName("OpenMC")
@@ -53,7 +51,7 @@ def run_app(user_args):
                        QtCore.Qt.AlignHCenter | QtCore.Qt.AlignBottom)
     app.processEvents()
     # load OpenMC model on another thread
-    openmc_args = {'threads': user_args.threads}
+    openmc_args = {'threads': user_args.threads, 'model_path': user_args.model_path}
     loader_thread = Thread(target=_openmcReload, kwargs=openmc_args)
     loader_thread.start()
     # while thread is working, process app events
@@ -67,7 +65,7 @@ def run_app(user_args):
 
     font_metric = QtGui.QFontMetrics(app.font())
     screen_size = app.primaryScreen().size()
-    mainWindow = MainWindow(font_metric, screen_size)
+    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()

openmc_plotter/main_window.py

@@ -1,5 +1,6 @@
 import copy
 from functools import partial
+from pathlib import Path
 import pickle
 from threading import Thread
 
@@ -26,7 +27,7 @@
 from .tools import ExportDataDialog
 
 
-def _openmcReload(threads=None):
+def _openmcReload(threads=None, model_path='.'):
     # reset OpenMC memory, instances
     openmc.lib.reset()
     openmc.lib.finalize()
@@ -35,18 +36,22 @@ def _openmcReload(threads=None):
     args = ["-c"]
     if threads is not None:
         args += ["-s", str(threads)]
+    args.append(model_path)
     openmc.lib.init(args)
     openmc.lib.settings.verbosity = 1
 
+
 class MainWindow(QMainWindow):
     def __init__(self,
                  font=QtGui.QFontMetrics(QtGui.QFont()),
-                 screen_size=QtCore.QSize()):
+                 screen_size=QtCore.QSize(),
+                 model_path='.'):
         super().__init__()
 
         self.screen = screen_size
         self.font_metric = font
         self.setWindowTitle('OpenMC Plot Explorer')
+        self.model_path = Path(model_path)
 
     def loadGui(self, use_settings_pkl=True):
 
@@ -449,7 +454,7 @@ def loadModel(self, reload=False, use_settings_pkl=True):
         if reload:
             self.resetModels()
         else:
-            self.model = PlotModel(use_settings_pkl)
+            self.model = PlotModel(use_settings_pkl, self.model_path)
 
             # update plot and model settings
             self.updateRelativeBases()
@@ -1152,7 +1157,7 @@ def saveSettings(self):
             self.model.statepoint.close()
 
         # get hashes for material.xml and geometry.xml at close
-        mat_xml_hash, geom_xml_hash = hash_model()
+        mat_xml_hash, geom_xml_hash = hash_model(self.model_path)
 
         pickle_data = {
             'version': self.model.version,
@@ -1161,7 +1166,11 @@ def saveSettings(self):
             'mat_xml_hash': mat_xml_hash,
             'geom_xml_hash': geom_xml_hash
         }
-        with open('plot_settings.pkl', 'wb') as file:
+        if self.model_path.is_file():
+            settings_pkl = self.model_path.with_name('plot_settings.pkl')
+        else:
+            settings_pkl = self.model_path / 'plot_settings.pkl'
+        with settings_pkl.open('wb') as file:
             pickle.dump(pickle_data, file)
 
     def exportTallyData(self):

openmc_plotter/plotmodel.py

@@ -61,10 +61,10 @@
                  'Rel. Error': 'rel_err'}
 
 
-def hash_file(filename):
+def hash_file(path):
     # return the md5 hash of a file
     h = hashlib.md5()
-    with open(filename,'rb') as file:
+    with path.open('rb') as file:
         chunk = 0
         while chunk != b'':
             # read 32768 bytes at a time
@@ -73,15 +73,17 @@ def hash_file(filename):
     return h.hexdigest()
 
 
-def hash_model():
+def hash_model(model_path):
     """Get hash values for materials.xml and geometry.xml (or model.xml)"""
-    # TODO: Add support for model in file other than model.xml
-    if Path('model.xml').is_file():
-        mat_xml_hash = hash_file('model.xml')
+    if model_path.is_file():
+        mat_xml_hash = hash_file(model_path)
+        geom_xml_hash = ""
+    elif (model_path / 'model.xml').exists():
+        mat_xml_hash = hash_file(model_path / 'model.xml')
         geom_xml_hash = ""
     else:
-        mat_xml_hash = hash_file('materials.xml')
-        geom_xml_hash = hash_file('geometry.xml')
+        mat_xml_hash = hash_file(model_path / 'materials.xml')
+        geom_xml_hash = hash_file(model_path / 'geometry.xml')
     return mat_xml_hash, geom_xml_hash
 
 
@@ -92,6 +94,8 @@ class PlotModel:
     ----------
     use_settings_pkl : bool
         If True, use plot_settings.pkl file to reload settings
+    model_path : pathlib.Path
+        Path to model XML file or directory
 
     Attributes
     ----------
@@ -128,7 +132,7 @@ class PlotModel:
         have unapplied changes
     """
 
-    def __init__(self, use_settings_pkl):
+    def __init__(self, use_settings_pkl, model_path):
         """ Initialize PlotModel class attributes """
 
         # Retrieve OpenMC Cells/Materials
@@ -160,8 +164,13 @@ def __init__(self, use_settings_pkl):
 
         self.defaultView = self.getDefaultView()
 
-        if use_settings_pkl and os.path.isfile('plot_settings.pkl'):
-            with open('plot_settings.pkl', 'rb') as file:
+        if model_path.is_file():
+            settings_pkl = model_path.with_name('plot_settings.pkl')
+        else:
+            settings_pkl = model_path / 'plot_settings.pkl'
+
+        if use_settings_pkl and settings_pkl.is_file():
+            with settings_pkl.open('rb') as file:
                 try:
                     data = pickle.load(file)
                 except AttributeError:
@@ -189,7 +198,7 @@ def __init__(self, use_settings_pkl):
 
                         # get materials.xml and geometry.xml hashes to
                         # restore additional settings if possible
-                        mat_xml_hash, geom_xml_hash = hash_model()
+                        mat_xml_hash, geom_xml_hash = hash_model(model_path)
                         if mat_xml_hash == data['mat_xml_hash'] and \
                             geom_xml_hash == data['geom_xml_hash']:
                             restore_domains = True