Merge pull request #58 from aceglia/master

Ploting reaction forces

bioviz/__init__.py

@@ -7,6 +7,7 @@
 import scipy
 import xarray as xr
 import pandas
+
 try:
     import biorbd
 except ImportError:
@@ -280,6 +281,7 @@ def __init__(
         show_analyses_panel=True,
         background_color=(0.5, 0.5, 0.5),
         force_wireframe=False,
+        experimental_forces_colors=(85, 78, 0),
         **kwargs,
     ):
         """
@@ -321,6 +323,7 @@ def __init__(
             contacts_size=contacts_size,
             segments_center_of_mass_size=segments_center_of_mass_size,
             force_wireframe=force_wireframe,
+            force_color=experimental_forces_colors,
         )
         self.vtk_model_markers: VtkModel = None
         self.is_executing = False
@@ -337,6 +340,11 @@ def __init__(
         self.show_experimental_markers = False
         self.experimental_markers = None
         self.experimental_markers_color = experimental_markers_color
+        self.show_experimental_forces = False
+        self.experimental_forces = None
+        self.segment_forces = []
+        self.experimental_forces_color = experimental_forces_colors
+        self.force_normalization_ratio = None
 
         self.show_contacts = show_contacts
         self.show_global_ref_frame = show_global_ref_frame
@@ -828,6 +836,9 @@ def __animate_from_slider(self):
         if self.show_experimental_markers:
             self.__set_experimental_markers_from_frame()
 
+        if self.show_experimental_forces:
+            self.__set_experimental_forces_from_frame()
+
         # Update graph of muscle analyses
         self.__update_muscle_analyses_graphs(True, True, True, True)
 
@@ -923,7 +934,12 @@ def __set_movement_slider(self):
         # Update the slider bar and frame count
         self.movement_slider[0].setEnabled(True)
         self.movement_slider[0].setMinimum(1)
-        experiment_shape = 0 if self.experimental_markers is None else self.experimental_markers.shape[2]
+        experiment_shape = 0
+        if self.experimental_forces is not None:
+            experiment_shape = self.experimental_forces.shape[2]
+        if self.experimental_markers is not None:
+            experiment_shape = max(self.experimental_markers.shape[2], experiment_shape)
+
         q_shape = 0 if self.animated_Q is None else self.animated_Q.shape[0]
         self.movement_slider[0].setMaximum(max(q_shape, experiment_shape))
         pal = QPalette()
@@ -937,9 +953,7 @@ def __load_experimental_data_from_button(self):
         # Load the actual movement
         options = QFileDialog.Options()
         options |= QFileDialog.DontUseNativeDialog
-        file_name = QFileDialog.getOpenFileName(
-            self.vtk_window, "Data to load", "", "C3D (*.c3d)", options=options
-        )
+        file_name = QFileDialog.getOpenFileName(self.vtk_window, "Data to load", "", "C3D (*.c3d)", options=options)
         if not file_name[0]:
             return
         self.load_experimental_markers(file_name[0])
@@ -958,7 +972,7 @@ def load_experimental_markers(self, data, auto_start=True, ignore_animation_warn
 
         else:
             raise RuntimeError(
-                f"Wrong type of experimental markers data ({type(data)}. "
+                f"Wrong type of experimental markers data ({type(data)}). "
                 f"Allowed type are numpy array (3xNxT), data array (3xNxT) or .c3d file (str)."
             )
 
@@ -974,14 +988,82 @@ def load_experimental_markers(self, data, auto_start=True, ignore_animation_warn
         if auto_start:
             self.__start_stop_animation()
 
+    def load_experimental_forces(
+        self, data, segments=None, normalization_ratio=0.2, auto_start=True, ignore_animation_warning=True
+    ):
+        if isinstance(data, (np.ndarray, xr.DataArray)):
+            self.experimental_forces = data if isinstance(data, xr.DataArray) else xr.DataArray(data)
+        else:
+            raise RuntimeError(
+                f"Wrong type of experimental force data ({type(data)}). "
+                f"Allowed type are numpy array (SxNxT), data array (SxNxT)."
+            )
+
+        if segments:
+            self.segment_forces = segments if isinstance(segments, (list, np.ndarray)) else [segments]
+        else:
+            self.segment_forces = ["ground"] * self.experimental_forces.shape[0]
+
+        if len(self.segment_forces) != self.experimental_forces.shape[0]:
+            raise RuntimeError(
+                "Number of segment must match number of experimental  forces. "
+                f"You have {len(segments)} and {self.experimental_forces.shape[0]}."
+            )
+
+        self.force_normalization_ratio = normalization_ratio
+
+        self.show_experimental_forces = True
+        self.__set_movement_slider()
+
+        if ignore_animation_warning:
+            self.animation_warning_already_shown = True
+        if auto_start:
+            self.__start_stop_animation()
+
     def __set_markers_from_q(self):
         self.markers[0:3, :, :] = self.Markers.get_data(Q=self.Q, compute_kin=False)
         self.vtk_model.update_markers(self.markers.isel(time=[0]))
 
     def __set_experimental_markers_from_frame(self):
         t_slider = self.movement_slider[0].value() - 1
         t = t_slider if t_slider < self.experimental_markers.shape[2] else self.experimental_markers.shape[2] - 1
-        self.vtk_model_markers.update_markers(self.experimental_markers.isel(time=t))
+        self.vtk_model_markers.update_markers(self.experimental_markers[:, :, t : t + 1].isel(time=[0]))
+
+    def __set_experimental_forces_from_frame(self):
+        segment_names = []
+        for i in range(self.model.nbSegment()):
+            segment_names.append(self.model.segment(i).name().to_string())
+        global_jcs = self.allGlobalJCS.get_data(Q=self.Q, compute_kin=False)
+        segment_jcs = []
+
+        for segment in self.segment_forces:
+            if isinstance(segment, str):
+                if segment == "ground":
+                    segment_jcs.append(np.identity(4))
+                else:
+                    segment_jcs.append(global_jcs[segment_names.index(segment)])
+            elif isinstance(segment, (float, int)):
+                segment_jcs.append(global_jcs[segment])
+            else:
+                raise RuntimeError("Wrong type of segment.")
+
+        max_forces = []
+        for i, forces in enumerate(self.experimental_forces):
+            max_forces.append(
+                max(
+                    np.sqrt(
+                        (forces[3, :] - forces[0, :]) ** 2
+                        + (forces[4, :] - forces[1, :]) ** 2
+                        + (forces[5, :] - forces[2, :]) ** 2
+                    )
+                )
+            )
+
+        t_slider = self.movement_slider[0].value() - 1
+        t = t_slider if t_slider < self.experimental_forces.shape[2] else self.experimental_forces.shape[2] - 1
+        self.vtk_model.update_force(
+            segment_jcs, self.experimental_forces[:, :, t : t + 1], max_forces, self.force_normalization_ratio
+        )
 
     def __set_contacts_from_q(self):
         self.contacts[0:3, :, :] = self.Contacts.get_data(Q=self.Q, compute_kin=False)

bioviz/biorbd_vtk.py

@@ -25,7 +25,15 @@
     vtkWindowToImageFilter,
     vtkPolygon,
     vtkExtractEdges,
+    vtkArrowSource,
+    vtkNamedColors,
+    vtkMath,
+    vtkMatrix4x4,
+    vtkMinimalStandardRandomSequence,
+    vtkTransform,
+    vtkTransformPolyDataFilter,
 )
+
 from vtk.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
 
 from pyomeca import Markers, Rototrans
@@ -168,6 +176,8 @@ def __init__(
         muscle_opacity=1.0,
         rt_length=0.1,
         rt_width=2,
+        force_color=(85, 78, 0),
+        force_opacity=1.0,
     ):
         """
         Creates a model that will holds things to plot
@@ -244,6 +254,14 @@ def __init__(
         self.wrapping_opacity = wrapping_opacity
         self.wrapping_actors = list()
 
+        self.all_forces = []
+        self.force_centers = []
+        self.max_forces = []
+        self.force_color = force_color
+        self.force_opacity = force_opacity
+        self.normalization_ratio = 0.2
+        self.force_actors = list()
+
     def set_markers_color(self, markers_color):
         """
         Dynamically change the color of the markers
@@ -1212,3 +1230,181 @@ def create_global_ref_frame(self):
 
         self.parent_window.ren.AddActor(actor)
         self.parent_window.ren.ResetCamera()
+
+    def set_force_color(self, force_color):
+        """
+        Dynamically change the color of the force
+        Parameters
+        ----------
+        force_color : tuple(int)
+            Color the force should be drawn (1 is max brightness)
+        """
+        self.force_color = force_color
+        self.update_force(self.force_centers, self.all_force, self.max_forces, self.normalization_ratio)
+
+    def set_force_opacity(self, force_opacity):
+        """
+        Dynamically change the opacity of the _force
+        Parameters
+        ----------
+        force_opacity : float
+            Opacity of the force (0.0 is completely transparent, 1.0 completely opaque)
+        Returns
+        -------
+
+        """
+        self.force_opacity = force_opacity
+        self.update_force(self.force_centers, self.all_force, self.max_forces, self.normalization_ratio)
+
+    def new_force_set(self, segment_jcs, all_forces, max_forces, normalization_ratio):
+        """
+        Define a new force set.
+        Parameters
+        ----------
+        segment_jcs : list
+            list of transformation for each segment
+        all_forces : np.ndarray
+            force array. Dims are : (1) segments, (2) application and magnitude point arrow coordinates (3) frame
+        max_forces : list
+            list of maximal force for each segment on all frames
+
+        """
+        # Arrow visualization parameters
+        arrow_source = vtkArrowSource()
+        arrow_source.SetTipResolution(15)
+        arrow_source.SetShaftResolution(8)
+        arrow_source.SetShaftRadius(0.015)
+        arrow_source.SetTipLength(0.2)
+        arrow_source.SetTipRadius(0.08)
+
+        self.arrow_source = arrow_source
+
+        transform = vtkTransform()
+
+        transform_polydata = vtkTransformPolyDataFilter()
+        transform_polydata.SetTransform(transform)
+        transform_polydata.SetInputConnection(arrow_source.GetOutputPort())
+
+        self.all_forces = all_forces
+        # Remove previous actors from the scene
+        for actor in self.force_actors:
+            self.parent_window.ren.RemoveActor(actor)
+        self.force_actors = list()
+
+        for i, forces in enumerate(all_forces):
+            # Create a mapper
+            mapper = vtkPolyDataMapper()
+            mapper.SetInputConnection(transform_polydata.GetOutputPort())
+
+            # Create an actor
+            self.force_actors.append(vtkActor())
+            self.force_actors[i].SetMapper(mapper)
+            self.force_actors[i].GetProperty().SetColor(self.force_color)
+            self.force_actors[i].GetProperty().SetOpacity(self.force_opacity)
+
+            self.parent_window.ren.AddActor(self.force_actors[i])
+            self.parent_window.ren.ResetCamera()
+
+        # Set rt orientations
+        self.n_force = len(all_forces)
+        self.update_force(segment_jcs, all_forces, max_forces, normalization_ratio)
+
+    @staticmethod
+    def compute_basis_force(application_point, magnitude_point):
+        """
+        Compute basis to plot vtk arrow object from two points.
+        Parameters
+        ----------
+        application_point : list
+            list of the 3 coordinates for the arrow starting point
+        magnitude_point: list
+            list of the 3 coordinates for the arrow ending point
+        Return
+        ----------
+            matrix of the transformation and the length of the arrow.
+        """
+        # Compute a basis
+        normalizedX = [0.0] * 3
+        normalizedY = [0.0] * 3
+        normalizedZ = [0.0] * 3
+
+        # The X axis is a vector from start to end
+        vtkMath.Subtract(magnitude_point, application_point, normalizedX)
+        length = vtkMath.Norm(normalizedX)
+        vtkMath.Normalize(normalizedX)
+
+        rng = vtkMinimalStandardRandomSequence()
+        rng.SetSeed(8775070)
+        max_r = 10.0
+
+        # The Z axis is an arbitrary vector cross X
+        arbitrary = [0.0] * 3
+        for i in range(0, 3):
+            arbitrary[i] = rng.GetRangeValue(-max_r, max_r)
+            rng.Next()
+        vtkMath.Cross(normalizedX, arbitrary, normalizedZ)
+        vtkMath.Normalize(normalizedZ)
+
+        # The Y axis is Z cross X
+        vtkMath.Cross(normalizedZ, normalizedX, normalizedY)
+        matrix = vtkMatrix4x4()
+
+        # Create the direction cosine matrix
+        matrix.Identity()
+        for i in range(0, 3):
+            matrix.SetElement(i, 0, normalizedX[i])
+            matrix.SetElement(i, 1, normalizedY[i])
+            matrix.SetElement(i, 2, normalizedZ[i])
+
+        return matrix, length
+
+    def update_force(self, segment_jcs, all_forces, max_forces, normalization_ratio):
+        """
+        Update force on the screen (but do not repaint)
+        Parameters
+        ----------
+        segment_jcs : list
+            list of roto-translation matrix for each segment
+        all_forces : np.ndarray
+            force array. Dims are : (1) nb of segments, (2) arrow start and end points coordinates (3) frame
+        max_forces : list
+            list of maximal force for each segment compute from all frames
+        normalization_ratio : float
+            ratio to normalize force for visualization
+        """
+        if self.all_forces == []:
+            self.new_force_set(segment_jcs, all_forces, max_forces, normalization_ratio)
+            return  # Prevent calling update_markers recursively
+
+        self.max_forces = max_forces
+        self.all_forces = all_forces
+        for i, forces in enumerate(all_forces):
+            # Express force from current segment basis to global basis
+            rot_seg = segment_jcs[i][:3, :3]
+            trans_seg = segment_jcs[i][:-1, 3:]
+            force_magnitude = np.dot(rot_seg, forces[3:, :])
+            force_magnitude = force_magnitude + trans_seg
+            force_application = np.dot(rot_seg, forces[:3, :])
+            force_application = force_application + trans_seg
+            application_point = [force_application[0, :], force_application[1, :], force_application[2, :]]
+            magnitude_point = [force_magnitude[0, :], force_magnitude[1, :], force_magnitude[2, :]]
+            # Compute a basis for the arrow scaling
+            matrix, length = self.compute_basis_force(application_point, magnitude_point)
+            # Normalize force for visualization
+            length = length * normalization_ratio / max_forces[i]
+            transform = vtkTransform()
+            transform.Translate(application_point)
+            transform.Concatenate(matrix)
+            transform.Scale(length, length, length)
+
+            # Create an actor
+            mapper = self.force_actors[i].GetMapper()
+            transform_polydata = vtkTransformPolyDataFilter()
+            transform_polydata.SetTransform(transform)
+            transform_polydata.SetInputConnection(self.arrow_source.GetOutputPort())
+
+            mapper.SetInputConnection(transform_polydata.GetOutputPort())
+
+            self.force_actors[i].SetMapper(mapper)
+            self.force_actors[i].GetProperty().SetColor(self.force_color)
+            self.force_actors[i].GetProperty().SetOpacity(self.force_opacity)

bioviz/mesh.py

@@ -3,12 +3,7 @@
 
 
 class Mesh(Markers):
-    def __new__(
-        cls,
-        vertex=None,
-        triangles=np.ndarray((3, 0)),
-        **kwargs,
-    ):
+    def __new__(cls, vertex=None, triangles=np.ndarray((3, 0)), **kwargs):
         if isinstance(triangles, list):
             triangles = np.array(triangles)