diff --git a/openmc_plotter/main_window.py b/openmc_plotter/main_window.py index 2ae8e20..0bc738f 100755 --- a/openmc_plotter/main_window.py +++ b/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() diff --git a/openmc_plotter/plotgui.py b/openmc_plotter/plotgui.py index b04b1b6..5375bb0 100644 --- a/openmc_plotter/plotgui.py +++ b/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, diff --git a/openmc_plotter/plotmodel.py b/openmc_plotter/plotmodel.py index a339fd2..91c8e12 100644 --- a/openmc_plotter/plotmodel.py +++ b/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):