Added possibility to plot the contact forces

src/hippopt/robot_planning/utilities/foot_contact_state_plotter.py

@@ -3,7 +3,7 @@
 import logging
 import math
 import multiprocessing
-from typing import TypeVar
+from typing import Callable, TypeVar
 
 import matplotlib.axes
 import matplotlib.pyplot as plt
@@ -21,25 +21,40 @@
 
 @dataclasses.dataclass
 class ContactPointStatePlotterSettings:
-    axes: list[matplotlib.axes.Axes] | None = dataclasses.field(default=None)
+    complementarity_axes: list[matplotlib.axes.Axes] | None = dataclasses.field(
+        default=None
+    )
+    force_axes: matplotlib.axes.Axes | None = dataclasses.field(default=None)
     terrain: TerrainDescriptor = dataclasses.field(default=None)
 
-    input_axes: dataclasses.InitVar[list[matplotlib.axes.Axes]] = dataclasses.field(
+    input_complementarity_axes: dataclasses.InitVar[
+        list[matplotlib.axes.Axes]
+    ] = dataclasses.field(default=None)
+    input_force_axes: dataclasses.InitVar[matplotlib.axes.Axes] = dataclasses.field(
         default=None
     )
     input_terrain: dataclasses.InitVar[TerrainDescriptor] = dataclasses.field(
         default=None
     )
 
     def __post_init__(
-        self, input_axes: list[matplotlib.axes.Axes], input_terrain: TerrainDescriptor
+        self,
+        input_complementarity_axes: list[matplotlib.axes.Axes],
+        input_force_axes: matplotlib.axes.Axes,
+        input_terrain: TerrainDescriptor,
     ):
-        self.axes = None
-        if isinstance(input_axes, list):
-            if len(input_axes) != 2:
+        self.complementarity_axes = None
+        if isinstance(input_complementarity_axes, list):
+            if len(input_complementarity_axes) != 2:
                 raise ValueError("input_axes must be a list of length 2.")
 
-            self.axes = input_axes
+            self.complementarity_axes = input_complementarity_axes
+
+        self.force_axes = (
+            input_force_axes
+            if isinstance(input_force_axes, matplotlib.axes.Axes)
+            else None
+        )
 
         self.terrain = (
             input_terrain
@@ -54,15 +69,17 @@ def __init__(
         settings: ContactPointStatePlotterSettings = ContactPointStatePlotterSettings(),
     ):
         self.settings = settings
-        self._axes = self.settings.axes
-        self._fig = None
+        self._complementarity_axes = self.settings.complementarity_axes
+        self._complementarity_fig = None
+        self._force_axes = self.settings.force_axes
+        self._force_fig = None
 
     def plot_complementarity(
         self,
         states: list[ContactPointState],
         time_s: float | list[float] | np.ndarray = None,
         title: str = "Contact Point Complementarity",
-    ):
+    ) -> None:
         _time_s = copy.deepcopy(time_s)
         if _time_s is None or isinstance(_time_s, float) or _time_s.size == 1:
             single_step = _time_s if _time_s is not None else 0.0
@@ -73,8 +90,10 @@ def plot_complementarity(
                 "timestep_s and foot_contact_states have different lengths."
             )
 
-        if self._axes is None:
-            self._fig, self._axes = plt.subplots(nrows=1, ncols=2)
+        if self._complementarity_axes is None:
+            self._complementarity_fig, self._complementarity_axes = plt.subplots(
+                nrows=1, ncols=2
+            )
             plt.tight_layout()
 
         height_function = self.settings.terrain.height_function()
@@ -83,21 +102,53 @@ def plot_complementarity(
         positions = np.array([height_function(s.p) for s in states]).flatten()
         forces = np.array([normal_direction_fun(s.p).T @ s.f for s in states]).flatten()
         complementarity_error = np.multiply(positions, forces)
-        self._axes[1].plot(_time_s, complementarity_error)
-        self._axes[1].set_ylabel("Complementarity Error [Nm]")
-        self._axes[1].set_xlabel("Time [s]")
-        self._axes[0].plot(_time_s, positions)
-        self._axes[0].set_ylabel("Height [m]", color="C0")
-        self._axes[0].tick_params(axis="y", color="C0", labelcolor="C0")
-        axes_force = self._axes[0].twinx()
+        self._complementarity_axes[1].plot(_time_s, complementarity_error)
+        self._complementarity_axes[1].set_ylabel("Complementarity Error [Nm]")
+        self._complementarity_axes[1].set_xlabel("Time [s]")
+        self._complementarity_axes[0].plot(_time_s, positions)
+        self._complementarity_axes[0].set_ylabel("Height [m]", color="C0")
+        self._complementarity_axes[0].tick_params(axis="y", color="C0", labelcolor="C0")
+        self._complementarity_axes[0].set_xlabel("Time [s]")
+        axes_force = self._complementarity_axes[0].twinx()
         axes_force.plot(_time_s, forces, "C1")
         axes_force.set_ylabel("Normal Force [N]", color="C1")
         axes_force.tick_params(axis="y", color="C1", labelcolor="C1")
         axes_force.spines["right"].set_color("C1")
         axes_force.spines["left"].set_color("C0")
 
-        if self._fig is not None:
-            self._fig.suptitle(title)
+        if self._complementarity_fig is not None:
+            self._complementarity_fig.suptitle(title)
+            plt.draw()
+            plt.pause(0.001)
+            plt.show()
+
+    def plot_forces(
+        self,
+        states: list[ContactPointState],
+        time_s: float | list[float] | np.ndarray = None,
+        title: str = "Contact Point Forces",
+    ) -> None:
+        _time_s = copy.deepcopy(time_s)
+        if _time_s is None or isinstance(_time_s, float) or _time_s.size == 1:
+            single_step = _time_s if _time_s is not None else 0.0
+            _time_s = np.linspace(0, len(states) * single_step, len(states))
+
+        if len(_time_s) != len(states):
+            raise ValueError(
+                "timestep_s and foot_contact_states have different lengths."
+            )
+
+        if self._force_axes is None:
+            self._force_fig, self._force_axes = plt.subplots(nrows=1, ncols=1)
+            plt.tight_layout()
+
+        forces = np.array([s.f for s in states])
+        self._force_axes.plot(_time_s, forces)
+        self._force_axes.set_ylabel("Force [N]")
+        self._force_axes.set_xlabel("Time [s]")
+
+        if self._force_fig is not None:
+            self._force_fig.suptitle(title)
             plt.draw()
             plt.pause(0.001)
             plt.show()
@@ -120,23 +171,20 @@ def __init__(
         settings: FootContactStatePlotterSettings = FootContactStatePlotterSettings(),
     ):
         self._settings = settings
-        self._ext_process = None
+        self._ext_proc_complementarity = None
+        self._ext_proc_forces = None
         self._logger = logging.getLogger("[hippopt::FootContactStatePlotter]")
 
-    def plot_complementarity(
+    def _plot(
         self,
         states: list[FootContactState],
-        time_s: float | list[float] | np.ndarray = None,
-        title: str = "Foot Contact Complementarity",
-        blocking: bool = False,
-    ):
-        if self._ext_process is not None:
-            self._logger.warning(
-                "A plot is already running. "
-                "Make sure to close the previous plot first."
-            )
-            self._ext_process.join()
-            self._ext_process = None
+        time_s: float | list[float] | np.ndarray,
+        title: str,
+        blocking: bool,
+        plot_function: Callable[
+            [list[FootContactState], np.ndarray, str, int, TerrainDescriptor], None
+        ],
+    ) -> multiprocessing.Process | None:
         _time_s = copy.deepcopy(time_s)
         _states = copy.deepcopy(states)
         _terrain = copy.deepcopy(self._settings.terrain)
@@ -150,10 +198,10 @@ def plot_complementarity(
             )
 
         if len(_states) == 0:
-            return
+            return None
 
-        self._ext_process = multiprocessing.Process(
-            target=FootContactStatePlotter._create_complementarity_plot,
+        process = multiprocessing.Process(
+            target=plot_function,
             args=(
                 _states,
                 _time_s,
@@ -162,10 +210,13 @@ def plot_complementarity(
                 _terrain,
             ),
         )
-        self._ext_process.start()
+        process.start()
 
         if blocking:
-            self._ext_process.join()
+            process.join()
+            process = None
+
+        return process
 
     @staticmethod
     def _create_complementarity_plot(
@@ -174,7 +225,7 @@ def _create_complementarity_plot(
         title: str,
         number_of_columns: int,
         terrain: TerrainDescriptor,
-    ):
+    ) -> None:
         number_of_points = len(states[0])
         number_of_plots = number_of_points + 1
         _number_of_columns = (
@@ -195,7 +246,7 @@ def _create_complementarity_plot(
         _point_plotters = [
             ContactPointStatePlotter(
                 ContactPointStatePlotterSettings(
-                    input_axes=[el, last_plot],
+                    input_complementarity_axes=[el, last_plot],
                     terrain=terrain,
                 )
             )
@@ -216,8 +267,104 @@ def _create_complementarity_plot(
         plt.pause(0.001)
         plt.show()
 
-    def close(self):
-        if self._ext_process is not None:
-            self._ext_process.terminate()
-            self._ext_process = None
+    @staticmethod
+    def _create_forces_plot(
+        states: list[FootContactState],
+        time_s: np.ndarray,
+        title: str,
+        number_of_columns: int,
+        _: TerrainDescriptor,
+    ) -> None:
+        number_of_points = len(states[0])
+        number_of_plots = number_of_points
+        _number_of_columns = (
+            math.floor(math.sqrt(number_of_plots))
+            if number_of_columns < 1
+            else number_of_columns
+        )
+        number_of_rows = math.ceil(number_of_plots / _number_of_columns)
+
+        _fig, axes_list = plt.subplots(
+            nrows=number_of_rows,
+            ncols=_number_of_columns,
+            squeeze=False,
+        )
+        plt.tight_layout()
+        last_plot_column = (
+            number_of_points - 1 - _number_of_columns * (number_of_rows - 1)
+        )
+        _point_plotters = [
+            ContactPointStatePlotter(
+                ContactPointStatePlotterSettings(
+                    input_force_axes=el,
+                )
+            )
+            for row in axes_list
+            for el in row
+        ]
+        for i in range(last_plot_column + 1, _number_of_columns):
+            axes_list[number_of_rows - 1][i].remove()
+
+        for p in range(number_of_points):
+            contact_states = [state[p] for state in states]
+            _point_plotters[p].plot_forces(states=contact_states, time_s=time_s)
+
+        _fig.suptitle(title)
+        plt.draw()
+        plt.pause(0.001)
+        plt.show()
+
+    def plot_complementarity(
+        self,
+        states: list[FootContactState],
+        time_s: float | list[float] | np.ndarray = None,
+        title: str = "Foot Contact Complementarity",
+        blocking: bool = False,
+    ) -> None:
+        if self._ext_proc_complementarity is not None:
+            self._logger.warning(
+                "A complementarity plot is already running. "
+                "Make sure to close the previous plot first."
+            )
+            self._ext_proc_complementarity.join()
+            self._ext_proc_complementarity = None
+
+        self._ext_proc_complementarity = self._plot(
+            states,
+            time_s,
+            title,
+            blocking,
+            self._create_complementarity_plot,
+        )
+
+    def plot_forces(
+        self,
+        states: list[FootContactState],
+        time_s: float | list[float] | np.ndarray = None,
+        title: str = "Foot Contact Forces",
+        blocking: bool = False,
+    ) -> None:
+        if self._ext_proc_forces is not None:
+            self._logger.warning(
+                "A force plot is already running. "
+                "Make sure to close the previous plot first."
+            )
+            self._ext_proc_forces.join()
+            self._ext_proc_forces = None
+
+        self._ext_proc_forces = self._plot(
+            states,
+            time_s,
+            title,
+            blocking,
+            self._create_forces_plot,
+        )
+
+    def close(self) -> None:
+        if self._ext_proc_complementarity is not None:
+            self._ext_proc_complementarity.terminate()
+            self._ext_proc_complementarity = None
+        if self._ext_proc_forces is not None:
+            self._ext_proc_forces.terminate()
+            self._ext_proc_forces = None
         plt.close("all")