smail alijagic
63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
import sys
|
|
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())
|
|
|
|
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["p.e.x"], w["p.e.y"],
|
|
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["p.e.x"], b["p.e.y"],
|
|
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=20)
|
|
print(f"Saved {sys.argv[2]}")
|
|
else:
|
|
plt.show() |