added arrows

visualization.py

@@ -1,52 +1,63 @@
 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)")
-
+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))
+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
+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)
-
-
+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]}")
+    print(f"Saved {sys.argv[2]}")
 else:
     plt.show()