Merge pull request #56 from Ipuch/chunks

Author: Ipuch

environment.yml

@@ -4,8 +4,8 @@ channels:
 - default
 dependencies:
 - ezc3d
-- rerun-sdk>=0.20.1
-- numpy
+- rerun-sdk=0.21.0
+- numpy=1.26.4
 - biorbd>=1.10.5
 - trimesh
 - pyomeca

pyorerun/biorbd_components/model_display_options.py

@@ -10,6 +10,9 @@ class DisplayModelOptions:
     markers_color: tuple[int, int, int] = (0, 0, 255)
     markers_radius: float = 0.01
 
+    centers_of_mass_color: tuple[int, int, int] = (0, 0, 0)
+    centers_of_mass_radius: float = 0.01
+
     ligaments_color: tuple[int, int, int] = (255, 255, 0)
     ligaments_radius: float = 0.01
 

pyorerun/biorbd_components/model_interface.py

@@ -1,9 +1,13 @@
+from functools import cached_property
+
 import biorbd
 import numpy as np
-from biorbd import GeneralizedCoordinates, segment_index
+from biorbd import GeneralizedCoordinates
 
 from .model_display_options import DisplayModelOptions
 
+MINIMAL_SEGMENT_MASS = 1e-08
+
 
 class BiorbdSegment:
     """
@@ -14,32 +18,36 @@ def __init__(self, segment, index):
         self.segment = segment
         self._index: int = index
 
-    @property
+    @cached_property
     def name(self) -> str:
         return self.segment.name().to_string()
 
-    @property
+    @cached_property
     def id(self) -> int:
         return self._index
 
-    @property
+    @cached_property
     def has_mesh(self) -> bool:
         has_mesh = self.segment.characteristics().mesh().hasMesh()
         if has_mesh:
             return not self.mesh_path.endswith("/")  # Avoid empty mesh path
         return has_mesh
 
-    @property
+    @cached_property
     def has_meshlines(self) -> bool:
         has_mesh = self.segment.characteristics().mesh().hasMesh()
         if has_mesh:
             return self.mesh_path.endswith("/")  # Avoid empty mesh path
         return has_mesh
 
-    @property
+    @cached_property
     def mesh_path(self) -> str:
         return self.segment.characteristics().mesh().path().absolutePath().to_string()
 
+    @cached_property
+    def mass(self) -> float:
+        return self.segment.characteristics().mass()
+
 
 class BiorbdModelNoMesh:
     """
@@ -55,30 +63,38 @@ def __init__(self, path: str, options=None):
     def from_biorbd_object(cls, model: biorbd.Model, options=None):
         return cls(model.path().absolutePath().to_string(), options)
 
-    @property
+    @cached_property
     def name(self):
         return self.path.split("/")[-1].split(".")[0]
 
-    @property
+    @cached_property
     def marker_names(self) -> tuple[str, ...]:
         return tuple([s.to_string() for s in self.model.markerNames()])
 
-    @property
+    @cached_property
     def nb_markers(self) -> int:
         return self.model.nbMarkers()
 
-    @property
+    @cached_property
     def segment_names(self) -> tuple[str, ...]:
         return tuple([s.name().to_string() for s in self.model.segments()])
 
-    @property
+    @cached_property
     def nb_segments(self) -> int:
         return self.model.nbSegment()
 
-    @property
+    @cached_property
     def segments(self) -> tuple[BiorbdSegment, ...]:
         return tuple(BiorbdSegment(s, i) for i, s in enumerate(self.model.segments()))
 
+    @cached_property
+    def segments_with_mass(self) -> tuple[BiorbdSegment, ...]:
+        return tuple([s for s in self.segments if s.mass > MINIMAL_SEGMENT_MASS])
+
+    @cached_property
+    def segment_names_with_mass(self) -> tuple[str, ...]:
+        return tuple([s.name for s in self.segments_with_mass])
+
     def segment_homogeneous_matrices_in_global(self, q: np.ndarray, segment_index: int) -> np.ndarray:
         """
         Returns a biorbd object containing the roto-translation matrix of the segment in the global reference frame.
@@ -91,24 +107,27 @@ def markers(self, q: np.ndarray) -> np.ndarray:
         """
         Returns a [N_markers x 3] array containing the position of each marker in the global reference frame
         """
-        return np.array(
-            [self.model.markers(GeneralizedCoordinates(q))[i].to_array() for i in range(self.model.nbMarkers())]
-        )
+        return np.array([self.model.markers(GeneralizedCoordinates(q))[i].to_array() for i in range(self.nb_markers)])
 
-    def center_of_mass(self, q: np.ndarray) -> np.ndarray:
+    def centers_of_mass(self, q: np.ndarray) -> np.ndarray:
         """
-        Returns the position of the center of mass in the global reference frame
+        Returns the position of the centers of mass in the global reference frame
         """
-        return self.model.CoM(GeneralizedCoordinates(q)).to_array()
+        all_com = np.array([com.to_array() for com in self.model.CoMbySegment(GeneralizedCoordinates(q))])
+        all_com_with_mass = np.zeros((len(self.segments_with_mass), 3))
+        for i, segment in enumerate(self.segments_with_mass):
+            all_com_with_mass[i, :] = all_com[segment.id, :]
 
-    @property
+        return all_com_with_mass
+
+    @cached_property
     def nb_ligaments(self) -> int:
         """
         Returns the number of ligaments
         """
         return self.model.nbLigaments()
 
-    @property
+    @cached_property
     def ligament_names(self) -> tuple[str, ...]:
         """
         Returns the names of the ligaments
@@ -129,14 +148,14 @@ def ligament_strips(self, q: np.ndarray) -> list[list[np.ndarray]]:
             ligaments.append(ligament_strip)
         return ligaments
 
-    @property
+    @cached_property
     def nb_muscles(self) -> int:
         """
         Returns the number of ligaments
         """
         return self.model.nbMuscles()
 
-    @property
+    @cached_property
     def muscle_names(self) -> tuple[str, ...]:
         """
         Returns the names of the ligaments
@@ -157,32 +176,32 @@ def muscle_strips(self, q: np.ndarray) -> list[list[np.ndarray]]:
             muscles.append(muscle_strip)
         return muscles
 
-    @property
+    @cached_property
     def nb_q(self) -> int:
         return self.model.nbQ()
 
-    @property
+    @cached_property
     def dof_names(self) -> tuple[str, ...]:
         return tuple(s.to_string() for s in self.model.nameDof())
 
-    @property
+    @cached_property
     def q_ranges(self) -> tuple[tuple[float, float], ...]:
         q_ranges = [q_range for segment in self.model.segments() for q_range in segment.QRanges()]
         return tuple((q_range.min(), q_range.max()) for q_range in q_ranges)
 
-    @property
+    @cached_property
     def gravity(self) -> np.ndarray:
         return self.model.getGravity().to_array()
 
-    @property
+    @cached_property
     def has_mesh(self) -> bool:
         return any([s.has_mesh for s in self.segments])
 
-    @property
+    @cached_property
     def has_meshlines(self) -> bool:
         return any([s.has_meshlines for s in self.segments])
 
-    @property
+    @cached_property
     def has_soft_contacts(self) -> bool:
         return self.model.nbSoftContacts() > 0
 
@@ -193,14 +212,14 @@ def soft_contacts(self, q: np.ndarray) -> np.ndarray:
         soft_contacts = self.model.softContacts(q, True)
         return np.array([soft_contacts[i].to_array() for i in range(self.model.nbSoftContacts())])
 
-    @property
+    @cached_property
     def soft_contacts_names(self) -> tuple[str, ...]:
         """
         Returns the names of the soft contacts
         """
         return tuple([s.to_string() for s in self.model.softContactNames()])
 
-    @property
+    @cached_property
     def soft_contact_radii(self) -> tuple[float, ...]:
         """
         Returns the radii of the soft contacts
@@ -227,7 +246,7 @@ def __init__(self, path, options=None):
     def segments(self) -> tuple[BiorbdSegment, ...]:
         return tuple([s for s in super().segments if s.has_mesh or s.has_meshlines])
 
-    @property
+    @cached_property
     def meshlines(self) -> list[np.ndarray]:
 
         meshes = []

pyorerun/biorbd_components/model_updapter.py

@@ -19,6 +19,7 @@ def __init__(self, name, model: BiorbdModelNoMesh | BiorbdModel):
         self.name = name
         self.model = model
         self.markers = self.create_markers_updater()
+        self.centers_of_mass = self.create_centers_of_mass_updater()
         self.soft_contacts = self.create_soft_contacts_updater()
         self.ligaments = self.create_ligaments_updater()
         self.segments = self.create_segments_updater()
@@ -73,6 +74,17 @@ def create_markers_updater(self):
             callable_markers=self.model.markers,
         )
 
+    def create_centers_of_mass_updater(self):
+        return MarkersUpdater(
+            self.name + "/centers_of_mass",
+            marker_properties=MarkerProperties(
+                markers_names=self.model.segment_names_with_mass,
+                color=np.array(self.model.options.centers_of_mass_color),
+                radius=self.model.options.centers_of_mass_radius,
+            ),
+            callable_markers=self.model.centers_of_mass,
+        )
+
     def create_soft_contacts_updater(self):
         if not self.model.has_soft_contacts:
             return EmptyUpdater(self.name + "/soft_contacts")
@@ -160,7 +172,14 @@ def components(self) -> list[Any]:
         all_segment_components = []
         for segment in self.segments:
             all_segment_components.extend(segment.components)
-        return [self.markers, self.soft_contacts, *all_segment_components, self.ligaments, self.muscles]
+        return [
+            self.markers,
+            self.centers_of_mass,
+            self.soft_contacts,
+            *all_segment_components,
+            self.ligaments,
+            self.muscles,
+        ]
 
     @property
     def component_names(self) -> list[str]:

pyproject.toml

@@ -21,8 +21,8 @@ classifiers = [
      ]
 dependencies = [
 #    "ezc3d", # Not yet available on pypi, use `conda install -c conda-forge ezc3d`
-    "numpy",
-    "rerun-sdk==0.20.1",
+    "numpy==1.26.4",
+    "rerun-sdk==0.21.0",
     "trimesh",
     "pyomeca",
     "tk",

setup.py

@@ -5,8 +5,8 @@
     version="1.2.5",
     install_requires=[
         "ezc3d",  # Not yet available on pypi, use `conda install -c conda-forge ezc3d`
-        "numpy",
-        "rerun-sdk>=0.20.1",
+        "numpy=1.26.4",
+        "rerun-sdk=0.21.0",
         "trimesh",
         "pyomeca",
         "tk",