Add sheet 1 and 2

Jaslo/sheet1.py

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

Jaslo/sheet2.py

+import numpy as np
+import matplotlib.pyplot as plt
+
+# Problem 2.1
+# a)
+def verlet(force, x0, p0, m, dt, N):
+    assert(x0.shape == p0.shape)
+
+    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 * force(x[i-1]) * dt
+        x[i] = x[i-1] + p[i] / m * dt
+        p[i] = p[i] + 1/2 * force(x[i]) * dt
+
+    return x, p
+
+r1_0 = np.array([1/3, 0, 0])
+r2_0 = -r1_0
+x0 = np.array([r1_0, r2_0]).flatten()
+
+p1_0 = np.array([-6/7, 3/7, -2/7])
+p2_0 = -p1_0
+p0 = np.array([p1_0, p2_0]).flatten()
+
+m = np.array([1, 1, 1, 1, 1, 1]) # ??
+a = 1 # ??
+epsilon = 1 # ??
+
+k = 1 # ??
+dt = 1e-2 # 1/omega_0
+t_max = 10 # 1/omeag_0
+t = np.arange(0, t_max, dt)
+
+def F_H(x):
+    r1 = x[:3]
+    r2 = x[3:]
+
+    return np.array([
+        -k * (r1 - r2),
+        k * (r1 - r2),
+    ]).flatten()
+
+x, p = verlet(F_H, x0, p0, m, dt, int(t_max / dt))
+r1 = x[:,:3]
+r2 = x[:,3:]
+p1 = p[:,:3]
+p2 = p[:,3:]
+
+plt.plot(r1[:,0], r1[:,1], label="Particle 1", lw=4)
+plt.plot(r2[:,0], r2[:,1], label="Particle 2")
+
+plt.xlabel(r"$x / a$")
+plt.ylabel(r"$y / a$")
+plt.title("Trajectories")
+plt.legend()
+plt.savefig("problem2.1a-trajectories.png")
+plt.show()
+
+p_tot = p1 + p2
+plt.plot(t, np.linalg.norm(p_tot - p_tot[0], axis=1))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$p / \sqrt{m \epsilon}$")
+plt.title("Momentum Deviations")
+plt.savefig("problem2.1a-dp.png")
+plt.show()
+
+J = np.cross(r1, p1, axis=1) + np.cross(r2, p2, axis=1)
+plt.plot(t, np.linalg.norm(J - J[0], axis=1))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
+plt.title("Angular Momentum Deviations")
+plt.savefig("problem2.1a-dJ.png")
+plt.show()
+
+plt.plot(t, 1 - np.inner(J, J[0]) / (np.linalg.norm(J, axis=1) * np.linalg.norm(J[0])))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
+plt.title("Momentum Deviations function")
+plt.savefig("problem2.1a-J?.png")
+plt.show()
+
+E = 1/2 * np.linalg.norm(p1, axis=1)**2 + 1/2 * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2
+plt.plot(t, E - E[0])
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$E / \epsilon$")
+plt.title("Energy Deviations")
+plt.savefig("problem2.1a-dE.png")
+plt.show()
+
+print(f"Maximum energy deviation: {np.max(np.abs((E - E[0]) / E[0]))}")
+
+# b)
+def FENE(x):
+    r1 = x[:3]
+    r2 = x[3:]
+
+    f1 = -epsilon * (r1 - r2) / (a**2 - np.linalg.norm(r1 - r2)**2)
+    f2 = -f1
+
+    return np.array([
+        f1,
+        f2,
+    ]).flatten()
+
+x_F, p_F = verlet(FENE, x0, p0, m, dt, int(t_max / dt))
+r1_F = x_F[:,:3]
+r2_F = x_F[:,3:]
+p1_F = p_F[:,:3]
+p2_F = p_F[:,3:]
+
+plt.plot(r1_F[:,0], r1_F[:,1], label="Particle 1")
+plt.plot(r2_F[:,0], r2_F[:,1], label="Particle 2")
+plt.xlabel(r"$x / a$")
+plt.ylabel(r"$y / a$")
+plt.title("Trajectories")
+plt.legend()
+plt.savefig("problem2.1b-trajectories.png")
+plt.show()
+
+p_totF = p1_F + p2_F
+plt.plot(t, np.linalg.norm(p_totF - p_totF[0], axis=1))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$p / \sqrt{m \epsilon}$")
+plt.title("Momentum Deviations")
+plt.savefig("problem2.1b-dp.png")
+plt.show()
+
+J_F = np.cross(r1_F, p1_F, axis=1) + np.cross(r2_F, p2_F, axis=1)
+plt.plot(t, np.linalg.norm(J_F - J_F[0], axis=1))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
+plt.title("Angular Momentum Deviations")
+plt.savefig("problem2.1b-dJ.png")
+plt.show()
+
+plt.plot(t, 1 - np.inner(J, J[0]) / (np.linalg.norm(J, axis=1) * np.linalg.norm(J[0])))
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$J / a \sqrt{m \epsilon}$")
+plt.title("Momentum Deviations function")
+plt.savefig("problem2.1a-J?.png")
+plt.show()
+
+E_F = 1/2 * np.linalg.norm(p1_F, axis=1)**2 + 1/2 * np.linalg.norm(p2_F, axis=1)**2 + 1/2 * k * np.linalg.norm(r1_F - r2_F, axis=1)**2
+plt.plot(t, E_F - E_F[0])
+plt.xlabel(r"$t / \omega_0^{-1}$")
+plt.ylabel(r"$E / \epsilon$")
+plt.title("Energy Deviations")
+plt.savefig("problem2.1b-dE.png")
+plt.show()
+
+print(f"Maximum energy deviation: {np.max(np.abs((E_F - E_F[0]) / E_F[0]))}")
+
+# c)
+lambda_c = 1/2 - (2 * np.sqrt(326) + 36)**(1/3) / 12 + 1 / (6 * (2 * np.sqrt(326) + 36)**(1/3))
+
+def verlet_BABAB(force, x0, p0, m, dt, N, lambda_):
+    assert(x0.shape == p0.shape)
+
+    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] + lambda_ * force(x[i-1]) * dt
+        x[i] = x[i-1] + 1/2 * p[i] / m * dt
+        p[i] = p[i] + (1 - 2 * lambda_) * force(x[i]) * dt
+        x[i] = x[i] + 1/2 * p[i] / m * dt
+        p[i] = p[i] + lambda_ * force(x[i]) * dt
+
+    return x, p
+
+for f in (F_H, FENE,):
+    for n in range(1, 5):
+        for j in range(1, 6):
+            dt = 4**(-j)
+            t = np.arange(0, t_max, dt)
+
+            x, p = verlet_BABAB(f, x0, p0, m, dt, len(t), lambda_c)
+            r1 = x[:,:3]
+            r2 = x[:,3:]
+            p1 = p[:,:3]
+            p2 = p[:,3:]
+
+            E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2
+
+            plt.plot(t, (E - E[0]) / dt**n, label="$k =" + str(j) + "$")
+
+        plt.xlabel(r"$t / \omega_0^{-1}$")
+        plt.ylabel(r"$\frac{E}{t^" + str(n) + r"} / \epsilon \omega_0^" + str(n) + r"$")
+        plt.title("$n =" + str(n) + "$")
+        plt.legend(loc="lower left")
+        plt.savefig(f"problem2.1c-dE-{f.__name__}-n={n}.png")
+        plt.show()
+
+for lambda_ in (1/3, 0, 1/2):
+    dt = 4**(-2)
+    t = np.arange(0, t_max, dt)
+
+    x, p = verlet_BABAB(f, x0, p0, m, dt, len(t), lambda_)
+    r1 = x[:,:3]
+    r2 = x[:,3:]
+    p1 = p[:,:3]
+    p2 = p[:,3:]
+
+    E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2
+
+    print(f"Maximum energy deviation for {lambda_}: {np.max(np.abs(E))}")
+
+# d)
+theta = 0.08398315262876693
+lambda_d = 0.6822365335719091
+rho = 0.2539785108410595
+mu = -0.03230286765269967
+
+def verlet_d(force, x0, p0, m, dt, N):
+    assert(x0.shape == p0.shape)
+
+    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] + theta * force(x[i-1]) * dt
+        x[i] = x[i-1] + rho * p[i] / m * dt
+        p[i] = p[i] + lambda_d * force(x[i]) * dt
+        x[i] = x[i] + mu * p[i] / m * dt
+        p[i] = p[i] + (1 - 2 * (lambda_d + theta)) / 2 * force(x[i]) * dt
+        x[i] = x[i] + (1 - 2 * (mu + rho)) * p[i] / m * dt
+        p[i] = p[i] + (1 - 2 * (lambda_d + theta)) / 2 * force(x[i]) * dt
+        x[i] = x[i] + mu * p[i] / m * dt
+        p[i] = p[i] + lambda_d * force(x[i]) * dt
+        x[i] = x[i] + rho * p[i] / m * dt
+        p[i] = p[i] + theta * force(x[i]) * dt
+
+    return x, p
+
+for f in (F_H, FENE):
+    for n in range(1, 5):
+        for j in range(1, 4):
+            dt = 5 * 2**(-j)
+            t = np.arange(0, t_max, dt)
+
+            x, p = verlet_d(f, x0, p0, m, dt, len(t))
+            r1 = x[:,:3]
+            r2 = x[:,3:]
+            p1 = p[:,:3]
+            p2 = p[:,3:]
+
+            E = 1/(2 * m[0]) * np.linalg.norm(p1, axis=1)**2 + 1/(2 * m[3]) * np.linalg.norm(p2, axis=1)**2 + 1/2 * k * np.linalg.norm(r1 - r2, axis=1)**2
+
+            plt.plot(t, (E - E[0]) / dt**n, label="$k =" + str(j) + "$")
+
+        plt.xlabel(r"$t / \omega_0^{-1}$")
+        plt.ylabel(r"$\frac{E}{t^" + str(n) + r"} / \epsilon \omega_0^" + str(n) + r"$")
+        plt.title("$n =" + str(n) + "$")
+        plt.legend(loc="lower left")
+        plt.savefig(f"problem2.1d-dE-{f.__name__}-n={n}.png")
+        plt.show()