bug fix: plot_total_traj support t color coding

track_analyzer/plotting.py

@@ -1316,6 +1316,7 @@ def plot_total_traj(data_dir, df, dim=3, plot_dir=None, plot_fn=None, plot_confi
         traj_length = traj.shape[0]
         first_frame = traj['frame'].min()
         x0, y0 = traj[traj['frame'] == first_frame][['x_scaled', 'y_scaled']].values[0]
+        t = traj['t'].values
         x = traj['x_scaled'].values
         y = traj['y_scaled'].values
         if dim == 3:
@@ -1324,9 +1325,11 @@ def plot_total_traj(data_dir, df, dim=3, plot_dir=None, plot_fn=None, plot_confi
             traj['x_scaled'] -= x0
             traj['y_scaled'] -= y0
 
-        # color
-        if color_code == "z" and dim == 3:
+        # color code
+        if color_code == "z":
             colors = tpr.get_cmap_color(z, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
+        elif color_code == "t":
+            colors = tpr.get_cmap_color(t, cmap, vmin=cmap_lim[0], vmax=cmap_lim[1])
         elif color_code == "random":
             colors = color_list[i % len(color_list)]
         elif color_code == "none":
@@ -1339,13 +1342,14 @@ def plot_total_traj(data_dir, df, dim=3, plot_dir=None, plot_fn=None, plot_confi
         else:
             colors = color_list[0]
 
-        if color_code == 'z' and dim == 3:
-            for j in range(1, traj_length):
-                ax.plot([x[j - 1], x[j]], [y[j - 1], y[j]], ls='-', color=colors[j])
-            ax.plot(x[-1], y[-1], marker='.', color=colors[-1])
-        else:
-            ax.plot(x, y, ls='-', color=colors)
-            ax.plot(x[-1], y[-1], marker='.', color=colors)
+        if traj_length > 1:  # dont plot single points
+            if color_code == 'z' or color_code == "t":
+                for j in range(1, traj_length):
+                    ax.plot([x[j - 1], x[j]], [y[j - 1], y[j]], ls='-', color=colors[j])
+                ax.plot(x[-1], y[-1], marker='.', color=colors[-1])
+            else:
+                ax.plot(x, y, ls='-', color=colors)
+                ax.plot(x[-1], y[-1], marker='.', color=colors)
 
         if hide_labels is False:
             s = '{}'.format(int(track))