diff --git a/src/dagmc_geometry_slice_plotter/utils.py b/src/dagmc_geometry_slice_plotter/utils.py index 62976f9..1f39453 100644 --- a/src/dagmc_geometry_slice_plotter/utils.py +++ b/src/dagmc_geometry_slice_plotter/utils.py @@ -37,21 +37,33 @@ def plot_axis_slice( if view_direction == "-z": plane_normal = [0, 0, -1] rotate_plot = 0 + x_label = "X [cm]" + y_label = "Y [cm]" if view_direction == "z": plane_normal = [0, 0, 1] rotate_plot = 0 + x_label = "X [cm]" + y_label = "Y [cm]" if view_direction == "-y": plane_normal = [0, -1, 0] rotate_plot = 90 + x_label = "X [cm]" + y_label = "Z [cm]" if view_direction == "y": plane_normal = [0, 1, 0] rotate_plot = 90 + x_label = "X [cm]" + y_label = "Z [cm]" if view_direction == "-x": plane_normal = [-1, 0, 0] rotate_plot = -90 + x_label = "Y [cm]" + y_label = "Z [cm]" if view_direction == "x": plane_normal = [1, 0, 0] rotate_plot = 90 + x_label = "Y [cm]" + y_label = "Z [cm]" slice = plot_slice( dagmc_file_or_trimesh_object=dagmc_file_or_trimesh_object, @@ -60,6 +72,9 @@ def plot_axis_slice( rotate_plot=rotate_plot, ) + slice.xlabel(x_label) + slice.ylabel(y_label) + return slice diff --git a/tests/test_plot_axis_slice.py b/tests/test_plot_axis_slice.py new file mode 100644 index 0000000..4ce24ed --- /dev/null +++ b/tests/test_plot_axis_slice.py @@ -0,0 +1,42 @@ +import os +import unittest +from pathlib import Path +import trimesh + +import matplotlib.pyplot as plt +from dagmc_geometry_slice_plotter import plot_axis_slice + + +class TestPlotSliceOfTrimeshObject(unittest.TestCase): + """Tests the neutronics utilities functionality and use cases""" + + def setUp(self): + + h5m_filename_smaller = "tests/dagmc.h5m" + + self.trimesh_mesh_object_smaller = trimesh.load_mesh( + h5m_filename_smaller, process=False + ) + + def test_create_default_plot(self): + """Tests returned object is a matplotlib plot""" + + plot = plot_axis_slice( + view_direction="x", + dagmc_file_or_trimesh_object=self.trimesh_mesh_object_smaller, + ) + + assert isinstance(plot, type(plt)) + + def test_create_default_plot_file(self): + """Tests output file creation""" + + os.system("rm test_plot.png") + + plot = plot_axis_slice( + view_direction="y", + dagmc_file_or_trimesh_object=self.trimesh_mesh_object_smaller, + ) + plot.savefig("test_plot.png") + + assert Path("test_plot.png").is_file()