# %%
f = plt.figure(figsize=(12, 12))
axes = f.subplot_mosaic("""
AAADDD
AAAEEE
AAAFFF
BBBGGG
CCCHHI
""")
f.tight_layout()
_ = draw.Hairpin(ax=axes["A"])
for n, bout in enumerate(bouts):
# draw 2d and linearized tracking
draw.Tracking(bout.x, bout.y, ax=axes["A"])
draw.Tracking(bout.linearized.x, bout.linearized.y, ax=axes["B"])
# draw speed accel and ang vel along track
axes["D"].plot(bout.linearized.x, bout.speed, color="k", alpha=0.5)
axes["E"].plot(bout.linearized.x, bout.acceleration, color="k", alpha=0.5)
axes["F"].plot(bout.linearized.x, bout.thetadot, color="k", alpha=0.5)
axes["G"].plot(bout.linearized.x, bout.thetadotdot, color="k", alpha=0.5)
# mark peak acceleration etc
# peak_accel = np.where(bout.acceleration > np.percentile(bout.acceleration, 80))[0]
# draw.Tracking.scatter(bout.x[peak_accel], bout.y[peak_accel],
# color=blue, zorder=100, ax=axes['A'], label='peak accel' if n == 0 else None)
center_line,
K,
angle_cost=1,
length_cost=0.01)
# ----------------------------------- PLOT ----------------------------------- #
f = plt.figure(figsize=(8, 12))
axes = f.subplot_mosaic("""
AAA
AAA
AAA
BBB
""")
# draw tracking 2d
draw.Tracking(center_line.x, center_line.y, ax=axes["A"])
draw.Tracking.scatter(center_line.x,
center_line.y,
ax=axes["A"],
color=blue_grey)
draw.Tracking(left_line.x, left_line.y, ls="--", alpha=0.5, ax=axes["A"])
draw.Tracking(right_line.x,
right_line.y,
ls="--",
alpha=0.5,
color="b",
ax=axes["A"])
for pts in control_points.values():
draw.Tracking.scatter(pts.x,
pts.y,
from myterial import pink
import draw
f, ax = plt.subplots(figsize=(7, 10))
# load and draw tracking data
from fcutils.path import files
for fp in files(
"/Users/federicoclaudi/Dropbox (UCL)/Rotation_vte/Locomotion/analysis/control/",
"*.h5",
):
tracking = pd.read_hdf(fp, key="hdf")
tracking.x = 20 - tracking.x + 20
# draw.Tracking.scatter(tracking.x, tracking.y, c=tracking.theta, vmin=0, vmax=360, cmap='bwr', lw=1, ec='k')
draw.Tracking(tracking.x, tracking.y, alpha=0.7)
# draw hairpin arena
draw.Hairpin(ax)
# draw waypoints
wps = Waypoints()
for wp in wps:
draw.Arrow(wp.x, wp.y, wp.theta, 2, width=4, color=pink, outline=True)
# fit dubin path
dubin = DubinPath(wps).fit()
draw.Tracking(dubin.x, dubin.y, lw=2, color="k")
plt.show()
body_only=False,
movement=False)
if tracking.empty:
logger.info(f'"{session}" - no tracking')
continue
body = tracking.loc[tracking.bpname == 'body'].iloc[0]
snout = tracking.loc[tracking.bpname == 'snout'].iloc[0]
paw = tracking.loc[tracking.bpname == 'right_fl'].iloc[0]
tail = tracking.loc[tracking.bpname == 'tail_base'].iloc[0]
f, ax = plt.subplots(figsize=(9, 12))
draw.Hairpin(ax=ax)
draw.Tracking(body.x, body.y, ax=ax)
for bp, color in zip((snout, paw, tail), (teal, salmon, indigo)):
draw.Tracking.scatter(bp.x[::30],
bp.y[::30],
color=color,
alpha=.5,
zorder=100,
ax=ax,
label=bp.bpname)
ax.set(title=f'{session} - {len(tracking.x)} frames')
ax.legend()
recorder.add_figure(f, session, svg=False)
def animate_one(ROI: str, crossing_id: int, FPS: int):
scale = 1 / 30 # to scale velocity vectors
save_folder = (paths.analysis_folder / "behavior" /
"roi_crossings_animations")
# ----------------------------------- prep ----------------------------------- #
# load tracking
bouts = pd.read_hdf(paths.analysis_folder / "behavior" / "roi_crossings" /
f"{ROI}_crossings.h5")
bout = bouts.sort_values("duration").iloc[crossing_id]
logger.info(f"Animating bout {crossing_id} - {round(bout.duration, 2)}s")
# tracking: dict = ROICrossing.get_crossing_tracking(bout.crossing_id)
# generate a path from tracking
path = Path(bout.x, bout.y)
# create fig and start camera
if ROI == "T4":
f = plt.figure(figsize=(12, 12))
elif "S" in ROI:
f = plt.figure(figsize=(8, 14))
else:
f = plt.figure(figsize=(5, 14))
axes = f.subplot_mosaic("""
AA
AA
BB
""")
camera = Camera(f)
# ----------------------------------- plot ----------------------------------- #
n_interpolated_frames = int(60 / FPS)
n_frames = len(bout.x) - 1
trajectory = GrowingPath()
for frame in track(range(n_frames), total=n_frames):
# repeat each frame N times interpolating between frames
for interpol in np.linspace(0, 1, n_interpolated_frames):
# get interpolated quantities
x = interpolate_at_frame(path.x, frame, interpol)
y = interpolate_at_frame(path.y, frame, interpol)
speed = interpolate_at_frame(path.speed, frame, interpol)
trajectory.update(x, y, speed=speed)
# time elapsed
_time = (np.arange(len(trajectory.speed)) / n_interpolated_frames /
60)
# plot the arena
draw.ROI(ROI, set_ax=True, shade=False, ax=axes["A"])
# plot tracking so far
draw.Tracking(trajectory.x, trajectory.y, lw=3, ax=axes["A"])
draw.Dot(x, y, s=50, ax=axes["A"])
# plot speed and velocity vectors
draw.Arrow(
x,
y,
interpolate_at_frame(path.velocity.angle,
frame,
interpol,
method="step"),
L=scale *
interpolate_at_frame(path.velocity.magnitude, frame, interpol),
color=colors.velocity,
outline=True,
width=2,
ax=axes["A"],
)
draw.Arrow(
x,
y,
path.acceleration.angle[frame],
L=4 * scale * path.acceleration.magnitude[frame],
color=colors.acceleration,
outline=True,
width=2,
ax=axes["A"],
zorder=200,
)
# plot speed trace
axes["B"].fill_between(_time,
0,
trajectory.speed,
alpha=0.5,
color=colors.speed)
axes["B"].plot(_time, trajectory.speed, lw=4, color=colors.speed)
axes["A"].set(title=f"{ROI} - crossing: {crossing_id}")
axes["B"].set(xlabel="time (s)", ylabel="speed (cm/s)")
camera.snap()
# ------------------------------- video creation ------------------------------- #
save_path = save_folder / f"{ROI}_{bout.crossing_id}.mp4"
logger.info(f"Saving animation @: '{save_path}'")
animation = camera.animate(interval=1000 / FPS)
animation.save(save_path, fps=FPS)
logger.info("Done!")
plt.cla()
plt.close(f)
del camera
del animation
# ----------------------------------- plot ----------------------------------- #
f = plot_roi_crossing(bout, step=4)
recorder.add_figures(svg=False)
plt.close(f)
# get controller
MODEL = KinematicsLQR(path)
# MODEL = PurePursuit(path)
animation_step = 5
frame = 0
for t in np.arange(0, maxTime, dt):
# ---------------------------------- control --------------------------------- #
target_speed = MODEL.step(t)
# ----------------------------------- plot ----------------------------------- #
if frame % animation_step == 0:
plt.cla()
# draw track
draw.Tracking(path.x, path.y, lw=3)
# draw car trajectory
draw.Tracking(
MODEL.trajectory.x, MODEL.trajectory.y, lw=1.5, color=salmon_dark
)
draw.ControlCar(
MODEL.vehicle.x,
MODEL.vehicle.y,
MODEL.vehicle.theta,
-MODEL.vehicle.steer,
)
# draw car velocity and acceleration vectors
if len(MODEL.trajectory.x) > 5:
velocity, _, _, acceleration, _, _ = va.compute_vectors(
axes = f.subplot_mosaic("""
AABC
AADE
""")
for ax in "ABCDE":
axes[ax].axis("equal")
draw.Hairpin(ax=axes[ax], set_ax=True if ax == "A" else False)
if ax != "A":
axes[ax].axis("equal")
axes[ax].axis("off")
# draw tracking traces
for bout in bouts:
draw.Tracking(bout.x + 0.5, bout.y + 0.5, ax=axes["A"])
# draw speed heatmap
draw.Tracking.heatmap(
X + 0.5,
Y + 0.5,
c=S,
gridsize=(40, 60),
ax=axes["B"],
vmin=0,
vmax=80,
mincnt=10,
colorbar=True,
cmap="inferno",
)
f, axes = plt.subplots(figsize=(18, 12), ncols=3, nrows=2)
axes = axes.flatten()
for i, bout in selected_bouts.iterrows():
# generate a path from tracking
path = Path(bout.x, bout.y)
ax = axes[i]
# draw arena and tracking
draw.ROI(ROI, set_ax=True, shade=False, ax=ax)
# draw.Tracking(bouts.x, bouts.y, alpha=0.5)
# draw bout tracking and velocity/acceleration vectors
# draw.Tracking.scatter(path.x, path.y, c=path.tangent.angle, cmap='bwr', vmin=-180, vmax=180)
draw.Tracking(path.x, path.y, color=[0.6, 0.6, 0.6], lw=5, ax=ax)
draw.Arrows(
path.x,
path.y,
path.velocity.angle,
L=path.velocity.magnitude,
color=blue,
step=2,
outline=True,
width=2,
ax=ax,
)
draw.Arrows(
path.x,
path.y,
def plot_roi_crossing(
crossing: pd.Series,
tracking: dict = None,
step: int = 5,
highlight: str = None,
arrow_scale=0.5,
) -> plt.Figure:
"""
Plots an entire ROI crossing, tracking, vectors and mouse
"""
if highlight == "velocity":
v_alpha, a_alpha = 1, 0.4
elif highlight == "acceleration":
v_alpha, a_alpha = 0.4, 1
else:
v_alpha, a_alpha = 1, 1
# path = Path(
# convolve_with_gaussian(crossing.x, kernel_width=11),
# convolve_with_gaussian(crossing.y, kernel_width=11))
path = Path(crossing.x, crossing.y)
f = plt.figure(figsize=(20, 10))
axes = f.subplot_mosaic(
"""
AACC
AADD
BBEE
"""
)
f._save_name = (
f"{crossing.roi}_{crossing.crossing_id}" + ""
if highlight is None
else highlight
)
# draw main crossing plot
draw.ROI(crossing.roi, ax=axes["A"], set_ax=True)
draw.Tracking(path.x, path.y, lw=0.5, color="k", ax=axes["A"])
# draw velocity and acceleartion vectors
draw.Arrows(
path.x,
path.y,
path.velocity.angle,
label="velocity",
L=arrow_scale * path.velocity.magnitude / 30,
step=step,
color=colors.velocity,
ax=axes["A"],
outline=True,
alpha=v_alpha,
)
draw.Arrows(
path.x,
path.y,
path.acceleration.angle,
label="acceleration",
L=arrow_scale * 6 * path.acceleration.magnitude / 30,
step=step,
color=colors.acceleration,
ax=axes["A"],
outline=True,
alpha=a_alpha,
)
draw.Tracking.scatter(
path.x[::step], path.y[::step], color="k", ax=axes["A"], zorder=200,
)
# draw speed traces
time = np.arange(len(path.speed)) / 60
axes["B"].fill_between(time, 0, path.speed, alpha=0.5, color=colors.speed)
axes["B"].plot(
time, path.speed, lw=4, color=colors.speed,
)
if tracking is not None:
for bp in PAWS:
axes["C"].plot(
tracking[bp]["bp_speed"], color=colors.bodyparts[bp], label=bp
)
axes["C"].plot(
tracking["body"]["bp_speed"],
color=colors.bodyparts["body"],
label="body",
)
axes["C"].plot(
np.mean(
np.vstack(
[tracking[bp]["bp_speed"] for bp in PAWS if "h" in bp]
),
0,
),
color="k",
label="mean",
)
# clean plots
axes["A"].legend()
axes["B"].set(xlabel="time (s)", ylabel="speed (cm/s)")
axes["C"].legend()
f.tight_layout()
return f
colors = make_palette(pink_dark, blue_dark, len(P))
traces = []
for n, p in enumerate(P):
trace, _ = tp.fit_best_trace(
control_points,
center_line,
K,
angle_cost=p,
length_cost=(1 - p) * 5e-3,
)
# draw best trace
X = convolve_with_gaussian(trace.x, kernel_width=11)
Y = convolve_with_gaussian(trace.y, kernel_width=11)
traces.append(Path(X, Y))
draw.Tracking(X, Y, color=colors[n], lw=3)
# %%
# ------------------------------- load tracking ------------------------------ #
_bouts = pd.read_hdf(paths.analysis_folder / "behavior" / "saved_data" /
f"complete_bouts.h5")
_bouts = _bouts.loc[_bouts.duration < 8]
print(f"Kept {len(_bouts)} bouts")
bouts = []
for i, bout in _bouts.iterrows():
bouts.append(LocomotionBout(bout))
# %%
# ----------------------------------- plot ----------------------------------- #
f, ax = plt.subplots(figsize=(8, 12))
f, ax = plt.subplots(figsize=(7, 10))
# load and draw tracking data
from fcutils.path import files
for n, fp in enumerate(
files(
"/Users/federicoclaudi/Dropbox (UCL)/Rotation_vte/Locomotion/analysis/control/",
"*.h5",
)
):
tracking = pd.read_hdf(fp, key="hdf")
tracking.x = 20 - tracking.x + 20
# draw.Tracking.scatter(tracking.x, tracking.y, c=tracking.speed, vmin=0, vmax=100, cmap='bwr', lw=1, ec='k')
draw.Tracking(tracking.x, tracking.y, alpha=0.7)
if n == 5:
wps = Waypoints.from_tracking(
tracking.x, tracking.y, tracking.speed, tracking.theta,
)
# draw hairpin arena
draw.Hairpin(ax)
# draw waypoints
for wp in wps:
draw.Arrow(wp.x, wp.y, wp.theta, 2, width=4, color="g")
# fit and animate quintic polinomial
# wps = [
import draw
# x = np.linspace(0, 30, 100)
# y = np.sin(x)
c = np.linspace(0, 1.5 * np.pi, 100)
x = c
y = np.sin(c)
path = Path(x, y)
downsampled = path.downsample(1)
downsampled_eu = path.downsample_euclidean(1)
upsampled = downsampled.interpolate(0.25)
draw.Tracking(path.x, path.y, color="k", alpha=0.5)
draw.Tracking.scatter(
path.x, path.y, color="k", label="original", alpha=0.5
)
draw.Tracking(downsampled.x, downsampled.y, color="r")
draw.Tracking.scatter(
downsampled.x, downsampled.y, color="r", label="downsampled"
)
draw.Tracking(upsampled.x, upsampled.y, color="b")
draw.Tracking.scatter(
upsampled.x, upsampled.y, color="b", label="resampled"
)
draw.Tracking(downsampled_eu.x, downsampled_eu.y, color="g")
axes = f.subplot_mosaic("""
AABC
AADE
""")
for ax in "ABCDE":
axes[ax].axis("equal")
draw.Hairpin(ax=axes[ax], set_ax=True if ax == "A" else False)
if ax != "A":
axes[ax].axis("equal")
axes[ax].axis("off")
# draw tracking traces
for bout in bouts:
draw.Tracking(bout.body.x + 0.5, bout.body.y + 0.5, ax=axes["A"])
# draw speed heatmap
draw.Tracking.heatmap(
X + 0.5,
Y + 0.5,
c=S,
gridsize=(40, 60),
ax=axes["B"],
vmin=0,
vmax=80,
mincnt=3,
colorbar=True,
cmap="inferno",
)
axes = f.subplot_mosaic("""
AAABB
AAACC
""")
# draw the curvature
_ = axes["B"].plot(center_line.comulative_distance,
center_line.curvature,
lw=4,
zorder=100)
axes["B"].set(xticks=[], ylabel="curvature")
# draw single bouts speeds
for bout, centered_bout in zip(bouts, centered_bouts):
# tracking
draw.Tracking(bout.body.x, bout.body.y, lw=0.5, ax=axes["A"])
# curvature
axes["B"].plot(
bout.body.comulative_distance,
bout.body.curvature,
color="k",
alpha=0.25,
)
# speed
axes["C"].plot(centered_bout.x, bout.body.speed, color="k", alpha=0.25)
# draw the v_max
for y_max in [2.5]:
vmax = np.sqrt((y_max * 60) / (center_line.curvature))
draw.ROI(ROI, ax=axes["A"])
SKIP = 0
N = 10 # len(_bouts)
S, T = np.full((80, N), np.nan), np.full((80, N), np.nan)
_S, _T = np.full((80, N), np.nan), np.full((80, N), np.nan)
simulated = []
for i in track(range(N)):
bout = paths[i]
trajectory, _time = MSD(
bout, skip=SKIP, start_frame=2, end_frame=-2
).simulate()
# plot
draw.Tracking(bout.x, bout.y, zorder=-1, ax=axes["A"], alpha=0.8)
draw.Tracking(
trajectory.x,
trajectory.y,
zorder=10,
ax=axes["A"],
color="red",
alpha=0.8,
)
simulated.append(trajectory)
# store tracking data
S[: len(bout), i] = bout.speed
# store simulation results
_S[_time[0] : _time[-1], i] = trajectory.speed
import draw
f, ax = plt.subplots(figsize=(7, 10))
# load and draw tracking data
from fcutils.path import files
for fp in files(
"/Users/federicoclaudi/Dropbox (UCL)/Rotation_vte/Locomotion/analysis/control/",
# r'D:\Dropbox (UCL)\Rotation_vte\Locomotion\analysis\control',
"*.h5",
):
tracking = pd.read_hdf(fp, key="hdf")
tracking.x = 20 - tracking.x + 20
draw.Tracking(tracking.x, tracking.y, alpha=0.7)
# draw hairpin arena
draw.Hairpin(ax)
# draw waypoints
wps = Waypoints(use="spline")
for wp in wps:
draw.Arrow(
wp.x,
wp.y,
wp.theta,
2.5,
width=4,
color=pink,
outline=True,
logger.info(f"Loaded {len(crossings)} crossings")
color = arena.ROIs_dict[crossings.iloc[0].roi].color
# draw tracking
for i, cross in crossings.iterrows():
if i % 30 == 0:
axes["A"].scatter(
cross.x[0],
cross.y[0],
zorder=100,
color=color,
lw=0.5,
ec=[0.3, 0.3, 0.3],
)
draw.Tracking(cross.x, cross.y, ax=axes["A"], alpha=0.7)
# draw heatmaps for speed and acceleration for fastest crossings
fast_crossings = crossings.sort_values('duration').iloc[:250]
visuals.plot_heatmap_2d(
fast_crossings,
gridsize=25,
key="speed",
ax=axes["B"],
)
visuals.plot_heatmap_2d(
fast_crossings,
gridsize=25,
key="acceleration",
ax=axes["D"],
vmin=-4,
f = plt.figure(figsize=(20, 12), constrained_layout=True)
f.suptitle(f'{unit.brain_region} - {unit.unit_id}')
axes = f.subplot_mosaic('''
AABBCCFF
AADDEEGG
''')
unit = units.loc[units.unit_id == 1036].iloc[0]
frate = convolve_with_gaussian(unit.firing_rate)
draw.ROI(roi, ax=axes['A'])
X, Y = [], []
for bout in bouts:
draw.Tracking(bout.x, bout.y - 2, ax=axes['A'])
axes['B'].plot(bout.velocity.magnitude, color='k', lw=.5)
axes['C'].plot(bout.thetadot, color='k', lw=.5)
axes['D'].plot(frate[bout.start_frame:bout.end_frame])
X.extend(list(frate[bout.start_frame + 2:bout.end_frame - 2]))
Y.extend(list(bout.speed))
sns.regplot(X, Y, ax=axes['E'])
starts = [b.start_frame for b in bouts]
ends = [b.end_frame for b in bouts]
draw.Raster(unit.spikes, starts, ends)
# %%