lbaumberger
52 lines
1.4 KiB
Python
52 lines
1.4 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.gif)")
|
|
|
|
|
|
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")
|
|
|
|
|
|
# Lets define the 2 groups
|
|
# WHITE = 1 (left -> right, drawn as filled black circles, like paper)
|
|
# BLACK = 0 (right -> left, drawn as empty circles)
|
|
s_white = ax.scatter([],[], s=30, facecolors = "black", edgecolor = "black", zorder = 2)
|
|
s_black = ax.scatter([],[], s=30, facecolors = "white", edgecolor = "black", zorder = 2)
|
|
|
|
|
|
|
|
# update our graphic from t to t+1
|
|
title = ax.set_title("")
|
|
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)
|
|
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() |