visualization.py

from enum import Enum

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv
from contourpy import LineType
from matplotlib.colors import LinearSegmentedColormap, Normalize
from matplotlib.tri import Triangulation

from geometry import RevoluteLink


class LineType(Enum):
    REGULAR = 0
    SENSOR = 1
    ACTUATED = 2


class PlotHandler:
    LINE_WIDTH = 3
    X_SIZE = 900
    Y_SIZE = 600
    SHOW_FRAME_NUMBER = True
    SHOW_DEGREES = True  # If false, shows radians
    PAUSE_BETWEEN_FRAMES_MS = 30  # milliseconds
    MECH_NAME = "Mechanism Visualization"
    ANG_DISP_NAME = "Angular Position Visualization"
    JACOBIAN_COLOUR_BAR_LIMITS = (0.0, 4.0)

    JACOBIAN_COLOUR_BAR_POINTS = [
        (0.0, "red"),  # value = 0
        (0.3333 / 4.0, "yellow"),  # normalize 0 into 0–1 range later
        (1 / 4.0, "green"),  # normalize 1 into 0–1 range later
        (3 / 4.0, "yellow"),  # normalize 2 into 0–1 range later
        (1.0, "red"),  # value = 4
    ]

    # Build colourmap
    JACOBIAN_COLOUR_MAP = LinearSegmentedColormap.from_list(
        "custom_step", JACOBIAN_COLOUR_BAR_POINTS
    )
    JACOBIAN_COLOUR_NORMALIZE = Normalize(vmin=0, vmax=4)

    @staticmethod
    def _add_observer(target, event_name, callback):
        if hasattr(target, "AddObserver"):  # VTK style
            return target.AddObserver(event_name, callback)
        if hasattr(target, "add_observer"):  # PyVista wrapper style
            return target.add_observer(event_name, callback)
        raise AttributeError("Interactor has no AddObserver/add_observer")

    @staticmethod
    def _create_repeating_timer(iren, ms):
        # Try common variants across VTK/PyVista wrappers
        for name in (
            "CreateRepeatingTimer",
            "CreateTimer",
            "create_repeating_timer",
            "create_timer",
        ):
            if hasattr(iren, name):
                fn = getattr(iren, name)
                try:
                    # Some variants accept (ms, repeating=True/False); others just (ms)
                    return fn(ms, True) if fn.__code__.co_argcount >= 2 else fn(ms)
                except TypeError:
                    return fn(ms)
        raise AttributeError("Interactor has no repeating timer method")

    def on_timer(self, _, __):
        if self.is_1D:
            if (
                self.increasing
                and self.current_data_index >= len(self.joint_positions) - 1
            ) or (self.current_data_index <= 0 and not self.increasing):
                self.increasing = not self.increasing

            if self.increasing:
                self.current_data_index = self.current_data_index + 1
            else:
                self.current_data_index = self.current_data_index - 1
        else:
            if self.current_data_index >= len(self.joint_positions) - 1:
                self.current_data_index = 0
                self.direction = (self.direction + 1) % 4
            else:
                self.current_data_index = self.current_data_index + 1

        data = self.get_data_for_current_index()

        for i in range(len(self.meshes)):
            self.meshes[i].points[:] = [data[i * 2][0:3], data[i * 2 + 1][0:3]]
            pass

        if self.SHOW_FRAME_NUMBER:
            self.text.SetText(
                1, str(self.direction) + " - " + str(self.current_data_index)
            )

        # Update link labels
        for i in range(len(self.link_names) - 1):
            self.link_label_locations[i].points[:] = (
                data[i * 2][0:3] + data[i * 2 + 1][0:3]
            ) / 2
            self.link_label_locations[i].Modified()
        self.pl.render()

    # Start the repeating timer *after* the first render so the interactor is live
    def start_timer_once(self, _=None):
        if getattr(self.pl, "_timer_started", False):
            return
        self.pl._timer_started = True
        iren = self.pl.iren  # vtkRenderWindowInteractor
        self._add_observer(iren, "TimerEvent", self.on_timer)
        self.pl._timer_id = self._create_repeating_timer(
            iren, self.PAUSE_BETWEEN_FRAMES_MS
        )

    def hover(self, event):
        """Hover event that generates annotations for plots"""
        # Hide all annotations by default
        any_visible = False
        for p in self.plots:
            p["annot"].set_visible(False)

        if event.inaxes is None:
            self.fig.canvas.draw_idle()
            return

        # Find which subplot the mouse is in
        for p in self.plots:
            if event.inaxes is p["ax"]:
                tri = p["tri"]
                z = p["z"]
                annot = p["annot"]
                trifinder = p["trifinder"]

                # event.xdata, event.ydata are in data coords of this axes
                if event.xdata is None or event.ydata is None:
                    break

                tri_index = trifinder(event.xdata, event.ydata)

                if tri_index == -1:
                    break  # mouse not over any triangle in this axes

                # For shading='flat', a single value per triangle
                tri_vertices = tri.triangles[tri_index]
                val = z[tri_vertices].mean()

                annot.xy = (event.xdata, event.ydata)
                annot.set_text(f"{val:.3f}")
                annot.set_visible(True)
                any_visible = True
                break  # we've handled this event

        if any_visible:
            self.fig.canvas.draw_idle()
        else:
            self.fig.canvas.draw_idle()

    def __init__(
        self,
        joint_positions,
        sensor_values,
        geometry_manager,
    ) -> None:
        # Setup window for mechanism visualization
        self.pl = pv.Plotter(
            window_size=(self.X_SIZE, self.Y_SIZE), title=self.MECH_NAME
        )
        self.pl.add_text(self.MECH_NAME, font_size=10)
        self.pl.show_grid()
        self.pl.add_axes()
        self.meshes = []
        self.link_names = geometry_manager.get_link_names()
        self.sensor_joint_ids = geometry_manager.get_sensor_joint_ids()
        self.actuated_joint_ids = geometry_manager.get_actuated_joint_ids()
        self.geometry_manager = geometry_manager

        self.joint_positions = joint_positions
        self.increasing = True
        self.current_data_index = 0
        self.plots = []
        self.fig = []

        matplotlib.use("Qt5Agg")

        # Line direction: ranges from 0 to 3
        self.direction = 0
        self.is_1D = len(self.actuated_joint_ids) == 1

        data = self.get_data_for_current_index()

        for i in range(0, len(data) - 1, 2):
            mesh = pv.Line(data[i][0:3], data[i + 1][0:3])

            if int(data[i][3]) == LineType.REGULAR.value:
                self.pl.add_mesh(mesh, color="black", line_width=self.LINE_WIDTH)
            elif int(data[i][3]) == LineType.SENSOR.value:
                self.pl.add_mesh(mesh, color="blue", line_width=self.LINE_WIDTH)
            elif int(data[i][3]) == LineType.ACTUATED.value:
                self.pl.add_mesh(mesh, color="green", line_width=self.LINE_WIDTH)

            self.meshes.append(mesh)

        i = 0
        self.link_labels = []
        self.link_label_locations = []
        # Add link labels
        for i in range(len(self.link_names) - 1):
            self.link_label_locations.append(
                pv.PolyData((data[i * 2][0:3] + data[i * 2 + 1][0:3]) / 2)
            )
            self.link_labels.append(
                self.pl.add_point_labels(
                    self.link_label_locations[i],
                    [self.link_names[i]],
                    font_size=10,
                    text_color="black",
                    shape_opacity=0.0,
                    show_points=False,
                    justification_horizontal="center",
                    justification_vertical="center",
                    always_visible=True,
                    name=f"label_{i}",
                )
            )

        self.text = self.pl.add_text(
            "",
            position="upper_left",
            font_size=12,
            color="black",
        )

        if len(self.actuated_joint_ids) == 1:
            plt.ion()
            # Plot angular displacements
            plt.subplot(1, 2, 1)
            angular_displacement_data = (
                self.get_data_for_angular_displacement_visualization(sensor_values)
            )
            if self.SHOW_DEGREES:
                angular_displacement_data[:, :] = np.degrees(
                    angular_displacement_data[:, :]
                )

            plt.plot(angular_displacement_data[:, 0], angular_displacement_data[:, 1])
            plt.xlabel(
                f"{geometry_manager.actuated_links[0].name} \n Input Joint Angle (deg)"
                if self.SHOW_DEGREES
                else f"{geometry_manager.actuated_links[0].name} \n Input Joint Angle (rad)"
            )
            plt.ylabel(
                f"{geometry_manager.sensors[0].name} Angular Position (deg)"
                if self.SHOW_DEGREES
                else f"{geometry_manager.sensors[0].name} Angular Position (rad)"
            )
            plt.title("Angular Position")
            plt.grid()
            plt.show(block=False)

            plt.subplot(1, 2, 2)

            # Plot Sensitivities
            plt.plot(angular_displacement_data[:, 0], angular_displacement_data[:, 2])
            plt.xlabel(
                f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (deg)"
                if self.SHOW_DEGREES
                else f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (rad)"
            )
            plt.ylabel(
                f"{geometry_manager.sensors[0].name} Change In Angle Per Input Angle( deg/deg)"
                if self.SHOW_DEGREES
                else f"{geometry_manager.sensors[0].name} Change In Angle Per Input Angle( rad/rad)"
            )
            plt.title("Angular Sensitivity")
            plt.grid()
            plt.show(block=False)

        else:
            # Plot positions

            plt.ion()
            self.fig, self.axs = plt.subplots(
                3, len(sensor_values[0]) - 2, figsize=(15, 8)
            )
            self.fig.tight_layout(pad=4.0)
            if self.SHOW_DEGREES:
                sensor_values[:, 2:] = np.degrees(sensor_values[:, 2:])
                if isinstance(geometry_manager.actuated_links[0], RevoluteLink):
                    sensor_values[:, 0] = np.degrees(sensor_values[:, 0])
                if isinstance(geometry_manager.actuated_links[1], RevoluteLink):
                    sensor_values[:, 1] = np.degrees(sensor_values[:, 1])

            for sensor_id in range(len(sensor_values[0]) - 2):
                ax = plt.subplot(3, len(sensor_values[0]) - 2, sensor_id + 1)
                tri = Triangulation(sensor_values[:, 0], sensor_values[:, 1])
                tpc = ax.tripcolor(
                    tri,
                    sensor_values[:, sensor_id + 2],
                    cmap="viridis",
                    edgecolor="none",
                    shading="gouraud",
                    label=f"Sensor {sensor_id}",
                )
                annot = ax.annotate(
                    "",
                    xy=(0, 0),
                    xytext=(10, 10),
                    textcoords="offset points",
                    bbox=dict(boxstyle="round", fc="w"),
                    ha="center",
                    visible=False,
                )
                trifinder = tri.get_trifinder()
                self.plots.append(
                    dict(
                        ax=ax,
                        tri=tri,
                        z=sensor_values[:, sensor_id + 2],
                        tpc=tpc,
                        annot=annot,
                        trifinder=trifinder,
                    )
                )
                plt.colorbar(
                    tpc,
                    label=(
                        f"Angular Position (deg)"
                        if self.SHOW_DEGREES
                        else f"Angular Position (rad)"
                    ),
                )
                if isinstance(geometry_manager.actuated_links[0], RevoluteLink):
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (rad)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (rad)"
                    )
                else:
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nDisplacement (mm)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nDisplacement (mm)"
                    )
                plt.title(
                    f"{self.ANG_DISP_NAME} - {geometry_manager.sensors[sensor_id].name}"
                )
                plt.grid()

            plt.show(block=False)

        # Jacobian visualization
        if len(self.actuated_joint_ids) == 2:
            # Rearrange x,y, output data for jacobian visualization
            x = np.unique(sensor_values[:, 0])
            y = np.unique(sensor_values[:, 1])
            z = [np.empty((len(x), len(y))), np.empty((len(x), len(y)))]
            index = 0
            for i in range(len(x)):
                for j in range(len(y)):
                    z[0][i][j] = sensor_values[index][2]
                    z[1][i][j] = sensor_values[index][3]
                    index += 1

            for sensor_id in range(len(sensor_values[0]) - 2):
                gradient_x, gradient_y = np.gradient(z[sensor_id], x, y)
                gradient_x = 1 / gradient_x
                gradient_y = 1 / gradient_y

                # Plot gradient for first revolute joint
                ax = plt.subplot(3, len(sensor_values[0]) - 2, sensor_id + 3)
                tri = Triangulation(sensor_values[:, 0], sensor_values[:, 1])
                zz = abs(gradient_x.ravel())
                if isinstance(geometry_manager.actuated_links[0], RevoluteLink):
                    tpc = ax.tripcolor(
                        tri,
                        zz,
                        cmap=self.JACOBIAN_COLOUR_MAP,
                        norm=self.JACOBIAN_COLOUR_NORMALIZE,
                        edgecolor="none",
                        shading="gouraud",
                        label=f"Sensor {sensor_id}",
                        clim=self.JACOBIAN_COLOUR_BAR_LIMITS,
                    )
                else:
                    tpc = ax.tripcolor(
                        tri,
                        zz,
                        edgecolor="none",
                        shading="gouraud",
                        label=f"Sensor {sensor_id}",
                    )
                annot = ax.annotate(
                    "",
                    xy=(0, 0),
                    xytext=(10, 10),
                    textcoords="offset points",
                    bbox=dict(boxstyle="round", fc="w"),
                    ha="center",
                    visible=False,
                )
                trifinder = tri.get_trifinder()
                self.plots.append(
                    dict(
                        ax=ax,
                        tri=tri,
                        z=zz,
                        tpc=tpc,
                        annot=annot,
                        trifinder=trifinder,
                    )
                )
                if isinstance(geometry_manager.actuated_links[0], RevoluteLink):
                    plt.colorbar(
                        tpc,
                        label=(
                            f"Change in input per degree \n{geometry_manager.sensors[sensor_id].name} output (deg/deg)"
                            if self.SHOW_DEGREES
                            else f"Change in input per radian \n{geometry_manager.sensors[sensor_id].name} output (rad/rad)"
                        ),
                    )
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (rad)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (rad)"
                    )
                else:
                    plt.colorbar(
                        tpc,
                        label=(
                            f"Change in input per degree \n{geometry_manager.sensors[sensor_id].name} output (deg/mm)"
                            if self.SHOW_DEGREES
                            else f"Change in input per radian \n{geometry_manager.sensors[sensor_id].name} output (rad/mm)"
                        ),
                    )
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nDisplacement (mm)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nDisplacement (mm)"
                    )
                plt.title(
                    f"Jacobian Partial Derivative - {geometry_manager.actuated_links[0].name} WRT {geometry_manager.sensors[sensor_id].name} Output"
                )
                plt.grid()

                # Plot gradient wrt second revolute joint
                ax = plt.subplot(
                    3,
                    len(sensor_values[0]) - 2,
                    sensor_id + 5,
                )
                tri = Triangulation(sensor_values[:, 0], sensor_values[:, 1])
                zz = abs(gradient_y.ravel())
                if isinstance(geometry_manager.actuated_links[1], RevoluteLink):
                    tpc = ax.tripcolor(
                        tri,
                        zz,
                        cmap=self.JACOBIAN_COLOUR_MAP,
                        norm=self.JACOBIAN_COLOUR_NORMALIZE,
                        edgecolor="none",
                        shading="gouraud",
                        label=f"Sensor {sensor_id}",
                        clim=self.JACOBIAN_COLOUR_BAR_LIMITS,
                    )
                else:
                    tpc = ax.tripcolor(
                        tri,
                        zz,
                        edgecolor="none",
                        shading="gouraud",
                        label=f"Sensor {sensor_id}",
                    )
                annot = ax.annotate(
                    "",
                    xy=(0, 0),
                    xytext=(10, 10),
                    textcoords="offset points",
                    bbox=dict(boxstyle="round", fc="w"),
                    ha="center",
                    visible=False,
                )
                trifinder = tri.get_trifinder()
                self.plots.append(
                    dict(
                        ax=ax,
                        tri=tri,
                        z=zz,
                        tpc=tpc,
                        annot=annot,
                        trifinder=trifinder,
                    )
                )
                if isinstance(geometry_manager.actuated_links[1], RevoluteLink):
                    plt.colorbar(
                        tpc,
                        label=(
                            f"Change in input per degree \n {geometry_manager.sensors[sensor_id].name} output (deg/deg)"
                            if self.SHOW_DEGREES
                            else f"Change in input per radian \n {geometry_manager.sensors[sensor_id].name} output (rad/rad)"
                        ),
                    )
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[0].name} \nInput Joint Angle (rad)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (deg)"
                        if self.SHOW_DEGREES
                        else f"{geometry_manager.actuated_links[1].name} \nInput Joint Angle (rad)"
                    )
                else:
                    plt.colorbar(
                        tpc,
                        label=(
                            f"Change in input per degree \n {geometry_manager.sensors[sensor_id].name} output (deg/mm)"
                            if self.SHOW_DEGREES
                            else f"Change in input per radian \n {geometry_manager.sensors[sensor_id].name} output (rad/mm)"
                        ),
                    )
                    plt.xlabel(
                        f"{geometry_manager.actuated_links[0].name} \nDisplacement (mm)"
                    )
                    plt.ylabel(
                        f"{geometry_manager.actuated_links[1].name} \nDisplacement (mm)"
                    )
                plt.title(
                    f"Jacobian Partial Derivative - {geometry_manager.actuated_links[1].name} WRT {geometry_manager.sensors[sensor_id].name} Output"
                )
                plt.grid()
                self.fig.canvas.mpl_connect("motion_notify_event", self.hover)

        # This runs right after the first frame is drawn
        self.pl.add_on_render_callback(self.start_timer_once)

    def get_data_for_angular_displacement_visualization(
        self, angular_displacement_data
    ):
        """Generates a matrix for angular displacement data

        Returns:
            np array [[float], [float], [float]]: [x, y, sensitivity]
        """
        x_data = np.column_stack(angular_displacement_data[0])[0]
        y_data = np.column_stack(angular_displacement_data[1])[0]
        sensitivity_data = np.gradient(y_data, x_data)
        data = np.column_stack([x_data, y_data, sensitivity_data])
        return data

    def get_data_for_current_index(self):
        """Creates a matrix for the current data index

        Returns:
            np array [[float],[float],[float],[float]]: [[x],[y],[z],[line type]]
        """
        x = 0
        y = 0
        if self.is_1D:
            joint_pos = self.joint_positions[self.current_data_index]
        else:
            # Loop round the four edges of the mechanism motion
            if self.direction == 0:
                joint_pos = self.joint_positions[0][self.current_data_index]
            elif self.direction == 1:
                joint_pos = self.joint_positions[self.current_data_index][-1]
            elif self.direction == 2:
                joint_pos = self.joint_positions[-1][-self.current_data_index - 1]
            else:
                joint_pos = self.joint_positions[-self.current_data_index - 1][0]

        x_data = []
        y_data = []
        z_data = []
        joint_type = []

        for link in self.geometry_manager.links:
            x_data.append(joint_pos[link.start.id * 3])
            y_data.append(joint_pos[link.start.id * 3 + 1])
            z_data.append(joint_pos[link.start.id * 3 + 2])
            x_data.append(joint_pos[link.end.id * 3])
            y_data.append(joint_pos[link.end.id * 3 + 1])
            z_data.append(joint_pos[link.end.id * 3 + 2])
            joint_type.append(LineType.REGULAR.value)
            joint_type.append(LineType.REGULAR.value)

        # Add sensors
        for sensor in self.geometry_manager.sensors:
            x_data.append(joint_pos[sensor.start.id * 3])
            y_data.append(joint_pos[sensor.start.id * 3 + 1])
            z_data.append(joint_pos[sensor.start.id * 3 + 2])
            x_data.append(joint_pos[sensor.end.id * 3])
            y_data.append(joint_pos[sensor.end.id * 3 + 1])
            z_data.append(joint_pos[sensor.end.id * 3 + 2])
            joint_type.append(LineType.SENSOR.value)
            joint_type.append(LineType.SENSOR.value)

        # Add revolute joints
        for actuated_link in self.geometry_manager.actuated_links:
            x_data.append(joint_pos[actuated_link.start.id * 3])
            y_data.append(joint_pos[actuated_link.start.id * 3 + 1])
            z_data.append(joint_pos[actuated_link.start.id * 3 + 2])
            x_data.append(joint_pos[actuated_link.end.id * 3])
            y_data.append(joint_pos[actuated_link.end.id * 3 + 1])
            z_data.append(joint_pos[actuated_link.end.id * 3 + 2])
            joint_type.append(LineType.ACTUATED.value)
            joint_type.append(LineType.ACTUATED.value)

        data = np.column_stack([x_data, y_data, z_data, joint_type])

        return data

    def show(self):
        self.pl.show()