From 1111d7163a923b2b02adde3290013b4ba8d64e80 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Tue, 28 Nov 2023 15:58:45 +0000
Subject: [PATCH] adding plotling backend option

---
 examples/plot_minimal_example.py        |   3 +-
 examples/plot_plotly_example.py         |  95 ++++++++++
 src/openmc_regular_mesh_plotter/core.py | 229 +++++++++++++++---------
 3 files changed, 245 insertions(+), 82 deletions(-)
 create mode 100644 examples/plot_plotly_example.py

diff --git a/examples/plot_minimal_example.py b/examples/plot_minimal_example.py
index 04ca4af..f5a02b0 100644
--- a/examples/plot_minimal_example.py
+++ b/examples/plot_minimal_example.py
@@ -78,5 +78,6 @@
     colorbar=False,
 )
 
-plot.figure.savefig("example_openmc_regular_mesh_plotter.png")
 plot.title.set_text("")
+plot.figure.savefig("example_openmc_regular_mesh_plotter.png")
+print('file created example_openmc_regular_mesh_plotter.png')
\ No newline at end of file
diff --git a/examples/plot_plotly_example.py b/examples/plot_plotly_example.py
new file mode 100644
index 0000000..1446bb6
--- /dev/null
+++ b/examples/plot_plotly_example.py
@@ -0,0 +1,95 @@
+import openmc
+from matplotlib.colors import LogNorm
+from openmc_regular_mesh_plotter import plot_mesh_tally
+
+
+# MATERIALS
+mat_1 = openmc.Material()
+mat_1.add_element("Li", 1)
+mat_1.set_density("g/cm3", 0.45)
+my_materials = openmc.Materials([mat_1])
+
+# GEOMETRY
+# surfaces
+inner_surface = openmc.Sphere(r=200)
+outer_surface = openmc.Sphere(r=400, boundary_type="vacuum")
+# regions
+inner_region = -inner_surface
+outer_region = -outer_surface & +inner_surface
+# cells
+inner_cell = openmc.Cell(region=inner_region)
+outer_cell = openmc.Cell(region=outer_region)
+outer_cell.fill = mat_1
+my_geometry = openmc.Geometry([inner_cell, outer_cell])
+
+# SIMULATION SETTINGS
+my_settings = openmc.Settings()
+my_settings.batches = 10
+my_settings.inactive = 0
+my_settings.particles = 5000
+my_settings.run_mode = "fixed source"
+# Create a DT point source
+try:
+    source = openmc.IndependentSource()
+except:
+    # work with older versions of openmc
+    source = openmc.Source()
+source.space = openmc.stats.Point((100, 0, 0))
+source.angle = openmc.stats.Isotropic()
+source.energy = openmc.stats.Discrete([14e6], [1])
+my_settings.source = source
+
+# Tallies
+my_tallies = openmc.Tallies()
+mesh = openmc.RegularMesh().from_domain(
+    my_geometry,  # the corners of the mesh are being set automatically to surround the geometry
+    dimension=[40, 40, 40],
+)
+mesh_filter = openmc.MeshFilter(mesh)
+mesh_tally_1 = openmc.Tally(name="mesh_tally")
+mesh_tally_1.filters = [mesh_filter]
+mesh_tally_1.scores = ["heating"]
+my_tallies.append(mesh_tally_1)
+
+model = openmc.model.Model(my_geometry, my_materials, my_settings, my_tallies)
+# sp_filename = model.run()
+
+sp_filename = '/home/jshimwell/reactor_csg_neutronics_model/statepoint.12.h5'
+
+# post process simulation result
+statepoint = openmc.StatePoint(sp_filename)
+
+print(statepoint.tallies)
+
+# extracts the mesh tally by name
+# my_mesh_tally = statepoint.get_tally(name="mesh_tally")
+# my_mesh_tally = statepoint.get_tally(name="neutron_H3-production_on_regular_xy_mesh")
+my_mesh_tally = statepoint.get_tally(name="neutron_H3-production_on_regular_xz_mesh")
+
+# default tally units for heating are in eV per source neutron
+# for this example plot we want Mega Joules per second per cm3 or Mjcm^-3s^-1
+neutrons_per_second = 1e21
+eV_to_joules = 1.60218e-19
+joules_to_mega_joules = 1e-6
+scaling_factor = neutrons_per_second * eV_to_joules * joules_to_mega_joules
+# note that volume_normalization is enabled so this will also change the units to divide by the volume of each mesh voxel
+# alternatively you could set volume_normalization to false and divide by the mesh.volume[0][0][0] in the scaling factor
+# in a regular mesh all the voxels have the same volume so the [0][0][0] just picks the first volume
+
+plot = plot_mesh_tally(
+    plotting_backend='plotly',
+    tally=my_mesh_tally,
+    outline=True,  # enables an outline around the geometry
+    geometry=my_geometry,  # needed for outline
+    norm=LogNorm(),  # log scale
+    colorbar=True,
+    scaling_factor=scaling_factor,
+    colorbar_kwargs={'title':'Heating [MJ cm-3s-1]'},
+    basis='xz',
+)
+
+# setting title of the plot
+plot.update_layout({'title': ' made with openmc_regular_mesh_plotter'})
+plot.show()
+# plot.write_html("example_openmc_regular_mesh_plotter.html")
+print('file created example_openmc_regular_mesh_plotter.html')
\ No newline at end of file
diff --git a/src/openmc_regular_mesh_plotter/core.py b/src/openmc_regular_mesh_plotter/core.py
index 6f219f8..6ab0e1e 100644
--- a/src/openmc_regular_mesh_plotter/core.py
+++ b/src/openmc_regular_mesh_plotter/core.py
@@ -45,6 +45,7 @@ def plot_mesh_tally(
     scaling_factor: typing.Optional[float] = None,
     colorbar_kwargs: dict = {},
     outline_kwargs: dict = _default_outline_kwargs,
+    plotting_backend: str = 'matplotlib',
     **kwargs,
 ) -> "matplotlib.image.AxesImage":
     """Display a slice plot of the mesh tally score.
@@ -175,90 +176,156 @@ def plot_mesh_tally(
             slice_index,
         )
 
-    im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **default_imshow_kwargs)
-
-    if colorbar:
-        fig.colorbar(im, **colorbar_kwargs)
-
-    if outline and geometry is not None:
-        import matplotlib.image as mpimg
-
-        # code to make sure geometry outline is in the middle of the mesh voxel
-        # two of the three dimensions are just in the center of the mesh
-        # but the slice can move one axis off the center so this needs calculating
-        x0, y0, z0 = mesh.lower_left
-        x1, y1, z1 = mesh.upper_right
-        nx, ny, nz = mesh.dimension
-        center_of_mesh = mesh.bounding_box.center
-
-        if basis == "xy":
-            zarr = np.linspace(z0, z1, nz + 1)
-            center_of_mesh_slice = [
-                center_of_mesh[0],
-                center_of_mesh[1],
-                (zarr[slice_index] + zarr[slice_index + 1]) / 2,
-            ]
-        if basis == "xz":
-            yarr = np.linspace(y0, y1, ny + 1)
-            center_of_mesh_slice = [
-                center_of_mesh[0],
-                (yarr[slice_index] + yarr[slice_index + 1]) / 2,
-                center_of_mesh[2],
-            ]
-        if basis == "yz":
-            xarr = np.linspace(x0, x1, nx + 1)
-            center_of_mesh_slice = [
-                (xarr[slice_index] + xarr[slice_index + 1]) / 2,
-                center_of_mesh[1],
-                center_of_mesh[2],
-            ]
-
-        model = openmc.Model()
-        model.geometry = geometry
-        plot = openmc.Plot()
-        plot.origin = center_of_mesh_slice
-        bb_width = mesh.bounding_box.extent[basis]
-        plot.width = (bb_width[0] - bb_width[1], bb_width[2] - bb_width[3])
-        aspect_ratio = (bb_width[0] - bb_width[1]) / (bb_width[2] - bb_width[3])
-        pixels_y = math.sqrt(pixels / aspect_ratio)
-        pixels = (int(pixels / pixels_y), int(pixels_y))
-        plot.pixels = pixels
-        plot.basis = basis
-        plot.color_by = outline_by
-        model.plots.append(plot)
-
-        with TemporaryDirectory() as tmpdir:
-            # Run OpenMC in geometry plotting mode
-            model.plot_geometry(False, cwd=tmpdir)
-
-            # 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")
-            img = mpimg.imread(str(img_path))
-
-        # Combine R, G, B values into a single int
-        rgb = (img * 256).astype(int)
-        image_value = (rgb[..., 0] << 16) + (rgb[..., 1] << 8) + (rgb[..., 2])
-
-        if basis == "xz":
-            image_value = np.rot90(image_value, 2)
-        elif basis == "yz":
-            image_value = np.rot90(image_value, 2)
-        else:  # basis == 'xy'
-            image_value = np.rot90(image_value, 2)
-
-        # Plot image and return the axes
-        axes.contour(
-            image_value,
-            origin="upper",
-            levels=np.unique(image_value),
-            extent=(x_min, x_max, y_min, y_max),
-            **outline_kwargs,
+    if plotting_backend == 'plotly':
+        import plotly.graph_objects as go
+        dx = abs(x_min - x_max)/(data.shape[0]-1)
+        dy = abs(y_min - y_max)/(data.shape[1]-1)
+
+        fig = go.Figure()
+        fig.add_trace(
+            go.Contour(
+                z=data,
+                x0=x_min,
+                dx=dx,
+                y0=y_min,
+                dy=dy,
+                showscale=colorbar,
+                colorbar=colorbar_kwargs,
+                colorscale='viridis',
+                line = {'width': 0}
+            ))
+        fig.update_layout(
+            xaxis_title=xlabel,
+            yaxis_title=ylabel,
+            # can be used to ensure aspect ratio but prevents resizing
+            # height=abs(y_min - y_max),
+            # width=abs(x_min - x_max)
         )
+        fig.update_yaxes(
+            scaleanchor="x",
+            scaleratio=1,
+        )
+
+        if outline and geometry is not None:
+            
+            image_value = get_outline(mesh, basis, slice_index, geometry, outline_by, pixels)
+            
+            dx = abs(x_min - x_max)/(image_value.shape[0]-1)
+            dy = abs(y_min - y_max)/(image_value.shape[1]-1)
+            fig.add_trace(
+              go.Contour(
+                z=image_value,
+                x0=x_min,
+                dx=dx,
+                y0=y_min,
+                dy=dy,
+                contours_coloring='lines',
+                colorscale=[[0, 'black'], [1.0, 'black']],
+                line = {'width': 1, 'color': 'rgb(50,50,50)'}
+        #         contours=dict(
+        #     start=0,
+        #     end=8,
+        #     size=2,
+        # ),
+                # ncontours=1
+                # line = {'width': 0}
+              )  
+            )
+        print()
+        return fig
+
+    elif plotting_backend == 'matplotlib':
+
+        im = axes.imshow(data, extent=(x_min, x_max, y_min, y_max), **default_imshow_kwargs)
 
-    return axes
+        if colorbar:
+            fig.colorbar(im, **colorbar_kwargs)
 
+        if outline and geometry is not None:
+
+            image_value = get_outline(mesh, basis, slice_index, geometry, outline_by, pixels)
+
+            # Plot image and return the axes
+            axes.contour(
+                image_value,
+                origin="upper",
+                levels=np.unique(image_value),
+                extent=(x_min, x_max, y_min, y_max),
+                **outline_kwargs,
+            )
+
+        return axes
+    else:
+        raise ValueError(f'plotting_backend must be either matplotlib or plotly, not {plotting_backend}')
+
+def get_outline(mesh, basis, slice_index, geometry, outline_by, pixels):
+    import matplotlib.image as mpimg
+
+    # code to make sure geometry outline is in the middle of the mesh voxel
+    # two of the three dimensions are just in the center of the mesh
+    # but the slice can move one axis off the center so this needs calculating
+    x0, y0, z0 = mesh.lower_left
+    x1, y1, z1 = mesh.upper_right
+    nx, ny, nz = mesh.dimension
+    center_of_mesh = mesh.bounding_box.center
+
+    if basis == "xy":
+        zarr = np.linspace(z0, z1, nz + 1)
+        center_of_mesh_slice = [
+            center_of_mesh[0],
+            center_of_mesh[1],
+            (zarr[slice_index] + zarr[slice_index + 1]) / 2,
+        ]
+    if basis == "xz":
+        yarr = np.linspace(y0, y1, ny + 1)
+        center_of_mesh_slice = [
+            center_of_mesh[0],
+            (yarr[slice_index] + yarr[slice_index + 1]) / 2,
+            center_of_mesh[2],
+        ]
+    if basis == "yz":
+        xarr = np.linspace(x0, x1, nx + 1)
+        center_of_mesh_slice = [
+            (xarr[slice_index] + xarr[slice_index + 1]) / 2,
+            center_of_mesh[1],
+            center_of_mesh[2],
+        ]
+
+    model = openmc.Model()
+    model.geometry = geometry
+    plot = openmc.Plot()
+    plot.origin = center_of_mesh_slice
+    bb_width = mesh.bounding_box.extent[basis]
+    plot.width = (bb_width[0] - bb_width[1], bb_width[2] - bb_width[3])
+    aspect_ratio = (bb_width[0] - bb_width[1]) / (bb_width[2] - bb_width[3])
+    pixels_y = math.sqrt(pixels / aspect_ratio)
+    pixels = (int(pixels / pixels_y), int(pixels_y))
+    plot.pixels = pixels
+    plot.basis = basis
+    plot.color_by = outline_by
+    model.plots.append(plot)
+
+    with TemporaryDirectory() as tmpdir:
+        # Run OpenMC in geometry plotting mode
+        model.plot_geometry(False, cwd=tmpdir)
+
+        # 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")
+        img = mpimg.imread(str(img_path))
+
+    # Combine R, G, B values into a single int
+    rgb = (img * 256).astype(int)
+    image_value = (rgb[..., 0] << 16) + (rgb[..., 1] << 8) + (rgb[..., 2])
+
+    if basis == "xz":
+        image_value = np.rot90(image_value, 2)
+    elif basis == "yz":
+        image_value = np.rot90(image_value, 2)
+    else:  # basis == 'xy'
+        image_value = np.rot90(image_value, 2)
+    return image_value
 
 # TODO currently we allow slice index, but this code will be useful if want to
 # allow slicing by axis values / coordinates.