diff options
| -rw-r--r-- | README.md | 4 | ||||
| -rw-r--r-- | pykepler.py | 176 |
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() |