StructurePlotter

Bases: VtkPlotter

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

Parameters:

Name	Type	Description	Default
structure	Structure	The pymatgen Structure object to plot.	required

Source code in src/baderkit/plotting/toolkit/structure.py

class StructurePlotter(VtkPlotter):
    """
    A convenience class for creating plots of crystal structures using
    pyvista's package for VTK.

    Parameters
    ----------
    structure : Structure
        The pymatgen Structure object to plot.

    """

    def __init__(
        self,
        structure: Structure,
        wrap_atoms=True,
        atom_metallicness=0.0,
        atom_roughness=0.0,
        radii_scale=0.3,
        atom_radii=None,
        atom_colors=None,
        **kwargs,
    ):

        # create initial class variables
        self._atom_opacities = np.array([1 for i in range(len(structure))], dtype=float)
        self._wrap_atoms = wrap_atoms
        self._atom_metallicness = atom_metallicness
        self._atom_roughness = atom_roughness
        self._radii_scale = radii_scale
        radii = []
        for s in structure:
            try:
                radius = s.specie.atomic_radius
            except:
                radius = 1.0
            radii.append(radius)
        self._atom_radii = np.array(radii)
        self._atom_colors = np.array(
            [ATOM_COLORS.get(s.specie.symbol, (1.00, 1.00, 1.00)) for s in structure]
        )

        # atom poly data
        self._map_wrapped_to_atoms = None
        self._atom_poly = None
        self._wrapped_atom_poly = None
        self._sphere_mesh = pv.Sphere(
            radius=1.0, theta_resolution=15, phi_resolution=15
        )
        # generate initial plotter
        super().__init__(structure=structure, **kwargs)
        if atom_radii is not None:
            try:
                self.atom_radii = atom_radii
            except:
                logging.info("Improper atom radii provided. Defaults will be used")
        if atom_colors is not None:
            try:
                self.atom_colors = atom_colors
            except:
                logging.info("Improper atom colors provided. Defaults will be used")

    ###########################################################################
    # Properties and Setters
    ###########################################################################
    @property
    def atom_opacities(self) -> NDArray[float]:
        """

        Returns
        -------
        NDArray[float]
            Whether or not each atom is visible. This is actually the opacity
            if desired.

        """
        return self._atom_opacities

    @atom_opacities.setter
    def atom_opacities(self, atom_opacities: NDArray[float]):

        # convert to array
        atom_opacities = np.array(atom_opacities, dtype=np.float64)

        # make sure we have the write shape
        assert (
            self.atom_opacities.shape == atom_opacities.shape
        ), "Length match the number of atoms"

        # set visible atoms
        self._atom_opacities = atom_opacities
        # update
        self._update_site_meshes()

    @property
    def wrap_atoms(self):
        return self._wrap_atoms

    @wrap_atoms.setter
    def wrap_atoms(self, wrap_atoms: bool):
        assert type(wrap_atoms) == bool, "wrap_atoms must be a bool"

        # update
        self._wrap_atoms = wrap_atoms
        self._update_site_meshes()

    @property
    def atom_radii(self) -> NDArray[float]:
        """

        Returns
        -------
        NDArray[float]
            The radius to display for each atom in the structure. The actual
            displayed radius will be radii_scale*radius.

        """
        return self._atom_radii

    @atom_radii.setter
    def atom_radii(self, atom_radii: NDArray[float]):
        # fix atom_radii to be a list and make any negative values == 0.01
        atom_radii = np.array(atom_radii, dtype=np.float64)
        assert (
            self.atom_radii.shape == atom_radii.shape
        ), "Length match the number of atoms"

        atom_radii[atom_radii <= 0.01] = 0.01
        self._atom_radii = atom_radii
        self._update_site_meshes()

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

        Returns
        -------
        float
            A constant to multiply atom radii by

        """
        return self._radii_scale

    @radii_scale.setter
    def radii_scale(self, radii_scale: float):
        # ensure scale is not at or below zero
        radii_scale = float(radii_scale)
        radii_scale = max(radii_scale, 0.1)
        # update
        self._radii_scale = radii_scale
        self._update_site_meshes()

    @property
    def atom_colors(self) -> list[str]:
        """

        Returns
        -------
        list[str]
            The atom_colors to use for each atom as hex codes.

        """
        return self._atom_colors

    @atom_colors.setter
    def atom_colors(self, atom_colors: NDArray[float | str]):
        # for each site, check if the radius has changed and if it has remove it
        # then remake
        assert len(self.atom_colors) == len(
            atom_colors
        ), "Length match the number of atoms"

        # convert provided colors to rgb
        atom_colors = np.array(
            [pv.Color(i).float_rgb for i in atom_colors], dtype=np.float64
        )
        self._atom_colors = atom_colors
        self._update_site_meshes()

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

        Returns
        -------
        float
            The amount of metallic character in the atom display.

        """
        return self._atom_metallicness

    @atom_metallicness.setter
    def atom_metallicness(self, atom_metallicness: float):
        # update all atoms
        actor = self.plotter.actors["atom_glyphs"]
        actor.prop.metallic = atom_metallicness
        self._atom_metallicness = atom_metallicness

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

        Returns
        -------
        float
            The amount of roughness in the atom display.

        """
        return self._atom_roughness

    @atom_roughness.setter
    def atom_roughness(self, atom_roughness: float):
        # update all atoms
        actor = self.plotter.actors["atom_glyphs"]
        actor.prop.roughness = atom_roughness
        self._atom_roughness = atom_roughness

    @property
    def atom_df(self) -> pd.DataFrame:
        """

        Returns
        -------
        atom_df : TYPE
            A dataframe summarizing the properties of the atom meshes.

        """
        # construct a pandas dataframe for each atom
        visible = self.atom_opacities > 0
        atom_df = pd.DataFrame(
            {
                "Label": self.structure.labels,
                "Visible": visible,
                "Color": self.atom_colors,
                "Radius": self.atom_radii,
            }
        )
        return atom_df

    @atom_df.setter
    def atom_df(self, atom_df: pd.DataFrame):
        # set each property from the dataframe
        self.atom_opacities = atom_df["Visible"]
        self.atom_colors = atom_df["Color"]
        self.atom_radii = atom_df["Radius"]

    def _create_atom_polydata(self):
        wrapped_atom_coords = []
        corresponding_sites = []

        for i in range(len(self.structure)):
            frac_coords = self.structure.frac_coords[i]
            # get wrapped atoms as well
            wrapped_coords, shifts = self._wrap_near_edge(frac_coords)
            cart_coords = wrapped_coords @ self.structure.lattice.matrix
            # add to lists
            wrapped_atom_coords.extend(cart_coords)
            for j, shift in enumerate(shifts):
                corresponding_sites.append(i)

        # save the map from wrapped points to their original site
        self._map_wrapped_to_atoms = np.array(corresponding_sites, dtype=int)
        # create poly data of points for both the atoms with and without wrapping
        self._atom_poly = pv.PolyData(self.structure.cart_coords)
        self._wrapped_atom_poly = pv.PolyData(wrapped_atom_coords)

    def _update_site_meshes(self):
        if self._wrapped_atom_poly is None or self._atom_poly is None:
            self._create_atom_polydata()

        # get appropriate poly data
        if self.wrap_atoms:
            # get poly data including wrapped atoms
            atoms = self._wrapped_atom_poly
            # get atom colors
            atom_colors = self.atom_colors[self._map_wrapped_to_atoms]
            # get alpha values
            alpha = self.atom_opacities[self._map_wrapped_to_atoms]
            # get radii
            radii = self.atom_radii[self._map_wrapped_to_atoms] * self.radii_scale
        else:
            # get poly data without wrapped atoms
            atoms = self._atom_poly
            # get atom colors
            atom_colors = self.atom_colors
            # get alpha values
            alpha = self.atom_opacities
            # get radii
            radii = self.atom_radii * self.radii_scale

        sphere = self._sphere_mesh
        # add alpha to colors
        colors = np.column_stack((atom_colors, alpha))

        # update poly data scalars
        atoms["atom_colors"] = colors
        atoms["atom_radii"] = radii

        # generate glyphs
        glyphs = atoms.glyph(geom=sphere, scale="atom_radii", orient=False)

        # add the atom glyphs to our plotter. This automatically overwrites any
        # previous meshes
        self.plotter.add_mesh(
            glyphs,
            scalars="atom_colors",
            rgb=True,
            name="atom_glyphs",
            pbr=self.pbr,
        )

    def _create_plot(self) -> 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.

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

        """
        plotter = super()._create_plot()

        # add site meshes
        self._update_site_meshes()

        return plotter

atom_colors property writable

Returns:

Type	Description
list[str]	The atom_colors to use for each atom as hex codes.

atom_df property writable

Returns:

Name	Type	Description
atom_df	TYPE	A dataframe summarizing the properties of the atom meshes.

atom_metallicness property writable

Returns:

Type	Description
float	The amount of metallic character in the atom display.

atom_opacities property writable

Returns:

Type	Description
NDArray[float]	Whether or not each atom is visible. This is actually the opacity if desired.

atom_radii property writable

Returns:

Type	Description
NDArray[float]	The radius to display for each atom in the structure. The actual displayed radius will be radii_scale*radius.

atom_roughness property writable

Returns:

Type	Description
float	The amount of roughness in the atom display.

background_color property writable

Returns:

Type	Description
str	The color of the plot background as a hex code, rgb array, or color string.

lattice_thickness property writable

Returns:

Type	Description
float	The thickness of the lines outlining the unit cell.

light_color property writable

Returns:

Type	Description
str	The color of the light shining on the scene as a hex code, rgb array, or color string.

light_intensity property writable

Returns:

Type	Description
str	The intensity of the light on the scene from 0-1

parallel_projection property writable

Returns:

Type	Description
bool	If True, a parallel projection scheme will be used rather than perspective.

pbr property writable

Returns:

Type	Description
bool	If True, physically based rendering will be used

radii_scale property writable

Returns:

Type	Description
float	A constant to multiply atom radii by

show_axes property writable

Returns:

Type	Description
bool	Whether or not to show the axis widget.

show_lattice property writable

Returns:

Type	Description
bool	Whether or not to display the outline of the unit cell.

get_plot_html()

Creates an html string representing the current state of the StructurePlotter class.

Returns:

Type	Description
str	The html string representing the current StructurePlotter class.

Source code in src/baderkit/plotting/base/base.py

def get_plot_html(self) -> str:
    """
    Creates an html string representing the current state of the StructurePlotter
    class.

    Returns
    -------
    str
        The html string representing the current StructurePlotter class.

    """
    if sys.platform == "win32":
        # We can return the html directly without opening a subprocess. And
        # we need to because the "fork" start method doesn't work
        html_plotter = self.plotter.export_html(filename=None)
        return html_plotter.read()
    # BUG-FIX: On Linux and maybe MacOS, pyvista's export_html must be run
    # as a main process. To do this within our streamlit apps, we use python's
    # multiprocess to run the process as is done in [stpyvista](https://github.com/edsaac/stpyvista/blob/main/src/stpyvista/trame_backend.py)
    queue = Queue(maxsize=1)
    process = Process(target=_export_html, args=(queue, self.plotter))
    process.start()
    html_plotter = queue.get().read()
    process.join()
    return html_plotter

get_plot_screenshot(filename=None, transparent_background=None, return_img=True, window_size=None, scale=None)

Creates a screenshot of the current state of the StructurePlotter class. This is a wraparound of pyvista's screenshot method

Parameters:

Name	Type	Description	Default
filename	str | Path | BytesIO	Location to write image to. If None, no image is written.	None
transparent_background	bool	Whether to make the background transparent. The default is looked up on the plotter’s theme.	None
return_img	bool	If True, a numpy.ndarray of the image will be returned. Defaults to True.	True
window_size	tuple[int, int]	Set the plotter’s size to this (width, height) before taking the screenshot.	None
scale	int	Set the factor to scale the window size to make a higher resolution image. If None this will use the image_scale property on this plotter which defaults to one.	None

Returns:

Type	Description
NDArray[float]	Array containing pixel RGB and alpha. Sized: [Window height x Window width x 3] if transparent_background is set to False. [Window height x Window width x 4] if transparent_background is set to True.

Source code in src/baderkit/plotting/base/base.py

def get_plot_screenshot(
    self,
    filename: str | Path | io.BytesIO = None,
    transparent_background: bool = None,
    return_img: bool = True,
    window_size: tuple[int, int] = None,
    scale: int = None,
) -> NDArray[float]:
    """
    Creates a screenshot of the current state of the StructurePlotter class.
    This is a wraparound of pyvista's screenshot method

    Parameters
    ----------
    filename: str | Path | io.BytesIO
        Location to write image to. If None, no image is written.

    transparent_background: bool
        Whether to make the background transparent.
        The default is looked up on the plotter’s theme.

    return_img: bool
        If True, a numpy.ndarray of the image will be returned. Defaults to
        True.

    window_size: tuple[int, int]
        Set the plotter’s size to this (width, height) before taking the
        screenshot.

    scale: int
        Set the factor to scale the window size to make a higher resolution image. If None this will use the image_scale property on this plotter which defaults to one.

    Returns
    -------
    NDArray[float]
        Array containing pixel RGB and alpha. Sized:

        [Window height x Window width x 3] if transparent_background is set to False.

        [Window height x Window width x 4] if transparent_background is set to True.

    """

    plotter = self.plotter

    # if our plotter is not currently rendered, we want to temporarily set
    # it to be off screen to take the screenshot, then set it back
    plotter.render()
    screenshot = plotter.screenshot(
        filename=filename,
        transparent_background=transparent_background,
        return_img=return_img,
        window_size=window_size,
        scale=scale,
    )

    return screenshot

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 StructurePlotter class.

Source code in src/baderkit/plotting/base/base.py

def rebuild(self) -> pv.Plotter:
    """
    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
        StructurePlotter class.

    """
    self.plotter = None
    plotter = self._create_plotter()
    plotter.suppress_rendering = True
    self._suppressing = True
    self._create_plot()
    plotter.suppress_rendering = False
    self._suppressing = True

show()

Renders the plot to a window. After closing the window, a new instance must be created to plot again. Pressing q pauses the rendering allowing changes to be made without fully exiting.

Returns:

Type	Description
None.

Source code in src/baderkit/plotting/base/base.py

def show(self):
    """
    Renders the plot to a window. After closing the window, a new instance
    must be created to plot again. Pressing q pauses the rendering allowing
    changes to be made without fully exiting.

    Returns
    -------
    None.

    """

    self.plotter.show()

soft_rebuild()

reuilds the current pyvista plotter object with current settings.

Returns:

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

Source code in src/baderkit/plotting/base/base.py

def soft_rebuild(self) -> pv.Plotter:
    """
    reuilds the current pyvista plotter object with current settings.

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

    """
    plotter = self.plotter
    if plotter is None:
        plotter = self._create_plotter()
    plotter.suppress_rendering = True
    self._suppressing = True
    self._create_plot()
    plotter.suppress_rendering = False
    self._suppressing = True