Add tally colorbar min/max rescaling options (#169)

Author: paulromano

openmc_plotter/docks.py

@@ -773,17 +773,12 @@ class ColorForm(QWidget):
         Selector for colormap
     dataIndicatorCheckBox : QCheckBox
         Inidcates whether or not the data indicator will appear on the colorbar
-    userMinMaxBox : QCheckBox
-        Indicates whether or not the user defined values in the min and max
-        will be used to set the bounds of the colorbar.
+    minMaxTypeBox : QComboBox
+        Dropdown to select min/max type: "Full data", "Visible data", or "Custom"
     maxBox : ScientificDoubleSpinBox
-        Max value of the colorbar. If the userMinMaxBox is checked, this will be
-        the user's input. If the userMinMaxBox is not checked, this box will
-        hold the max value of the visible data.
+        Max value of the colorbar. Only visible when minMaxTypeBox is set to "Custom".
     minBox : ScientificDoubleSpinBox
-        Min value of the colorbar. If the userMinMaxBox is checked, this will be
-        the user's input. If the userMinMaxBox is not checked, this box will
-        hold the max value of the visible data.
+        Min value of the colorbar. Only visible when minMaxTypeBox is set to "Custom".
     scaleBox : QCheckBox
         Indicates whether or not the data is displayed on a log or linear
         scale
@@ -844,10 +839,13 @@ def __init__(self, model, main_window, field, colormaps=None):
         self.dataIndicatorCheckBox.stateChanged.connect(
             data_indicator_connector)
 
-        # User specified min/max check box
-        self.userMinMaxBox = QCheckBox()
-        minmax_connector = partial(main_window.toggleTallyDataUserMinMax)
-        self.userMinMaxBox.stateChanged.connect(minmax_connector)
+        # Min/max type dropdown
+        self.minMaxTypeBox = QComboBox()
+        self.minMaxTypeBox.addItem("Full data")
+        self.minMaxTypeBox.addItem("Visible data")
+        self.minMaxTypeBox.addItem("Custom")
+        minmax_type_connector = partial(main_window.setTallyMinMaxType)
+        self.minMaxTypeBox.currentIndexChanged.connect(minmax_type_connector)
 
         # Data min spin box
         self.minBox = ScientificDoubleSpinBox()
@@ -861,6 +859,10 @@ def __init__(self, model, main_window, field, colormaps=None):
         max_connector = partial(main_window.editTallyDataMax)
         self.maxBox.valueChanged.connect(max_connector)
 
+        # Labels for min/max (so we can show/hide them)
+        self.minLabel = QLabel("Min: ")
+        self.maxLabel = QLabel("Max: ")
+
         # Linear/Log scaling check box
         self.scaleBox = QCheckBox()
         scale_connector = partial(main_window.toggleTallyLogScale)
@@ -894,9 +896,9 @@ def __init__(self, model, main_window, field, colormaps=None):
         self.layout.addRow("Colormap: ", self.colormapBox)
         self.layout.addRow("Reverse colormap: ", self.reverseCmapBox)
         self.layout.addRow("Data Indicator: ", self.dataIndicatorCheckBox)
-        self.layout.addRow("Custom Min/Max: ", self.userMinMaxBox)
-        self.layout.addRow("Min: ", self.minBox)
-        self.layout.addRow("Max: ", self.maxBox)
+        self.layout.addRow("Min/max: ", self.minMaxTypeBox)
+        self.layout.addRow(self.minLabel, self.minBox)
+        self.layout.addRow(self.maxLabel, self.maxBox)
         self.layout.addRow("Log Scale: ", self.scaleBox)
         self.layout.addRow("Clip Data: ", self.clipDataBox)
         self.layout.addRow("Mask Zeros: ", self.maskZeroBox)
@@ -914,16 +916,26 @@ def updateDataIndicator(self):
         cv = self.model.currentView
         self.dataIndicatorCheckBox.setChecked(cv.tallyDataIndicator)
 
-    def setMinMaxEnabled(self, enable):
-        enable = bool(enable)
-        self.minBox.setEnabled(enable)
-        self.maxBox.setEnabled(enable)
+    def updateMinMaxType(self):
+        """Update the min/max type dropdown and show/hide min/max inputs."""
+        cv = self.model.currentView
+        type_map = {'full': 0, 'visible': 1, 'custom': 2}
+        idx = type_map.get(cv.tallyDataMinMaxType, 0)
+        self.minMaxTypeBox.blockSignals(True)
+        self.minMaxTypeBox.setCurrentIndex(idx)
+        self.minMaxTypeBox.blockSignals(False)
+        # Show/hide min/max inputs based on whether custom is selected
+        show_custom = (cv.tallyDataMinMaxType == 'custom')
+        self.minLabel.setVisible(show_custom)
+        self.minBox.setVisible(show_custom)
+        self.maxLabel.setVisible(show_custom)
+        self.maxBox.setVisible(show_custom)
 
     def updateMinMax(self):
         cv = self.model.currentView
         self.minBox.setValue(cv.tallyDataMin)
         self.maxBox.setValue(cv.tallyDataMax)
-        self.setMinMaxEnabled(cv.tallyDataUserMinMax)
+        self.updateMinMaxType()
 
     def updateTallyVisibility(self):
         cv = self.model.currentView
@@ -951,7 +963,6 @@ def update(self):
 
         self.alphaBox.setValue(cv.tallyDataAlpha)
         self.visibilityBox.setChecked(cv.tallyDataVisible)
-        self.userMinMaxBox.setChecked(cv.tallyDataUserMinMax)
         self.scaleBox.setChecked(cv.tallyDataLogScale)
 
         self.updateMinMax()

openmc_plotter/main_window.py

@@ -1037,10 +1037,29 @@ def toggleTallyDataClip(self, state):
         av = self.model.activeView
         av.clipTallyData = bool(state)
 
-    def toggleTallyDataUserMinMax(self, state, apply=False):
+    def setTallyMinMaxType(self, index, apply=False):
+        """Set the min/max type for tally data.
+
+        Parameters
+        ----------
+        index : int
+            Index of the selected option: 0='full', 1='visible', 2='custom'
+        apply : bool
+            Whether to apply changes immediately
+        """
         av = self.model.activeView
-        av.tallyDataUserMinMax = bool(state)
-        self.tallyDock.tallyColorForm.setMinMaxEnabled(bool(state))
+        type_map = {0: 'full', 1: 'visible', 2: 'custom'}
+        new_type = type_map.get(index, 'full')
+        av.tallyDataMinMaxType = new_type
+
+        # Immediately update visibility of min/max fields based on selection
+        show_custom = (new_type == 'custom')
+        form = self.tallyDock.tallyColorForm
+        form.minLabel.setVisible(show_custom)
+        form.minBox.setVisible(show_custom)
+        form.maxLabel.setVisible(show_custom)
+        form.maxBox.setVisible(show_custom)
+
         if apply:
             self.applyChanges()
 

openmc_plotter/plotgui.py

@@ -567,7 +567,7 @@ def updatePixmap(self):
         # draw tally image
         if image_data is not None:
 
-            if not cv.tallyDataUserMinMax:
+            if cv.tallyDataMinMaxType != 'custom':
                 cv.tallyDataMin = data_min
                 cv.tallyDataMax = data_max
             else:
@@ -577,6 +577,9 @@ def updatePixmap(self):
             # always mask out negative values
             image_mask = image_data < 0.0
 
+            # mask out invalid values (NaN/Inf)
+            image_mask |= ~np.isfinite(image_data)
+
             if cv.clipTallyData:
                 image_mask |= image_data < data_min
                 image_mask |= image_data > data_max
@@ -587,6 +590,21 @@ def updatePixmap(self):
             # mask out invalid values
             image_data = np.ma.masked_where(image_mask, image_data)
 
+            # auto-rescale based on displayed (imshow) data
+            if cv.tallyDataMinMaxType == 'visible' and not cv.tallyContours:
+                displayed = image_data.compressed()
+                if displayed.size:
+                    visible_min = float(displayed.min())
+                    visible_max = float(displayed.max())
+                    # Fall back to full data range if visible range is invalid for log scale
+                    if cv.tallyDataLogScale and visible_min <= 0:
+                        pass  # keep the full data range
+                    else:
+                        data_min = visible_min
+                        data_max = visible_max
+                        cv.tallyDataMin = data_min
+                        cv.tallyDataMax = data_max
+
             if extents is None:
                 extents = data_bounds
 

openmc_plotter/plotmodel.py

@@ -937,6 +937,8 @@ class PlotViewIndependent:
         Minimum scale value for tally data
     tallyDataLogScale : bool
         Indicator of logarithmic scale for tally data
+    tallyDataMinMaxType : str
+        Type of min/max scaling for tally data: 'full', 'visible', or 'custom'
     tallyMaskZeroValues : bool
         Indicates whether or not zero values in tally data should be masked
     clipTallyData: bool
@@ -980,7 +982,7 @@ def __init__(self):
         self.tallyDataVisible = True
         self.tallyDataAlpha = 1.0
         self.tallyDataIndicator = False
-        self.tallyDataUserMinMax = False
+        self.tallyDataMinMaxType = 'full'  # 'full', 'visible', or 'custom'
         self.tallyDataMin = 0.0
         self.tallyDataMax = np.inf
         self.tallyDataLogScale = False
@@ -999,6 +1001,14 @@ def __setstate__(self, state):
             self.outlinesCell = False
         if not hasattr(self, 'outlinesMat'):
             self.outlinesMat = False
+        # Migrate old boolean attributes to new tallyDataMinMaxType
+        if not hasattr(self, 'tallyDataMinMaxType'):
+            if getattr(self, 'tallyDataUserMinMax', False):
+                self.tallyDataMinMaxType = 'custom'
+            else:
+                self.tallyDataMinMaxType = 'full'
+        # Remove old attributes if present
+        self.__dict__.pop('tallyDataUserMinMax', None)
 
     def getDataLimits(self):
         return self.data_minmax