Commit 652caff8 authored by amichaut's avatar amichaut
Browse files

bug fix: plot_total_traj support t color coding

parent 32b1eaee
This diff is collapsed.
This diff is collapsed.
......@@ -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))
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment