#!/usr/bin/env python3 """ Two-body problem in 3D — the invariable plane, from skew initial velocities. Setup (non-dimensional: G = 1, m1 = m2 = 1/2, so mu = G(m1+m2) = 1) ------------------------------------------------------------------- R1(0) = (+d, 0, 0) V1(0) = (0, v, 0) -- velocity along +y R2(0) = (-d, 0, 0) V2(0) = (0, 0, v) -- velocity along +z with d = 1/2 (initial separation |R2 - R1| = 1) and v = 1/2. The two velocities are *skew*: they do not lie in a common plane with the line joining the bodies in any obvious way, and the total linear momentum P = m1 V1 + m2 V2 = (v/2) (yhat + zhat) =/= 0 does not vanish, so the barycentre drifts. Nevertheless the *relative* motion is planar, because h = r x rdot = 2 d v (yhat + zhat) is constant, and h happens to be parallel to P for these initial conditions. The plane containing the two bodies therefore translates along its own normal: in the inertial frame the bodies trace two interlocking helices, while in the barycentric frame they trace two similar ellipses in one fixed plane. With r(0) . rdot(0) = 0 the initial point is an apsis, and h = 1/sqrt(2), eps = -3/4, a = 2/3, e = 1/2, T = 2 pi sqrt(a^3/mu) = 3.4201 (the start is apoapsis, since the initial relative speed v*sqrt(2) = 0.707 is below the local circular speed 1). Integration ----------- Runge-Kutta-Fehlberg 7(8) (Fehlberg, NASA TR R-287, 1968), implemented from the tableau below with adaptive step-size control; the 8th-order solution is propagated (local extrapolation) and the 7(8) difference drives the step size. Accepted steps are then interpolated onto a uniform time grid with cubic Hermite interpolation, using the exact velocities as derivatives. Usage ----- python two-body-3d-plane.py --verify # tableau + invariants check python two-body-3d-plane.py --save img/two-body-3d-plane.gif """ from __future__ import annotations import argparse import os from fractions import Fraction as F import numpy as np # ── Runge–Kutta–Fehlberg 7(8) tableau ──────────────────────────────────────── # Coefficients are entered as exact rationals and converted once, so that the # order conditions are satisfied to machine precision. _C = [F(0), F(2, 27), F(1, 9), F(1, 6), F(5, 12), F(1, 2), F(5, 6), F(1, 6), F(2, 3), F(1, 3), F(1), F(0), F(1)] _A = [ [], [F(2, 27)], [F(1, 36), F(1, 12)], [F(1, 24), F(0), F(1, 8)], [F(5, 12), F(0), F(-25, 16), F(25, 16)], [F(1, 20), F(0), F(0), F(1, 4), F(1, 5)], [F(-25, 108), F(0), F(0), F(125, 108), F(-65, 27), F(125, 54)], [F(31, 300), F(0), F(0), F(0), F(61, 225), F(-2, 9), F(13, 900)], [F(2), F(0), F(0), F(-53, 6), F(704, 45), F(-107, 9), F(67, 90), F(3)], [F(-91, 108), F(0), F(0), F(23, 108), F(-976, 135), F(311, 54), F(-19, 60), F(17, 6), F(-1, 12)], [F(2383, 4100), F(0), F(0), F(-341, 164), F(4496, 1025), F(-301, 82), F(2133, 4100), F(45, 82), F(45, 164), F(18, 41)], [F(3, 205), F(0), F(0), F(0), F(0), F(-6, 41), F(-3, 205), F(-3, 41), F(3, 41), F(6, 41), F(0)], [F(-1777, 4100), F(0), F(0), F(-341, 164), F(4496, 1025), F(-289, 82), F(2193, 4100), F(51, 82), F(33, 164), F(12, 41), F(0), F(1)], ] # 7th-order weights _B7 = [F(41, 840), F(0), F(0), F(0), F(0), F(34, 105), F(9, 35), F(9, 35), F(9, 280), F(9, 280), F(41, 840), F(0), F(0)] # 8th-order weights _B8 = [F(0), F(0), F(0), F(0), F(0), F(34, 105), F(9, 35), F(9, 35), F(9, 280), F(9, 280), F(0), F(41, 840), F(41, 840)] C = np.array([float(c) for c in _C]) A = np.zeros((13, 13)) for i, row in enumerate(_A): for j, aij in enumerate(row): A[i, j] = float(aij) B7 = np.array([float(b) for b in _B7]) B8 = np.array([float(b) for b in _B8]) STAGES = 13 ORDER = 8 def rkf78_step(f, t, y, h): """One RKF7(8) step. Returns (y8, error estimate vector, stage slopes).""" k = np.empty((STAGES, y.size)) k[0] = f(t, y) for i in range(1, STAGES): k[i] = f(t + C[i] * h, y + h * (A[i, :i] @ k[:i])) y7 = y + h * (B7 @ k) y8 = y + h * (B8 @ k) return y8, y8 - y7, k def integrate(f, t0, y0, tf, rtol=1e-12, atol=1e-14, h0=1e-3, safety=0.9): """Adaptive RKF7(8). Returns (t, y, dy) at the accepted steps.""" t, y = float(t0), np.asarray(y0, dtype=float) h = float(h0) ts, ys, ds = [t], [y.copy()], [f(t, y)] while t < tf: h = min(h, tf - t) y_new, err, k = rkf78_step(f, t, y, h) scale = atol + rtol * np.maximum(np.abs(y), np.abs(y_new)) e = np.max(np.abs(err) / scale) if e <= 1.0: # accept t += h y = y_new ts.append(t) ys.append(y.copy()) ds.append(f(t, y)) factor = safety * (1.0 / e) ** (1.0 / ORDER) if e > 0 else 5.0 h *= min(5.0, max(0.1, factor)) return np.array(ts), np.array(ys), np.array(ds) def hermite_resample(ts, ys, ds, t_out): """Cubic Hermite interpolation of the step history onto a uniform grid.""" idx = np.clip(np.searchsorted(ts, t_out, side='right') - 1, 0, len(ts) - 2) h = (ts[idx + 1] - ts[idx])[:, None] s = ((t_out - ts[idx]) / h[:, 0])[:, None] y0, y1, d0, d1 = ys[idx], ys[idx + 1], ds[idx], ds[idx + 1] h00 = 2 * s**3 - 3 * s**2 + 1 h10 = s**3 - 2 * s**2 + s h01 = -2 * s**3 + 3 * s**2 h11 = s**3 - s**2 return h00 * y0 + h10 * h * d0 + h01 * y1 + h11 * h * d1 # ── Dynamics ───────────────────────────────────────────────────────────────── G = 1.0 M1 = M2 = 0.5 M_TOT = M1 + M2 MU = G * M_TOT D0 = 0.5 # each body's initial distance from the origin V0 = 0.5 # magnitude of both initial velocities def rhs(t, y): """y = [R1, R2, V1, V2] (12 components).""" r1, r2, v1, v2 = y[0:3], y[3:6], y[6:9], y[9:12] d = r2 - r1 s = np.dot(d, d) ** 1.5 a1 = G * M2 * d / s a2 = -G * M1 * d / s return np.concatenate([v1, v2, a1, a2]) def initial_state(): r1 = np.array([+D0, 0.0, 0.0]) r2 = np.array([-D0, 0.0, 0.0]) v1 = np.array([0.0, V0, 0.0]) # body 1 along +y v2 = np.array([0.0, 0.0, V0]) # body 2 along +z return np.concatenate([r1, r2, v1, v2]) def elements(y): """Relative-orbit invariants from a state vector.""" r = y[3:6] - y[0:3] v = y[9:12] - y[6:9] h = np.cross(r, v) eps = 0.5 * v @ v - MU / np.linalg.norm(r) e_vec = np.cross(v, h) / MU - r / np.linalg.norm(r) a = -MU / (2 * eps) return dict(h=h, hmag=np.linalg.norm(h), eps=eps, e_vec=e_vec, e=np.linalg.norm(e_vec), a=a, T=2 * np.pi * np.sqrt(a**3 / MU) if a > 0 else np.inf) def simulate(n_orbits=2.0, n_frames=300, rtol=1e-12): y0 = initial_state() el = elements(y0) tf = n_orbits * el['T'] ts, ys, ds = integrate(rhs, 0.0, y0, tf, rtol=rtol) t_out = np.linspace(0.0, tf, n_frames) Y = hermite_resample(ts, ys, ds, t_out) R1, R2 = Y[:, 0:3], Y[:, 3:6] V1, V2 = Y[:, 6:9], Y[:, 9:12] Rcm = (M1 * R1 + M2 * R2) / M_TOT Vcm = (M1 * V1 + M2 * V2) / M_TOT return dict(t=t_out, R1=R1, R2=R2, V1=V1, V2=V2, Rcm=Rcm, Vcm=Vcm, el=el, steps=len(ts), y_end=ys[-1]) # ── Verification ───────────────────────────────────────────────────────────── def verify(): """Order of the tableau, then the invariants of the actual run.""" print('RKF7(8) — order verification on y\' = y cos t, y(0) = 1\n') f = lambda t, y: y * np.cos(t) exact = lambda t: np.exp(np.sin(t)) prev = None for n in (10, 20, 40, 80): h = 2.0 / n t, y = 0.0, np.array([1.0]) for _ in range(n): y, _, _ = rkf78_step(f, t, y, h) t += h err = abs(y[0] - exact(2.0)) rate = '' if prev is None else f' observed order = {np.log2(prev/err):5.2f}' print(f' h = {h:7.4f} |error| = {err:.3e}{rate}') prev = err print('\nTwo-body run — conservation of the integrals\n') s = simulate() el = s['el'] print(f" accepted RKF7(8) steps : {s['steps']}") print(f" mu = G(m1+m2) : {MU:.6f}") print(f" h = |r x rdot| : {el['hmag']:.10f}") print(f" eps : {el['eps']:.10f}") print(f" a, e : {el['a']:.10f}, {el['e']:.10f}") print(f" T : {el['T']:.10f}") print(f" h direction (unit) : {el['h']/el['hmag']}") print(f" V_cm : {s['Vcm'][0]}") cos = (s['Vcm'][0] @ el['h']) / (np.linalg.norm(s['Vcm'][0]) * el['hmag']) print(f" cos angle(V_cm, h) : {cos:.12f}") el_end = elements(s['y_end']) print(f"\n drift in |h| over 2 periods : {abs(el_end['hmag']-el['hmag']):.3e}") print(f" drift in eps over 2 periods : {abs(el_end['eps']-el['eps']):.3e}") print(f" drift in e over 2 periods : {abs(el_end['e']-el['e']):.3e}") nhat = el['h'] / el['hmag'] dev1 = np.abs((s['R1'] - s['Rcm']) @ nhat).max() dev2 = np.abs((s['R2'] - s['Rcm']) @ nhat).max() print(f"\n max |(R1 - Rcm).nhat| : {dev1:.3e} (planarity)") print(f" max |(R2 - Rcm).nhat| : {dev2:.3e} (planarity)") drift = s['Rcm'] - (s['Rcm'][0] + np.outer(s['t'], s['Vcm'][0])) print(f" max |Rcm - (Rcm0 + Vcm t)| : {np.abs(drift).max():.3e} " '(uniform barycentre drift)') # ── Rendering ──────────────────────────────────────────────────────────────── CREAM, INK, DIM = '#fcfbf7', '#2f2a24', '#8b857c' MAROON, BLUE, GREEN, AMBER, PURPLE = ('#6b1f1f', '#2f5ea8', '#2d6a2d', '#b45309', '#7e22ce') def _plane_quad(centre, nhat, half): """Corners of a square patch of the plane through `centre` normal to `nhat`.""" u = np.array([1.0, 0.0, 0.0]) u = u - (u @ nhat) * nhat u /= np.linalg.norm(u) w = np.cross(nhat, u) return np.array([centre + a * half * u + b * half * w for a, b in ((1, 1), (-1, 1), (-1, -1), (1, -1))]) def render(path, n_orbits=2.0, n_frames=260, fps=25, dpi=100): import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt from matplotlib import animation from mpl_toolkits.mplot3d.art3d import Poly3DCollection plt.rcParams.update({ 'figure.facecolor': CREAM, 'savefig.facecolor': CREAM, 'axes.facecolor': CREAM, 'text.color': INK, 'font.size': 11, 'axes.labelcolor': INK, 'xtick.color': INK, 'ytick.color': INK, }) s = simulate(n_orbits=n_orbits, n_frames=n_frames) el = s['el'] nhat = el['h'] / el['hmag'] n = n_frames R1, R2, Rcm = s['R1'], s['R2'], s['Rcm'] r1b, r2b = R1 - Rcm, R2 - Rcm # barycentric frame fig = plt.figure(figsize=(12.0, 5.6)) axL = fig.add_subplot(1, 2, 1, projection='3d') axR = fig.add_subplot(1, 2, 2, projection='3d') for ax in (axL, axR): ax.set_facecolor(CREAM) for pane in (ax.xaxis, ax.yaxis, ax.zaxis): pane.set_pane_color((0.988, 0.984, 0.969, 1.0)) pane._axinfo['grid'].update(color='#ded9cf', linewidth=0.6) ax.set_xlabel('$x$', labelpad=-6) ax.set_ylabel('$y$', labelpad=-6) ax.set_zlabel('$z$', labelpad=-6) ax.tick_params(labelsize=7, colors=DIM, pad=-2) for axis in (ax.xaxis, ax.yaxis, ax.zaxis): axis.set_major_locator(plt.MaxNLocator(4)) # limits — equal aspect, left panel must contain the whole drift allpts = np.vstack([R1, R2]) ctr = allpts.mean(axis=0) span = np.abs(allpts - ctr).max() * 1.12 axL.set_xlim(ctr[0] - span, ctr[0] + span) axL.set_ylim(ctr[1] - span, ctr[1] + span) axL.set_zlim(ctr[2] - span, ctr[2] + span) spanB = np.abs(np.vstack([r1b, r2b])).max() * 1.30 for lim in (axR.set_xlim, axR.set_ylim, axR.set_zlim): lim(-spanB, spanB) axL.set_box_aspect((1, 1, 1)) axR.set_box_aspect((1, 1, 1)) fig.text(0.26, 0.905, 'inertial frame — the plane drifts along $\\mathbf{h}$', ha='center', va='center', color=MAROON, fontsize=12.5) fig.text(0.76, 0.905, 'barycentric frame — one fixed plane', ha='center', va='center', color=MAROON, fontsize=12.5) # ── static artists ── axL.plot(*Rcm.T, color=AMBER, lw=1.0, ls=(0, (5, 4)), zorder=1) # barycentric orbits, drawn faintly in full axR.plot(*r1b.T, color=MAROON, lw=0.9, alpha=0.30) axR.plot(*r2b.T, color=BLUE, lw=0.9, alpha=0.30) quadR = _plane_quad(np.zeros(3), nhat, spanB * 0.95) axR.add_collection3d(Poly3DCollection([quadR], facecolor=BLUE, alpha=0.10, edgecolor=BLUE, linewidths=0.8)) # ── dynamic artists ── trailL1, = axL.plot([], [], [], color=MAROON, lw=1.6) trailL2, = axL.plot([], [], [], color=BLUE, lw=1.6) dotL1, = axL.plot([], [], [], 'o', color=MAROON, ms=8) dotL2, = axL.plot([], [], [], 'o', color=BLUE, ms=8) dotLc, = axL.plot([], [], [], 'o', color=AMBER, ms=5) segL, = axL.plot([], [], [], color=INK, lw=0.9, alpha=0.55) planeL = Poly3DCollection([_plane_quad(Rcm[0], nhat, span * 0.55)], facecolor=BLUE, alpha=0.11, edgecolor=BLUE, linewidths=0.8) axL.add_collection3d(planeL) hL, = axL.plot([], [], [], color=GREEN, lw=2.0) trailR1, = axR.plot([], [], [], color=MAROON, lw=2.0) trailR2, = axR.plot([], [], [], color=BLUE, lw=2.0) dotR1, = axR.plot([], [], [], 'o', color=MAROON, ms=9) dotR2, = axR.plot([], [], [], 'o', color=BLUE, ms=9) axR.plot([0], [0], [0], 'o', color=AMBER, ms=6) segR, = axR.plot([], [], [], color=INK, lw=1.0, alpha=0.6) hlen = spanB * 0.85 axR.plot([0, nhat[0] * hlen], [0, nhat[1] * hlen], [0, nhat[2] * hlen], color=GREEN, lw=2.4) axR.text(nhat[0] * hlen + 0.10 * spanB, nhat[1] * hlen, nhat[2] * hlen * 1.10, r'$\mathbf{h}$', color=GREEN, fontsize=15) from matplotlib.lines import Line2D handles = [ Line2D([], [], color=MAROON, marker='o', ls='-', label='$m_1 = 1/2$'), Line2D([], [], color=BLUE, marker='o', ls='-', label='$m_2 = 1/2$'), Line2D([], [], color=AMBER, marker='o', ls=(0, (5, 4)), label='barycentre'), Line2D([], [], color=GREEN, lw=2.4, label=r'$\mathbf{h}$ (constant)'), Line2D([], [], color=BLUE, lw=6, alpha=0.25, label='orbital plane $\\perp\\,\\mathbf{h}$'), ] fig.legend(handles=handles, loc='lower center', ncol=5, frameon=False, fontsize=10, bbox_to_anchor=(0.5, 0.005)) fig.text(0.5, 0.965, r'$m_1 = m_2 = \frac{1}{2}$, $\mathbf{R}_{1,2}(0) = (\pm\frac{1}{2},0,0)$, ' r'$\mathbf{V}_1(0) = \frac{1}{2}\hat{\mathbf{y}}$, ' r'$\mathbf{V}_2(0) = \frac{1}{2}\hat{\mathbf{z}}$ ' rf"$\Rightarrow$ $e = {el['e']:.2f}$, $a = {el['a']:.3f}$, " rf"$h = {el['hmag']:.3f}$", ha='center', va='top', fontsize=12, color=INK) fig.subplots_adjust(left=0.0, right=1.0, bottom=0.055, top=0.99, wspace=0.0) # Camera. The orbital plane is normal to n = (yhat + zhat)/sqrt(2), so the # two panels want different viewpoints: the left one looks nearly *along* # the plane (edge-on) to expose the drift and the helices, the right one # looks nearly *down* the normal so the ellipses open out. trail = max(1, n // 2) def frame(k): k = min(k, n - 1) lo = max(0, k - trail) sl = slice(lo, k + 1) trailL1.set_data_3d(*R1[sl].T) trailL2.set_data_3d(*R2[sl].T) dotL1.set_data_3d(*R1[k:k + 1].T) dotL2.set_data_3d(*R2[k:k + 1].T) dotLc.set_data_3d(*Rcm[k:k + 1].T) segL.set_data_3d(*np.vstack([R1[k], R2[k]]).T) planeL.set_verts([_plane_quad(Rcm[k], nhat, span * 0.55)]) hL.set_data_3d(*np.vstack([Rcm[k], Rcm[k] + nhat * span * 0.45]).T) trailR1.set_data_3d(*r1b[:k + 1].T) trailR2.set_data_3d(*r2b[:k + 1].T) dotR1.set_data_3d(*r1b[k:k + 1].T) dotR2.set_data_3d(*r2b[k:k + 1].T) segR.set_data_3d(*np.vstack([r1b[k], r2b[k]]).T) sweep = 40.0 * k / n axL.view_init(elev=18, azim=-50.0 + sweep) axR.view_init(elev=26, azim=15.0 + sweep) return () hold = 25 ani = animation.FuncAnimation(fig, frame, frames=n + hold, blit=False) os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) writer = 'ffmpeg' if path.endswith('.mp4') else 'pillow' print(f'Writing {path} ({n + hold} frames) …') ani.save(path, writer=writer, fps=fps, dpi=dpi, savefig_kwargs=dict(facecolor=CREAM)) print('done:', path, f'{os.path.getsize(path)/1e6:.2f} MB') def main(): ap = argparse.ArgumentParser(description=__doc__.split('\n')[1]) ap.add_argument('--verify', action='store_true', help='check the tableau order and the conserved quantities') ap.add_argument('--save', metavar='PATH', default=None, help='render the animation to PATH (.gif or .mp4)') ap.add_argument('--frames', type=int, default=260) ap.add_argument('--orbits', type=float, default=2.0) ap.add_argument('--fps', type=int, default=25) ap.add_argument('--dpi', type=int, default=100) args = ap.parse_args() if args.verify: verify() if args.save: render(args.save, n_orbits=args.orbits, n_frames=args.frames, fps=args.fps, dpi=args.dpi) if not args.verify and not args.save: ap.print_help() if __name__ == '__main__': main()