Adding the option to specify a universe level when coloring by cell.

Author: pshriwise

openmc_plotter/docks.py

@@ -145,6 +145,14 @@ def _createOptionsBox(self):
         self.colorbyBox.currentTextChanged[str].connect(
             self.main_window.editColorBy)
 
+        # Universe level (applies to cell coloring)
+        self.universeLevelBox = QComboBox(self)
+        self.universeLevelBox.addItem('lowest')
+        for i in range(self.model.max_universe_levels):
+            self.universeLevelBox.addItem(str(i))
+        self.universeLevelBox.currentTextChanged[str].connect(
+            self.main_window.editUniverseLevel)
+
         # Alpha
         self.domainAlphaBox = QDoubleSpinBox(self)
         self.domainAlphaBox.setValue(self.model.activeView.domainAlpha)
@@ -180,6 +188,7 @@ def _createOptionsBox(self):
         self.opLayout.addRow('Height:', self.heightBox)
         self.opLayout.addRow('Basis:', self.basisBox)
         self.opLayout.addRow('Color By:', self.colorbyBox)
+        self.opLayout.addRow('Universe Level', self.universeLevelBox)
         self.opLayout.addRow('Plot alpha:', self.domainAlphaBox)
         self.opLayout.addRow('Visible:', self.visibilityBox)
         self.opLayout.addRow('Outlines:', self.outlinesBox)
@@ -250,6 +259,13 @@ def updateHeight(self):
 
     def updateColorBy(self):
         self.colorbyBox.setCurrentText(self.model.activeView.colorby)
+        if self.model.activeView.colorby != 'cell':
+            self.universeLevelBox.setEnabled(False)
+        else:
+            self.universeLevelBox.setEnabled(True)
+
+    def updateUniverseLevel(self):
+        self.universeLevelBox.setCurrentIndex(self.model.activeView.level + 1)
 
     def updatePlotAlpha(self):
         self.domainAlphaBox.setValue(self.model.activeView.domainAlpha)

openmc_plotter/main_window.py

@@ -26,6 +26,7 @@
 from .overlays import ShortcutsOverlay
 from .tools import ExportDataDialog
 
+_COORD_LEVELS = 0
 
 def _openmcReload():
     # reset OpenMC memory, instances
@@ -35,7 +36,6 @@ def _openmcReload():
     openmc.lib.settings.output_summary = False
     openmc.lib.init(["-c"])
 
-
 class MainWindow(QMainWindow):
     def __init__(self, font=QtGui.QFontMetrics(QtGui.QFont()), screen_size=QtCore.QSize()):
         super().__init__()
@@ -627,6 +627,15 @@ def editColorBy(self, domain_kind, apply=False):
         if apply:
             self.applyChanges()
 
+    def editUniverseLevel(self, level, apply=False):
+        if level == 'lowest':
+            self.model.activeView.level = -1
+        else:
+            self.model.activeView.level = int(level)
+        self.dock.updateUniverseLevel()
+        if apply:
+            self.applyChanges()
+
     def toggleOverlaps(self, state, apply=False):
         self.model.activeView.color_overlaps = bool(state)
         self.colorDialog.updateOverlap()

openmc_plotter/plotmodel.py

@@ -96,6 +96,7 @@ def __init__(self):
         # Retrieve OpenMC Cells/Materials
         self.modelCells = openmc.lib.cells
         self.modelMaterials = openmc.lib.materials
+        self.max_universe_levels = openmc.lib._coord_levels()
 
         # Cell/Material ID by coordinates
         self.ids = None
@@ -170,7 +171,9 @@ def getDefaultView(self):
         else:
             zcenter = 0.00
 
-        default = PlotView([xcenter, ycenter, zcenter], width, height)
+        level = self.max_universe_levels
+
+        default = PlotView([xcenter, ycenter, zcenter], width, height, level)
         return default
 
     def resetColors(self):
@@ -653,13 +656,13 @@ class PlotView(openmc.lib.plot._PlotBase):
         Label of the currently selected tally
     """
 
-    def __init__(self, origin, width, height):
+    def __init__(self, origin, width, height, level):
         """ Initialize PlotView attributes """
 
         super().__init__()
 
         # View Parameters
-        self.level = -1
+        self.level = level
         self.origin = origin
         self.width = width
         self.height = height
@@ -791,7 +794,7 @@ def urc(self):
         else:
             x = self.origin[0] + self.width / 2.0
             y = self.origin[1]
-            z = self.origin[2] + height / 2.0
+            z = self.origin[2] + self.height / 2.0
         return x, y, z
 
     def adopt_plotbase(self, view):