COrrect typose in lolliplot and change world2body transform to be faster

doc/source/tutorials/02a-orientation-background.ipynb

@@ -525,7 +525,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.5"
+   "version": "3.6.3"
   }
  },
  "nbformat": 4,

navipy/tools/plots.py

@@ -35,7 +35,7 @@ corresponding color-map (scalar mappable)
     """
 
     frame_series = pd.Series(index=np.arange(
-        min(frame_range), max(frame_range)+1))  # +1 include the last frame
+        min(frame_range), max(frame_range) + 1))  # +1 include the last frame
     frame_series[:] = np.linspace(0, 1, frame_series.shape[0])
     return get_color_dataframe(frame_series, cmap=cmap)
 
@@ -70,9 +70,10 @@ index as in series.
     # substract offset from values, so the smalles one is zero
     normalised_values -= normalised_values.min()
     # all values are zero and have the same value
-    assert normalised_values.max() != 0, 'series.values are constant'
+    if normalised_values.max() == 0:
+        raise ValueError('series.values are constant')
     # actually normalize the values
-    normalised_values /= normalised_values.max()
+    normalised_values = normalised_values / normalised_values.max()
     colors = cmap(normalised_values)
     # create the dataframe from color
     df_colors = pd.DataFrame(
@@ -117,9 +118,9 @@ def draw_frame(frame,
     origin = frame[:, 3]
     for (i, color) in enumerate(colors):
         v = frame[:, i]
-        xs = [origin[0], origin[0]+v[0]]
-        ys = [origin[1], origin[1]+v[1]]
-        zs = [origin[2], origin[2]+v[2]]
+        xs = [origin[0], origin[0] + v[0]]
+        ys = [origin[1], origin[1] + v[1]]
+        zs = [origin[2], origin[2] + v[2]]
         a = Arrow3D(xs, ys, zs,
                     mutation_scale=mutation_scale,
                     lw=lw,

navipy/trajectories/__init__.py

@@ -10,6 +10,8 @@ import navipy.maths.homogeneous_transformations as htf
 from navipy.errorprop import propagate_error
 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
 
 
 def _markers2position(x, kwargs):
@@ -434,7 +436,7 @@ class Trajectory(pd.DataFrame):
             self.loc[index_i, self.rotation_mode] = orientation
         return self
 
-    def world2body(self, markers):
+    def world2body(self, markers, indeces=None):
         """ Transform markers in world coordinate to body coordinate
         """
         if not isinstance(markers, pd.Series):
@@ -446,28 +448,28 @@ class Trajectory(pd.DataFrame):
             msg += ' (i,"x"), (i,"y"),(i,"z")\n'
             msg += 'here i is the index of the marker'
             raise TypeError(msg)
-        transformed_markers = pd.DataFrame(index=self.index,
-                                           columns=markers.index,
-                                           dtype=float)
-        # Looping through each time point along the trajectory
+        if indeces is None:
+            # Looping through each time point along the trajectory
+            indeces = self.index
+        # Looping throught the indeces
         # to get the homogeneous transformation from the position orientation
         # and then apply the transformed to the marker position
-        for index_i, row in self.iterrows():
+        transformed_markers = pd.DataFrame(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
+        for index_i in indeces:
+            row = self.loc[index_i]
             angles = row.loc[self.rotation_mode].values
             translate = row.loc['location'].values
             trans_mat = htf.compose_matrix(angles=angles,
                                            translate=translate,
                                            axes=self.rotation_mode)
-            # More than one marker may be transformed
-            # The marker are assume to be a multiIndex dataframe
-            for marki in markers.index.levels[0]:
-                markpos = [markers.loc[marki, 'x'],
-                           markers.loc[marki, 'y'],
-                           markers.loc[marki, 'z'],
-                           1]
-                markpos = trans_mat.dot(markpos)[:3]
-                transformed_markers.loc[index_i,
-                                        (marki, ['x', 'y', 'z'])] = markpos
+            tmarker = trans_mat.dot(homogeneous_markers.transpose())[:3]
+            transformed_markers.loc[index_i, :] = tmarker.transpose().flatten()
 
         return transformed_markers
 
@@ -503,10 +505,10 @@ class Trajectory(pd.DataFrame):
                                    'dalpha_2']] = rot.squeeze()
         return velocity
 
-    def lollipops(self, ax,
+    def lollipops(self, ax=None,
                   colors=None, step_lollipop=1,
                   offset_lollipop=0, lollipop_marker='o',
-                  linewidth=1, lollipop_tail_width=1,
+                  linewidth=1, lollipop_tail_width=1, lollipop_tail_length=1,
                   lollipop_head_size=1, stickdir='backward'
                   ):
         """ lollipops plot
@@ -531,19 +533,27 @@ class Trajectory(pd.DataFrame):
         :param step_lollipop: number of frames between two lollipops
         :param offset_lollipop: the first lollipop to be plotted
         :param lollipop_marker: the head of the lollipop
-        :param linewidth:
-        :param lollipop_tail_width:
-        :param lollipop_head_size:
-        :param stickdir: The direction of the stick of the animal (backward or forward) 
+        :param linewidth: The width of the line connecting lollipops
+        :param lollipop_tail_width: The width of the lollipop stick
+        :param lollipop_tail_length: The length of the lollipop stick
+        :param lollipop_head_size: The size of the lollipop
+        :param stickdir: The direction of the stick of the animal \
+(backward or forward)
         """
+        # import time
+        # t_start = time.time()
+        if ax is None:
+            fig = plt.figure()
+            ax = fig.add_subplot(111, projection='3d')
         direction = self.facing_direction()
         if colors is None:
             timeseries = pd.Series(data=self.index,
                                    index=self.index)
             colors, sm = get_color_dataframe(timeseries)
-        # Create a continuous index from colors
-        frames = np.arange(colors.index.min(), colors.index.max() + 1)
-
+        # Create a continuous index from trajectory
+        frames = np.arange(self.index.min(), self.index.max() + 1)
+        # Select indeces to save time
+        indeces = frames[offset_lollipop::step_lollipop]
         # Calculate agent tail direction
         # use function in trajectory to get any point bodyref to worldref
         tailmarker = pd.Series(data=0,
@@ -551,10 +561,10 @@ class Trajectory(pd.DataFrame):
                                    [[0],
                                     ['x', 'y', 'z']]))
         if stickdir == 'forward':
-            tailmarker.loc[0, 'x'] = -1
+            tailmarker.loc[0, 'x'] = -lollipop_tail_length
         else:
-            tailmarker.loc[0, 'x'] = 1
-        tail = self.world2body(tailmarker)
+            tailmarker.loc[0, 'x'] = lollipop_tail_length
+        tail = self.world2body(tailmarker, indeces=indeces)
         tail = tail.loc[:, 0]
         # Plot the agent trajectory
         # - loop through consecutive point
@@ -563,36 +573,49 @@ class Trajectory(pd.DataFrame):
         x = self.loc[:, ('location', 'x')]
         y = self.loc[:, ('location', 'y')]
         z = self.loc[:, ('location', 'z')]
-
-        for frame_i, frame_j in zip(frames[:-1], frames[1:]):
-            # Frames may be missing in trajectory,
-            # and therefore can not be plotted
-            if (frame_i in self.index) and \
-                    (frame_j in self.index) and \
-                    (frame_i in self.index):
-                color = [colors.r[frame_i], colors.g[frame_i],
-                         colors.b[frame_i], colors.a[frame_i]]
-                # Create the line to plot
-                line_x = [x[frame_i], x[frame_j]]
-                line_y = [y[frame_i], y[frame_j]]
-                line_z = [z[frame_i], z[frame_j]]
-                # Actual plot command
-                ax.plot(xs=line_x, ys=line_y, zs=line_z,
-                        color=color, linewidth=linewidth)
+        # 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
+            # in the trajectory, a rather long process
+            for frame_i, frame_j in zip(frames[:-1], frames[1:]):
+                # Frames may be missing in trajectory,
+                # and therefore can not be plotted
+                if (frame_i in self.index) and \
+                   (frame_j in self.index) and \
+                   (frame_i in self.index):
+                    color = [colors.r[frame_i], colors.g[frame_i],
+                             colors.b[frame_i], colors.a[frame_i]]
+                    # Create the line to plot
+                    line_x = [x[frame_i], x[frame_j]]
+                    line_y = [y[frame_i], y[frame_j]]
+                    line_z = [z[frame_i], z[frame_j]]
+                    # Actual plot command
+                    ax.plot(xs=line_x, ys=line_y, zs=line_z,
+                            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()
         # Plot the lollipop
         # - loop through the frames with a step of step_lollipop
         # - The lollipop is colored with the color of this frame
         # - Each lollipop is composed of a marker,
         # a point on the agent trajectory
         #   and a line representing the body (anti facing direction)
-        for frame_i in frames[offset_lollipop::step_lollipop]:
+        for frame_i in indeces:
             # Frames may be missing in trajectory,
             # and therefore can not be plotted
             if (frame_i in self.index) and  \
-                    (frame_i in direction.index) and \
-                    (frame_i in colors.index):
-                color = [colors.r[frame_i], colors.g[frame_i],
-                         colors.b[frame_i], colors.a[frame_i]]
+                    (frame_i in direction.index):
+                if isinstance(colors, pd.DataFrame):
+                    if frame_i not in colors.index:
+                        continue
+                    color = [colors.r[frame_i], colors.g[frame_i],
+                             colors.b[frame_i], colors.a[frame_i]]
+                else:
+                    color = colors
                 # Create the line to plot
                 line_x = [self.x[frame_i],
                           tail.x[frame_i]]
@@ -609,3 +632,4 @@ class Trajectory(pd.DataFrame):
                         color=color,
                         marker=lollipop_marker,
                         markersize=lollipop_head_size)
+        # print(time.time() - t_start)