From 18022c0dc2ef9c13c8d12d5dd5ca882a66b1de2a Mon Sep 17 00:00:00 2001 From: shimwell Date: Sun, 18 Feb 2024 01:21:30 +0000 Subject: [PATCH 1/2] plotly colors working --- src/openmc_geometry_plot/app.py | 102 ++++--------- src/openmc_geometry_plot/core.py | 249 +++++++++++++++++++++++++++++++ 2 files changed, 280 insertions(+), 71 deletions(-) diff --git a/src/openmc_geometry_plot/app.py b/src/openmc_geometry_plot/app.py index bbc4ac7..f87329a 100644 --- a/src/openmc_geometry_plot/app.py +++ b/src/openmc_geometry_plot/app.py @@ -355,26 +355,27 @@ def main(): and isinstance(plot_bottom, float) ): + if basis == 'xy': + origin=( + (plot_left+plot_right)/2, + (plot_top+plot_bottom)/2, + slice_value, + ) + elif basis == 'yz': + origin=( + slice_value, + (plot_left+plot_right)/2, + (plot_top+plot_bottom)/2, + ) + elif basis == 'xz': + origin=( + (plot_left+plot_right)/2, + slice_value, + (plot_top+plot_bottom)/2, + ) + if backend == "matplotlib": - if basis == 'xy': - origin=( - (plot_left+plot_right)/2, - (plot_top+plot_bottom)/2, - slice_value, - ) - elif basis == 'yz': - origin=( - slice_value, - (plot_left+plot_right)/2, - (plot_top+plot_bottom)/2, - ) - elif basis == 'xz': - origin=( - (plot_left+plot_right)/2, - slice_value, - (plot_top+plot_bottom)/2, - ) width_x=plot_left-plot_right width_y=plot_top-plot_bottom print('origin',origin) @@ -394,7 +395,6 @@ def main(): plt.title(title) plt.savefig("openmc_plot_geometry_image.png") st.pyplot(plt) - # st.image("openmc_plot_geometry_image.png", use_column_width="always") with open("openmc_plot_geometry_image.png", "rb") as file: st.sidebar.download_button( @@ -404,58 +404,18 @@ def main(): mime="image/png", ) else: - data = [ - go.Heatmap( - z=color_data_slice, - showscale=False, - colorscale="viridis", - x0=plot_left, - dx=abs(plot_left - plot_right) / (len(color_data_slice[0]) - 1), - y0=plot_bottom, - dy=abs(plot_bottom - plot_top) / (len(color_data_slice) - 1), - # colorbar=dict(title=dict(side="right", text=cbar_label)), - # text = material_ids, - # hovertemplate= - # # 'material ID = %{z}
'+ - # "Cell ID = %{z}
" + - # # '
%{text}
'+ - # xlabel[:2].title() - # + ": %{x} cm
" - # + ylabel[:2].title() - # + ": %{y} cm
", - ) - ] - - if outline is not None: - data.append( - go.Contour( - z=outline_data_slice, - x0=plot_left, - dx=abs(plot_left - plot_right) - / (len(outline_data_slice[0]) - 1), - y0=plot_bottom, - dy=abs(plot_bottom - plot_top) - / (len(outline_data_slice) - 1), - contours_coloring="lines", - line_width=1, - colorscale=[[0, "rgb(0, 0, 0)"], [1.0, "rgb(0, 0, 0)"]], - showscale=False, - ) - ) - - plot = go.Figure(data=data) - - plot.update_layout( - xaxis={"title": xlabel}, - # reversed autorange is required to avoid image needing rotation/flipping in plotly - yaxis={"title": ylabel, "autorange": "reversed"}, - title=title, - autosize=False, - height=800, - ) - plot.update_yaxes( - scaleanchor="x", - scaleratio=1, + from openmc_geometry_plot import plot_plotly + plot = plot_plotly( + my_geometry, + origin=origin, + # width=[width_x,width_y], + pixels=pixels, + basis=basis, + color_by=color_by, + colors=my_colors, + legend=legend, + axis_units=axis_units, + outline=outline, ) plot.write_html("openmc_plot_geometry_image.html") diff --git a/src/openmc_geometry_plot/core.py b/src/openmc_geometry_plot/core.py index 195cefc..1cf7a4c 100644 --- a/src/openmc_geometry_plot/core.py +++ b/src/openmc_geometry_plot/core.py @@ -5,6 +5,8 @@ from tempfile import mkdtemp from pathlib import Path import matplotlib.pyplot as plt +import math +from tempfile import TemporaryDirectory from PIL import Image @@ -523,6 +525,253 @@ def get_slice_of_cell_ids( return trimmed_image_value +def plot_plotly( + geometry, + origin=None, + width=None, + pixels=40000, + basis='xy', + color_by='cell', + colors=None, + seed=None, + openmc_exec='openmc', + axes=None, + legend=False, + axis_units='cm', + # legend_kwargs=_default_legend_kwargs, + outline=False, + **kwargs +): + """Display a slice plot of the universe. + + Parameters + ---------- + origin : iterable of float + Coordinates at the origin of the plot. If left as None, + universe.bounding_box.center will be used to attempt to ascertain + the origin with infinite values being replaced by 0. + width : iterable of float + Width of the plot in each basis direction. If left as none then the + universe.bounding_box.width() will be used to attempt to + ascertain the plot width. Defaults to (10, 10) if the bounding_box + contains inf values + pixels : Iterable of int or int + If iterable of ints provided then this directly sets the number of + pixels to use in each basis direction. If int provided then this + sets the total number of pixels in the plot and the number of + pixels in each basis direction is calculated from this total and + the image aspect ratio. + basis : {'xy', 'xz', 'yz'} + The basis directions for the plot + color_by : {'cell', 'material'} + Indicate whether the plot should be colored by cell or by material + colors : dict + Assigns colors to specific materials or cells. Keys are instances of + :class:`Cell` or :class:`Material` and values are RGB 3-tuples, RGBA + 4-tuples, or strings indicating SVG color names. Red, green, blue, + and alpha should all be floats in the range [0.0, 1.0], for example: + + .. code-block:: python + + # Make water blue + water = openmc.Cell(fill=h2o) + universe.plot(..., colors={water: (0., 0., 1.)) + seed : int + Seed for the random number generator + openmc_exec : str + Path to OpenMC executable. + axes : matplotlib.Axes + Axes to draw to + + .. versionadded:: 0.13.1 + legend : bool + Whether a legend showing material or cell names should be drawn + + .. versionadded:: 0.14.0 + legend_kwargs : dict + Keyword arguments passed to :func:`matplotlib.pyplot.legend`. + + .. versionadded:: 0.14.0 + outline : bool + Whether outlines between color boundaries should be drawn + + .. versionadded:: 0.14.0 + axis_units : {'km', 'm', 'cm', 'mm'} + Units used on the plot axis + + .. versionadded:: 0.14.0 + **kwargs + Keyword arguments passed to :func:`matplotlib.pyplot.imshow` + + Returns + ------- + matplotlib.axes.Axes + Axes containing resulting image + + """ + import matplotlib.image as mpimg + # import matplotlib.patches as mpatches + import matplotlib.pyplot as plt + import plotly.graph_objects as go + + # Determine extents of plot + if basis == 'xy': + x, y = 0, 1 + xlabel, ylabel = f'x [{axis_units}]', f'y [{axis_units}]' + elif basis == 'yz': + x, y = 1, 2 + xlabel, ylabel = f'y [{axis_units}]', f'z [{axis_units}]' + elif basis == 'xz': + x, y = 0, 2 + xlabel, ylabel = f'x [{axis_units}]', f'z [{axis_units}]' + + bb = geometry.bounding_box + # checks to see if bounding box contains -inf or inf values + if np.isinf(bb.extent[basis]).any(): + if origin is None: + origin = (0, 0, 0) + if width is None: + width = (10, 10) + else: + if origin is None: + # if nan values in the bb.center they get replaced with 0.0 + # this happens when the bounding_box contains inf values + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + origin = np.nan_to_num(bb.center) + if width is None: + bb_width = bb.width + x_width = bb_width['xyz'.index(basis[0])] + y_width = bb_width['xyz'.index(basis[1])] + width = (x_width, y_width) + + if isinstance(pixels, int): + aspect_ratio = width[0] / width[1] + pixels_y = math.sqrt(pixels / aspect_ratio) + pixels = (int(pixels / pixels_y), int(pixels_y)) + + axis_scaling_factor = {'km': 0.00001, 'm': 0.01, 'cm': 1, 'mm': 10} + + x_min = (origin[x] - 0.5*width[0]) * axis_scaling_factor[axis_units] + x_max = (origin[x] + 0.5*width[0]) * axis_scaling_factor[axis_units] + y_min = (origin[y] - 0.5*width[1]) * axis_scaling_factor[axis_units] + y_max = (origin[y] + 0.5*width[1]) * axis_scaling_factor[axis_units] + + with TemporaryDirectory() as tmpdir: + model = openmc.Model() + model.geometry = geometry + if seed is not None: + model.settings.plot_seed = seed + + # Determine whether any materials contains macroscopic data and if + # so, set energy mode accordingly + for mat in geometry.get_all_materials().values(): + if mat._macroscopic is not None: + model.settings.energy_mode = 'multi-group' + break + + # Create plot object matching passed arguments + plot = openmc.Plot() + plot.origin = origin + plot.width = width + plot.pixels = pixels + plot.basis = basis + plot.color_by = color_by + + if colors is not None: + + colors_based_on_ids = {} + for key, value in colors.items(): + colors_based_on_ids[key] = get_rgb_from_int(key.id) + plot.colors = colors_based_on_ids + + model.plots.append(plot) + + # Run OpenMC in geometry plotting mode + model.plot_geometry(False, cwd=tmpdir, openmc_exec=openmc_exec) + + # Read image from file + img_path = Path(tmpdir) / f'plot_{plot.id}.png' + if not img_path.is_file(): + img_path = img_path.with_suffix('.ppm') + # todo see if we can just read in image once + img = mpimg.imread(str(img_path)) + image_values = Image.open(img_path) + image_values = np.asarray(image_values) + image_values = [ + [get_int_from_rgb(inner_entry) for inner_entry in outer_entry] + for outer_entry in image_values + ] + + image_values = np.array(image_values) + image_values[image_values == 16777215] = 0 + + data = [] + + if outline: + # Combine R, G, B values into a single int + rgb = (img * 256).astype(int) + image_value = (rgb[..., 0] << 16) + \ + (rgb[..., 1] << 8) + (rgb[..., 2]) + + data.append( + go.Contour( + z=image_value, + contours_coloring='none', + # colorscale=dcolorsc, + showscale=False, + x0=x_min, + dx=abs(x_min - x_max) / (img.shape[0] - 1), + y0=y_min, + dy=abs(y_min - y_max) / (img.shape[1] - 1), + ) + ) + + dcolorsc=[ + [0, 'white'], + ] + + rgb_cols = [f'rgb({c[0]},{c[1]},{c[2]})' for c in list(colors.values())] + mat_ids=[mat.id for mat in colors.keys()] + highest_mat_id = max(mat_ids) + for rgb_col, mat_id in zip(rgb_cols, mat_ids): + dcolorsc.append(((1/highest_mat_id)*mat_id,rgb_col)) + + data.append( + go.Heatmap( + z=image_values, + # showscale=True, + colorscale=dcolorsc, + x0=x_min, + dx=abs(x_min - x_max) / (img.shape[0] - 1), + y0=y_min, + dy=abs(y_min - y_max) / (img.shape[1] - 1), + colorbar= dict( + tick0= 0, + xref="container", + tickmode= 'array', + tickvals= mat_ids, + ticktext= mat_ids, # TODO add material names + title=f'{color_by.title()} IDs' + ) + ) + ) + plot = go.Figure(data=data) + + plot.update_layout( + xaxis={"title": xlabel}, + # reversed autorange is required to avoid image needing rotation/flipping in plotly + yaxis={"title": ylabel, "autorange": "reversed"}, + # title=title, + autosize=False, + height=800, + ) + plot.update_yaxes( + scaleanchor="x", + scaleratio=1, + ) + return plot + # patching openmc openmc.Geometry.get_dagmc_universe = get_dagmc_universe From 70d8fb7472a577b1136b68b36bd460319813dba3 Mon Sep 17 00:00:00 2001 From: Jonathan Shimwell Date: Sun, 18 Feb 2024 21:37:19 +0000 Subject: [PATCH 2/2] improvements to plotly and sidebar --- src/openmc_geometry_plot/app.py | 53 +++++++++++++++++++------------- src/openmc_geometry_plot/core.py | 53 +++++++++++++++++++------------- 2 files changed, 63 insertions(+), 43 deletions(-) diff --git a/src/openmc_geometry_plot/app.py b/src/openmc_geometry_plot/app.py index f87329a..b79394b 100644 --- a/src/openmc_geometry_plot/app.py +++ b/src/openmc_geometry_plot/app.py @@ -47,8 +47,8 @@ def main(): 👉 Create your ```openmc.Geometry()``` and export the geometry xml file using ```export_to_xml()```. Not got a geometry.xml file handy, right mouse 🖱️ click and save these links - [ example 1 ](https://fusion-energy.github.io/openmc_geometry_plot/examples/csg_tokamak/geometry.xml), - [ example 2 ](https://fusion-energy.github.io/openmc_geometry_plot/examples/csg_cylinder_box/geometry.xml) + [ example 1 ](https://raw.githubusercontent.com/fusion-energy/openmc_geometry_plot/31be0556f3f34c102cab3de094df08f48acad5ca/examples/csg_tokamak/geometry.xml), + [ example 2 ](https://raw.githubusercontent.com/fusion-energy/openmc_geometry_plot/31be0556f3f34c102cab3de094df08f48acad5ca/examples/csg_cylinder_box/geometry.xml) """ ) @@ -57,7 +57,7 @@ def main(): 👉 Create your DAGMC h5m file using tools like [CAD-to-h5m](https://github.com/fusion-energy/cad_to_dagmc) Not got a DAGMC h5m file handy, right mouse 🖱️ click and save these links - [ example 1 ](https://fusion-energy.github.io/openmc_geometry_plot/examples/dagmc_tokamak/dagmc_180_tokamak.h5m) + [ example 1 ](https://github.com/fusion-energy/openmc_geometry_plot/raw/31be0556f3f34c102cab3de094df08f48acad5ca/examples/dagmc_tokamak/dagmc_180_tokamak.h5m) """ ) @@ -91,6 +91,10 @@ def main(): mat_ids = range(0, len(dag_universe.material_names) + 1) # mat_names = dag_universe.material_names + all_cells_ids = range(0, dag_universe.n_cells + 1) + all_cells ={} + for cell_id in all_cells_ids: + all_cells[cell_id] =openmc.Cell(cell_id=cell_id) if len(mat_ids) >= 1: set_mat_ids = set(mat_ids) @@ -113,6 +117,11 @@ def main(): # find all material names mat_ids = range(0, len(dag_universe.material_names) + 1) + all_cells_ids = range(0, dag_universe.n_cells + 1) + all_cells ={} + for cell_id in all_cells_ids: + all_cells[cell_id] =openmc.Cell(cell_id=cell_id) + if len(mat_ids) >= 1: set_mat_ids = set(mat_ids) else: @@ -172,18 +181,11 @@ def main(): ) backend = st.sidebar.selectbox( label="Ploting backend", - options=("matplotlib", "plotly"), + options=("plotly", "matplotlib"), index=0, key="geometry_plotting_backend", help="Create png images with MatPlotLib or HTML plots with Plotly", ) - outline = st.sidebar.selectbox( - label="Outline", - options=("materials", "cells", None), - index=0, - key="outline", - help="Allows an outline to be drawn around the cells or materials, select None for no outline", - ) legend = st.sidebar.selectbox( label="Legend", options=(True, False), @@ -249,6 +251,15 @@ def main(): slice_min = float(bb[0][slice_index]) slice_max = float(bb[1][slice_index]) + if isinstance(slice_min, float) and isinstance(slice_max, float): + slice_value = st.sidebar.slider( + label="Slice value", + min_value=slice_min, + max_value=slice_max, + value=(slice_min + slice_max) / 2, + key="slice_slider", + help="Set the value of the slice axis", + ) if isinstance(x_min, float) and isinstance(x_max, float): plot_right, plot_left = st.sidebar.slider( label="Left and right values for the horizontal axis", @@ -268,15 +279,6 @@ def main(): key="bottom_top_slider", help="Set the lowest visible value and highest visible value on the vertical axis", ) - if isinstance(slice_min, float) and isinstance(slice_max, float): - slice_value = st.sidebar.slider( - label="Slice value", - min_value=slice_min, - max_value=slice_max, - value=(slice_min + slice_max) / 2, - key="slice_slider", - help="Set the value of the slice axis", - ) color_by = st.sidebar.selectbox( label="Color by", @@ -285,7 +287,13 @@ def main(): key="color_by", help="Should the plot be colored by material or by cell", ) - + outline = st.sidebar.selectbox( + label="Outline", + options=(True, False), + index=0, + key="outline", + help="Allows an outline to be drawn around the cells or materials, select None for no outline", + ) selected_color_map = st.sidebar.selectbox( label="Color map", options=colormaps(), index=82 ) # index 82 is tab20c @@ -415,7 +423,8 @@ def main(): colors=my_colors, legend=legend, axis_units=axis_units, - outline=outline, + outline=outline, + title=title ) plot.write_html("openmc_plot_geometry_image.html") diff --git a/src/openmc_geometry_plot/core.py b/src/openmc_geometry_plot/core.py index 1cf7a4c..d661dbf 100644 --- a/src/openmc_geometry_plot/core.py +++ b/src/openmc_geometry_plot/core.py @@ -7,9 +7,12 @@ import matplotlib.pyplot as plt import math from tempfile import TemporaryDirectory - +import warnings from PIL import Image - +import matplotlib.image as mpimg +# import matplotlib.patches as mpatches +import matplotlib.pyplot as plt +import plotly.graph_objects as go from numpy import asarray @@ -26,6 +29,9 @@ def get_int_from_rgb(rgb: typing.Tuple[int, int, int]) -> int: blue = rgb[2] return (red << 16) + (green << 8) + blue +def get_hover_text_from_id(id, color_by): + return f'{color_by} {id}' + def get_plot_extent( self, plot_left, plot_right, plot_bottom, plot_top, slice_value, bb, view_direction @@ -540,6 +546,7 @@ def plot_plotly( axis_units='cm', # legend_kwargs=_default_legend_kwargs, outline=False, + title='', **kwargs ): """Display a slice plot of the universe. @@ -609,10 +616,7 @@ def plot_plotly( Axes containing resulting image """ - import matplotlib.image as mpimg - # import matplotlib.patches as mpatches - import matplotlib.pyplot as plt - import plotly.graph_objects as go + # Determine extents of plot if basis == 'xy': @@ -643,7 +647,7 @@ def plot_plotly( bb_width = bb.width x_width = bb_width['xyz'.index(basis[0])] y_width = bb_width['xyz'.index(basis[1])] - width = (x_width, y_width) + width = (x_width+3, y_width+3) # makes width a bit bigger so that the edges don't get cut if isinstance(pixels, int): aspect_ratio = width[0] / width[1] @@ -709,14 +713,9 @@ def plot_plotly( data = [] if outline: - # Combine R, G, B values into a single int - rgb = (img * 256).astype(int) - image_value = (rgb[..., 0] << 16) + \ - (rgb[..., 1] << 8) + (rgb[..., 2]) - data.append( go.Contour( - z=image_value, + z=image_values, contours_coloring='none', # colorscale=dcolorsc, showscale=False, @@ -737,6 +736,20 @@ def plot_plotly( for rgb_col, mat_id in zip(rgb_cols, mat_ids): dcolorsc.append(((1/highest_mat_id)*mat_id,rgb_col)) + cbar = dict( + tick0= 0, + xref="container", + tickmode= 'array', + tickvals= mat_ids, + ticktext= mat_ids, # TODO add material names + title=f'{color_by.title()} IDs', + ) + + hovertext = [ + [get_hover_text_from_id(inner_entry, color_by) for inner_entry in outer_entry] + for outer_entry in image_values + ] + data.append( go.Heatmap( z=image_values, @@ -746,17 +759,14 @@ def plot_plotly( dx=abs(x_min - x_max) / (img.shape[0] - 1), y0=y_min, dy=abs(y_min - y_max) / (img.shape[1] - 1), - colorbar= dict( - tick0= 0, - xref="container", - tickmode= 'array', - tickvals= mat_ids, - ticktext= mat_ids, # TODO add material names - title=f'{color_by.title()} IDs' - ) + colorbar= cbar, + showscale=legend, + hoverinfo='text', + text=hovertext ) ) plot = go.Figure(data=data) + plot.update_layout( xaxis={"title": xlabel}, @@ -765,6 +775,7 @@ def plot_plotly( # title=title, autosize=False, height=800, + title=title ) plot.update_yaxes( scaleanchor="x",