import numpy as np import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation import ffmpeg class nbody_system: def __init__(self, masses, x, v, unit_system = "Solar"): self.n = len(masses) self.m = masses assert(x.shape[0]==self.n and v.shape[0]==self.n) self.dim = x.shape[1] assert(v.shape[1] == self.dim) self.x = x self.v = v if(unit_system == "SI"): #SI units self.G = 6.674e-11 elif(unit_system == "solar" or unit_system == "Solar"): #mass unit M_sun, length unit AU, time unit year self.G = 4.*np.pi**2 else: print("unsupported unit system:"+unit_system) exit() self.min_dsq = 1.e99 for i in range(self.n-1): for j in range(i+1, self.n): dsq = np.sum((self.x[i, :]-self.x[j, :])**2) if(dsq*1.e-6 < self.min_dsq): self.min_dsq = dsq*1.e-6 #prepare other constants for integrator self.c41 = 1.0/(2.0 - 2.0**(1.0/3.0)) self.c40 = 1.0 - 2.0*self.c41 self.c61 = -1.177679984178871007 self.c62 = 0.235573213359358134 self.c63 = 0.784513610477557264 self.c60 = 1.0-2.0*(self.c61+self.c62+self.c63) def total_energy(self): totk = 0.0 for i in range(self.n): totk += np.sum(self.v[i, :]**2)*self.m[i] totp = 0. for i in range(self.n-1): for j in range(i+1, self.n): totp += self.m[i]*self.m[j]/np.sqrt(np.sum((self.x[i, :]-self.x[j, :])**2)) return totk/2. - totp*self.G def total_angular_momentum(self): if(self.dim == 2): #return a scalar L_z return np.sum((self.x[:, 0]*self.v[:, 1]-self.x[:, 1]*self.v[:, 0])*self.m) elif(self.dim == 3): #return a vector [L_x, L_y, L_z] Lz = np.sum((self.x[:, 0]*self.v[:, 1]-self.x[:, 1]*self.v[:, 0])*self.m) Lx = np.sum((self.x[:, 1]*self.v[:, 2]-self.x[:, 2]*self.v[:, 1])*self.m) Ly = np.sum((self.x[:, 2]*self.v[:, 0]-self.x[:, 0]*self.v[:, 2])*self.m) return np.array([Lx, Ly, Lz]) else: print("Angular momentum is not defined for dimension:"+str(self.dim)) exit() def evolve_x(self, dt): self.x += self.v * dt def evolve_v(self, dt): fac = self.G * dt for i in range(self.n-1): for j in range(i+1, self.n): f = fac*(self.x[i, :]-self.x[j, :])/max(np.sum((self.x[i, :]-self.x[j, :])**2), self.min_dsq)**1.5 self.v[i, :] -= f*self.m[j] self.v[j, :] += f*self.m[i] def symplectic_o2step(self, dt, c1, c2): self.evolve_v(c1*dt) self.evolve_x((c1+c2)*dt/2.0) def symplectic_2nd(self, dt, nsteps): self.evolve_v(dt/2.0) for j in range(nsteps-1): self.evolve_x(dt) self.evolve_v(dt) self.evolve_x(dt) self.evolve_v(dt/2.0) def symplectic_4th(self, dt, nsteps): self.evolve_x(self.c41 * dt/2.0) for j in range(nsteps-1): self.symplectic_o2step(dt, self.c41, self.c40) self.symplectic_o2step(dt, self.c40, self.c41) self.symplectic_o2step(dt, self.c41, self.c41) self.symplectic_o2step(dt, self.c41, self.c40) self.symplectic_o2step(dt, self.c40, self.c41) self.symplectic_o2step(dt, self.c41, 0.0) def symplectic_6th(self, dt, nsteps): self.evolve_x(self.c63*dt/2.0) for j in range(nsteps-1): self.symplectic_o2step(dt, self.c63, self.c62) self.symplectic_o2step(dt, self.c62, self.c61) self.symplectic_o2step(dt, self.c61, self.c60) self.symplectic_o2step(dt, self.c60, self.c61) self.symplectic_o2step(dt, self.c61, self.c62) self.symplectic_o2step(dt, self.c62, self.c63) self.symplectic_o2step(dt, self.c63, self.c63) self.symplectic_o2step(dt, self.c63, self.c62) self.symplectic_o2step(dt, self.c62, self.c61) self.symplectic_o2step(dt, self.c61, self.c60) self.symplectic_o2step(dt, self.c60, self.c61) self.symplectic_o2step(dt, self.c61, self.c62) self.symplectic_o2step(dt, self.c62, self.c63) self.symplectic_o2step(dt, self.c63, 0.0) #demo code for 2D nbody simulation #set up n-body system n = int(input(r"enter number of particles (00 and n<=20) dim = 2 masses = np.random.rand(n) x = np.random.normal(size=(n, dim)) v = np.random.normal(size=(n, dim)) for i in range(n): v[i, :] = v[i, :] - np.dot(v[i, :], x[i, :])/np.sqrt(np.sum(x[i, :]**2)) masses[0] += 2. #make one big mass xmean = np.zeros(dim) vmean = np.zeros(dim) for i in range(n): xmean += x[i, :]*masses[i] vmean += v[i, :]*masses[i] summass = np.sum(masses) xmean /= summass vmean /= summass for i in range(n): x[i, :] -= xmean v[i, :] -= vmean nb = nbody_system(masses = masses, x = x, v = v, unit_system = "Solar") print(nb.m) print(nb.x) print(nb.v) #evolution scheme t_total = 1. #evolve 1 year nsteps = 365 #output every day n_fine = 24*60 #update every minutes dt = t_total/nsteps dt_fine = dt/n_fine initial_E = nb.total_energy() initial_L = nb.total_angular_momentum() positions = np.empty((nsteps, n, dim)) print("time, energy error, angular momentum error") for i in range(nsteps): nb.symplectic_2nd(dt_fine, n_fine) positions[i, :, :] = nb.x print(i*dt, nb.total_energy()/initial_E - 1., nb.total_angular_momentum()/initial_L-1.) #below we make a movie from the saved positions fig, ax = plt.subplots(figsize=(10, 8)) ax.set_xlim(-10, 10) ax.set_ylim(-10, 10) ax.set_xlabel('X Position') ax.set_ylabel('Y Position') ax.set_title('Animated Particle Trajectories', fontsize=14) ax.grid(True, alpha=0.3) ax.set_aspect('equal', 'box') trajectories = [] particles = [] colors = plt.cm.Set2(np.linspace(0, 1, n)) for i in range(n): line, = ax.plot([], [], color=colors[i], alpha=0.6, linewidth=1.5) trajectories.append(line) particle, = ax.plot([], [], 'o', color=colors[i], markersize=10, markeredgecolor='white', markeredgewidth=2) particles.append(particle) trail_length = 5 def update(frame): for i in range(n): start_idx = max(0, frame - trail_length) x_data = positions[start_idx:frame+1, i, 0] y_data = positions[start_idx:frame+1, i, 1] trajectories[i].set_data(x_data, y_data) if frame < nsteps: particles[i].set_data([positions[frame, i, 0]], [positions[frame, i, 1]]) # ax.set_title(f'Particle Trajectories - Frame {frame}/{nsteps}', fontsize=14) return trajectories + particles ani = FuncAnimation(fig, update, frames=nsteps, interval=50, blit=True) ani.save('nbody.mp4', writer='ffmpeg', fps=20, dpi=150) #ani.save('particle_trajectories.gif', writer='pillow', fps=20)