brown/particles.py

58 lines
1.6 KiB
Python
Raw Normal View History

2019-07-12 19:52:12 +00:00
from brown.interaction import UFuncWrapper
from brown.brown import BrownIterator
import numpy as np
from collections import deque
from copy import copy
import matplotlib.pyplot as plt
import matplotlib.animation as ani
c = np.array([5, 10, 20, 30, 0, 0, 0, 1, -20, 0, -2, -0.1, 2, 0, 0, 0, 0, 0, 0], dtype=np.float16)
#force_function = UFuncWrapper(0, c)
#interaction2D = UFuncWrapper(1, c)
borders_x = [-100, 100]
borders_y = [-100, 100]
n_particles = 6
frames = 1000
x_coords = np.random.uniform(borders_x[0] / 2, borders_x[1] / 2, n_particles).astype(np.float16)
y_coords = np.random.uniform(borders_y[0] / 2, borders_y[1] / 2, n_particles).astype(np.float16)
x_momenta = np.zeros(n_particles, dtype=np.float16)
y_momenta = np.zeros(n_particles, dtype=np.float16)
fig = plt.figure(figsize=(7, 7))
ax = fig.add_axes([0, 0, 1, 1], frameon=False)
ax.set_xlim(*borders_x)
ax.set_xticks([])
ax.set_ylim(*borders_y)
ax.set_yticks([])
plot, = ax.plot(x_coords, y_coords, "b.")
center_of_mass, = ax.plot(x_coords.mean(), y_coords.mean(), "r-")
center_of_mass_history_x = deque([x_coords.mean()])
center_of_mass_history_y = deque([y_coords.mean()])
brown = BrownIterator(-1, c
, x_coords, y_coords
, y_momenta, y_momenta
, borders_x, borders_y
, border_dampening=1
, dt=0.001)
u = iter(brown)
def update(i):
data = next(u)
center_of_mass_history_x.append(x_coords.mean())
center_of_mass_history_y.append(y_coords.mean())
plot.set_data(*data)
center_of_mass.set_data(center_of_mass_history_x, center_of_mass_history_y)
animation = ani.FuncAnimation(fig, update, range(frames), interval=1)
plt.show()