From 6927e8cf767f7b30c3a3363f200702593047e1a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Kn=C3=BCttel?= <daniel.knuettel@daknuett.eu>
Date: Wed, 31 Jul 2019 12:05:33 +0200
Subject: [PATCH] added option to draw arrows

---
 particles.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/particles.py b/particles.py
index c56ce87..c312300 100644
--- a/particles.py
+++ b/particles.py
@@ -16,17 +16,20 @@ from coefficients import c
 # the BrownIterator).
 borders_x = [-100, 100]
 borders_y = [-100, 100]
-n_particles = 600
+n_particles = 60
 # Idk, seems to not do anyting.
 frames = 100
 # Only spawn in 1/x of the borders.
-spawn_restriction = 3
+spawn_restriction = 2
 # Time resolution. Note that setting this to a too
 # high value (i.e. low resolution) will lead to 
 # erratic behaviour, because potentials can be skipped.
-dt = 0.1
+dt = 0.01
 c[-1] = dt
 
+# Draw arrows, makes the simulation fucking slow.
+draw_arrows = True
+
 # Initial positions.
 x_coords = np.random.uniform(borders_x[0] / spawn_restriction, borders_x[1] / spawn_restriction, n_particles).astype(np.float32)
 y_coords = np.random.uniform(borders_y[0] / spawn_restriction, borders_y[1] / spawn_restriction, n_particles).astype(np.float32)
@@ -49,6 +52,12 @@ ax.set_yticks([])
 # Plot the initial values.
 plot, = ax.plot(x_coords, y_coords, "b.")
 center_of_mass, = ax.plot(x_coords.mean(), y_coords.mean(), "r-")
+if(draw_arrows):
+	arrows = [ plt.Arrow(x, y, dx, dy) for x, y, dx, dy in zip(x_coords, y_coords, x_momenta, y_momenta)]
+	for arrow in arrows:
+		ax.add_patch(arrow)
+
+
 # Keep track of the center of mass.
 center_of_mass_history_x = deque([x_coords.mean()])
 center_of_mass_history_y = deque([y_coords.mean()])
@@ -57,9 +66,9 @@ brown = BrownIterator(-1, c # Max iterations, simulation parameters.
 			, x_coords, y_coords
 			, y_momenta, y_momenta
 			# The boundary condition: reflect at the borders,
-			#, borders_x, borders_y
+			, borders_x, borders_y
 			# or just let propagate to infinity.
-			, [], []
+			#, [], []
 			# Let the border dampen the system, border_dampening < 1 => energy is absorbed.
 			, border_dampening=1
 			, dt=dt)
@@ -67,6 +76,7 @@ brown = BrownIterator(-1, c # Max iterations, simulation parameters.
 u = iter(brown)
 
 def update(i):
+	global arrows
 	# Get the next set of positions.
 	data = next(u)
 	center_of_mass_history_x.append(x_coords.mean())
@@ -74,6 +84,12 @@ def update(i):
 
 	plot.set_data(*data)
 	center_of_mass.set_data(center_of_mass_history_x, center_of_mass_history_y)
+	if(draw_arrows):
+		for arrow in arrows:
+			ax.patches.remove(arrow)
+		arrows = [plt.Arrow(x, y, dx, dy) for x, y, dx, dy in zip(data[0], data[1], u.px, u.py)]
+		for arrow in arrows:
+			ax.add_patch(arrow)
 
 animation = ani.FuncAnimation(fig, update, range(frames), interval=1)
 plt.show()