Display gravity vector

bioviz/__init__.py

@@ -301,6 +301,7 @@ def __init__(
         loaded_model=None,
         show_meshes=True,
         show_global_center_of_mass=True,
+        show_gravity_vector=True,
         show_segments_center_of_mass=True,
         segments_center_of_mass_size=0.005,
         show_global_ref_frame=True,
@@ -399,6 +400,7 @@ def __init__(
         self.soft_contacts_color = soft_contacts_color
         self.show_global_ref_frame = show_global_ref_frame
         self.show_global_center_of_mass = show_global_center_of_mass
+        self.show_gravity_vector = show_gravity_vector
         self.show_segments_center_of_mass = show_segments_center_of_mass
         self.show_local_ref_frame = show_local_ref_frame
         self.biorbd_compiled_with_muscles = hasattr(biorbd.Model, "nbMuscles")
@@ -572,6 +574,8 @@ def set_q(self, Q, refresh_window=True):
             self.__set_meshes_from_q()
         if self.show_global_center_of_mass:
             self.__set_global_center_of_mass_from_q()
+        if self.show_gravity_vector:
+            self.__set_gravity_vector()
         if self.show_segments_center_of_mass:
             self.__set_segments_center_of_mass_from_q()
         if self.show_markers:
@@ -1135,6 +1139,16 @@ def __set_global_center_of_mass_from_q(self):
         self.global_center_of_mass.loc[{"channel": 0, "time": 0}] = com.squeeze()
         self.vtk_model.update_global_center_of_mass(self.global_center_of_mass.isel(time=[0]))
 
+    def __set_gravity_vector(self):
+        start = [0, 0, 0]
+        magnitude = self.model.getGravity().to_array()
+        gravity = np.concatenate((start, magnitude))
+        length = np.linalg.norm(gravity)
+        id_matrix = np.identity(4)
+        self.vtk_model.new_gravity_vector(
+            id_matrix, gravity, length, normalization_ratio=0.3, vector_color=(0, 0, 0)
+        )
+
     def __set_segments_center_of_mass_from_q(self):
         coms = self.CoMbySegment.get_data(Q=self.Q, compute_kin=False)
         for k, com in enumerate(coms):
@@ -1148,7 +1162,6 @@ def __set_meshes_from_q(self):
 
     def __set_muscles_from_q(self):
         muscles = self.musclesPointsInGlobal.get_data(Q=self.Q)
-
         idx = 0
         cmp = 0
         for group_idx in range(self.model.nbMuscleGroups()):

bioviz/biorbd_vtk.py

@@ -271,6 +271,8 @@ def __init__(
         self.normalization_ratio = 0.2
         self.force_actors = list()
 
+        self.gravity_actors = list()
+
     def set_markers_color(self, markers_color):
         """
         Dynamically change the color of the markers
@@ -1346,7 +1348,6 @@ def set_force_color(self, force_color):
             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):
         """
@@ -1360,7 +1361,6 @@ def set_force_opacity(self, force_opacity):
 
         """
         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):
         """
@@ -1514,3 +1514,70 @@ def update_force(self, segment_jcs, all_forces, max_forces, normalization_ratio)
             self.force_actors[i].SetMapper(mapper)
             self.force_actors[i].GetProperty().SetColor(self.force_color)
             self.force_actors[i].GetProperty().SetOpacity(self.force_opacity)
+
+    def new_gravity_vector(self, segment_rt, gravity, length, normalization_ratio, vector_color):
+        """
+        Define a new gravity vector.
+        Parameters
+        ----------
+        segment_rt : np.ndarray
+            homogeneous matrix in which coordinates are applied
+        gravity : np.ndarray
+            gravity array with 3 application coordinates and 3 magnitude coordinates
+        length : float
+            absolut length of gravity vector to scale arrow
+        normalization_ratio: float
+            ratio to scale arrow
+        vector_color: tuple
+            tuple of RGB for vector color
+
+        """
+        # 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
+
+        # Arrow orientation
+        transform = vtkTransform()
+        transform_polydata = vtkTransformPolyDataFilter()
+        transform_polydata.SetTransform(transform)
+        transform_polydata.SetInputConnection(arrow_source.GetOutputPort())
+
+        rot_seg = segment_rt[:3, :3]
+        trans_seg = segment_rt[:-1, 3:]
+        force_magnitude = np.dot(rot_seg, gravity[3:])
+        force_magnitude = force_magnitude + trans_seg.reshape(3)
+        force_application = np.dot(rot_seg, gravity[:3])
+        force_application = force_application + trans_seg.reshape(3)
+        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 / length
+        transform = vtkTransform()
+        transform.Translate(application_point)
+        transform.Concatenate(matrix)
+        transform.Scale(length, length, length)
+
+        # Create an actor
+        mapper = vtkPolyDataMapper()
+
+        transform_polydata = vtkTransformPolyDataFilter()
+        transform_polydata.SetTransform(transform)
+        transform_polydata.SetInputConnection(self.arrow_source.GetOutputPort())
+        mapper.SetInputConnection(transform_polydata.GetOutputPort())
+
+        self.gravity_actors = vtkActor()
+        self.gravity_actors.SetMapper(mapper)
+        self.gravity_actors.GetProperty().SetColor(vector_color)
+
+        self.parent_window.ren.AddActor(self.gravity_actors)
+        self.parent_window.ren.ResetCamera()