|
| 1 | +# pylint: disable=too-many-locals, too-many-statements, no-member |
| 2 | +""" |
| 3 | +The Autonomous Cooperative Consensus Orbit Determination (ACCORD) framework. |
| 4 | +Author: Beth Probert |
| 5 | +Email: beth.probert@strath.ac.uk |
| 6 | +
|
| 7 | +Copyright (C) 2025 Applied Space Technology Laboratory |
| 8 | +
|
| 9 | +This program is free software: you can redistribute it and/or modify |
| 10 | +it under the terms of the GNU General Public License as published by |
| 11 | +the Free Software Foundation, either version 3 of the License, or |
| 12 | +(at your option) any later version. |
| 13 | +
|
| 14 | +This program is distributed in the hope that it will be useful, |
| 15 | +but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 16 | +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 17 | +GNU General Public License for more details. |
| 18 | +
|
| 19 | +You should have received a copy of the GNU General Public License |
| 20 | +along with this program. If not, see <http://www.gnu.org/licenses/>. |
| 21 | +
|
| 22 | +""" |
| 23 | + |
| 24 | +import os |
| 25 | +import numpy as np |
| 26 | +import matplotlib.pyplot as plt |
| 27 | +from matplotlib.animation import FuncAnimation, PillowWriter |
| 28 | + |
| 29 | + |
| 30 | +def visualise_orbits(data_path='sim_data/sim_results.npz', |
| 31 | + output_gif='images/orbits.gif', frames=200, interval=50) -> None: |
| 32 | + """ |
| 33 | + Creates an animated GIF of satellite orbits around Earth. |
| 34 | +
|
| 35 | + Args: |
| 36 | + data_path: Path to the .npz file containing 'truth' data. |
| 37 | + output_gif: Path to save the resulting GIF. |
| 38 | + frames: Number of frames to animate (subsamples the data if less than total steps). |
| 39 | + interval: Delay between frames in milliseconds. Default 50ms for smoother motion. |
| 40 | +
|
| 41 | + Returns: |
| 42 | + None. Saves the animation as a GIF at the specified location. |
| 43 | + """ |
| 44 | + if not os.path.exists(data_path): |
| 45 | + print(f"Error: Data file {data_path} not found.") |
| 46 | + return |
| 47 | + |
| 48 | + # Load data |
| 49 | + print(f"Loading data from {data_path}...") |
| 50 | + data = np.load(data_path, allow_pickle=True) |
| 51 | + truth = data['truth'] # Shape: (steps, 6*N) |
| 52 | + |
| 53 | + steps, state_size = truth.shape |
| 54 | + n_sats = state_size // 6 |
| 55 | + print(f"Loaded {n_sats} satellites over {steps} time steps.") |
| 56 | + |
| 57 | + # Subsample frame indices |
| 58 | + frames = min(frames, steps) |
| 59 | + |
| 60 | + indices = np.linspace(0, steps - 1, frames, dtype=int) |
| 61 | + |
| 62 | + # Setup 3D Plot |
| 63 | + fig = plt.figure(figsize=(12, 10)) |
| 64 | + ax = fig.add_subplot(111, projection='3d') |
| 65 | + |
| 66 | + # Earth model |
| 67 | + r_e = 6378e3 # Earth radius in meters |
| 68 | + u = np.linspace(0, 2 * np.pi, 50) |
| 69 | + v = np.linspace(0, np.pi, 50) |
| 70 | + x_earth = r_e * np.outer(np.cos(u), np.sin(v)) |
| 71 | + y_earth = r_e * np.outer(np.sin(u), np.sin(v)) |
| 72 | + z_earth = r_e * np.outer(np.ones(np.size(u)), np.cos(v)) |
| 73 | + ax.plot_surface(x_earth, y_earth, z_earth, color='blue', # type: ignore [attr-defined] |
| 74 | + alpha=0.1, rstride=2, cstride=2) # type: ignore [attr-defined] |
| 75 | + |
| 76 | + # Initialise satellite markers |
| 77 | + dots = [] |
| 78 | + trails = [] |
| 79 | + trail_len = 30 # Number of animation steps to show as a tail |
| 80 | + |
| 81 | + # Randomly select up to 50 satellites to visualise if N is too large for clear animation |
| 82 | + viz_n = min(n_sats, 50) |
| 83 | + viz_indices = np.random.choice(range(n_sats), viz_n, replace=False) |
| 84 | + |
| 85 | + colors = plt.cm.viridis(np.linspace(0, 1, viz_n)) # type: ignore [attr-defined] |
| 86 | + |
| 87 | + for i in range(viz_n): |
| 88 | + dot, = ax.plot([], [], [], 'o', color=colors[i], markersize=4) |
| 89 | + dots.append(dot) |
| 90 | + trail, = ax.plot([], [], [], '-', color=colors[i], alpha=0.4, linewidth=1.5) |
| 91 | + trails.append(trail) |
| 92 | + |
| 93 | + # Add text overlays |
| 94 | + _ = ax.text2D(0.05, 0.95, f"Satellites shown: {viz_n} / {n_sats}", # type: ignore [attr-defined] |
| 95 | + transform=ax.transAxes, fontsize=12, fontweight='bold') |
| 96 | + step_text = ax.text2D(0.05, 0.90, "", transform=ax.transAxes, # type: ignore [attr-defined] |
| 97 | + fontsize=12) # type: ignore [attr-defined] |
| 98 | + |
| 99 | + # Set plot limits |
| 100 | + # Calculate limits from first step positions to ensure Earth and orbits are visible |
| 101 | + max_dist = np.max(np.linalg.norm(truth[0, :3], axis=0)) |
| 102 | + limit = max_dist * 1.2 if max_dist > 0 else r_e * 2 |
| 103 | + |
| 104 | + ax.set_xlim(-limit, limit) |
| 105 | + ax.set_ylim(-limit, limit) |
| 106 | + ax.set_zlim(-limit, limit) # type: ignore [attr-defined] |
| 107 | + ax.set_xlabel('X (m)') |
| 108 | + ax.set_ylabel('Y (m)') |
| 109 | + ax.set_zlabel('Z (m)') # type: ignore [attr-defined] |
| 110 | + ax.set_title('Satellite Constellation Animation', fontsize=16) |
| 111 | + |
| 112 | + def init(): |
| 113 | + for dot, trail in zip(dots, trails): |
| 114 | + dot.set_data([], []) |
| 115 | + dot.set_3d_properties([]) |
| 116 | + trail.set_data([], []) |
| 117 | + trail.set_3d_properties([]) |
| 118 | + step_text.set_text("") |
| 119 | + return dots + trails + [step_text] |
| 120 | + |
| 121 | + def update(frame): |
| 122 | + # Update each satellite |
| 123 | + actual_step = indices[frame] |
| 124 | + |
| 125 | + # Calculate start index for trail in terms of ACTUAL steps to ensure smoothness |
| 126 | + start_frame_idx = max(0, frame - trail_len) |
| 127 | + start_step = indices[start_frame_idx] |
| 128 | + |
| 129 | + for i, sat_idx in enumerate(viz_indices): |
| 130 | + # Current position |
| 131 | + pos = truth[actual_step, sat_idx*6 : sat_idx*6+3] |
| 132 | + dots[i].set_data([pos[0]], [pos[1]]) |
| 133 | + dots[i].set_3d_properties([pos[2]]) |
| 134 | + |
| 135 | + # Trail using full resolution data between start_step and actual_step |
| 136 | + trail_pos = truth[start_step : actual_step + 1, sat_idx*6 : sat_idx*6+3] |
| 137 | + if len(trail_pos) > 0: |
| 138 | + trails[i].set_data(trail_pos[:, 0], trail_pos[:, 1]) |
| 139 | + trails[i].set_3d_properties(trail_pos[:, 2]) |
| 140 | + |
| 141 | + # Update step counter |
| 142 | + step_text.set_text(f"Timestep: {actual_step}") |
| 143 | + |
| 144 | + return dots + trails + [step_text] |
| 145 | + |
| 146 | + print("Creating animation...") |
| 147 | + ani = FuncAnimation(fig, update, frames=frames, init_func=init, blit=True, interval=interval) |
| 148 | + |
| 149 | + print(f"Saving to {output_gif}...") |
| 150 | + writer = PillowWriter(fps=1000 // interval) |
| 151 | + ani.save(output_gif, writer=writer) |
| 152 | + print("Done!") |
| 153 | + |
| 154 | +if __name__ == "__main__": |
| 155 | + # Check for existing data files |
| 156 | + potential_files = [ |
| 157 | + 'sim_data/sim_results.npz', |
| 158 | + 'sim_data/ekf_simulation_results.npz' |
| 159 | + ] |
| 160 | + |
| 161 | + CHOSEN_FILE = None |
| 162 | + for f in potential_files: |
| 163 | + if os.path.exists(f): |
| 164 | + CHOSEN_FILE = f |
| 165 | + break |
| 166 | + |
| 167 | + if CHOSEN_FILE: |
| 168 | + visualise_orbits(data_path=CHOSEN_FILE) |
| 169 | + else: |
| 170 | + print("Could not find any .npz simulation data in sim_data/.") |
0 commit comments