Source code for geomfum.wrap.polyscope

"""Wraps polyscope functions."""

import polyscope as ps

import geomstats.backend as gs
from geomfum.plot import ShapePlotter




class PsMeshPlotter(ShapePlotter):
    """Plotting object to display meshes."""

    # NB: for now assumes only one mesh is plotted

    def __init__(self, colormap="viridis", backend=""):
        super().__init__()

        self.colormap = colormap

        self._plotter = ps
        self._name = "Mymesh"

        self._plotter.init(backend)



    def add_mesh(self, mesh):
        """Add mesh to plot.

        Parameters
        ----------
        mesh : TriangleMesh
            Mesh to be plotted.
        """
        self._plotter.register_surface_mesh(self._name, gs.to_numpy(mesh.vertices), gs.to_numpy( mesh.faces))
        return self




    def set_vertex_scalars(self, scalars, name="scalars"):
        """Set vertex scalars on mesh.

        Parameters
        ----------
        scalars : array-like
            Value at each vertex.
        name : str
            Scalar field name.
        """
        ps.get_surface_mesh(self._name).add_scalar_quantity(
            name,
            scalars,
            defined_on="vertices",
            cmap=self.colormap,
            enabled=True,
        )
        return self




    def highlight_vertices(self, coords, color=(1.0, 0.0, 0.0), size=0.01,):
        """
        Highlight vertices on a mesh using Polyscope by adding a point cloud.

        Parameters
        ----------
        coords : array-like, shape = [n_vertices, 3]
            Coordinates of vertices to highlight.
        color : tuple
            Color of the highlighted vertices (e.g., (1.0, 0.0, 0.0)).
        radius : float
            Radius of the rendered points (visual size).
        """
        name = 'Highlighted_points'
        self._plotter.register_point_cloud(name, coords, radius = size, color = color)
        return self




    def show(self):
        """Display plot."""
        self._plotter.show()