diff --git a/bioviz/__init__.py b/bioviz/__init__.py index 59f679e..58afacd 100644 --- a/bioviz/__init__.py +++ b/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() diff --git a/bioviz/biorbd_vtk.py b/bioviz/biorbd_vtk.py index 05df66f..bb2e3ec 100644 --- a/bioviz/biorbd_vtk.py +++ b/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 diff --git a/environment.yml b/environment.yml index 9d70a58..a636bff 100644 --- a/environment.yml +++ b/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