Initial commitHEADmaster

2 files changed, 180 insertions, 0 deletions

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..d2ca960
--- /dev/null
+++ b/README.md
@@ -0,0 +1,4 @@
+# Pykepler
+
+Numerically solve the restricted 3-body problem and animate the movement of the
+bodies
diff --git a/pykepler.py b/pykepler.py
new file mode 100644
index 0000000..d6241c3
--- /dev/null
+++ b/pykepler.py
@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+
+# pykepler: Numerically solve the restricted 3-body problem and animate the
+# movement of the bodies
+#
+# Copyright (C) 2025 Thomas Guillermo Albers Raviola
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from functools import partial
+
+import numpy as np
+from scipy.integrate import solve_ivp
+
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from matplotlib import rcParams
+from matplotlib.animation import FuncAnimation
+
+def unpack_r(rv):
+    return rv[0:2], rv[2:4], rv[4:6]
+
+def unpack_v(rv):
+    return rv[6:8], rv[8:10], rv[10:12]
+
+def distance(x1, x2):
+    return np.sqrt(np.sum((x1 - x2)**2))
+
+def grav_force(r1, r2):
+    d = distance(r1, r2)
+    return (r2 - r1) / d**3
+
+#
+# Initial value problem has form
+# x' = F(x, t)
+#
+# Set x = [r, v], where r is position of body and v = r'
+# Thus, x' = [r', v'] = [v, a] = [v, -G*M*m/abs(d)**2 * e]
+#
+
+def func(t, rv, mu):
+    # Unpack components
+    r1, r2, r3 = unpack_r(rv)
+    v1, v2, v3 = unpack_v(rv)
+
+    f12 = grav_force(r1, r2)
+    f13 = grav_force(r1, r3)
+    f23 = grav_force(r2, r3)
+
+    return np.concatenate((v1, v2, v3,
+                           mu[1] * f12 + mu[2] * f13,
+                           -mu[0] * f12 + mu[2] * f23,
+                           -mu[1] * f23 - mu[0] * f13))
+
+
+def angle(x1, x2):
+    ux1 = x1 / np.linalg.norm(x1)
+    ux2 = x2 / np.linalg.norm(x2)
+    return 180.0 * np.arccos(np.clip(np.dot(ux1, ux2), -1.0, 1.0)) / np.pi
+
+
+def update(sol, text_artist, tri_artist, pnt_artists, pos_artists, frame):
+    '''Update plot of the IVP up to the current timestamp'''
+
+    if frame > 0:
+        r1, r2, r3 = unpack_r(sol.y)
+
+        r1 = r1[:, frame-1]
+        r2 = r2[:, frame-1]
+        r3 = r3[:, frame-1]
+
+        tri_artist.set_data([r1[0], r2[0], r3[0], r1[0]],
+                            [r1[1], r2[1], r3[1], r1[1]])
+
+        alpha = angle(r2 - r1, r3 - r1)
+        beta = angle(r1 - r2, r3 - r2)
+        gamma = angle(r1 - r3, r2 - r3)
+
+        a = distance(r2, r3)
+        b = distance(r3, r1)
+        c = distance(r2, r1)
+
+        fmt = 'α:{:.2f}, β:{:.2f}, γ:{:.2f}\na: {:.2f}, b:{:.2f}, c:{:.2f}'
+        text_artist.set_text(fmt.format(alpha, beta, gamma, a, b, c))
+
+
+    for artist, r in zip(pos_artists, unpack_r(sol.y)):
+        artist.set_data(r[0, 0:frame], r[1, 0:frame])
+
+    for artist, r in zip(pnt_artists, unpack_r(sol.y)):
+        artist.set_data([r[0, frame-1]], [r[1, frame-1]])
+
+    return [text_artist, tri_artist] + pos_artists + pnt_artists
+
+
+def animate_ivp(sol, fps=20):
+    def progress(current_frame, total_frames):
+        str = '\r  Rendering frame [{}/{}] ...'
+        print(str.format(current_frame + 1, total_frames), end='', flush=True)
+
+    fig, ax = plt.subplots(figsize=(10,10), dpi=300)
+    ax.set_xlabel('x')
+    ax.set_ylabel('y')
+    ax.set_xlim(-1, 1)
+    ax.set_ylim(-1, 1)
+
+    pos_artists = []
+    pnt_artists = []
+
+    tri_artist, = ax.plot([], [], ls='--', color='black')
+    for i in range(3):
+        pos_artists.append(ax.plot([], [])[0])
+        pnt_artists.append(ax.plot([], [], '.', label='m{}'.format(i + 1))[0])
+
+    text_artist = ax.text(0.05, 0.95, '',
+                          backgroundcolor='#7F7F7F33',
+                          transform=ax.transAxes,
+                          verticalalignment='top') # bbox=props
+
+    ax.legend()
+
+    frames = np.arange(0, len(sol.t))
+    animation = FuncAnimation(fig,
+                              partial(update, sol,
+                                      text_artist, tri_artist,
+                                      pnt_artists, pos_artists),
+                              frames=frames,
+                              interval=sol.t[-1], blit=True, repeat=False)
+
+    animation.save('three-body.mp4', fps=fps, progress_callback=progress)
+    # Reserve newline for progress
+    print('')
+
+
+def main():
+    mu1 = .7
+    mu2 = 1 - mu1
+    mu = np.array([mu1, mu2, 1e-12])
+    omega = 1
+
+    ir1 = mu[1] * np.array([-1, 0])
+    ir2 = mu[0] * np.array([1, 0])
+    # Trojan satelite at L4 (ahead)
+    ir3 = np.array([.5 - mu2, np.sqrt(3) / 2])
+
+    iv1 = mu[1] * np.array([0, -1])
+    iv2 = mu[0] * np.array([0, 1])
+    # yz - zy, zx - xz
+    iv3 = np.array([- ir3[1], ir3[0]])
+
+    irv = np.concatenate((ir1, ir2, ir3, iv1, iv2, iv3))
+
+    t_span = (0.0, 4)
+    t_eval = np.linspace(t_span[0], t_span[1], 1024)
+
+    sol = solve_ivp(func, t_span, irv,
+                    t_eval=t_eval, method='DOP853',
+                    rtol=1e-10, atol=1e-20,
+                    args=(mu,))
+
+    animate_ivp(sol)
+
+
+if __name__ == '__main__':
+    main()