Skip to content
Snippets Groups Projects
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()