extra checks for pixels (#4009)

Co-authored-by: Jonathan Shimwell <jon@proximafusion.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: 3 people

openmc/model/model.py

@@ -34,6 +34,15 @@ def __call__(self, val: float, **kwargs: Any) -> None:
         ...
 
 
+def _check_pixels(pixels: int | Sequence[int]) -> None:
+    if isinstance(pixels, Integral):
+        check_greater_than('pixels', pixels, 0)
+    else:
+        check_length('pixels', pixels, 2)
+        for p in pixels:
+            check_greater_than('pixels', p, 0)
+
+
 class Model:
     """Model container.
 
@@ -1051,12 +1060,7 @@ def _set_plot_defaults(
         pixels: int | Sequence[int],
         basis: str
     ):
-        if isinstance(pixels, int):
-            check_greater_than('pixels', pixels, 0)
-        else:
-            check_length('pixels', pixels, 2)
-            for p in pixels:
-                check_greater_than('pixels', p, 0)
+        _check_pixels(pixels)
 
         x, y, _ = _BASIS_INDICES[basis]
 
@@ -1220,6 +1224,8 @@ def slice_data(
         """
         import openmc.lib
 
+        _check_pixels(pixels)
+
         if width is not None and (u_span is not None or v_span is not None):
             raise ValueError("width is mutually exclusive with u_span/v_span.")
 

openmc/plots.py

@@ -383,36 +383,36 @@ def id_map_to_rgb(
     """
     # Initialize RGB array with white background (values between 0 and 1 for matplotlib)
     img = np.ones(id_map.shape, dtype=float)
-    
+
     # Get the appropriate index based on color_by
     if color_by == 'cell':
         id_index = 0  # Cell IDs are in the first channel
     elif color_by == 'material':
         id_index = 2  # Material IDs are in the third channel
     else:
         raise ValueError("color_by must be either 'cell' or 'material'")
-    
+
     # Get all unique IDs in the plot
     unique_ids = np.unique(id_map[:, :, id_index])
-    
+
     # Generate default colors if not provided
     if colors is None:
         colors = {}
-    
+
     # Convert colors dict to use IDs as keys
     color_map = {}
     for key, color in colors.items():
         if isinstance(key, (openmc.Cell, openmc.Material)):
             color_map[key.id] = color
         else:
             color_map[key] = color
-    
+
     # Generate random colors for IDs not in color_map
     rng = np.random.RandomState(1)
     for uid in unique_ids:
         if uid > 0 and uid not in color_map:
             color_map[uid] = rng.randint(0, 256, (3,))
-    
+
     # Apply colors to each pixel
     for uid in unique_ids:
         if uid == -1:  # Background/void
@@ -432,7 +432,7 @@ def id_map_to_rgb(
                 rgb = color
             mask = id_map[:, :, id_index] == uid
             img[mask] = np.array(rgb) / 255.0
-    
+
     return img
 
 class PlotBase(IDManagerMixin):

tests/unit_tests/test_model.py

@@ -679,6 +679,8 @@ def test_model_plot_invalid_inputs():
         model.plot(pixels=(0, 100))
     with pytest.raises(ValueError):
         model.plot(pixels=(100,))
+    with pytest.raises(ValueError):
+        model.slice_data(u_span=(2, 0, 0), v_span=(0, 2, 0), pixels=-1)
 
 
 def test_model_id_map_initialization(run_in_tmpdir):