Merge pull request #60 from Ipuch/Dev_Contact_spheres

Displaying SoftContactSphere

bioviz/__init__.py

@@ -126,6 +126,27 @@ def _get_data_from_casadi(self, Q=None, compute_kin=True):
             if self.m.nbContacts():
                 self.data[:, :, 0] = np.array(self.contacts(Q))
 
+    class SoftContacts(BiorbdFunc):
+        def __init__(self, model):
+            super().__init__(model)
+            self.data = np.ndarray((3, self.m.nbSoftContacts(), 1))
+
+        def _prepare_function_for_casadi(self):
+            q_sym = casadi.MX.sym("Q", self.m.nbQ(), 1)
+            self.soft_contacts = biorbd.to_casadi_func("SoftContacts", self.m.softContacts, q_sym, True)
+
+        def _get_data_from_eigen(self, Q=None, compute_kin=True):
+            if compute_kin:
+                soft_contacts = self.m.softContacts(Q, True)
+            else:
+                soft_contacts = self.m.softContacts(Q, False)
+            for i in range(self.m.nbSoftContacts()):
+                self.data[:, i, 0] = soft_contacts[i].to_array()
+
+        def _get_data_from_casadi(self, Q=None, compute_kin=True):
+            if self.m.nbContacts():
+                self.data[:, :, 0] = np.array(self.soft_contacts(Q))
+
     class CoM(BiorbdFunc):
         def __init__(self, model):
             super().__init__(model)
@@ -276,12 +297,14 @@ def __init__(
         markers_size=0.010,
         show_contacts=True,
         contacts_size=0.010,
+        show_soft_contacts=True,
+        soft_contacts_color=(0.11, 0.63, 0.95),
         show_muscles=True,
         show_wrappings=True,
         show_analyses_panel=True,
         background_color=(0.5, 0.5, 0.5),
         force_wireframe=False,
-        experimental_forces_colors=(85, 78, 0),
+        experimental_forces_color=(85, 78, 0),
         **kwargs,
     ):
         """
@@ -311,10 +334,20 @@ def __init__(
             show_contacts = False
             show_muscles = False
             show_wrappings = False
+            show_soft_contacts = False
 
         # Create the plot
         self.vtk_window = VtkWindow(background_color=background_color)
         self.vtk_markers_size = markers_size
+
+        # soft_contact sphere sizes
+        radius = []
+        for i in range(self.model.nbSoftContacts()):
+            c = self.model.softContact(i)
+            if c.typeOfNode() == biorbd.SOFT_CONTACT_SPHERE:
+                radius.append(biorbd.SoftContactSphere(self.model.softContact(i)).radius())
+        soft_contacts_size = radius
+
         self.vtk_model = VtkModel(
             self.vtk_window,
             markers_color=(0, 0, 1),
@@ -323,7 +356,9 @@ 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,
+            force_color=experimental_forces_color,
+            soft_contacts_size=soft_contacts_size,
+            soft_contacts_color=soft_contacts_color,
         )
         self.vtk_model_markers: VtkModel = None
         self.is_executing = False
@@ -343,10 +378,12 @@ def __init__(
         self.show_experimental_forces = False
         self.experimental_forces = None
         self.segment_forces = []
-        self.experimental_forces_color = experimental_forces_colors
+        self.experimental_forces_color = experimental_forces_color
         self.force_normalization_ratio = None
 
         self.show_contacts = show_contacts
+        self.show_soft_contacts = show_soft_contacts
+        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_segments_center_of_mass = show_segments_center_of_mass
@@ -375,6 +412,9 @@ def __init__(
         if self.show_contacts:
             self.Contacts = InterfacesCollections.Contact(self.model)
             self.contacts = Markers(np.ndarray((3, self.model.nbContacts(), 1)))
+        if self.show_soft_contacts:
+            self.SoftContacts = InterfacesCollections.SoftContacts(self.model)
+            self.soft_contacts = Markers(np.ndarray((3, self.model.nbSoftContacts(), 1)))
         if self.show_global_center_of_mass:
             self.CoM = InterfacesCollections.CoM(self.model)
             self.global_center_of_mass = Markers(np.ndarray((3, 1, 1)))
@@ -523,8 +563,12 @@ def set_q(self, Q, refresh_window=True):
             self.__set_segments_center_of_mass_from_q()
         if self.show_markers:
             self.__set_markers_from_q()
+        if self.show_markers:
+            self.__set_markers_from_q()
         if self.show_contacts:
             self.__set_contacts_from_q()
+        if self.show_soft_contacts:
+            self.__set_soft_contacts_from_q()
         if self.show_wrappings:
             self.__set_wrapping_from_q()
 
@@ -1069,6 +1113,10 @@ def __set_contacts_from_q(self):
         self.contacts[0:3, :, :] = self.Contacts.get_data(Q=self.Q, compute_kin=False)
         self.vtk_model.update_contacts(self.contacts.isel(time=[0]))
 
+    def __set_soft_contacts_from_q(self):
+        self.soft_contacts[0:3, :, :] = self.SoftContacts.get_data(Q=self.Q, compute_kin=False)
+        self.vtk_model.update_soft_contacts(self.soft_contacts.isel(time=[0]))
+
     def __set_global_center_of_mass_from_q(self):
         com = self.CoM.get_data(Q=self.Q, compute_kin=False)
         self.global_center_of_mass.loc[{"channel": 0, "time": 0}] = com.squeeze()

bioviz/biorbd_vtk.py

@@ -158,6 +158,9 @@ def __init__(
         contacts_color=(0, 1, 0),
         contacts_size=0.01,
         contacts_opacity=1.0,
+        soft_contacts_color=(1, 0.35, 0),
+        soft_contacts_size=0.1,
+        soft_contacts_opacity=0.35,
         global_ref_frame_length=0.15,
         global_ref_frame_width=5,
         global_center_of_mass_size=0.0075,
@@ -214,6 +217,12 @@ def __init__(
         self.contacts_opacity = contacts_opacity
         self.contacts_actors = list()
 
+        self.soft_contacts = Markers()
+        self.soft_contacts_size = soft_contacts_size
+        self.soft_contacts_color = soft_contacts_color
+        self.soft_contacts_opacity = soft_contacts_opacity
+        self.soft_contacts_actors = list()
+
         self.has_global_ref_frame = False
         self.global_ref_frame_length = global_ref_frame_length
         self.global_ref_frame_width = global_ref_frame_width
@@ -456,6 +465,103 @@ def update_contacts(self, contacts):
             source.SetRadius(self.contacts_size)
             mapper.SetInputConnection(source.GetOutputPort())
 
+    def set_soft_contacts_color(self, soft_contacts_color):
+        """
+        Dynamically change the color of the soft_contacts
+        Parameters
+        ----------
+        soft_contacts_color : tuple(int)
+            Color the soft_contacts should be drawn (1 is max brightness)
+        """
+        self.soft_contacts_color = soft_contacts_color
+        self.update_soft_contacts(self.soft_contacts)
+
+    def set_soft_contacts_size(self, soft_contacts_size):
+        """
+        Dynamically change the size of the soft_contacts
+        Parameters
+        ----------
+        soft_contacts_size : float
+            Size the soft_contacts should be drawn
+        """
+        self.soft_contacts_size = soft_contacts_size
+        self.update_soft_contacts(self.soft_contacts)
+
+    def set_soft_contacts_opacity(self, soft_contacts_opacity):
+        """
+        Dynamically change the opacity of the soft_contacts
+        Parameters
+        ----------
+        soft_contacts_opacity : float
+            Opacity of the soft_contacts (0.0 is completely transparent, 1.0 completely opaque)
+        Returns
+        -------
+
+        """
+        self.soft_contacts_opacity = soft_contacts_opacity
+        self.update_soft_contacts(self.soft_contacts)
+
+    def new_soft_contacts_set(self, soft_contacts):
+        """
+        Define a new marker set. This function must be called each time the number of soft_contacts change
+        Parameters
+        ----------
+        soft_contacts : Markers3d
+            One frame of soft_contacts
+
+        """
+        if soft_contacts.time.size != 1:
+            raise IndexError("soft_contacts should be from one frame only")
+        self.soft_contacts = soft_contacts
+
+        # Remove previous actors from the scene
+        for actor in self.soft_contacts_actors:
+            self.parent_window.ren.RemoveActor(actor)
+        self.soft_contacts_actors = list()
+
+        # Create the geometry of a point (the coordinate) points = vtk.vtkPoints()
+        for i in range(soft_contacts.channel.size):
+            # Create a mapper
+            mapper = vtkPolyDataMapper()
+
+            # Create an actor
+            self.soft_contacts_actors.append(vtkActor())
+            self.soft_contacts_actors[i].SetMapper(mapper)
+
+            self.parent_window.ren.AddActor(self.soft_contacts_actors[i])
+            self.parent_window.ren.ResetCamera()
+
+        # Update marker position
+        self.update_soft_contacts(self.soft_contacts)
+
+    def update_soft_contacts(self, soft_contacts):
+        """
+        Update position of the soft_contacts on the screen (but do not repaint)
+        Parameters
+        ----------
+        soft_contacts : Markers3d
+            One frame of soft_contacts
+
+        """
+
+        if soft_contacts.time.size != 1:
+            raise IndexError("soft_contacts should be from one frame only")
+        if soft_contacts.channel.size != self.soft_contacts.channel.size:
+            self.new_soft_contacts_set(soft_contacts)
+            return  # Prevent calling update_soft_contacts recursively
+        self.soft_contacts = soft_contacts
+        soft_contacts = np.array(soft_contacts)
+
+        for i, actor in enumerate(self.soft_contacts_actors):
+            # mapper = actors.GetNextActor().GetMapper()
+            mapper = actor.GetMapper()
+            self.soft_contacts_actors[i].GetProperty().SetColor(self.soft_contacts_color)
+            self.soft_contacts_actors[i].GetProperty().SetOpacity(self.soft_contacts_opacity)
+            source = vtkSphereSource()
+            source.SetCenter(soft_contacts[0:3, i])
+            source.SetRadius(self.soft_contacts_size[i])
+            mapper.SetInputConnection(source.GetOutputPort())
+
     def set_global_center_of_mass_color(self, global_center_of_mass_color):
         """
         Dynamically change the color of the global center of mass

environment.yml

@@ -3,7 +3,7 @@ name: bioptim
 channels:
   - conda-forge
 dependencies:
-  - biorbd >=1.6.1
+  - biorbd >=1.8.1
   - python >=3.9
   - numpy
   - matplotlib