Merge pull request #32 from fusion-energy/develop

general update

README.md

@@ -91,12 +91,44 @@ plot.savefig('example_4_slice.png')
 
 DAGMC files can also be made using packages like [cad_to_dagmc](https://github.com/fusion-energy/cad_to_dagmc) or [stl_to_h5m](https://github.com/fusion-energy/stl_to_h5m). This example uses CadQuery to make an STL file then converts the STL file to a DAGMC h5m file and plots a slice through the geometry  
 
+# Custom plots
+
+You can also use the package to access the the coordinates and lines that make
+up the outline and then plot these lines with your own MatplotLib script or
+other plotting package.
+
+This example makes use of low level functionality in the package to extract the
+coordinates and lines then plot them in MatPlotLib. This offers users full
+customization of how the plots appear and also allows the plots to be combined
+with other types of plots such as 
+
+```python
+from dagmc_geometry_slice_plotter import get_slice_coordinates
+import matplotlib.pyplot as plt
+
+data = get_slice_coordinates(
+    dagmc_file_or_trimesh_object="dagmc.h5m",
+    plane_origin=[0, 0, 200],
+    plane_normal=[0, 0, 1],
+)
+
+plt.axes().set_aspect("equal")  # scales the x y axis 1:1
+
+for xy in data:
+    plt.plot(*xy, color="black", linewidth=1)
+
+plt.savefig("example_7_slice.png")
+```
 
 # Related packages
 
-This package is used by the [openmc_plot](https://github.com/fusion-energy/openmc_plot) Python package which has a web deployed version at [xsplot.com](http://www.xsplot.com)
+- This package is used by the [openmc_plot](https://github.com/fusion-energy/openmc_plot) Python package which has a web deployed version at [xsplot.com](http://www.xsplot.com)
+
+
+- This package can also be used together with the [openmc_geometry_plot](https://github.com/fusion-energy/openmc_geometry_plot) Python package to combine outline slice plots of DAGMC geometry with colored areas for materials or cells
 
-This package is used by the [regular_mesh_plotter](https://github.com/fusion-energy/regular_mesh_plotter) Python package to combine slice plots with regular mesh tally results and produce images like below.
+- This package can also be used together with the [regular_mesh_plotter](https://github.com/fusion-energy/regular_mesh_plotter) Python package to combine outline
+slice plots with regular mesh tally results and produce images like below.
 
 ![paramak plot openmc regular mesh tally](https://user-images.githubusercontent.com/8583900/138322007-daf1eb6f-ca42-4d9c-9581-8dbc9da94fe5.png)
 

examples/example_7_from_readme.py

@@ -0,0 +1,18 @@
+# this example makes use of the dagmc_geometry_slice_ploter to extract the
+# dagmc geometry outline.
+
+from dagmc_geometry_slice_plotter import get_slice_coordinates
+import matplotlib.pyplot as plt
+
+data = get_slice_coordinates(
+    dagmc_file_or_trimesh_object="dagmc.h5m",
+    plane_origin=[0, 0, 200],
+    plane_normal=[0, 0, 1],
+)
+
+plt.axes().set_aspect("equal")
+
+for xy in data:
+    plt.plot(*xy, color="black", linewidth=1)
+
+plt.savefig("example_7_slice.png")

src/dagmc_geometry_slice_plotter/__init__.py

@@ -1,10 +1,19 @@
-from importlib.metadata import version
+try:
+    # this works for python 3.7 and lower
+    from importlib.metadata import version, PackageNotFoundError
+except (ModuleNotFoundError, ImportError):
+    # this works for python 3.8 and higher
+    from importlib_metadata import version, PackageNotFoundError
+try:
+    __version__ = version("dagmc_geometry_slice_plotter")
+except PackageNotFoundError:
+    from setuptools_scm import get_version
 
-__version__ = version("dagmc_geometry_slice_plotter")
+    __version__ = get_version(root="..", relative_to=__file__)
 
 __all__ = ["__version__"]
 
-from .utils import plot_slice_of_dagmc_file
-from .utils import plot_slice_of_trimesh_object
-from .utils import plot_slice
-from .utils import plot_axis_slice
+from .core import get_slice_coordinates
+from .core import plot_slice_of_trimesh_object
+from .core import plot_slice
+from .core import plot_axis_slice

src/dagmc_geometry_slice_plotter/core.py

@@ -0,0 +1,215 @@
+from pathlib import Path
+from typing import Tuple, Union
+
+import matplotlib.pyplot as plt
+import trimesh
+from matplotlib import transforms
+
+from .utils import view_direction_to_plane_normal
+from .utils import view_direction_to_rotate_plot
+from .utils import view_direction_to_x_y_label
+from .utils import plane_normal_to_x_y_label
+
+
+def plot_axis_slice(
+    dagmc_file_or_trimesh_object: Union[str, trimesh.Trimesh, trimesh.Scene],
+    view_direction: str,
+    plane_origin: Tuple[float, float, float] = None,
+):
+    """Axis aligned plane normal slices through a 3D DAGMC geometry file
+    (h5m format) and produces a matplotlib plot of the slice.
+
+    Args:
+        dagmc_file_or_trimesh_object: the filename of the DAGMC h5m file or a
+            A trimesh mesh object. This can be created from a DAGMC h5m file in
+            the following way 'trimesh_mesh_object = trimesh.load_mesh(
+            dagmc_filename, process=False)'.
+        view_direction: the axis to view the geometry from. Acceptable options
+            are '-z', 'z', '-x', 'x', '-y', 'y' (just like Paraview)
+        plane_origin: the origin of the plain, if None then the centroid of
+            the mesh will be used.
+
+    Return:
+        A matplotlib.pyplot object
+    """
+
+    acceptable_values = ["-z", "z", "-x", "x", "-y", "y"]
+    if view_direction not in acceptable_values:
+        raise ValueError(
+            f"view_direction must be one of the following {acceptable_values}"
+        )
+
+    plane_normal = view_direction_to_plane_normal(view_direction)
+    rotate_plot = view_direction_to_rotate_plot(view_direction)
+
+    slice = plot_slice(
+        dagmc_file_or_trimesh_object=dagmc_file_or_trimesh_object,
+        plane_origin=plane_origin,
+        plane_normal=plane_normal,
+        rotate_plot=rotate_plot,
+    )
+
+    return slice
+
+
+def plot_slice(
+    dagmc_file_or_trimesh_object: Union[str, trimesh.Trimesh, trimesh.Scene],
+    plane_origin: Tuple[float, float, float] = None,
+    plane_normal: Union[Tuple[float, float, float], str] = [0, 0, 1],
+    rotate_plot: float = 0,
+):
+    """Arbitrary plane normal slices through a 3D DAGMC geometry file
+    (h5m format) and produces a matplotlib plot of the slice.
+
+    Args:
+        dagmc_file_or_trimesh_object: the filename of the DAGMC h5m file or a
+            A trimesh mesh object. This can be created from a DAGMC h5m file in
+            the following way 'trimesh_mesh_object = trimesh.load_mesh(
+            dagmc_filename, process=False)'.
+        plane_origin: the origin of the plain, if None then the centroid of
+            the mesh will be used.
+        plane_normal: the plane to slice the geometry on. Defaults to slicing
+            along the Z plane which is input as [0, 0, 1].
+        rotate_plot: the angle in degrees to rotate the plot by. Useful when
+            used in conjunction with changing plane_normal to orientate the
+            plot correctly.
+
+    Return:
+        A matplotlib.pyplot object
+    """
+
+    if isinstance(dagmc_file_or_trimesh_object, str):
+        if not Path(dagmc_file_or_trimesh_object).is_file():
+            raise FileNotFoundError(f"file {dagmc_file_or_trimesh_object} not found.")
+
+        trimesh_mesh_object = trimesh.load_mesh(
+            dagmc_file_or_trimesh_object, process=False
+        )
+    else:
+        trimesh_mesh_object = dagmc_file_or_trimesh_object
+
+    slice = plot_slice_of_trimesh_object(
+        trimesh_mesh_object=trimesh_mesh_object,
+        plane_origin=plane_origin,
+        plane_normal=plane_normal,
+        rotate_plot=rotate_plot,
+    )
+
+    return slice
+
+
+def plot_slice_of_trimesh_object(
+    trimesh_mesh_object: Union[trimesh.Trimesh, trimesh.Scene],
+    plane_origin: Tuple[float, float, float] = None,
+    plane_normal: Tuple[float, float, float] = [0, 0, 1],
+    rotate_plot: float = 0,
+) -> plt:
+    """Slices through a trimesh_mesh_object and produces a matplotlib plot of
+    the slice. Accepts a trimesh_mesh_object which avoids reloading the mesh
+    file each time a new plot is required.
+
+    Args:
+        dagmc_file_or_trimesh_object: the filename of the DAGMC h5m file or a
+            A trimesh mesh object. This can be created from a DAGMC h5m file in
+            the following way 'trimesh_mesh_object = trimesh.load_mesh(
+            dagmc_filename, process=False)'.
+        plane_origin: the origin of the plain, if None then the centroid of
+            the mesh will be used.
+        plane_normal: the plane to slice the geometry on. Defaults to slicing
+            along the Z plane which is input as [0, 0, 1].
+        rotate_plot: the angle in degrees to rotate the plot by. Useful when
+            used in conjunction with changing plane_normal to orientate the
+            plot correctly.
+
+    Return:
+        A matplotlib.pyplot object
+    """
+
+    plt.close()
+
+    # keep plot axis scaled the same
+    plt.axes().set_aspect("equal")  # an option to increase box size "datalim"
+
+    data = get_slice_coordinates(
+        trimesh_mesh_object,
+        plane_origin,
+        plane_normal,
+    )
+
+    if rotate_plot != 0:
+        base = plt.gca().transData
+        rot = transforms.Affine2D().rotate_deg(rotate_plot)
+
+        for i in data:
+            plt.plot(*i, color="black", linewidth=1, transform=rot + base)
+
+    else:
+        for i in data:
+            plt.plot(*i, color="black", linewidth=1)
+
+    x_label, y_label = plane_normal_to_x_y_label(plane_normal)
+    plt.xlabel(x_label)
+    plt.ylabel(y_label)
+
+    return plt
+
+
+def get_slice_coordinates(
+    dagmc_file_or_trimesh_object: Union[str, trimesh.Trimesh, trimesh.Scene],
+    plane_origin: Tuple[float, float, float] = None,
+    plane_normal: Tuple[float, float, float] = [0, 0, 1],
+):
+    """returns the outline x,y coordinates for each line in the slice. Can be
+    plotted by iterating through the lines and expanding them with *
+
+    Args:
+        trimesh_mesh_object: A trimesh mesh object. This can be created from a
+            DAGMC h5m file in the following way 'trimesh_mesh_object =
+            trimesh.load_mesh(dagmc_filename, process=False)'
+        plane_origin: the origin of the plain, if None then the centroid of
+            the mesh will be used.
+        plane_normal: the plane to slice the geometry on. Defaults to slicing
+            along the Z plane which is input as [0, 0, 1].
+        rotate_plot: the angle in degrees to rotate the plot by. Useful when
+            used in conjunction with changing plane_normal to orientate the
+            plot correctly.
+
+    Return:
+        A matplotlib.pyplot object
+    """
+
+    if isinstance(dagmc_file_or_trimesh_object, str):
+        if not Path(dagmc_file_or_trimesh_object).is_file():
+            raise FileNotFoundError(f"file {dagmc_file_or_trimesh_object} not found.")
+
+        trimesh_mesh_object = trimesh.load_mesh(
+            dagmc_file_or_trimesh_object, process=False
+        )
+    else:
+        trimesh_mesh_object = dagmc_file_or_trimesh_object
+
+    if plane_origin is None:
+        plane_origin = trimesh_mesh_object.centroid
+    slice = trimesh_mesh_object.section(
+        plane_origin=plane_origin,
+        plane_normal=plane_normal,
+    )
+
+    if slice is None:
+        msg = (
+            "The geometry slice returned None which means no geometry was "
+            "intersected, try changing the plane_origin or plane_normal"
+        )
+        raise ValueError(msg)
+
+    to_2D = trimesh.geometry.align_vectors(plane_normal, [0, 0, -1])
+
+    slice_2D, to_3D = slice.to_planar(to_2D=to_2D)
+
+    lines = []
+    for entity in slice_2D.entities:
+        discrete = entity.discrete(slice_2D.vertices)
+
+        lines.append(discrete.T)
+
+    return lines