Merge pull request #49 from Ipuch/chunks

refactor: minor faster ligament
pyomeca · Jan 16, 2025 · 5fb5ac3 · 5fb5ac3
2 parents 383f0bc + cc3b91b
commit 5fb5ac3
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 21 deletions.
diff --git a/examples/biorbd/msk_model.py b/examples/biorbd/msk_model.py
@@ -27,12 +27,12 @@ def main():
     # viz.add_animated_model(model, q, display_q=True)
     viz.add_animated_model(model, q, display_q=False)
     tic = time.time()
-    viz.rerun("msk_model")
+    viz.rerun("msk_model with chunks")
     toc = time.time()
     print(f"Time to run: {toc - tic}")
 
     tic = time.time()
-    viz.rerun_by_frame("msk_model with chunk")
+    # viz.rerun_by_frame("msk_model frame by frame")
     toc = time.time()
     print(f"Time to run with chunks: {toc - tic}")
 

diff --git a/pyorerun/biorbd_components/ligaments.py b/pyorerun/biorbd_components/ligaments.py
@@ -38,24 +38,20 @@ def compute_strips(self, q: np.ndarray) -> list:
 
     def to_chunk(self, q: np.ndarray) -> dict[str, list]:
         nb_frames = q.shape[1]
-        # partition = [self.nb_strips for _ in range(nb_frames)]
 
         strips_by_frame = self.compute_strips(q)
-        strips_output = {}
-        for s in range(self.nb_strips):
-            strips_output.update(
-                {
-                    f"{self.name}_{s}": [
-                        rr.LineStrips3D.indicator(),
-                        rr.components.LineStrip3DBatch([strips_by_frame[f][s] for f in range(nb_frames)]),
-                        rr.components.ColorBatch([self.properties.color for _ in range(nb_frames)]),
-                        rr.components.RadiusBatch([self.properties.radius for _ in range(nb_frames)]),
-                        rr.components.TextBatch([self.properties.strip_names[s] for _ in range(nb_frames)]),
-                    ]
-                }
-            )
-
-        return strips_output
+        colors = [self.properties.color for _ in range(nb_frames)]
+        radii = [self.properties.radius for _ in range(nb_frames)]
+        return {
+            f"{self.name}_{s}": [
+                rr.LineStrips3D.indicator(),
+                rr.components.LineStrip3DBatch([strips_by_frame[f][s] for f in range(nb_frames)]),
+                rr.components.ColorBatch(colors),
+                rr.components.RadiusBatch(radii),
+                rr.components.TextBatch([self.properties.strip_names[s] for _ in range(nb_frames)]),
+            ]
+            for s in range(self.nb_strips)
+        }
 
 
 class LigamentsUpdater(LineStripUpdater):
@@ -122,7 +118,6 @@ def to_chunk(self, q: np.ndarray, markers: np.ndarray) -> dict[str, list]:
 
         strips_output = {}
         for s in range(self.nb_strips):
-            print(self.name, s)
             strips_output.update(
                 {
                     f"{self.name}_{s}": [

diff --git a/pyorerun/rrc3d.py b/pyorerun/rrc3d.py
@@ -114,8 +114,7 @@ def rrc3d(
             phase_reruns[-1].add_video(vid_name, np.array(vid, dtype=np.uint8))
 
     multi_phase_rerun = MultiFrameRatePhaseRerun(phase_reruns)
-    # multi_phase_rerun.rerun(filename, notebook=notebook)
-    multi_phase_rerun.rerun_with_chunks(filename, notebook=notebook)
+    multi_phase_rerun.rerun(filename, notebook=notebook)
 
     if show_events:
         try: