|
| 1 | +import numpy as np |
| 2 | +import matplotlib.pyplot as plt |
| 3 | +from matplotlib.animation import FuncAnimation |
| 4 | + |
| 5 | +NUM_PARTICLES = 150 |
| 6 | +NUM_GROUPS = 4 |
| 7 | +G = 1.0 |
| 8 | +BETA = 0.9 |
| 9 | +BOUNDS = 100 |
| 10 | + |
| 11 | +particles = np.zeros(NUM_PARTICLES, dtype=[ |
| 12 | + ('position', float, 2), |
| 13 | + ('velocity', float, 2), |
| 14 | + ('force', float, 2), |
| 15 | + ('group', int, 1) |
| 16 | +]) |
| 17 | + |
| 18 | +particles['position'] = np.random.rand(NUM_PARTICLES, 2) * BOUNDS |
| 19 | +particles['velocity'] = np.zeros((NUM_PARTICLES, 2)) |
| 20 | +particles['group'] = np.random.randint(0, NUM_GROUPS, size=NUM_PARTICLES) |
| 21 | + |
| 22 | +fig, ax = plt.subplots() |
| 23 | +ax.set_facecolor('black') |
| 24 | +fig.set_size_inches(8, 8) |
| 25 | +ax.set_xticks([]) |
| 26 | +ax.set_yticks([]) |
| 27 | +plt.xlim(0, BOUNDS) |
| 28 | +plt.ylim(0, BOUNDS) |
| 29 | + |
| 30 | +scatter = ax.scatter( |
| 31 | + particles['position'][:, 0], |
| 32 | + particles['position'][:, 1], |
| 33 | + c=particles['group'], |
| 34 | + cmap='viridis', |
| 35 | + s=10 |
| 36 | +) |
| 37 | + |
| 38 | +def update(frame): |
| 39 | + particles['force'] = 0.0 |
| 40 | + |
| 41 | + for i in range(NUM_PARTICLES): |
| 42 | + for j in range(NUM_PARTICLES): |
| 43 | + if i == j: |
| 44 | + continue |
| 45 | + |
| 46 | + delta_pos = particles['position'][j] - particles['position'][i] |
| 47 | + dist_sq = max(np.sum(delta_pos**2), 1.0) |
| 48 | + |
| 49 | + force_mag = G / dist_sq |
| 50 | + |
| 51 | + force_vec = delta_pos * force_mag / np.sqrt(dist_sq) |
| 52 | + |
| 53 | + if particles['group'][i] == particles['group'][j]: |
| 54 | + particles['force'][i] -= force_vec |
| 55 | + else: |
| 56 | + particles['force'][i] += force_vec |
| 57 | + |
| 58 | + particles['velocity'] = (particles['velocity'] + particles['force']) * BETA |
| 59 | + particles['position'] += particles['velocity'] |
| 60 | + |
| 61 | + particles['position'] %= BOUNDS |
| 62 | + |
| 63 | + scatter.set_offsets(particles['position']) |
| 64 | + |
| 65 | + return scatter, |
| 66 | + |
| 67 | +ani = FuncAnimation(fig, update, frames=200, interval=10, blit=True) |
| 68 | + |
| 69 | +plt.show() |
0 commit comments