Framework for simulations

visualization.py

@@ -0,0 +1,52 @@
+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()