Select Git revision
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
sheet1.py 7.72 KiB
import math
import numpy as np
from matplotlib import pyplot as plt
from scipy.integrate import RK45
def verlet(F, x0, p0, m, dt, N):
if isinstance(p0, float):
x = np.zeros(N)
p = np.zeros(N)
else:
x = np.zeros((N, *x0.shape))
p = np.zeros((N, *p0.shape))
x[0] = x0
p[0] = p0
for i in range(1, N):
p[i] = p[i-1] + 1/2 * F(x[i-1]) * dt
x[i] = x[i-1] + 1/m * p[i] * dt
p[i] = p[i] + 1/2 * F(x[i]) * dt
return x, p
def verlet_inv(F, x0, p0, m, dt, N):
if isinstance(p0, float):
x = np.zeros(N)
p = np.zeros(N)
else:
x = np.zeros((N, *x0.shape))
p = np.zeros((N, *p0.shape))
x[0] = x0
p[0] = p0
for i in range(1, N):
p[i] = p[i-1] - 1/2 * F(x[i-1]) * dt
x[i] = x[i-1] - 1/m * p[i] * dt
p[i] = p[i] - 1/2 * F(x[i]) * dt
return x, p
def heun(F, r0, p0, m, dt, N):
p = np.zeros(N)
p[0] = p0
r = np.zeros(N)
r[0] = r0
for i in range(1, N):
rp = r[i-1] + 1/2 * p[i-1] * dt / m # predictor step
r[i] = rp + 1/m * (p[i-1] / 2 + F(r[i-1]) * dt / 2) * dt
p[i] = p[i-1] + F(rp) * dt
return r, p
def heun_inv(F, r0, p0, m, dt, N):
p = np.zeros(N)
p[0] = p0
r = np.zeros(N)
r[0] = r0
for i in range(1, N):
rp = r[i-1] - 1/2 * p[i-1] * dt / m # predictor step
r[i] = rp - 1/m * (p[i-1] / 2 - F(r[i-1]) * dt / 2) * dt
p[i] = p[i-1] - F(rp) * dt
return r, p
# Problem 1.1
# a)
dt = 1e-3 # s
t = np.arange(0, 10, dt) # s
x0 = 0 # m
p0 = 1e-3 # kg m / s
m = 1e-3 # kg
k = 0.1 # N / m
f = lambda x: -k * x
#r1, p1 = verlet_inv(f, x0, p0, m, dt, len(t))
r2, p2 = heun_inv(f, x0, p0, m, dt, len(t))
#plt.plot(r1, p1, label="Verlet")
plt.plot(r2, p2, label="Heun")
plt.legend()
plt.xlabel(r"$x(t)$")
plt.ylabel(r"$p(t)$")
plt.savefig("problem1.1c.png")
plt.show()
# b)
E_kin1 = p1**2 / (2 * m)
E_pot1 = 1/2 * k * r1**2
E1 = E_pot1 + E_kin1
dE1 = E1 - E1[0]
E_kin2 = p2**2 / (2 * m)
E_pot2 = 1/2 * k * r2**2
E2 = E_pot2 + E_kin2
dE2 = E2 - E2[0]
plt.plot(t, E1, label="Verlet")
plt.plot(t, E2, label="Heun")
plt.legend()
plt.title(r"$E$")
plt.xlabel(r"$t$")
plt.ylabel(r"$E$")
plt.show()
plt.plot(t, dE1, label="Verlet")
plt.plot(t, dE2, label="Heun")
plt.xlabel(r"$t$")
plt.ylabel(r"$E$")
plt.title(r"$\Delta E$")
plt.legend()
plt.show()
plt.plot(t, dE1 / dt, label="Verlet")
plt.plot(t, dE2 / dt, label="Heun")
plt.xlabel(r"$t$")
plt.ylabel(r"$E$")
plt.title(r"$\Delta E / \Delta t$")
plt.legend()
plt.show()
plt.plot(t, dE1 / dt**2, label="Verlet")
plt.plot(t, dE2 / dt**2, label="Heun")
plt.xlabel(r"$t$")
plt.ylabel(r"$E$")
plt.title(r"$\Delta E / (\Delta t)^2$")
plt.legend()
plt.show()
# Problem 1.2
# a)
R_c = 1 # AU
R_L = 2.57e-3 * R_c
m_s = 1 # solar mass
m_c = 3.00e-6 * m_s
m_L = 3.69e-8 * m_s
T_c = 1 # yr
T_L = 27.3 / 365.3 * T_c
G = 4 * np.pi**2 # AU^3 / m / yr^2
m0 = m_s
m1 = m_c
m2 = m_L
x0 = np.array([
R_c,
0,
R_c + R_L,
0,
0,
m_c * R_c * 2 * np.pi / T_c,
0,
m_L / m_c * m_c * R_c * 2 * np.pi / T_c + m_L * R_L * 2 * np.pi / T_L,
])
def f(t, x):
r1 = x[:2]
r2 = x[2:4]
p1 = x[4:6]
p2 = x[6:]
return np.array([
p1 / m1,
p2 / m2,
-G * m1 * (m0 * r1 / np.linalg.norm(r1)**3 + m2 * (r1 - r2) / np.linalg.norm(r1 - r2)**3),
-G * m2 * (m0 * r2 / np.linalg.norm(r2)**3 + m1 * (r2 - r1) / np.linalg.norm(r2 - r1)**3),
]).flatten()
n = 5
T = 0.3 # yr
dt = 4**(-n) # yr
N = int(T / dt)
t = np.arange(N) * dt
rk = RK45(f, 0, x0, T, first_step=dt, max_step=dt, atol=1e3, rtol=1e3)
x = np.zeros((N, len(x0)))
x[0] = x0
for i in range(1, N):
rk.step()
x[i] = rk.y
r1 = x[:,0:2]
r2 = x[:,2:4]
p1 = x[:,4:6]
p2 = x[:,6:8]
plt.plot(r1[:,0], r1[:,1])
plt.title(r"$r_1$, Runge-Kutta")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2a-r_1-n={n}.png")
plt.show()
plt.plot(r2[:,0] - r1[:,0], r2[:,1] - r1[:,1])
plt.title(r"$r_2 - r_1$, Runge-Kutta")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2a-r_2-r_1-n={n}.png")
plt.show()
H = np.linalg.norm(p1, axis=1)**2 / (2 * m1) + np.linalg.norm(p2, axis=1)**2 / (2 * m2) + \
-G * m0 * m1 / np.linalg.norm(r1, axis=1) - G * m0 * m2 / np.linalg.norm(r2, axis=1) + \
-G * m1 * m2 / np.linalg.norm(r2 - r1, axis=1)
L = np.abs(np.cross(r1, p1, axis=1) + np.cross(r2, p2, axis=1))
plt.plot(t, H)
plt.title(r"$H(t)$")
plt.xlabel(r"$t / \mathrm{yr}$")
plt.ylabel(r"$L / \frac{m_\odot \mathrm{AU}^2}{\mathrm{yr}^2}$")
plt.savefig(f"problem1.2a-Hn={n}.png")
plt.show()
plt.plot(t, L)
plt.title(r"$L(t)$")
plt.xlabel(r"$t / \mathrm{yr}$")
plt.ylabel(r"$L / \frac{m_\odot \mathrm{AU}^2}{\mathrm{yr}}$")
plt.savefig(f"problem1.2a-L-n={n}.png")
plt.show()
plt.plot(t, H - H[0])
plt.title(r"$H(t) - H(0)$")
plt.xlabel(r"$t / \mathrm{yr}$")
plt.ylabel(r"$L / \frac{m_\odot \mathrm{AU}^2}{\mathrm{yr}^2}$")
plt.savefig(f"problem1.2a-H-H0-n={n}.png")
plt.show()
plt.plot(t, L - L[0])
plt.title(r"$L(t) - L(0)$")
plt.xlabel(r"$t / \mathrm{yr}$")
plt.ylabel(r"$L / \frac{m_\odot \mathrm{AU}^2}{\mathrm{yr}}$")
plt.savefig(f"problem1.2a-L-L0-n={n}.png")
plt.show()
# b)
n_v = 5
dt_v = 4**(-n_v) # yr
def F(x):
r1 = x[0:2]
r2 = x[2:4]
return np.array([
-G * m1 * (m0 * r1 / np.linalg.norm(r1)**3 + m2 * (r1 - r2) / np.linalg.norm(r1 - r2)**3),
-G * m2 * (m0 * r2 / np.linalg.norm(r2)**3 + m1 * (r2 - r1) / np.linalg.norm(r2 - r1)**3),
]).flatten()
rv, pv = verlet(F, x0[:4], x0[4:], np.array([m1, m1, m2, m2]), dt_v, math.ceil(T / dt_v))
rv1 = rv[:,:2]
rv2 = rv[:,2:]
pv1 = pv[:,:2]
pv2 = pv[:,2:]
plt.plot(r1[:,0], r1[:,1], label="Runge-Kutta")
plt.plot(rv1[:,0], rv1[:,1], label="Verlet")
plt.legend()
plt.title(r"$r_1$, Verlet")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2b-r_v1-n={n_v}.png")
plt.show()
plt.plot(r2[:,0] - r1[:,0], r2[:,1] - r1[:,1], label="Runge-Kutta")
plt.plot(rv2[:,0] - rv1[:,0], rv2[:,1] - rv1[:,1], label="Verlet")
plt.legend()
plt.title(r"$r_2 - r_1$, Verlet")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2b-r_v2-r_v1-n={n_v}.png")
plt.show()
# c)
xT = np.array([
r1[-1],
r2[-1],
-p1[-1],
-p2[-1],
]).flatten()
n = 5
T = 0.3 # yr
dt = 4**(-n) # yr
N = int(T / dt)
t = np.arange(N) * dt
rk = RK45(f, 0, xT, T, first_step=dt, max_step=dt, atol=1e3, rtol=1e3)
x = np.zeros((N, len(x0)))
x[0] = xT
for i in range(1, N):
rk.step()
x[i] = rk.y
rb1 = x[:,0:2]
rb2 = x[:,2:4]
pb1 = x[:,4:6]
pb2 = x[:,6:8]
plt.plot(r1[:,0], r1[:,1], label="Forwards")
plt.plot(rb1[:,0], rb1[:,1], label="Backwards")
plt.legend()
plt.title(r"$r_1$, Runge-Kutta")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2c-r_1.png")
plt.show()
plt.plot(r2[:,0] - r1[:,0], r2[:,1] - r1[:,1], label="Forwards")
plt.plot(rb2[:,0] - rb1[:,0], rb2[:,1] - rb1[:,1], label="Backwards")
plt.legend()
plt.title(r"$r_2 - r_1$, Runge-Kutta")
plt.xlabel(r"$x / \mathrm{AU}$")
plt.ylabel(r"$y / \mathrm{AU}$")
plt.savefig(f"problem1.2c-r_2-r_1.png")
plt.show()
rvb, pvb = verlet(F, xT[:4], xT[4:], np.array([m1, m1, m2, m2]), dt_v, math.ceil(T / dt_v))
rvb1 = rvb[:,:2]
rvb2 = rvb[:,2:]
pvb1 = pvb[:,:2]
pvb2 = pvb[:,2:]
plt.plot(rv1[:,0], rv1[:,1], label="Forwards")
plt.plot(rvb1[:,0], rvb1[:,1], label="Backwards")
plt.legend()
plt.title(r"$r_1$, Verlet")
plt.xlabel(r"$x$")
plt.ylabel(r"$y$")
plt.savefig(f"problem1.2c-r_v1.png")
plt.show()
plt.plot(rv2[:,0] - rv1[:,0], rv2[:,1] - rv1[:,1], label="Forwards")
plt.plot(rvb2[:,0] - rvb1[:,0], rvb2[:,1] - rvb1[:,1], label="Backwards")
plt.legend()
plt.title(r"$r_2 - r_1$, Verlet")
plt.xlabel(r"$x$")
plt.ylabel(r"$y$")
plt.savefig(f"problem1.2c-r_v2-r_v1.png")
plt.show()