N-body simulation in 2D
October 16, 2022
A few days ago I implemented a n-body simulation in Python, which is used to generate the following animation:
Vectorized n-body implementation #
It uses a vectorized form of calculating the acceleration matrix, which is used to update the positions of the bodies. The velocity is integrated with leapfrog integration, which unlike Euler integration, is stable for oscillatory motion. It is pretty efficient. The acceleration matrix is calculated for 2D, but is easily extended to work in 3D as well.
Anyway, here is the code:
1n = 10
2gravity = 0.0001
3step_size = 20/1000
4half_step_size = 0.5 * step_size
5softening = 0.1
6
7position = 4 * (np.random.rand(n, 2) - 0.5)
8velocity = 1 * (np.random.rand(n, 2) - 0.5)
9acceleration = np.zeros((n, 2))
10mass = 500 * np.ones((n, 1))
11
12com_p = np.sum(np.multiply(mass, position), axis=0) / np.sum(mass, axis=0)
13com_v = np.sum(np.multiply(mass, velocity), axis=0) / np.sum(mass, axis=0)
14for p in position: p -= com_p
15for v in velocity: v -= com_v
16
17position[0] = np.array([0, 0])
18velocity[0] = np.array([0, 0])
19mass[0] = 10000
20
21fig = plt.figure()
22fig.patch.set_facecolor('white')
23ax = plt.axes(xlim=(-2.5, 2.5), ylim=(-2.5, 2.5))
24ax.set_aspect(1)
25scatter = ax.scatter(position[:,0], position[:,1], np.sqrt(mass),
26 color="#56c474", edgecolors='black', lw=1)
27
28
29def calculate_acceleration_matrix(P, M, gravity, softening):
30 x = P[:,0:1]
31 y = P[:,1:2]
32 dx = x.T - x
33 dy = y.T - y
34 inv_r3 = (dx**2 + dy**2 + softening**2) ** (-1.5)
35 ax = gravity * (dx * inv_r3) @ M
36 ay = gravity * (dy * inv_r3) @ M
37 return np.hstack((ax, ay))
38
39
40def animate(i):
41 global acceleration, velocity, position, softening
42 global step_size, half_step_size
43
44 velocity += acceleration * half_step_size
45 position += velocity * step_size
46
47 acceleration = calculate_acceleration_matrix(
48 position, mass, gravity, softening)
49
50 velocity += acceleration * half_step_size
51 scatter.set_offsets(position)
52
53
54anim = animation.FuncAnimation(
55 fig, animate, frames=400, interval=20, blit=False)
56
57anim.save('filename.gif', writer='imagemagick')