Grid

Bases: StructurePlotter

Source code in src/baderkit/plotting/core/plotter.py

class GridPlotter(StructurePlotter):
    def __init__(
        self,
        grid: Grid,
        off_screen: bool = False,
        # downscale: int | None = 400,
    ):
        """
        A convenience class for creating plots of crystal structures and isosurfaces
        using pyvista's package for VTK.

        Parameters
        ----------
        grid : Grid
            The Grid object to use for isosurfaces. The structure will be pulled
            from this grid.
        off_screen : bool, optional
            Whether or not the plotter should be in offline mode. The default is False.

        Returns
        -------
        None.

        """
        # apply StructurePlotter kwargs
        structure = grid.structure
        super().__init__(structure=structure, off_screen=off_screen)

        # Grid specific items
        # if downscale is not None:
        #     if grid.voxel_resolution > downscale:
        #         # downscale the grid for speed
        #         logging.info("Grid is above desired resolution. Downscaling.")
        #         grid = grid.regrid(downscale)
        self.grid = grid
        self._show_surface = True
        self._show_caps = True
        self._surface_opacity = 0.8
        self._cap_opacity = 0.8
        self._colormap = "viridis"
        self._use_solid_surface_color = False
        self._use_solid_cap_color = False
        self._surface_color = "#BA8E23"
        self._cap_color = "#BA8E23"

        # wrap values around to get one extra voxel on the far side of each axis.
        values = np.pad(grid.total, pad_width=((0, 1), (0, 1), (0, 1)), mode="wrap")
        self.shape = values.shape
        self.values = values.ravel(order="F")
        self.min_val = self.values.min()
        # make min val slightly above 0
        self.min_val += +0.0000001 * self.min_val
        self.max_val = self.values.max()
        # determine default iso if not provided
        self._iso_val = self.min_val  # np.mean(grid.total)
        # generate the structured grid
        indices = np.indices(self.shape).reshape(3, -1, order="F").T
        self.points = grid.get_cart_coords_from_vox(indices)
        self.structured_grid = self._make_structured_grid(self.values)
        # generate the surface
        self.surface = self.structured_grid.extract_surface()
        # update plotter
        self.plotter = self._create_grid_plot(off_screen)

    def _make_structured_grid(self, values: NDArray[float]) -> pv.StructuredGrid:
        """
        Creates a pyvista StructuredGrid object for making isosurfaces. This
        should generally not be called directly.

        Parameters
        ----------
        values : NDArray[float]
            A 3xN array of values representing the data in the structured grid.
            These should be raveled/reshaped using Fortran's conventions (order='F'')

        Returns
        -------
        structured_grid : pv.StructuredGrid
            A pyvista StructuredGrid with values representing the grid data.

        """
        structured_grid = pv.StructuredGrid()
        structured_grid.points = self.points
        structured_grid.dimensions = self.shape
        structured_grid["values"] = values
        return structured_grid

    @property
    def show_surface(self) -> bool:
        """

        Returns
        -------
        bool
            whether or not to display the isosurface.

        """
        return self._show_surface

    @show_surface.setter
    def show_surface(self, show_surface: bool):
        if "iso" in self.plotter.actors.keys():
            actor = self.plotter.actors["iso"]
            actor.visibility = show_surface
        self._show_surface = show_surface

    @property
    def show_caps(self) -> bool:
        """

        Returns
        -------
        bool
            Whether or not to display caps on the isosurface.

        """
        return self._show_caps

    @show_caps.setter
    def show_caps(self, show_caps: bool):
        if "cap" in self.plotter.actors.keys():
            actor = self.plotter.actors["cap"]
            actor.visibility = show_caps
        self._show_caps = show_caps

    @property
    def surface_opacity(self) -> float:
        """

        Returns
        -------
        float
            Opacity of the isosurface.

        """
        return self._surface_opacity

    @surface_opacity.setter
    def surface_opacity(self, surface_opacity: float):
        if "iso" in self.plotter.actors.keys():
            actor = self.plotter.actors["iso"]
            actor.prop.opacity = surface_opacity
        self._surface_opacity = surface_opacity

    @property
    def cap_opacity(self) -> float:
        """

        Returns
        -------
        float
            Opacity of the caps.

        """
        return self._cap_opacity

    @cap_opacity.setter
    def cap_opacity(self, cap_opacity: float):
        if "cap" in self.plotter.actors.keys():
            actor = self.plotter.actors["cap"]
            actor.prop.opacity = cap_opacity
        self._cap_opacity = cap_opacity

    @property
    def colormap(self) -> str:
        """

        Returns
        -------
        str
            The colormap for the caps and isosurface. This is ignored when the
            surface or caps are set to use solid colors. Valid options are those
            available in matplotlib.

        """
        return self._colormap

    @colormap.setter
    def colormap(self, colormap: str):
        # update settings
        self._colormap = colormap
        if not self.use_solid_surface_color:
            self._add_iso_mesh()
        if not self.use_solid_cap_color:
            self._add_cap_mesh()

    @property
    def use_solid_surface_color(self) -> bool:
        """

        Returns
        -------
        bool
            whether or not to use a solid color for the isosurface.
        """
        return self._use_solid_surface_color

    # TODO: Figure out a way to set the cmap without remaking the surface?
    @use_solid_surface_color.setter
    def use_solid_surface_color(self, use_solid_surface_color: bool):
        # update property
        self._use_solid_surface_color = use_solid_surface_color
        # remove surface and add it back with new color/cmap
        self._add_iso_mesh()

    @property
    def use_solid_cap_color(self) -> bool:
        """

        Returns
        -------
        bool
            whether or not to use a solid color for the caps.
        """
        return self._use_solid_cap_color

    @use_solid_cap_color.setter
    def use_solid_cap_color(self, use_solid_cap_color: bool):
        # update property
        self._use_solid_cap_color = use_solid_cap_color
        # remove cap and add it back with new color/cmap
        self._add_cap_mesh()

    @property
    def surface_color(self) -> str:
        """

        Returns
        -------
        str
            The color to use for the surface as a hex string. This is ignored if
            the surface is not set to use solid colors.

        """
        return self._surface_color

    @surface_color.setter
    def surface_color(self, surface_color: str):
        self._surface_color = surface_color
        if self.use_solid_surface_color:
            self._add_iso_mesh()

    @property
    def cap_color(self):
        """

        Returns
        -------
        str
            The color to use for the caps as a hex string. This is ignored if
            the caps are not set to use solid colors.

        """
        return self._cap_color

    @cap_color.setter
    def cap_color(self, cap_color: str):
        self._cap_color = cap_color
        if self.use_solid_cap_color:
            self._add_cap_mesh()

    @property
    def iso_val(self) -> float:
        """

        Returns
        -------
        float
            The value to set the isosurface to.

        """
        return self._iso_val

    @iso_val.setter
    def iso_val(self, iso_val: float):
        # make sure iso value is within range
        iso_val = max(self.min_val, min(iso_val, self.max_val))
        self._update_surface_mesh(iso_val)
        self._add_iso_mesh()
        self._add_cap_mesh()

    def _update_surface_mesh(self, iso_value: float):
        """
        Updates the surface meshes to the provided iso_value

        Parameters
        ----------
        iso_value : float
            The value to update the surface meshes to

        Returns
        -------
        None.

        """
        self.iso = self.structured_grid.contour([iso_value])
        self.cap = self.surface.contour_banded(
            2, rng=[iso_value, self.max_val], generate_contour_edges=False
        )

    def _get_surface_kwargs(self) -> dict:
        """
        Generates the keyword arguments to use when adding the surface to
        the plotter. We need this because setting a solid color vs. a colormap
        requires different keywords

        Returns
        -------
        dict
            The keyword arguments for setting the surface mesh in the plotter.

        """
        kwargs = {
            "opacity": self.surface_opacity,
            "pbr": True,
            "name": "iso",
        }
        kwargs["color"] = self.surface_color
        if self.use_solid_surface_color:
            kwargs["color"] = self.surface_color
        else:
            kwargs["colormap"] = self.colormap
            kwargs["scalars"] = "values"
            kwargs["clim"] = [self.min_val, self.max_val]
            kwargs["show_scalar_bar"] = False
        return kwargs

    def _get_cap_kwargs(self) -> dict:
        """
        Generates the keyword arguments to use when adding the caps to
        the plotter. We need this because setting a solid color vs. a colormap
        requires different keywords

        Returns
        -------
        dict
            The keyword arguments for setting the caps mesh in the plotter.

        """
        kwargs = {
            "opacity": self.cap_opacity,
            "pbr": True,
            "name": "cap",
        }
        if self.use_solid_cap_color:
            kwargs["color"] = self.cap_color
        else:
            kwargs["cmap"] = self.colormap
            kwargs["scalars"] = "values"
            kwargs["clim"] = [self.min_val, self.max_val]
            kwargs["show_scalar_bar"] = False
        return kwargs

    def _add_iso_mesh(self):
        """
        Removes the current isosurface mesh than adds a new one.

        Returns
        -------
        None.

        """
        if self.show_surface:
            if "iso" in self.plotter.actors.keys():
                self.plotter.remove_actor("iso")
            if len(self.iso["values"]) > 0:
                self.plotter.add_mesh(self.iso, **self._get_surface_kwargs())

    def _add_cap_mesh(self) -> dict:
        """
        Removes the current cap mesh than adds a new one.

        Returns
        -------
        None.

        """
        if self.show_caps:
            if "cap" in self.plotter.actors.keys():
                self.plotter.remove_actor("cap")
            if len(self.iso["values"]) > 0:
                self.plotter.add_mesh(self.cap, **self._get_cap_kwargs())

    def _create_grid_plot(self, off_screen) -> pv.Plotter():
        """
        Generates a pyvista.Plotter object from the current class variables.
        This is called when the class is first instanced and generally shouldn't
        be called again.

        Parameters
        ----------
        off_screen : bool
            Whether or not the plotter should run in off_screen mode.

        Returns
        -------
        plotter : pv.Plotter
            A pyvista Plotter object representing the provided Structure object.

        """
        # get initial plotter with structure
        plotter = self._create_structure_plot(off_screen=off_screen)
        # generate initial surface meshes
        self._update_surface_mesh(self.iso_val)
        # Add iso mesh
        if len(self.iso["values"]) > 0:
            plotter.add_mesh(self.iso, **self._get_surface_kwargs())
        # Add cap mesh
        if len(self.cap["values"]) > 0:
            plotter.add_mesh(self.cap, **self._get_cap_kwargs())
        return plotter

    def rebuild(self):
        """
        Builds a new pyvista plotter object representing the current state of
        the Plotter class.

        Returns
        -------
        pv.Plotter
            A pyvista Plotter object representing the current state of the
            GridPlotter class.

        """
        return self._create_grid_plot(self._off_screen)

cap_color property writable

Returns:

Type	Description
str	The color to use for the caps as a hex string. This is ignored if the caps are not set to use solid colors.

cap_opacity property writable

Returns:

Type	Description
float	Opacity of the caps.

colormap property writable

Returns:

Type	Description
str	The colormap for the caps and isosurface. This is ignored when the surface or caps are set to use solid colors. Valid options are those available in matplotlib.

iso_val property writable

Returns:

Type	Description
float	The value to set the isosurface to.

show_caps property writable

Returns:

Type	Description
bool	Whether or not to display caps on the isosurface.

show_surface property writable

Returns:

Type	Description
bool	whether or not to display the isosurface.

surface_color property writable

Returns:

Type	Description
str	The color to use for the surface as a hex string. This is ignored if the surface is not set to use solid colors.

surface_opacity property writable

Returns:

Type	Description
float	Opacity of the isosurface.

use_solid_cap_color property writable

Returns:

Type	Description
bool	whether or not to use a solid color for the caps.

use_solid_surface_color property writable

Returns:

Type	Description
bool	whether or not to use a solid color for the isosurface.

__init__(grid, off_screen=False)

A convenience class for creating plots of crystal structures and isosurfaces using pyvista's package for VTK.

Parameters:

Name	Type	Description	Default
grid	Grid	The Grid object to use for isosurfaces. The structure will be pulled from this grid.	required
off_screen	bool	Whether or not the plotter should be in offline mode. The default is False.	False

Returns:

Type	Description
None.

Source code in src/baderkit/plotting/core/plotter.py

def __init__(
    self,
    grid: Grid,
    off_screen: bool = False,
    # downscale: int | None = 400,
):
    """
    A convenience class for creating plots of crystal structures and isosurfaces
    using pyvista's package for VTK.

    Parameters
    ----------
    grid : Grid
        The Grid object to use for isosurfaces. The structure will be pulled
        from this grid.
    off_screen : bool, optional
        Whether or not the plotter should be in offline mode. The default is False.

    Returns
    -------
    None.

    """
    # apply StructurePlotter kwargs
    structure = grid.structure
    super().__init__(structure=structure, off_screen=off_screen)

    # Grid specific items
    # if downscale is not None:
    #     if grid.voxel_resolution > downscale:
    #         # downscale the grid for speed
    #         logging.info("Grid is above desired resolution. Downscaling.")
    #         grid = grid.regrid(downscale)
    self.grid = grid
    self._show_surface = True
    self._show_caps = True
    self._surface_opacity = 0.8
    self._cap_opacity = 0.8
    self._colormap = "viridis"
    self._use_solid_surface_color = False
    self._use_solid_cap_color = False
    self._surface_color = "#BA8E23"
    self._cap_color = "#BA8E23"

    # wrap values around to get one extra voxel on the far side of each axis.
    values = np.pad(grid.total, pad_width=((0, 1), (0, 1), (0, 1)), mode="wrap")
    self.shape = values.shape
    self.values = values.ravel(order="F")
    self.min_val = self.values.min()
    # make min val slightly above 0
    self.min_val += +0.0000001 * self.min_val
    self.max_val = self.values.max()
    # determine default iso if not provided
    self._iso_val = self.min_val  # np.mean(grid.total)
    # generate the structured grid
    indices = np.indices(self.shape).reshape(3, -1, order="F").T
    self.points = grid.get_cart_coords_from_vox(indices)
    self.structured_grid = self._make_structured_grid(self.values)
    # generate the surface
    self.surface = self.structured_grid.extract_surface()
    # update plotter
    self.plotter = self._create_grid_plot(off_screen)

rebuild()

Builds a new pyvista plotter object representing the current state of the Plotter class.

Returns:

Type	Description
Plotter	A pyvista Plotter object representing the current state of the GridPlotter class.

Source code in src/baderkit/plotting/core/plotter.py

def rebuild(self):
    """
    Builds a new pyvista plotter object representing the current state of
    the Plotter class.

    Returns
    -------
    pv.Plotter
        A pyvista Plotter object representing the current state of the
        GridPlotter class.

    """
    return self._create_grid_plot(self._off_screen)