Merge pull request #38 from fusion-energy/remove_base_plot

Remove base plot

README.md

@@ -50,15 +50,16 @@ Example 1 shows a OpenMC tally plotted
 ```python
 
 import regular_mesh_plotter as rmp
+import matplotlib.pyplot as plt
 
-my_plot = rmp.plot_regular_mesh_tally(
+rmp.plot_regular_mesh_tally(
     tally=my_tally,
     std_dev_or_tally_value="tally_value",
     x_label="X [cm]",
     y_label="Y [cm]",
 )
 
-my_plot.save_fig('openmc_mesh_tally_plot.png')
+plt.savefig('openmc_mesh_tally_plot.png')
 ```
 
 Example 4 shows a OpenMC tally plotted with an underlying DAGMC geometry

examples/examples_from_readme.py

@@ -2,7 +2,7 @@
 
 from dagmc_geometry_slice_plotter import plot_slice_of_dagmc_geometry
 from stl_to_h5m import stl_to_h5m
-
+import matplotlib.pyplot as plt
 import paramak
 
 my_reactor = paramak.SubmersionTokamak(
@@ -45,29 +45,29 @@
 )
 
 
-plot = plot_slice_of_dagmc_geometry(
+plot_slice_of_dagmc_geometry(
     dagmc_file_or_trimesh_object="dagmc.h5m",
     plane_normal=[0, 0, 1],
     output_filename="my_plot1.png",
 )
 
 
-plot = plot_slice_of_dagmc_geometry(
+plot_slice_of_dagmc_geometry(
     dagmc_file_or_trimesh_object="dagmc.h5m",
     plane_origin=[0, 0, 300],
     plane_normal=[0, 0, 1],
     output_filename="my_plot2.png",
 )
 
 
-plot = plot_slice_of_dagmc_geometry(
+plot_slice_of_dagmc_geometry(
     dagmc_file_or_trimesh_object="dagmc.h5m",
     plane_normal=[0, 1, 0],
     rotate_plot=45,
     output_filename="my_plot3.png",
 )
 
-plot.savefig("big.png", dpi=600)
+plt.savefig("big.png", dpi=600)
 
 plot_slice_of_dagmc_geometry(
     dagmc_file_or_trimesh_object="dagmc.h5m",

examples/plot_regular_mesh_dose_tally.py

@@ -1,5 +1,6 @@
 import regular_mesh_plotter as rmp
 import openmc
+import matplotlib.pyplot as plt
 
 # loads in the statepoint file containing tallies
 statepoint = openmc.StatePoint(filepath="statepoint.2.h5")
@@ -8,7 +9,7 @@
 my_tally = statepoint.get_tally(name="neutron_effective_dose_on_2D_mesh_xy")
 
 # creates a plot of the mesh tally
-my_plot = rmp.plot_regular_mesh_dose_tally(
+rmp.plot_regular_mesh_dose_tally(
     tally=my_tally,  # the openmc tally object to plot, must be a 2d mesh tally
     filename="plot_regular_mesh_dose_tally.png",  # the filename of the picture file saved
     x_label="X [cm]",
@@ -20,4 +21,4 @@
 )
 
 # displays the plot
-my_plot.show()
+plt.show()

examples/plot_regular_mesh_dose_tally_with_geometry.py

@@ -1,5 +1,6 @@
 import regular_mesh_plotter as rmp
 import openmc
+import matplotlib.pyplot as plt
 
 # loads in the statepoint file containing tallies
 statepoint = openmc.StatePoint(filepath="statepoint.2.h5")
@@ -8,7 +9,7 @@
 my_tally = statepoint.get_tally(name="neutron_effective_dose_on_2D_mesh_xy")
 
 # creates a plot of the mesh
-my_plot = rmp.plot_regular_mesh_dose_tally_with_geometry(
+rmp.plot_regular_mesh_dose_tally_with_geometry(
     tally=my_tally,
     dagmc_file_or_trimesh_object="dagmc.h5m",
     filename="plot_regular_mesh_dose_tally_with_geometry.png",
@@ -23,4 +24,4 @@
     source_strength=1e20,
 )
 
-my_plot.show()
+plt.show()

examples/plot_regular_mesh_tally.py

@@ -12,7 +12,6 @@
     tally=my_tally,
     filename="neutron_effective_dose_on_2D_mesh_xy.png",
     label="",
-    base_plt=None,
     x_label="X [cm]",
     y_label="Y [cm]",
     # rotate_plot: float = 0,

examples/plot_regular_mesh_tally_with_geometry.py

@@ -1,14 +1,16 @@
 import regular_mesh_plotter as rmp
 import openmc
 
+import matplotlib.pyplot as plt
+
 # loads in the statepoint file containing tallies
 statepoint = openmc.StatePoint(filepath="statepoint.2.h5")
 
 # gets one tally from the available tallies
 my_tally = statepoint.get_tally(name="neutron_effective_dose_on_2D_mesh_xy")
 
 # creates a plot of the mesh
-my_plot = rmp.plot_regular_mesh_tally_with_geometry(
+rmp.plot_regular_mesh_tally_with_geometry(
     tally=my_tally,
     dagmc_file_or_trimesh_object="dagmc.h5m",
     std_dev_or_tally_value="tally_value",
@@ -24,4 +26,4 @@
     source_strength=None,
 )
 
-my_plot.show()
+plt.show()

regular_mesh_plotter/core.py

@@ -19,21 +19,13 @@ def plot_regular_mesh_values(
     filename: Optional[str] = None,
     label="",
     title=None,
-    base_plt=None,
     extent=None,
     x_label="X [cm]",
     y_label="Y [cm]",
     rotate_plot: float = 0,
     **kwargs
 ):
 
-    if base_plt:
-        plt = base_plt
-    else:
-        import matplotlib.pyplot as plt
-
-        plt.plot()
-
     if rotate_plot != 0:
         x_center = sum(extent[:2]) / 2
         y_center = sum(extent[2:]) / 2
@@ -53,9 +45,8 @@ def plot_regular_mesh_values(
 
     # image_map = fig.imshow(values, norm=scale, vmin=vmin)
     plt.colorbar(image_map, label=label)
-    if filename:
+    if filename:  # TODO should we not let the users do that?
         plt.savefig(filename, dpi=300)
-    return plt
 
 
 def plot_regular_mesh_values_with_geometry(
@@ -81,21 +72,18 @@ def plot_regular_mesh_values_with_geometry(
         rotate_plot=rotate_geometry,
     )
 
-    both = plot_regular_mesh_values(
+    plot_regular_mesh_values(
         values=values,
         filename=filename,
         label=label,
         title=title,
-        base_plt=slice,
         extent=extent,
         x_label=x_label,
         y_label=y_label,
         rotate_plot=rotate_mesh,
         **kwargs
     )
 
-    return both
-
 
 def plot_regular_mesh_tally_with_geometry(
     tally,
@@ -126,35 +114,31 @@ def plot_regular_mesh_tally_with_geometry(
 
     extent = get_tally_extent(tally)
 
-    base_plt = dgsp.plot_slice_of_dagmc_geometry(
+    dgsp.plot_slice_of_dagmc_geometry(
         dagmc_file_or_trimesh_object=dagmc_file_or_trimesh_object,
         plane_origin=plane_origin,
         plane_normal=plane_normal,
         rotate_plot=rotate_geometry,
     )
 
-    plot = plot_regular_mesh_values(
+    plot_regular_mesh_values(
         values=value,
         filename=filename,
         label=label,
         title=title,
-        base_plt=base_plt,
         extent=extent,
         x_label=x_label,
         y_label=y_label,
         rotate_plot=rotate_mesh,
         **kwargs
     )
 
-    return plot
-
 
 def plot_regular_mesh_tally(
     tally,
     filename: Optional[str] = None,
     label="",
     title=None,
-    base_plt=None,
     x_label="X [cm]",
     y_label="Y [cm]",
     rotate_plot: float = 0,
@@ -177,28 +161,24 @@ def plot_regular_mesh_tally(
 
     extent = get_tally_extent(tally)
 
-    plot = plot_regular_mesh_values(
+    plot_regular_mesh_values(
         values=value,
         filename=filename,
         label=label,
         title=title,
-        base_plt=base_plt,
         extent=extent,
         x_label=x_label,
         y_label=y_label,
         rotate_plot=rotate_plot,
         **kwargs
     )
 
-    return plot
-
 
 def plot_regular_mesh_dose_tally(
     tally,
     filename: Optional[str] = None,
     label="",
     title=None,
-    base_plt=None,
     x_label="X [cm]",
     y_label="Y [cm]",
     rotate_plot: float = 0,
@@ -221,21 +201,18 @@ def plot_regular_mesh_dose_tally(
 
     extent = get_tally_extent(tally)
 
-    plot = plot_regular_mesh_values(
+    plot_regular_mesh_values(
         values=value,
         filename=filename,
         label=label,
         title=title,
-        base_plt=base_plt,
         extent=extent,
         x_label=x_label,
         y_label=y_label,
         rotate_plot=rotate_plot,
         **kwargs
     )
 
-    return plot
-
 
 def plot_regular_mesh_dose_tally_with_geometry(
     tally,
@@ -255,19 +232,18 @@ def plot_regular_mesh_dose_tally_with_geometry(
     **kwargs
 ):
 
-    slice = dgsp.plot_slice_of_dagmc_geometry(
+    dgsp.plot_slice_of_dagmc_geometry(
         dagmc_file_or_trimesh_object=dagmc_file_or_trimesh_object,
         plane_origin=plane_origin,
         plane_normal=plane_normal,
         rotate_plot=rotate_geometry,
     )
 
-    both = plot_regular_mesh_dose_tally(
+    plot_regular_mesh_dose_tally(
         tally=tally,
         filename=filename,
         label=label,
         title=title,
-        base_plt=slice,
         x_label=x_label,
         y_label=y_label,
         rotate_plot=rotate_mesh,
@@ -276,5 +252,3 @@ def plot_regular_mesh_dose_tally_with_geometry(
         std_dev_or_tally_value=std_dev_or_tally_value,
         **kwargs
     )
-
-    return both

tests/test_plot_regular_mesh_values.py

@@ -2,7 +2,6 @@
 import unittest
 from pathlib import Path
 
-import matplotlib
 import numpy as np
 from regular_mesh_plotter import plot_regular_mesh_values
 
@@ -77,9 +76,7 @@ def setUp(self):
 
     def test_plot_regular_mesh_values(self):
 
-        test_plot = plot_regular_mesh_values(values=self.values)
-
-        assert isinstance(test_plot, type(matplotlib.pyplot))
+        plot_regular_mesh_values(values=self.values)
 
     def test_plot_regular_mesh_values_with_output(self):
 

tests/test_plot_regular_mesh_values_with_geometry.py

@@ -77,20 +77,7 @@ def setUp(self):
 
     def test_plot_regular_mesh_values_with_geometry(self):
 
-        test_plot = plot_regular_mesh_values_with_geometry(
+        plot_regular_mesh_values_with_geometry(
             values=self.values,
             dagmc_file_or_trimesh_object="tests/example.stl",  # this could be a h5m file
         )
-
-        assert isinstance(test_plot, type(matplotlib.pyplot))
-
-    def test_plot_regular_mesh_values_with_geometry(self):
-
-        test_plot = plot_regular_mesh_values_with_geometry(
-            values=self.values,
-            dagmc_file_or_trimesh_object="tests/example.stl",  # this could be a h5m file
-            filename="test.png",
-        )
-
-        assert isinstance(test_plot, type(matplotlib.pyplot))
-        assert Path("test.png").is_file()