Speed up the transformation with multiprocess

navipy/trajectories/__init__.py

@@ -12,6 +12,9 @@ from .transformations import markers2translate, markers2euler
 from navipy.tools.plots import get_color_dataframe
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D  # noqa F401
+from multiprocessing import Pool
+from functools import partial
+import time
 
 
 def _markers2position(x, kwargs):
@@ -34,6 +37,18 @@ def _markers2angles(x, kwargs):
                          triangle_mode, euler_axes)
 
 
+def _markerstransform(index_i, trajectory,
+                      homogeneous_markers, rotation_mode):
+    row = trajectory.loc[index_i]
+    angles = row.loc[rotation_mode].values
+    translate = row.loc['location'].values
+    trans_mat = htf.compose_matrix(angles=angles,
+                                   translate=translate,
+                                   axes=rotation_mode)
+    tmarker = trans_mat.dot(homogeneous_markers)[:3]
+    return tmarker.transpose().flatten()
+
+
 class Trajectory(pd.DataFrame):
     def __init__(self, rotconv='rzyx', indeces=np.arange(1)):
         columns = self.__build_columns(rotconv)
@@ -451,16 +466,27 @@ class Trajectory(pd.DataFrame):
         if indeces is None:
             # Looping through each time point along the trajectory
             indeces = self.index
+        # More than one marker may be transformed
+        # The marker are assume to be a multiIndex dataframe
+        homogeneous_markers = markers.unstack()
+        homogeneous_markers['w'] = 1
+        homogeneous_markers = homogeneous_markers.transpose()
         # Looping throught the indeces
         # to get the homogeneous transformation from the position orientation
         # and then apply the transformed to the marker position
-        transformed_markers = pd.DataFrame(index=indeces,
+        with Pool() as p:
+            result = p.map(
+                partial(_markerstransform,
+                        trajectory=self,
+                        homogeneous_markers=homogeneous_markers,
+                        rotation_mode=self.rotation_mode),
+                indeces)
+        transformed_markers = pd.DataFrame(data=result,
+                                           index=indeces,
                                            columns=markers.index,
                                            dtype=float)
-        # More than one marker may be transformed
-        # The marker are assume to be a multiIndex dataframe
-        homogeneous_markers = markers.unstack()
-        homogeneous_markers['w'] = 1
+        return transformed_markers
+
         for index_i in indeces:
             row = self.loc[index_i]
             angles = row.loc[self.rotation_mode].values
@@ -541,7 +567,7 @@ class Trajectory(pd.DataFrame):
 (backward or forward)
         """
         # import time
-        # t_start = time.time()
+        t_start = time.time()
         if ax is None:
             fig = plt.figure()
             ax = fig.add_subplot(111, projection='3d')
@@ -573,8 +599,8 @@ class Trajectory(pd.DataFrame):
         x = self.loc[:, ('location', 'x')]
         y = self.loc[:, ('location', 'y')]
         z = self.loc[:, ('location', 'z')]
-        # print(time.time() - t_start)
-        # t_start = time.time()
+        print(time.time() - t_start)
+        t_start = time.time()
         if isinstance(colors, pd.DataFrame):
             # Each segment will be plotted with a different color
             # we therefore need to loop through all points
@@ -596,8 +622,8 @@ class Trajectory(pd.DataFrame):
                             color=color, linewidth=linewidth)
         else:
             ax.plot(xs=x, ys=y, zs=z, color=colors, linewidth=linewidth)
-        # print(time.time() - t_start)
-        # t_start = time.time()
+        print(time.time() - t_start)
+        t_start = time.time()
         # Plot the lollipop
         # - loop through the frames with a step of step_lollipop
         # - The lollipop is colored with the color of this frame
@@ -632,4 +658,4 @@ class Trajectory(pd.DataFrame):
                         color=color,
                         marker=lollipop_marker,
                         markersize=lollipop_head_size)
-        # print(time.time() - t_start)
+        print(time.time() - t_start)