From 6b03287d52aa8b55320ee99f63d0d993b93eb375 Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Thu, 13 Apr 2023 16:09:07 -0400 Subject: [PATCH 1/9] add auto legend to universe.plot --- openmc/universe.py | 50 +++++++++++++++++++++++++++++++++++++++------- 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index 2ed96f4bcb..a75af3b0de 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -87,7 +87,8 @@ class UniverseBase(ABC, IDManagerMixin): self._volume = volume_calc.volumes[self.id].n self._atoms = volume_calc.atoms[self.id] else: - raise ValueError('No volume information found for this universe.') + raise ValueError( + 'No volume information found for this universe.') else: raise ValueError('No volume information found for this universe.') @@ -168,7 +169,7 @@ class UniverseBase(ABC, IDManagerMixin): clone._cells = OrderedDict() for cell in self._cells.values(): clone.add_cell(cell.clone(clone_materials, clone_regions, - memo)) + memo)) # Memoize the clone memo[self] = clone @@ -295,7 +296,7 @@ class Universe(UniverseBase): def plot(self, origin=(0., 0., 0.), width=(1., 1.), pixels=(200, 200), basis='xy', color_by='cell', colors=None, seed=None, - openmc_exec='openmc', axes=None, **kwargs): + openmc_exec='openmc', axes=None, legend=False, **kwargs): """Display a slice plot of the universe. Parameters @@ -329,6 +330,8 @@ class Universe(UniverseBase): Axes to draw to .. versionadded:: 0.13.1 + legend : bool + whether a legend showing material or cell names should be drawn **kwargs Keyword arguments passed to :func:`matplotlib.pyplot.imshow` @@ -339,6 +342,7 @@ class Universe(UniverseBase): """ import matplotlib.image as mpimg + import matplotlib.patches as mpatches import matplotlib.pyplot as plt # Determine extents of plot @@ -396,6 +400,36 @@ class Universe(UniverseBase): height = pixels[0]*px/(params.top - params.bottom) fig.set_size_inches(width, height) + # add legend showing which colors represent which material + # or cell if that was requested + if legend: + if plot.colors is None: + raise Exception("Must set plot color dictionary if you would" + " like to have an automatic legend.") + + if color_by == "cell": + expected_key_type = openmc.Cell + elif color_by == "material": + expected_key_type = openmc.Material + else: + raise Exception("wtf?") + + patches = [] + for key, color in plot.colors.items(): + + if isinstance(key, int): + raise Exception( + "Cannot use IDs in colors dict for auto legend.") + elif not isinstance(key, expected_key_type): + raise Exception( + "Color dict key type does not match color_by") + + # this works whether we're doing cells or materials + key_patch = mpatches.Patch(color=color, label=key.name) + patches.append(key_patch) + + axes.legend(handles=patches) + # Plot image and return the axes return axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs) @@ -738,8 +772,9 @@ class DAGMCUniverse(UniverseBase): @property def material_names(self): dagmc_file_contents = h5py.File(self.filename) - material_tags_hex=dagmc_file_contents['/tstt/tags/NAME'].get('values') - material_tags_ascii=[] + material_tags_hex = dagmc_file_contents['/tstt/tags/NAME'].get( + 'values') + material_tags_ascii = [] for tag in material_tags_hex: candidate_tag = tag.tobytes().decode().replace('\x00', '') # tags might be for temperature or reflective surfaces @@ -905,7 +940,8 @@ class DAGMCUniverse(UniverseBase): openmc.Universe Universe instance """ - bounding_cell = openmc.Cell(fill=self, cell_id=bounding_cell_id, region=self.bounding_region(**kwargs)) + bounding_cell = openmc.Cell( + fill=self, cell_id=bounding_cell_id, region=self.bounding_region(**kwargs)) return openmc.Universe(cells=[bounding_cell]) @classmethod @@ -972,4 +1008,4 @@ class DAGMCUniverse(UniverseBase): clone.volume = self.volume clone.auto_geom_ids = self.auto_geom_ids clone.auto_mat_ids = self.auto_mat_ids - return clone \ No newline at end of file + return clone From 8bc434a67d0579b8af89cb5560289b2416d411cd Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Fri, 14 Apr 2023 12:36:24 -0400 Subject: [PATCH 2/9] move legend outside image Co-authored-by: Jonathan Shimwell --- openmc/universe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/openmc/universe.py b/openmc/universe.py index a75af3b0de..a7f9e3d6e4 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -428,7 +428,7 @@ class Universe(UniverseBase): key_patch = mpatches.Patch(color=color, label=key.name) patches.append(key_patch) - axes.legend(handles=patches) + axes.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0) # Plot image and return the axes return axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs) From 3880d1cec86f12b8289a7f67c1ef3b8364777e43 Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Mon, 17 Apr 2023 12:14:52 -0400 Subject: [PATCH 3/9] line length suggestion Co-authored-by: Jonathan Shimwell --- openmc/universe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index a7f9e3d6e4..32768c305f 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -404,8 +404,8 @@ class Universe(UniverseBase): # or cell if that was requested if legend: if plot.colors is None: - raise Exception("Must set plot color dictionary if you would" - " like to have an automatic legend.") + raise Exception("Must set plot color dictionary if you " + "would like to have an automatic legend.") if color_by == "cell": expected_key_type = openmc.Cell From e4c332360175b66c8261c9d41643913cb9958e65 Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Mon, 17 Apr 2023 12:15:29 -0400 Subject: [PATCH 4/9] remove unnecessary error checking Co-authored-by: Jonathan Shimwell --- openmc/universe.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index 32768c305f..a4efb5f732 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -409,10 +409,8 @@ class Universe(UniverseBase): if color_by == "cell": expected_key_type = openmc.Cell - elif color_by == "material": - expected_key_type = openmc.Material else: - raise Exception("wtf?") + expected_key_type = openmc.Material patches = [] for key, color in plot.colors.items(): From f65de85f9fba7e63e8c4db12157e0fa3f89f2d3d Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Mon, 17 Apr 2023 12:17:09 -0400 Subject: [PATCH 5/9] add versionadded for legend kwarg Co-authored-by: Jonathan Shimwell --- openmc/universe.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index a4efb5f732..ed693b84ac 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -330,11 +330,12 @@ class Universe(UniverseBase): Axes to draw to .. versionadded:: 0.13.1 - legend : bool - whether a legend showing material or cell names should be drawn **kwargs Keyword arguments passed to :func:`matplotlib.pyplot.imshow` + .. versionadded:: 0.13.4 + legend : bool + Whether a legend showing material or cell names should be drawn Returns ------- matplotlib.image.AxesImage From ad33834e0f659f90efe3ae1154b30a254b61ecf4 Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Mon, 17 Apr 2023 12:18:03 -0400 Subject: [PATCH 6/9] default to legend label using ID if no name given Co-authored-by: Jonathan Shimwell --- openmc/universe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/openmc/universe.py b/openmc/universe.py index ed693b84ac..ec95075315 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -424,7 +424,8 @@ class Universe(UniverseBase): "Color dict key type does not match color_by") # this works whether we're doing cells or materials - key_patch = mpatches.Patch(color=color, label=key.name) + label = key.name if key.name != '' else key.id + key_patch = mpatches.Patch(color=color, label=label) patches.append(key_patch) axes.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0) From a02d583ac18925927e02b831efba0b2fcf96ffae Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Tue, 18 Apr 2023 11:24:20 -0400 Subject: [PATCH 7/9] review comments addressed --- openmc/universe.py | 32 ++++++++++++++++++++++++-------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index ec95075315..b85d36895b 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -6,6 +6,7 @@ from numbers import Integral, Real from pathlib import Path from tempfile import TemporaryDirectory from xml.etree import ElementTree as ET +from warnings import warn import h5py import numpy as np @@ -294,9 +295,15 @@ class Universe(UniverseBase): return [self, cell] + cell.fill.find(p) return [] + # default kwargs that are passed to plt.legend in the plot method below. + # If you change this, be sure to change it in the docstring for plot below. + _default_legend_kwargs = {'bbox_to_anchor': ( + 1.05, 1), 'loc': 2, 'borderaxespad': 0.0} + def plot(self, origin=(0., 0., 0.), width=(1., 1.), pixels=(200, 200), basis='xy', color_by='cell', colors=None, seed=None, - openmc_exec='openmc', axes=None, legend=False, **kwargs): + openmc_exec='openmc', axes=None, legend=False, + legend_kwargs=_default_legend_kwargs, **kwargs): """Display a slice plot of the universe. Parameters @@ -330,12 +337,21 @@ class Universe(UniverseBase): Axes to draw to .. versionadded:: 0.13.1 + legend : bool + Whether a legend showing material or cell names should be drawn + + .. versionadded:: 0.13.4 + legend_kwargs : dict + Keyword arguments passed to :func:`matplotlib.pyplot.legend`. + The default is: + + {'bbox_to_anchor': (1.05, 1), 'loc':2, 'borderaxespad': 0.0} + + .. versionadded:: 0.13.4 **kwargs Keyword arguments passed to :func:`matplotlib.pyplot.imshow` .. versionadded:: 0.13.4 - legend : bool - Whether a legend showing material or cell names should be drawn Returns ------- matplotlib.image.AxesImage @@ -405,8 +421,8 @@ class Universe(UniverseBase): # or cell if that was requested if legend: if plot.colors is None: - raise Exception("Must set plot color dictionary if you " - "would like to have an automatic legend.") + raise ValueError("Must pass 'colors' dictionary if you " + "are adding a legend via legend=True.") if color_by == "cell": expected_key_type = openmc.Cell @@ -417,10 +433,10 @@ class Universe(UniverseBase): for key, color in plot.colors.items(): if isinstance(key, int): - raise Exception( + raise TypeError( "Cannot use IDs in colors dict for auto legend.") elif not isinstance(key, expected_key_type): - raise Exception( + raise TypeError( "Color dict key type does not match color_by") # this works whether we're doing cells or materials @@ -428,7 +444,7 @@ class Universe(UniverseBase): key_patch = mpatches.Patch(color=color, label=label) patches.append(key_patch) - axes.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0) + axes.legend(handles=patches, **legend_kwargs) # Plot image and return the axes return axes.imshow(img, extent=(x_min, x_max, y_min, y_max), **kwargs) From 2f87ce7f2789b5a5a75ae70099cfcb2d12bae30c Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Wed, 19 Apr 2023 12:22:29 -0400 Subject: [PATCH 8/9] Update openmc/universe.py Co-authored-by: Paul Romano --- openmc/universe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/openmc/universe.py b/openmc/universe.py index b85d36895b..3f49825472 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -296,7 +296,6 @@ class Universe(UniverseBase): return [] # default kwargs that are passed to plt.legend in the plot method below. - # If you change this, be sure to change it in the docstring for plot below. _default_legend_kwargs = {'bbox_to_anchor': ( 1.05, 1), 'loc': 2, 'borderaxespad': 0.0} From 22f4fc0902b0ed708a245c6b82a341dbbfee99b5 Mon Sep 17 00:00:00 2001 From: Gavin Ridley Date: Wed, 19 Apr 2023 12:22:39 -0400 Subject: [PATCH 9/9] Update openmc/universe.py Co-authored-by: Paul Romano --- openmc/universe.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/openmc/universe.py b/openmc/universe.py index 3f49825472..d5ae077178 100644 --- a/openmc/universe.py +++ b/openmc/universe.py @@ -342,9 +342,6 @@ class Universe(UniverseBase): .. versionadded:: 0.13.4 legend_kwargs : dict Keyword arguments passed to :func:`matplotlib.pyplot.legend`. - The default is: - - {'bbox_to_anchor': (1.05, 1), 'loc':2, 'borderaxespad': 0.0} .. versionadded:: 0.13.4 **kwargs