Created
April 11, 2023 18:33
-
-
Save Aragroth/23e7403c4187c0454246c11817f9cfa5 to your computer and use it in GitHub Desktop.
Time solved orbit display, including basic oblateness effect of Earth
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import mpl_toolkits.mplot3d.art3d as art3d | |
| from matplotlib.animation import FuncAnimation | |
| from mpl_toolkits.mplot3d import Axes3D | |
| import matplotlib.pyplot as plt | |
| from scipy import optimize | |
| import numpy as np | |
| import datetime | |
| import sys | |
| class OrbitStateVector: | |
| def __init__(self, r, v): | |
| self.radius = np.array(r) | |
| self.velocity = np.array(v) | |
| self.mu = 398_600 | |
| self.angular_momentum = self.calculate_angular_momentum() | |
| self.inclination = self.calculate_inclination() | |
| self.node = self.calculate_node() | |
| self.right_ascention = self.calculate_right_ascention() | |
| self.eccentricity = self.calculate_eccentricity() | |
| self.argument_of_perigee = self.calculate_argument_of_perigee() | |
| self.true_anomaly = self.calculate_true_anomaly() | |
| def calculate_angular_momentum(self): | |
| return np.cross(self.radius, self.velocity) | |
| def calculate_inclination(self): | |
| return np.arccos(self.angular_momentum[2] / self.angular_momentum_module) | |
| def calculate_node(self): | |
| return np.cross([0, 0, 1], self.angular_momentum) | |
| def calculate_right_ascention(self): | |
| res = np.arccos(self.node[0] / self.node_module) | |
| return res if self.node[1] >= 0 else (2 * np.pi - res) | |
| def calculate_eccentricity(self): | |
| res_r = (self.velocity_module ** 2 - self.mu / | |
| self.radius_module) * self.radius | |
| res_v = self.radius_module * self.radial_velocity * self.velocity | |
| return (1 / self.mu) * (res_r - res_v) | |
| def calculate_argument_of_perigee(self): | |
| res = np.arccos(self.node @ self.eccentricity / | |
| (self.node_module * self.eccentricity_module)) | |
| return res if self.eccentricity[2] >= 0 else (2 * np.pi - res) | |
| def calculate_true_anomaly(self): | |
| res = np.arccos(self.eccentricity @ self.radius / | |
| (self.eccentricity_module * self.radius_module)) | |
| return res if self.radial_velocity >= 0 else (2 * np.pi - res) | |
| @property | |
| def angular_momentum_module(self): | |
| return np.linalg.norm(self.angular_momentum) | |
| @property | |
| def radius_module(self): | |
| return np.linalg.norm(self.radius) | |
| @property | |
| def velocity_module(self): | |
| return np.linalg.norm(self.velocity) | |
| @property | |
| def radial_velocity(self): | |
| return self.radius @ self.velocity.T / self.radius_module | |
| @property | |
| def node_module(self): | |
| return np.linalg.norm(self.node) | |
| @property | |
| def eccentricity_module(self): | |
| return np.linalg.norm(self.eccentricity) | |
| class OblatenessTimeSolver: | |
| def __init__(self, initial: OrbitStateVector): | |
| self.initial = initial | |
| self.mu = 398_600 | |
| self.j2 = 0.00108263 | |
| self.earth_radius = 6378 | |
| self.e = self.initial.eccentricity_module | |
| self.radius_coeff = self.initial.angular_momentum_module ** 2 / self.mu | |
| self.semimajor = self.calculate_semimajor() | |
| self.period = self.calculate_period() | |
| self.mean_motion = 2 * np.pi / self.period | |
| self.eccentric_anomaly = self.calculate_eccentric_anomaly() | |
| self.initial_epoch = self.calculate_initial_epoch() | |
| def calculate_semimajor(self): | |
| return self.radius_coeff * (1 / (1 - self.e ** 2)) | |
| def calculate_period(self): | |
| return 2 * np.pi / np.sqrt(self.mu) * self.semimajor ** (3 / 2) | |
| def calculate_eccentric_anomaly(self): | |
| coeff = np.sqrt((1 - self.e) / (1 + self.e)) | |
| print('init', self.e) | |
| return 2 * np.arctan(coeff * np.tan(self.initial.true_anomaly / 2)) | |
| def calculate_initial_epoch(self): | |
| return (self.eccentric_anomaly - self.e * np.sin(self.eccentric_anomaly)) / self.mean_motion | |
| def state_after(self, seconds: int): | |
| time_final = self.initial_epoch + seconds | |
| periods_num = time_final / self.period | |
| time_position = self.period * (periods_num - int(periods_num)) | |
| mean_anomaly = self.mean_motion * time_position | |
| def kepler_equation(E): | |
| return E - self.e * np.sin(E) - mean_anomaly | |
| eccentric_position = optimize.root(kepler_equation, 2 * np.pi).x[0] | |
| coeff = np.sqrt((1 + self.e) / (1 - self.e)) | |
| # np.mod - чтобы сделать угол от 0 до 2 * pi | |
| true_anomaly = np.mod(2 * np.arctan(coeff * np.tan(eccentric_position / 2)), 2 * np.pi) | |
| radius_peri_module = self.radius_coeff / (1 + self.e * np.cos(true_anomaly)) | |
| radius_peri = radius_peri_module * np.array([np.cos(true_anomaly), np.sin(true_anomaly), 0]) | |
| vel_coeff = self.mu / self.initial.angular_momentum_module | |
| velocity_peri = vel_coeff * np.array([-np.sin(true_anomaly), self.e + np.cos(true_anomaly), 0]) | |
| numerator = - 3 / 2 * (np.sqrt(self.mu) * self.j2 * self.earth_radius ** 2) * np.cos(self.initial.inclination) | |
| ascending_rate = numerator / ((1 - self.e ** 2) ** 2 * self.semimajor ** (7/2)) | |
| final_ascending = self.initial.right_ascention + ascending_rate * seconds | |
| periapsis_rate = ascending_rate * (5 / 2 * np.sin(self.initial.inclination) ** 2 - 2) / np.cos(self.initial.inclination) | |
| final_periapsis = self.initial.argument_of_perigee + periapsis_rate * seconds | |
| incl = self.initial.inclination | |
| first_tr = np.array([ | |
| [np.cos(final_periapsis), np.sin(final_periapsis), 0], | |
| [-np.sin(final_periapsis), np.cos(final_periapsis), 0], | |
| [0, 0, 1], | |
| ]) | |
| second_tr = np.array([ | |
| [1, 0 , 0], | |
| [0, np.cos(incl), np.sin(incl)], | |
| [0, -np.sin(incl), np.cos(incl)], | |
| ]) | |
| third_tr = np.array([ | |
| [np.cos(final_ascending), np.sin(final_ascending), 0], | |
| [-np.sin(final_ascending), np.cos(final_ascending), 0], | |
| [0, 0, 1], | |
| ]) | |
| transition_m = (first_tr @ second_tr @ third_tr).T | |
| return OrbitStateVector(transition_m @ radius_peri, transition_m @ velocity_peri) | |
| initial_state = OrbitStateVector([-5813, 2359, -2612.9], [ 0.071575724, -5.59614959 * 1.15, -5.230611 * 1.15]) | |
| if (initial_state.eccentricity_module > 1.0): | |
| print(f"orbit must be elliptic. Current ecc.: {initial_state.eccentricity_module}") | |
| sys.exit(1) | |
| solver = OblatenessTimeSolver(initial_state) | |
| final_state = solver.state_after(96 * 60 * 60) | |
| fig = plt.figure(figsize=(8, 8)) | |
| ax = fig.add_subplot(111, projection='3d') | |
| t_min, t_max = 0, 96 * 1 * 3600 | |
| dt = 300 | |
| t = np.arange(t_min, t_max, dt) | |
| data = np.zeros((3, int(t_max / dt) + 1)) | |
| def update(i): | |
| ax.cla() | |
| r = solver.earth_radius | |
| u, v = np.mgrid[0: 2 * np.pi: 25j, 0: np.pi: 25j] | |
| x = r * np.cos(u) * np.sin(v) | |
| y = r * np.sin(u) * np.sin(v) | |
| z = r * np.cos(v) | |
| # Plot sphere | |
| ax.plot_surface(x, y, z, color='b', zorder=1) | |
| art3d.pathpatch_2d_to_3d(plt.Circle((0, 0), r), z=0, zdir='z') | |
| time_str = datetime.datetime.utcfromtimestamp(i * dt).strftime('day, %H:%M:%S') | |
| day = datetime.datetime.utcfromtimestamp(i * dt).strftime('%d') | |
| fig.suptitle(f"Time elapsed: {int(day) - 1} {time_str}") | |
| ax.set_xlim([-11000, 11000]) | |
| ax.set_ylim([-11000, 11000]) | |
| ax.set_zlim([-11000, 11000]) | |
| ax.set_box_aspect([1, 1, 1]) | |
| ticks = np.arange(-9000, 11000, 9000) | |
| ax.set_xticks(ticks) | |
| ax.set_yticks(ticks) | |
| ax.set_zticks(ticks) | |
| ax.set_xlabel('X') | |
| ax.set_ylabel('Y') | |
| ax.set_zlabel('Z') | |
| ax.set_title('3D Animation') | |
| r = solver.state_after(dt * i).radius | |
| for j in range(3): | |
| data[j][i] = r[j] | |
| ax.plot(data[0][:i + 1], data[1][:i + 1], data[2][:i + 1], c='r', zorder=15) | |
| ax.scatter(r[0], r[1], r[2], c='black', marker='o', zorder=30) | |
| ani = FuncAnimation(fig, update, frames=len(t), interval=0.001) | |
| plt.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment