lbaumberger
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
import sys
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.animation import FuncAnimation
|
|
|
|
if len(sys.argv) < 2:
|
|
sys.exit("Use it like the following: python visualization.py positions.csv (out.mp4)")
|
|
|
|
df = pd.read_csv(sys.argv[1])
|
|
frames = sorted(df["time"].unique())
|
|
|
|
# Replace p.e with direction-of-next-hop: where the ped actually moves to next frame.
|
|
# Falls back to p.e for the last frame and for teleport jumps.
|
|
df = df.sort_values(["id", "time"]).reset_index(drop=True)
|
|
df["dx"] = df.groupby("id")["x"].shift(-1) - df["x"]
|
|
df["dy"] = df.groupby("id")["y"].shift(-1) - df["y"]
|
|
mag = np.sqrt(df["dx"]**2 + df["dy"]**2)
|
|
TELEPORT_THRESHOLD = 5.0 # any per-step jump larger than this is the teleport, not motion
|
|
fallback = (mag > TELEPORT_THRESHOLD) | mag.isna() | (mag < 1e-9)
|
|
df.loc[fallback, "dx"] = df.loc[fallback, "p.e.x"]
|
|
df.loc[fallback, "dy"] = df.loc[fallback, "p.e.y"]
|
|
mag = np.sqrt(df["dx"]**2 + df["dy"]**2).replace(0, 1.0)
|
|
df["dx"] = df["dx"] / mag
|
|
df["dy"] = df["dy"] / mag
|
|
|
|
fig, ax = plt.subplots(figsize=(10, 2.2))
|
|
ax.set_xlim(0, 50); ax.set_ylim(0, 10)
|
|
ax.set_aspect("equal"); ax.set_facecolor("white")
|
|
ax.set_xticks([]); ax.set_yticks([])
|
|
for s in ("top", "right", "left", "bottom"):
|
|
ax.spines[s].set_color("gray")
|
|
|
|
s_white = ax.scatter([], [], s=30, facecolors="black", edgecolor="black", zorder=2)
|
|
s_black = ax.scatter([], [], s=30, facecolors="white", edgecolor="black", zorder=2)
|
|
title = ax.set_title("")
|
|
|
|
# Keep a reference to quiver objects so we can remove them each frame
|
|
quivers = []
|
|
|
|
def update(t):
|
|
sub = df[df["time"] == t]
|
|
w = sub[sub["color"] == 1]
|
|
b = sub[sub["color"] == 0]
|
|
|
|
s_white.set_offsets(w[["x", "y"]].values)
|
|
s_black.set_offsets(b[["x", "y"]].values)
|
|
|
|
# Remove old quiver arrows
|
|
for q in quivers:
|
|
q.remove()
|
|
quivers.clear()
|
|
|
|
# Draw fresh arrows each frame
|
|
arrow_scale = 1.5 # Tune this to change arrow length visually
|
|
if len(w):
|
|
q_white = ax.quiver(w["x"], w["y"], w["dx"], w["dy"],
|
|
color="black", angles="xy", scale_units="xy",
|
|
scale=1/arrow_scale, width=0.003,
|
|
headwidth=4, headlength=5, zorder=3)
|
|
quivers.append(q_white)
|
|
if len(b):
|
|
q_black = ax.quiver(b["x"], b["y"], b["dx"], b["dy"],
|
|
color="gray", angles="xy", scale_units="xy",
|
|
scale=1/arrow_scale, width=0.003,
|
|
headwidth=4, headlength=5, zorder=3)
|
|
quivers.append(q_black)
|
|
|
|
title.set_text(f"t = {t}")
|
|
return s_white, s_black, title
|
|
|
|
anim = FuncAnimation(fig, update, frames=frames, interval=50, blit=False)
|
|
|
|
if len(sys.argv) >= 3:
|
|
anim.save(sys.argv[2], writer="ffmpeg", fps=5)
|
|
print(f"Saved {sys.argv[2]}")
|
|
else:
|
|
plt.show() |