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()