Created
July 5, 2023 16:05
-
-
Save void4/1894f4eaf8c6345d647005b386e9cd48 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # stolen from here: https://amor.cms.hu-berlin.de/~rodrigus/Jupyter%20Notebooks/Solar_System_N-Body_Simulation.html | |
| import numpy as np | |
| import scipy as sp | |
| import time | |
| import scipy.constants as cs | |
| import matplotlib.pyplot as plt | |
| from numba import jit | |
| from numba import cuda | |
| from mpl_toolkits.mplot3d import Axes3D | |
| from astropy.time import Time | |
| from astroquery.jplhorizons import Horizons | |
| from copy import deepcopy | |
| from random import random | |
| #Vectorial acceleration function | |
| def a_t(r, m, epsilon): | |
| """ | |
| Function of matrices returning the gravitational acceleration | |
| between N-bodies within a system with positions r and masses m | |
| ------------------------------------------------- | |
| r is a N x 3 matrix of object positions | |
| m is a N x 1 vector of object masses | |
| epsilon is the softening factor to prevent numerical errors | |
| a is a N x 3 matrix of accelerations | |
| ------------------------------------------------- | |
| """ | |
| G = cs.gravitational_constant | |
| # positions r = [x,y,z] for all bodies in the N-Body System | |
| x = r[:,0:1] | |
| y = r[:,1:2] | |
| z = r[:,2:3] | |
| # matrices that store each pairwise body separation for each [x,y,z] direction: r_j - r_i | |
| dx = x.T - x | |
| dy = y.T - y | |
| dz = z.T - z | |
| #matrix 1/r^3 for the absolute value of all pairwise body separations together and | |
| #resulting acceleration components in each [x,y,z] direction | |
| inv_r3 = (dx**2 + dy**2 + dz**2 + epsilon**2)**(-1.5) | |
| ax = G * (dx * inv_r3) @ m | |
| ay = G * (dy * inv_r3) @ m | |
| az = G * (dz * inv_r3) @ m | |
| # pack together the three acceleration components | |
| a = np.hstack((ax,ay,az)) | |
| return a | |
| def omega(N, t_max, dt): | |
| """ | |
| Dummy acceleration function to give an estimate of the total memory | |
| consumption for a simulation with N bodies, total simulation time | |
| t_max and timestep dt | |
| ------------------------------------------------- | |
| N is the amount of bodies in the simulation | |
| t_max is the total simulation time in years | |
| dt is the timestep for each integration in days | |
| ------------------------------------------------- | |
| """ | |
| epsilon = 1 | |
| G = 1 | |
| r = np.ones((N,3)) | |
| v = np.ones((N,3)) | |
| """although the velocities aren't actually taken into account for computing the | |
| acceleration, they will be stored as Nx3 matrices for exactly as many iterations | |
| in the integration loop itself and take up memory accordingly""" | |
| m = np.ones((N,1)) | |
| # positions r = [x,y,z] for all bodies in the N-Body System | |
| x = r[:,0:1] | |
| y = r[:,1:2] | |
| z = r[:,2:3] | |
| # matrices that store each pairwise body separation for each [x,y,z] direction: r_j - r_i | |
| dx = x.T - x | |
| dy = y.T - y | |
| dz = z.T - z | |
| #matrix 1/r^3 for the absolute value of all pairwise body separations together and | |
| #resulting acceleration components in each [x,y,z] direction | |
| inv_r3 = (dx**2 + dy**2 + dz**2 + epsilon**2)**(-1.5) | |
| ax = G * (dx * inv_r3) @ m | |
| ay = G * (dy * inv_r3) @ m | |
| az = G * (dz * inv_r3) @ m | |
| # pack together the three acceleration components | |
| a = np.hstack((ax,ay,az)) | |
| # sum the memory usage of each matrix storing the positions, distances and accelerations | |
| memory_usage_per_iteration = r.nbytes + v.nbytes + x.nbytes + y.nbytes + z.nbytes + dx.nbytes + dy.nbytes + dz.nbytes + ax.nbytes + ay.nbytes + az.nbytes + inv_r3.nbytes + a.nbytes | |
| total_memory_usage = memory_usage_per_iteration * (t_max)/(dt*1e6) | |
| return total_memory_usage #in megabytes | |
| #Example for the estimated memory usage from a simulation with N = 13 bodies, t_max = 250 years and dt = 1 day | |
| year = 365*24*60*60 | |
| day = 24*60*60 | |
| omega(13,250*year,1*day) | |
| #Conversion Units | |
| AU = 149597870700 | |
| D = 24*60*60 | |
| #Simulation starting date in Y/M/D | |
| t_0 = "2018-10-26" | |
| #Get Starting Parameters for Sun-Pluto from Nasa Horizons | |
| jpl_ids = [0,1,2,399,301,4,5,6,7,8,9] | |
| NUMNATURAL = len(jpl_ids) | |
| r_list = [] | |
| v_list = [] | |
| m_list = [[1.989e30],[3.285e23],[4.867e24],[5.972e24],[7.34767e22],[6.39e23],[1.8989e27],[5.683e26],[8.681e25],[1.024e26],[1.309e22]] #Object masses for Sun-Pluto | |
| plot_colors = ['green','brown','orange','blue','gray','red','red','orange','cyan','blue','brown'] | |
| plot_labels = ['Barycenter Drift','Mercury Orbit','Venus Orbit','Earth Orbit','Moon Orbit','Mars Orbit','Jupiter Orbit','Saturn Orbit','Uranus Orbit','Neptune Orbit','Pluto Orbit'] | |
| for i, jpl_id in enumerate(jpl_ids): | |
| obj = Horizons(id=jpl_id, location="@sun", epochs=Time(t_0).jd, id_type='id').vectors() | |
| r_obj = [obj['x'][0], obj['y'][0], obj['z'][0]] | |
| v_obj = [obj['vx'][0], obj['vy'][0], obj['vz'][0]] | |
| r_list.append(r_obj) | |
| v_list.append(v_obj) | |
| print(plot_labels[i], np.linalg.norm(v_obj)*AU/D, "m/s")#np.linalg.norm(v_obj), | |
| #Get Starting Parameters for any extra object to add into the simulation with input for id/mass/plot_color/plot_label | |
| MAXSPEEDDIFF = 5000/AU*D | |
| print(MAXSPEEDDIFF) | |
| def sample_spherical_many(npoints=1, ndim=3): | |
| vec = np.random.randn(ndim, npoints) | |
| vec /= np.linalg.norm(vec, axis=0) | |
| return vec | |
| def sample_spherical(): | |
| return list(zip(*sample_spherical_many()))[0] | |
| def add_simulation_object(Id_obj,m_obj, plot_color, plot_label): | |
| ori = Horizons(id=Id_obj, location="@sun", epochs=Time(t_0).jd, id_type='id').vectors() | |
| print(ori) | |
| for i in range(10): | |
| obj = deepcopy(ori) | |
| r_obj = [obj['x'][0], obj['y'][0], obj['z'][0]] | |
| v_obj = [obj['vx'][0], obj['vy'][0], obj['vz'][0]] | |
| #print(v_obj) | |
| speed_before = np.linalg.norm(v_obj)*AU/D | |
| random_vector = sample_spherical() | |
| deltas = [] | |
| for j in range(3): | |
| delta = MAXSPEEDDIFF*random_vector[j]*np.random.uniform(0.5, 1) | |
| deltas.append(delta) | |
| v_obj[j] += delta | |
| speed_after = np.linalg.norm(v_obj)*AU/D | |
| print(f"Added: {plot_label+str(i)}\tSpeed: {speed_after:.2f} m/s\tDelta: {speed_after-speed_before:.2f} m/s\tApplied delta: {np.linalg.norm(deltas)*AU/D}") | |
| r_list.append(r_obj) | |
| v_list.append(v_obj) | |
| m_list.append([m_obj]) | |
| plot_colors.append(plot_color) | |
| plot_labels.append(plot_label+str(i)) | |
| #Sample ID for Neowise | |
| id_neowise = '90004475' | |
| m_neowise = 5e13 | |
| #Sample ID for Space X's Starman | |
| id_starman = 'SpaceX Roadster' | |
| m_starman = 1300 | |
| #Add NEOWISE and Starman to the simulation | |
| #add_simulation_object('90004479', 5e13, 'black','Comet Neowise') | |
| #add_simulation_object('SpaceX Roadster',1300, 'pink','Starman') | |
| #add_simulation_object('90000033',1e14, 'gray','Comet Halley') | |
| add_simulation_object("-48", 11110, "black", "Hubble") | |
| #Convert object staring value lists to numpy | |
| r_i = np.array(r_list)*AU | |
| v_i = np.array(v_list)*AU/D | |
| m_i = np.array(m_list) | |
| #pack together as list for the simulation function | |
| horizons_data = [r_i,v_i,m_i] | |
| earth_radius = 6371*1000#m | |
| def simulate_solar_system(N,dN,starting_values): # | |
| t0_sim_start = time.time() | |
| t = 0 | |
| t_max = 365*24*60*60*N #N year simulation time | |
| dt = 60*60*24*dN #dN day time step | |
| epsilon_s = 0.01 #softening default value | |
| r_i = starting_values[0] | |
| v_i = starting_values[1] | |
| m_i = starting_values[2] | |
| a_i = a_t(r_i, m_i, epsilon_s) | |
| ram_usage_estimate = omega(len(r_i), t_max, dt) #returns the estimated ram usage for the simulation | |
| # Simulation Main Loop using a Leapfrog Kick-Drift-Kick Algorithm | |
| k = int(t_max/dt) | |
| r_save = np.zeros((r_i.shape[0],3,k+1)) | |
| r_save[:,:,0] = r_i | |
| for i in range(k): | |
| if i % 1000 == 0: | |
| print(f"{i}/{k}") | |
| earth_position = r_i[3] | |
| for idx, sat_position in enumerate(r_i[NUMNATURAL:]): | |
| dist = np.linalg.norm(earth_position-sat_position) | |
| print(idx, dist)#"earth:", earth_position, "sat:", sat_position, | |
| if dist < earth_radius: | |
| print("crashed:", idx) | |
| # (1/2) kick | |
| v_i += a_i * dt/2.0 | |
| # drift | |
| r_i += v_i * dt | |
| # update accelerations | |
| a_i = a_t(r_i, m_i, epsilon_s) | |
| # (2/2) kick | |
| v_i += a_i * dt/2.0 | |
| # update time | |
| t += dt | |
| #update list | |
| r_save[:,:,i+1] = r_i | |
| sim_time = time.time()-t0_sim_start | |
| print('The required computation time for the N-Body Simulation was', round(sim_time,3), 'seconds. The estimated memory usage was', round(ram_usage_estimate,3), 'megabytes of RAM.') | |
| return r_save | |
| STEP = 0.01 | |
| YEARS = 1 | |
| print(f"DURATION: {YEARS} years\tSTEP: {STEP*60*60*24} seconds") | |
| #Run simulation for 250 years at a 1 day time-step | |
| r_save = simulate_solar_system(YEARS, STEP, horizons_data) | |
| #Plot for the outer planets | |
| fig2 = plt.figure() | |
| ax = plt.axes(projection='3d') | |
| ax.set_xlim3d(-50e11,50e11) | |
| ax.set_ylim3d(-50e11,50e11) | |
| ax.set_zlim3d(-50e11,50e11) | |
| #Input sim. data | |
| for i in range(0,10): #Plots the outer planets | |
| ax.plot3D(r_save[i,0,:],r_save[i,1,:],r_save[i,2,:], plot_colors[i],label=plot_labels[i]) | |
| if len(r_i)>=10: | |
| for i in range(10,len(r_i)): #Plots any additional objects | |
| ax.plot3D(r_save[i,0,:],r_save[i,1,:],r_save[i,2,:], plot_colors[i],label=plot_labels[i]) | |
| ax.legend(loc = 'upper right', prop={'size': 6.5}) | |
| #plt.savefig('Plots/3D_Plots_Outer_Planets.pdf',dpi=1200) | |
| plt.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment