From 6a6034fb6c056d640d02b5c5607006870d2c81cd Mon Sep 17 00:00:00 2001
From: brandhoj95 <brandhoj95@zedat.fu-berlin.de>
Date: Fri, 13 Dec 2024 15:38:16 +0100
Subject: [PATCH] added colored tracks
---
src/robofish/io/file.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index 1ce2f1e..1ea891f 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -928,7 +928,7 @@ class File(h5py.File):
else:
step_size = poses.shape[1]
- cmap = matplotlib.cm.get_cmap(cmap)
+ cmap = plt.get_cmap(cmap)
x_world, y_world = self.world_size
if figsize is None:
@@ -951,6 +951,10 @@ class File(h5py.File):
this_c = c
timesteps = poses.shape[1] - 1
+ colours = plt.get_cmap(cmap)(np.linspace(0, 1, 1 + timesteps // step_size))
+
+ if not lw_distances:
+ colours[:] = this_c
for t in range(0, timesteps, step_size):
if lw_distances:
lw = np.mean(line_width[t : t + step_size + 1])
@@ -958,7 +962,7 @@ class File(h5py.File):
ax.plot(
poses[fish_id, t : t + step_size + 1, 0],
poses[fish_id, t : t + step_size + 1, 1],
- c=this_c,
+ c=colours[t//step_size],
lw=lw,
)
# Plotting outside of the figure to have the label
--
GitLab