def escapeVelocity(n=20, T=100):
"""
Code to plot different inital velocities in the v_y direction
Args:
n: (int) Number of orbits/escapes
T: (float/int) Total time to run over
"""
v = np.linspace(2 * np.pi, 3 * np.pi, n, endpoint=True)
T = 100
dt = -4
N = np.log10(T) - dt
orbits = []
for i in range(n):
system_dict = {
"Sun": [0, 0, 0, 0, 0, 0],
"Earth": [1, 0, 0, 0, v[i], 0]
} # Store the different velocities
setInitialConditions(f"escape_init_{i}.dat", system_dict)
simulate(N=N,
dt=dt,
Nwrite=1000,
sys=f"escape_init_{i}.dat",
out=f"escape_{i}.dat",
fixSun=True,
quiet=True)
system = read_data_file(f"escape_{i}.dat")
orbits.append(system["Earth"].r)
NoOfColors = n
colors = list(Color("cyan").range_to(Color("orange"), NoOfColors))
colors = [color.get_rgb() for color in colors]
vticks = [round(v_i, 2) for v_i in v]
fig, ax = plt.subplots(1, 1)
ax.set_facecolor('black')
ax.set_xlabel("x [AU]", fontsize=15)
ax.set_ylabel("y [AU]", fontsize=15)
ax.set(xlim=(-50, 2), ylim=(-26, 26))
for i, r in enumerate(orbits):
ax.plot(r[0], r[1], color=colors[i], alpha=0.7)
cmap = mpl.colors.ListedColormap(colors)
norm = mpl.colors.Normalize(vmin=v[0], vmax=v[-1])
cbar = ax.figure.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap),
ax=ax,
orientation='vertical',
ticks=vticks)
cbar.ax.set_ylabel(r' $v_y$', rotation=0, fontsize=17)
cbar.ax.tick_params(labelsize=13)
ax.tick_params(axis='both', which='major', labelsize=12)
ax.scatter(0, 0, c="yellow", s=2)
ax.text(0, 15.5, ' $v_{esc}$', color="black", fontsize=17)
ax.text(-49, 6.5, "$v_{esc}=8.89$ [AU/yr]", color="white", fontsize=14)
ax.arrow(-43, 8.5, 4, 4, ec="white", fc="white", head_width=1)
plt.show()
def forward_simulation(dt, N, system_dict, GR=True):
# forwards the simulation according to dt and N
# only initial and final datapoints are stored
setInitialConditions("precession_init.dat", system_dict)
simulate(N,
dt,
Nwrite=2,
GR=GR,
sys="precession_init.dat",
out="precession.dat",
quiet=True)
def run_simulation(dt, T, v):
# runs a single escape simulation and returns bool whether Earth escaped
N = np.log10(T) - dt # log10 of N
system_dict = {"Sun": [0, 0, 0, 0, 0, 0], "Earth": [1, 0, 0, 0, v, 0]}
setInitialConditions("escape_init.dat", system_dict)
simulate(N=N,
dt=dt,
Nwrite=2,
sys="escape_init.dat",
out="escape.dat",
fixSun=True,
quiet=True)
return check_escape()
def findPrecession(dt, GR=False, newsim=False):
# Runs 100 one-year simulations adding up to 100 years forward, and finds the final perihelion angle.
start = time.time() # timer
# setting initial conditions for Mercury-Sun system with perihelion along x-axis
system_dict = {
"Sun": [0, 0, 0, 0, 0, 0],
"Mercury": [0.3075, 0, 0, 0, 12.44, 0]
}
N = -dt # log10 of N, results in one-year simulation
if newsim: # simulates from scratch
for i in range(100):
# simulates 1 years at a time
forward_simulation(dt, N, system_dict, GR=GR)
print(dt, i, GR) # just to keep track of simulations
system_dict = read_final_state()
# finally simulates just over one orbit, saving all the points to find the perihelion angle
setInitialConditions("precession_orbit_init.dat", system_dict)
N = np.log10(
0.3) - dt # simulates 0.3 of a year, which is just over one orbit
print("final")
GR_string = {True: "GR", False: "CLASSIC"}[GR]
simulate(N,
dt,
Nwrite=int(1e6),
GR=GR,
sys="precession_orbit_init.dat",
out=f"prec/final_{dt}_{GR_string}.dat",
quiet=True)
# Gets the precession of the century simulated and prints
precession_per_century = getPerihelionAngle(dt, GR=GR)
print(
f"The precession over one century (GR: {GR}) was {precession_per_century} arcseconds."
)
print(f"Took {time.time()-start:.1f} seconds to run")
return precession_per_century
def fullSystem():
N = 7
dt = np.log10(10**(-5) * 25)
Nwrite = int(1e4)
bodynames = [
"Sun", "Mercury", "Venus", "Earth", "Mars", "Jupiter", "Saturn",
"Uranus", "Neptune", "Pluto"
]
bodycolors = [
"yellow", "#86511D", "#F7C3F4", "#0EEB58", "red", "orange", "#FCDB0A",
"aqua", "blue", "grey"
]
gic("fullsystem.dat", bodies=bodynames, fixedCoM=True, date="2018-03-14")
simulate(N=N,
dt=dt,
Nwrite=Nwrite,
sys="fullsystem.dat",
out="fullsystem.out")
system = read_data_file("fullsystem.out")
with plt.style.context("seaborn-darkgrid"):
fig = plt.figure(figsize=(13, 8), dpi=130)
ax = fig.gca(projection='3d')
for i, bodyname in enumerate(bodynames):
body = system[bodyname]
r = body.r
v = body.v
h = np.cross(r, v, axis=-2)
inc = np.mean(
np.rad2deg(np.arccos(h[-1, :] / np.linalg.norm(h, axis=-2))))
if i != 0:
print(f"Inclination {bodyname}: {round(inc,3)} degr")
x, y, z = body.r[0, :], body.r[1, :], body.r[2, :]
ax.plot(x, y, z, lw=1, color=bodycolors[i], alpha=0.65, zorder=-1)
ax.scatter(x[-1],
y[-1],
z[-1],
color=bodycolors[i],
zorder=10,
s=18)
ax.text(x[-1],
y[-1],
z[-1],
bodyname,
color=Color(bodycolors[i], luminance=0.95).rgb,
fontsize=12,
zorder=1000)
ax.set_axis_off()
ax.set_facecolor("black")
ax.set_title("All bodies over 25 years", color="white")
lim = 20 # plot dimensions (AU)
zscale = 20 # scaling of z-axis to show inclination
ax.set_xlim(-lim, lim)
ax.set_ylim(-lim, lim)
ax.set_zlim(-lim / zscale, lim / zscale)
plt.show()
def radialDistance(dt=-4, T=20):
"""
Plot the radial deviation from Earths orbit without Jupiter
Args:
dt: (float) step length in years
T: (float/int) total time to run the simulation over
"""
scale_j_mass = [1, 10, 100,
1000] # Hard code in the different scaled masses
N = np.log10(T) - dt
N_write = 10000
# First run without Jupiter present
initFilenameNoJ = "No_jupiter.dat"
outFilenameNoJ = f"No_jupiter_{T}_{-dt}.dat"
if not initFilenameNoJ in os.listdir("initData"):
getInitialCondition(initFilenameNoJ, ["Sun", "Earth"],
date="2018-03-14")
if not outFilenameNoJ in os.listdir("data"):
simulate(N=N,
dt=dt,
Nwrite=N_write,
sys=initFilenameNoJ,
out=outFilenameNoJ,
fixSun=True,
quiet=True)
system = read_data_file(outFilenameNoJ)
rE_no_j = np.linalg.norm(
system["Earth"].r,
axis=0) # Earths distance from the sun without Jupiter
initFilenames = [f"SEJ_{scale}.dat" for scale in scale_j_mass]
outFilenames = [f"SEJ_{T}_{scale}_{N_write}.dat" for scale in scale_j_mass]
rE = [
] # To store the distances with Jupiter present at different scaled masses
for i in range(len(scale_j_mass)):
if not initFilenames[i] in os.listdir("initData"):
getInitialCondition(initFilenames[i], ["Sun", "Earth", "Jupiter"],
scaled_mass={"Jupiter": scale_j_mass[i]})
if not outFilenames[i] in os.listdir("data"):
simulate(N=N,
dt=dt,
Nwrite=N_write,
sys=initFilenames[i],
out=outFilenames[i],
fixSun=True,
quiet=True)
system = read_data_file(outFilenames[i])
R = np.linalg.norm(system["Earth"].r, axis=0)
rE.append(R)
t = np.linspace(0, T, N_write, endpoint=True)
with sns.axes_style("darkgrid"):
fig, ax = plt.subplots()
for R, lab in zip(rE, scale_j_mass):
diff = rE_no_j - R
if lab == 1:
ax.plot(t, diff, label="$M_J$")
else:
ax.plot(t, diff, label=f"{lab}$M_J$")
diff = np.abs(diff)
avg_diff = np.mean(diff)
std_diff = np.mean(diff)
print(
f"Jupiter mass = {lab:5}Mj, Avrage deviation = {avg_diff:.5E}, std = {std_diff:.5E}"
)
ax.tick_params(axis='both', which='major', labelsize=13)
ax.set_xlabel("Time [yr]", fontsize=14)
ax.set_ylabel("Deviation from orbit without Jupiter [AU]", fontsize=14)
ax.legend(loc='upper center',
bbox_to_anchor=(0.5, 1.16),
ncol=4,
fancybox=True,
shadow=True,
fontsize=14)
plt.show()
def sunEarthJupiter(dt=-4, T=15, jupiter_scale=1000):
"""
Plot the Sun-Earth-Jupiter system
Args:
dt: (float) step length in AU/yr
T: (float/int) how long to run the simulation in years
jupiter_scale: (int/float) In case you want to scale Jupiters mass up/down
"""
N = np.log10(T) - dt
N_write = 10000
initFilename = f"SEJ_{jupiter_scale}.dat"
outFilename = f"SEJ_{T}_{jupiter_scale}_{N_write}.dat"
bodies = ["Sun", "Earth", "Jupiter"]
scaled_mass = {"Jupiter": jupiter_scale}
if not initFilename in os.listdir("initData"):
getInitialCondition(initFilename,
bodies,
fixedCoM=False,
scaled_mass=scaled_mass,
date="2018-03-14")
if not outFilename in os.listdir("data"):
simulate(N=N,
dt=dt,
Nwrite=1000,
sys=initFilename,
out=outFilename,
fixSun=True,
quiet=True)
system = read_data_file(outFilename)
rS = system["Sun"].r
rE = system["Earth"].r - rS
rJ = system["Jupiter"].r - rS
l = np.max(rJ) + np.max(
rJ) / 10 # Make plot just a little bigger than Jupiters orbit
bodynames = ["Earth", "Jupiter"]
bodycolors = ["#0EEB58", "orange"]
fig = plt.figure()
fig.set_facecolor('black')
ax = fig.add_subplot(111, projection='3d')
ax.set_facecolor('black')
ax.grid(False)
ax.w_xaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
ax.w_yaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
ax.w_zaxis.set_pane_color((0.0, 0.0, 0.0, 0.0))
ax.set(xlim=(-l, l), ylim=(-l, l), zlim=(-l, l))
ax.scatter(0, 0, c="yellow", label="Sun")
ax.plot(rE[0, 1:], rE[1, 1:], rE[2, 1:], c="#0EEB58", label="Earth")
ax.plot(rJ[0], rJ[1], rJ[2], c="orange", label="Jupiter")
ax.legend(loc='upper center',
bbox_to_anchor=(0.5, 1.1),
ncol=3,
fancybox=True,
shadow=True,
fontsize=15)
plt.show()
def circularOrbit(dt=0.0001, T_end=1, method="verlet", N_write=1000):
"""
Makes plot of stable Sun/Earth orbit
Sun at origin no init vel, earth x = 1 AU vx = 2pi * AU/yr
Args:
dt: (float) time step in years
T_end: (int/float) end time in years
method: (string) "euler" or "verlet"
N_write: (int) number of points to write to file
"""
N = int(T_end/dt)
if N_write == None:
N_write = N
body_dict = {"Sun": [0,0,0,0,0,0],
"Earth": [1,0,0,0,2*np.pi,0]} # Initial conditions
initFilename = "SunEarthStable_init.dat"
outFilename = "SunEarthStable_" + "_".join([str(method), str(T_end), str(N), str(N_write)]) + ".dat" # Make filenames
check_init(initFilename, body_dict)
setInitialConditions(initFilename, body_dict)
exists = has_data(outFilename)
N = np.log10(N)
dt = np.log10(dt)
if not exists: # If the datafiles are missing, make them
simulate(N=N, dt = dt, method=method, Nwrite=N_write, sys=initFilename, out=outFilename, fixSun=True, quiet=True)
system = read_data_file(outFilename) # Reads the data
r = system["Earth"].r
Ek, Ep = energy(system["Earth"]) # Calculate energy
Ek_std, Ep_std = np.std(Ek), np.std(Ep)
Ek_mean, Ep_mean = np.mean(Ek), np.mean(Ep)
if system["method"] == 0:
method = "euler"
if system["method"] == 1:
method = "verlet"
print(f"Method = {method}, dt = {dt:E}, N={N:E}, T_end = {T_end}")
print(f"Ek initial = {Ek[0]:.4f}, Ep initial = {Ep[0]:.4f}")
print(f"Ek (mean) = {Ek_mean:.4f}, std = {Ek_std:.4E}, relative = {(Ek_std/Ek_mean):.4E}")
print(f"Ep (mean) = {Ep_mean:.4f}, std = {Ep_std:.4E}, relative = {(-Ep_std/Ep_mean):.4E}")
T = np.linspace(0, T_end, N_write, endpoint=True)
with sns.axes_style("darkgrid"):
fig, ax = plt.subplots()
l = 1.5
ax.set(xlim=(-l,l), ylim=(-l,l))
ax.tick_params(axis='both', which='major', labelsize=13)
ax.axis("equal")
ax.set_xlabel("x [AU]", fontsize=13)
ax.set_ylabel("y [AU]", fontsize=13)
ax.scatter(0,0, c="r", label="Sun")
ax.plot(r[0], r[1], c="b", label="Earth")
ax.legend(fontsize=13)
plt.show()
def error(dt_start=3, dt_end=7, n_tests=20):
"""
Making a plot of the relative error in the total energy of the two algorithms
Args:
dt_start: (int) -log10 of what dt to start at
dt_end: (int) -log10 of what dt to end at
n_tests: (int) number of tests to preform between dt_start and dt_end
"""
dt = np.linspace(dt_start, dt_end, n_tests, endpoint=True)*-1
N = -dt
methods = ["euler", "verlet"]
initFilename = "SunEarthStable_init.dat"
outFilenames = []
for method in methods:
for i in range(n_tests):
outFilenames.append(f"SunEarthStable_{method}_{dt_start}_{dt_end}_{i+1}.dat")
body_dict = {"Sun": [0,0,0,0,0,0],
"Earth": [1,0,0,0,2*np.pi,0]}
setInitialConditions(initFilename, body_dict)
check_init(initFilename, body_dict)
exists = has_data(outFilenames)
i = 0
tot = 2*n_tests
if not exists: # If files are missing, create them
for method in methods:
for N_, dt_ in zip(N,dt):
simulate(N=N_, dt = dt_, method=method, Nwrite=2, sys=initFilename, out=outFilenames[i], fixSun=True, quiet=True)
i += 1
print(f"{method}, log10(dt)={dt_:.2f}, log10(N)={N_:.2f}, {(i*100/tot):.1f}%") # Again, migth take a while
N = 10**N
dt = 10**dt
systems = []
eulerError = []
verletError = []
for outfile in outFilenames: # Sort through files and place the error in euler and verlet
system = read_data_file(outfile)
Ek, Ep = energy(system["Earth"])
E = Ek+Ep
error = np.abs((E[0]-E[-1])/E[0])
if system["method"] == 0:
eulerError.append(error)
if system["method"] == 1:
verletError.append(error)
with sns.axes_style("darkgrid"):
fix, ax = plt.subplots()
ax.invert_xaxis()
ax.tick_params(axis='both', which='major', labelsize=13)
ax.set_xlabel("$h$ [yr]", fontsize=14)
ax.set_ylabel("Relative error of $E_{tot}$", fontsize=14)
ax.set(xscale="log", yscale="log")
ax.scatter(dt, eulerError, label="$E_{tot}$ Euler")
ax.scatter(dt, verletError, label="$E_{tot}$ Verlet")
slope, const, r_value, p_value, std_err = stats.linregress(np.log10(dt), np.log10(eulerError))
ax.plot(dt, 10**const * dt**slope, c="k" ,linestyle="dashed",\
label=f"s = {slope:.3f}$\pm${std_err:.3f} \n$R^2$ = {(r_value**2):.5f}")
ax.legend(fontsize=13, loc=6)
plt.show()
def benchmark(N_start=2, N_end=7, n_tests=50):
"""
Preform a timing of the two algorithms and make a log-log plot
Args:
N_start: (int) log10 of the number of integration steps to start at
N_end: (int) log10 of the number of integration steps to end at
n_tests: (int) number of test between N_start and N_end
"""
N = np.log10(np.logspace(N_start, N_end, n_tests, endpoint=True, dtype=int))
methods = ["euler", "verlet"]
initFilename = "SunEarthStable_init.dat"
outFilenames = []
for method in methods:
for i in range(n_tests):
outFilenames.append(f"SunEarthStable_{method}_{N_start}_{N_end}_{i+1}.dat")
body_dict = {"Sun": [0,0,0,0,0,0],
"Earth": [1,0,0,0,2*np.pi,0]}
check_init(initFilename, body_dict)
exists = has_data(outFilenames)
i = 0
tot = 2*n_tests
if not exists: # If datafiles are missing, make them
for method in methods:
for n_ in N:
dt_ = -n_
simulate(N=n_, dt = dt_, method=method, Nwrite=2, sys=initFilename, out=outFilenames[i], fixSun=True, quiet=True)
i += 1
print(f"Done {method}, dt = {n_}, {(i*100/tot):.2f}%") # It takes a while... good to know where your at
eulerTime = []
verletTime = []
for outfile in outFilenames: # Sort through files and place the time in euler and verlet
system = read_data_file(outfile)
if system["method"] == 0:
eulerTime.append(system["time"])
if system["method"] == 1:
verletTime.append(system["time"])
eulerTime = np.array(eulerTime) # For easier handling
verletTime = np.array(verletTime)
with sns.axes_style("darkgrid"):
fig, ax = plt.subplots()
ax.set(xscale="log", yscale="log")
ax.tick_params(axis='both', which='major', labelsize=13)
ax.set_xlabel(xlabel="N", fontsize=15)
ax.set_ylabel(ylabel="Time taken [s]", fontsize=15)
ax.scatter(N, eulerTime, label="Euler")
ax.scatter(N, verletTime, label="Verlet")
# Only linfit if standard params (as some points have to be excluded)
if all([val is benchmark.__defaults__[i] for i, val in enumerate([N_start,N_end,n_tests])]):
sE, eE = 25, len(N) # What points to fit Euler
sV, eV = 25, len(N) # What point to fit Verlet
slopeE, constE, r_valueE, p_valueE, std_errE = stats.linregress(np.log10(N[sE:eE]), np.log10(eulerTime[sE:eE]))
# Linfit for Verlet
slopeV, constV, r_valueV, p_valueV, std_errV = \
stats.linregress(np.log10(N[sV:eV]), np.log10(verletTime[sV:eV]))
# Plot Euler and Verlet linfit
ax.plot(N[sE:eE], 10**constE * N[sE:eE]**slopeE, c="red",\
label=f"Slope Euler = {slopeE:.3f}$\pm${std_errE:.3f} \n $R^2$ = {(r_valueE**2):.4f}", markersize=3)
ax.plot(N[sV:eV], 10**constV * N[sV:eV]**slopeV, c="k" ,\
label=f"Slope Verlet = {slopeV:.3f}$\pm${std_errV:.3f} \n $R^2$ = {(r_valueV**2):.4f}", markersize=3)
ax.legend(fontsize=13)
plt.show()
def ellipticalOrbits(dt=0.0001, T_end=10, n_v = 4, method="verlet", N_write=10000):
"""
Makes plot of different elliptical orbits
Args:
dt: (float) Step length in yr
T_end: (int/float) Total time to run the simulation
n_v: (int) Number of orbits in the range v_y = [pi, 5pi/2]
method: (string) Verlet or Euler
N_write: (int) How many points to write
"""
N = int(T_end/dt)
if N_write == None:
N_write = N
VY = np.linspace(np.pi, 5*np.pi/2, n_v, endpoint=True) # Array holding different initial velocities in the y-direction
filenames = [f"SunEarthEllip_{n_v}_{i}_{T_end}.dat" for i in range(n_v)] # Storing filenames
for i in range(n_v):
body_dict = {"Sun": [0,0,0,0,0,0],
"Earth": [1,0,0,0,VY[i],0]}
check_init(filenames[i], body_dict)
exists = has_data(filenames)
dt = np.log10(dt)
N = np.log10(N)
if not exists:
for i in range(n_v):
simulate(N=N, dt = dt, method=method, Nwrite=N_write, sys=filenames[i], out=filenames[i], fixSun=True, quiet=True)
labs = ["$\pi$", "$3\pi/2$", "$2\pi$", "$5\pi/2$"] # Only true if standard parameters are used
with sns.axes_style("darkgrid"):
fig, ax = plt.subplots()
ax.set_xlabel("x [AU]", fontsize=14)
ax.set_ylabel("y [AU]", fontsize=14)
ax.axis("equal")
for i in range(n_v):
system = read_data_file(filenames[i])
rE = system["Earth"].r
ax.plot(rE[0], rE[1], label=f"$v_y=${labs[i]} AU/yr")
ax.scatter(0,0, c="r")
ax.tick_params(axis='both', which='major', labelsize=13)
ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.22),
ncol=2, fancybox=True, shadow=True, fontsize=13)
plt.show()
with sns.axes_style("darkgrid"):
fig, ax = plt.subplots()
T = np.linspace(0, T_end, N_write, endpoint=True)
ax.set_xlabel("Time [yr]", fontsize=14)
ax.set_ylabel("Relative error of |$\ell$|", fontsize=14)
for i in range(n_v):
system = read_data_file(filenames[i])
L = angular_momentum(system["Earth"])
L = np.abs((L[0]-L)/L[0])
ax.plot(T, L, label=f"$v_y=${labs[i]}")
ax.tick_params(axis='both', which='major', labelsize=13)
ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.18),
ncol=2, fancybox=True, shadow=True, fontsize=13)
plt.show()