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 check_init(filename, body_dict):
"""
Check if init file exsists, in case not make it
Args:
Filename: (string) What the init file should be named
Body_dict: (dictionary) Holding the initial conditions for each body
"""
if not filename in os.listdir("initData"):
setInitialConditions(filename, body_dict)
else:
pass
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 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()