Use new slice_plot function from OpenMC

openmc_plotter/main_window.py

@@ -1224,7 +1224,7 @@ def requestPlotUpdate(self, view=None):
             view_params = self.model.view_params_payload(view_snapshot)
             self.plot_manager.set_latest_view_params(view_params)
             self.plot_manager.clear_pending()
-            self.model.makePlot(view_snapshot, self.model.ids_map, self.model.properties)
+            self.model.makePlot(view_snapshot, self.model.geom_data, self.model.property_data)
             self.resetModels()
             self.showCurrentView()
             if not self.plot_manager.is_busy:
@@ -1244,10 +1244,10 @@ def _on_plot_started(self):
     def _on_plot_queued(self):
         self.plotIm.showUpdatingOverlay("Generating Plot... (update queued)")
 
-    def _on_plot_finished(self, view_snapshot, view_params, ids_map, properties):
+    def _on_plot_finished(self, view_snapshot, view_params, geom_data, property_data):
         if view_params != self.plot_manager.latest_view_params:
             return
-        self.model.makePlot(view_snapshot, ids_map, properties)
+        self.model.makePlot(view_snapshot, geom_data, property_data)
         self.resetModels()
         self.showCurrentView()
 

openmc_plotter/plotgui.py

@@ -255,8 +255,8 @@ def getIDinfo(self, event):
            and 0 <= xPos and xPos < self.model.currentView.h_res:
             id = self.model.ids[yPos, xPos]
             instance = self.model.instances[yPos, xPos]
-            temp = "{:g}".format(self.model.properties[yPos, xPos, 0])
-            density = "{:g}".format(self.model.properties[yPos, xPos, 1])
+            temp = "{:g}".format(self.model.property_data[yPos, xPos, 0])
+            density = "{:g}".format(self.model.property_data[yPos, xPos, 1])
         else:
             id = _NOT_FOUND
             instance = _NOT_FOUND
@@ -579,7 +579,7 @@ def updatePixmap(self):
             norm = SymLogNorm(
                 1E-10) if cv.color_scale_log[cv.colorby] else None
 
-            data = self.model.properties[:, :, idx]
+            data = self.model.property_data[:, :, idx]
             self.image = self.figure.subplots().imshow(data,
                                                        cmap=cmap,
                                                        norm=norm,

openmc_plotter/plotmodel.py

@@ -44,7 +44,8 @@
                     openmc.CellFilter,
                     openmc.DistribcellFilter,
                     openmc.CellInstanceFilter,
-                    openmc.MeshFilter)
+                    openmc.MeshFilter,
+                    openmc.MeshMaterialFilter)
 
 _PRODUCTIONS = ('delayed-nu-fission', 'prompt-nu-fission', 'nu-fission',
                'nu-scatter', 'H1-production', 'H2-production',
@@ -122,16 +123,24 @@ class PlotWorker(QObject):
     def generate_maps(self, work_item: PlotWorkItem):
         try:
             params = work_item.view_params
-            view_param = ViewParam(params["origin"], params["width"],
-                                   params["height"], params["h_res"])
-            view_param.h_res = params["h_res"]
-            view_param.v_res = params["v_res"]
-            view_param.basis = params["basis"]
-            view_param.level = params["level"]
-            view_param.color_overlaps = params["color_overlaps"]
-            ids_map = openmc.lib.id_map(view_param)
-            properties = openmc.lib.property_map(view_param)
-            self.finished.emit(work_item.view_params, ids_map, properties)
+
+            # Determine if we need filter bins for MeshMaterialFilter tally
+            filter_cpp = None
+            if params.get("filter_id") is not None:
+                filter_cpp = openmc.lib.filters[params["filter_id"]]
+
+            # Single call replaces id_map + property_map + get_plot_bins
+            geom_data, property_data = openmc.lib.slice_plot(
+                origin=params["origin"],
+                width=(params["width"], params["height"]),
+                basis=params["basis"],
+                pixels=(params["h_res"], params["v_res"]),
+                color_overlaps=params["color_overlaps"],
+                level=params["level"],
+                filter=filter_cpp,
+            )
+
+            self.finished.emit(work_item.view_params, geom_data, property_data)
         except Exception as exc:
             self.error.emit(str(exc))
 
@@ -224,12 +233,12 @@ def _start_next(self):
         self.work_requested.emit(work_item)
 
     @Slot(object, object, object)
-    def _on_worker_finished(self, view_params, ids_map, properties):
+    def _on_worker_finished(self, view_params, geom_data, property_data):
         request = self._in_flight_request
         self._in_flight_request = None
         if request is not None:
             self.plot_finished.emit(request.view_snapshot,
-                                    view_params, ids_map, properties)
+                                    view_params, geom_data, property_data)
         if self._pending_request is not None:
             self._start_next()
         else:
@@ -266,10 +275,10 @@ class PlotModel:
         Dictionary mapping material IDs to openmc.Material instances
     ids : NumPy int array (v_res, h_res, 1)
         Mapping of plot coordinates to cell/material ID by pixel
-    ids_map : NumPy int32 array (v_res, h_res, 3)
-        Mapping of cell and material ids
-    properties : Numpy float array (v_res, h_res, 3)
-        Mapping of cell temperatures and material densities
+    geom_data : NumPy int32 array (v_res, h_res, 3) or (v_res, h_res, 4)
+        Geometry data with cell IDs, instances, material IDs, and optionally filter bins
+    property_data : Numpy float array (v_res, h_res, 2)
+        Property data with cell temperatures and material densities
     image : NumPy int array (v_res, h_res, 3)
         The current RGB image data
     statepoint : StatePointModel
@@ -310,9 +319,9 @@ def __init__(self, use_settings_pkl, model_path, default_res):
         # Cell/Material ID by coordinates
         self.ids = None
 
-        # Return values from id_map and property_map
-        self.ids_map = None
-        self.properties = None
+        # Return values from slice_plot
+        self.geom_data = None
+        self.property_data = None
         self.map_view_params = None
 
         self.version = __version__
@@ -473,34 +482,57 @@ def view_params_payload(self, view: "PlotView"):
             "basis": str(vp.basis),
             "level": int(vp.level),
             "color_overlaps": bool(vp.color_overlaps),
+            "filter_id": self.get_active_mesh_material_filter_id(view),
         }
 
+    def get_active_mesh_material_filter_id(self, view: "PlotView") -> Optional[int]:
+        """Return the filter ID if displaying a MeshMaterialFilter tally, else None."""
+        if self._statepoint is None:
+            return None
+        if not view.tallyDataVisible or view.selectedTally is None:
+            return None
+
+        tally = self._statepoint.tallies[view.selectedTally]
+        if tally.contains_filter(openmc.MeshMaterialFilter):
+            filter = tally.find_filter(openmc.MeshMaterialFilter)
+            return filter.id
+        return None
+
     def can_reuse_maps(self, view: "PlotView"):
-        if self.ids_map is None or self.properties is None:
+        if self.geom_data is None or self.property_data is None:
             return False
         return self.map_view_params == self.view_params_payload(view)
 
     def makePlot(self, view: Optional["PlotView"] = None,
-                 ids_map=None, properties=None):
-        """ Generate new plot image from active view settings
-
-        Creates corresponding .xml files from user-chosen settings.
-        Runs OpenMC in plot mode to generate new plot image.
-        """
+                 geom_data=None, property_data=None):
+        """Generate new plot image from active view settings"""
         if view is None:
             view = self.activeView
         # update/call maps under 2 circumstances
-        #   1. this is the intial plot (ids_map/properties are None)
+        #   1. this is the intial plot (geom_data/property_data are None)
         #   2. The active (desired) view differs from the current view parameters
-        if ids_map is None or properties is None:
+        if geom_data is None or property_data is None:
             if (self.currentView.view_params != view.view_params) or \
-                (self.ids_map is None) or (self.properties is None):
-                self.ids_map = openmc.lib.id_map(view.view_params)
-                self.properties = openmc.lib.property_map(view.view_params)
+                (self.geom_data is None) or (self.property_data is None):
+                # Determine if we need filter bins for MeshMaterialFilter tally
+                filter_cpp = None
+                filter_id = self.get_active_mesh_material_filter_id(view)
+                if filter_id is not None:
+                    filter_cpp = openmc.lib.filters[filter_id]
+
+                self.geom_data, self.property_data = openmc.lib.slice_plot(
+                    origin=view.origin,
+                    width=(view.width, view.height),
+                    basis=view.basis,
+                    pixels=(view.h_res, view.v_res),
+                    color_overlaps=view.color_overlaps,
+                    level=view.level,
+                    filter=filter_cpp,
+                )
             self.map_view_params = self.view_params_payload(view)
         else:
-            self.ids_map = ids_map
-            self.properties = properties
+            self.geom_data = geom_data
+            self.property_data = property_data
             self.map_view_params = self.view_params_payload(view)
 
         # update current view
@@ -548,15 +580,15 @@ def makePlot(self, view: Optional["PlotView"] = None,
         # tally data
         self.tally_data = None
 
-        self.properties[self.properties < 0.0] = np.nan
+        self.property_data[self.property_data < 0.0] = np.nan
 
-        self.temperatures = self.properties[..., _PROPERTY_INDICES['temperature']]
-        self.densities = self.properties[..., _PROPERTY_INDICES['density']]
+        self.temperatures = self.property_data[..., _PROPERTY_INDICES['temperature']]
+        self.densities = self.property_data[..., _PROPERTY_INDICES['density']]
 
         minmax = {}
         for prop in _MODEL_PROPERTIES:
             idx = _PROPERTY_INDICES[prop]
-            prop_data = self.properties[:, :, idx]
+            prop_data = self.property_data[:, :, idx]
             minmax[prop] = (np.min(np.nan_to_num(prop_data)),
                             np.max(np.nan_to_num(prop_data)))
 
@@ -672,6 +704,13 @@ def create_tally_image(self, view: Optional[PlotView] = None):
                                                       nuclides,
                                                       view)
                 return image + (units_out,)
+
+        elif tally.contains_filter(openmc.MeshMaterialFilter):
+            image = self._create_tally_filter_image(
+                tally, tally_value, openmc.MeshMaterialFilter, scores, nuclides, view)
+            return image + (units_out,)
+
+
         elif contains_distribcell or contains_cellinstance:
             if tally_value == 'rel_err':
                 mean_data = self._create_distribcell_image(
@@ -944,17 +983,93 @@ def _do_op(array, tally_value, ax=0):
 
         return image_data, None, data_min, data_max
 
+    def _create_tally_filter_image(
+            self, tally: openmc.Tally, tally_value: TallyValueType, filter_class,
+            scores: Tuple[str], nuclides: Tuple[str], view: PlotView = None
+        ):
+        # some variables used throughout
+        if view is None:
+            view = self.currentView
+
+        def _do_op(array, tally_value, ax=0):
+            if tally_value == 'mean':
+                return np.sum(array, axis=ax)
+            elif tally_value == 'std_dev':
+                return np.sqrt(np.sum(array**2, axis=ax))
+
+        # start with reshaped data
+        data = tally.get_reshaped_data(tally_value)
+
+        # move mesh axes to the end of the filters
+        filter_idx = [type(filter) for filter in tally.filters].index(filter_class)
+        data = np.moveaxis(data, filter_idx, -1)
+
+        # sum over the rest of the tally filters
+        for tally_filter in tally.filters:
+            if type(tally_filter) is filter_class:
+                continue
+
+            selected_bins = self.appliedFilters[tally_filter]
+            if selected_bins:
+                # sum filter data for the selected bins
+                data = data[np.array(selected_bins)].sum(axis=0)
+            else:
+                # if the filter is completely unselected,
+                # set all of its data to zero and remove the axis
+                data[:] = 0.0
+                data = _do_op(data, tally_value)
+
+        # filter by selected nuclides
+        if not nuclides:
+            data = 0.0
+
+        selected_nuclides = []
+        for idx, nuclide in enumerate(tally.nuclides):
+            if nuclide in nuclides:
+                selected_nuclides.append(idx)
+        data = _do_op(data[np.array(selected_nuclides)], tally_value)
+
+        # filter by selected scores
+        if not scores:
+            data = 0.0
+
+        selected_scores = []
+        for idx, score in enumerate(tally.scores):
+            if score in scores:
+                selected_scores.append(idx)
+        data = _do_op(data[np.array(selected_scores)], tally_value)
+
+        # Extract filter bins from geom_data (computed during slice_plot call)
+        # geom_data has shape (v_res, h_res, 4) when filter was included
+        if self.geom_data.shape[2] < 4:
+            raise RuntimeError(
+                "Filter bins not available. Ensure slice_plot was called with "
+                "the appropriate filter for MeshMaterialFilter tallies."
+            )
+        bins = self.geom_data[:, :, 3]
+
+        # set image data
+        image_data = np.full_like(self.ids, np.nan, dtype=float)
+        mask = (bins >= 0)
+        image_data[mask] = data[bins[mask]]
+
+        # get dataset's min/max
+        data_min = np.min(data)
+        data_max = np.max(data)
+
+        return image_data, None, data_min, data_max
+
     @property
     def cell_ids(self):
-        return self.ids_map[:, :, 0]
+        return self.geom_data[:, :, 0]
 
     @property
     def instances(self):
-        return self.ids_map[:, :, 1]
+        return self.geom_data[:, :, 1]
 
     @property
     def mat_ids(self):
-        return self.ids_map[:, :, 2]
+        return self.geom_data[:, :, 2]
 
 
 class ViewParam(openmc.lib.plot._PlotBase):