""" Symbolic, auto differentiated implementation of the STTS in [1] References: [1] https://arc.aiaa.org/doi/full/10.2514/1.G003897 """ import itertools import numpy as np from sympy import Matrix, lambdify, diff, exp, symbols from typing import List, Tuple, Dict, Callable, Union from sympy import cos, sin, pi from sympy import symbols, Matrix, exp import time #################################################################### # EQUATION OF MOTION #################################################################### # Declare symbolic variables for position, velocity, and constants x, y, z, vx, vy, vz, t = symbols("x y z vx vy vz t", real=True) GM_Earth, GM_Moon, CD, S, m, den, rho0, R_Earth, r0, J2 = symbols("GM_Earth GM_Moon CD S m den rho0 R_Earth r0 J2") # Symbolic position and velocity vectors r = Matrix([x, y, z]) v = Matrix([vx, vy, vz]) # Density function (symbolic) rho_sym = rho0 * exp(-(r.norm() - r0) / den) # Position of Moon, simple circular orbit (symbolic) pos_moon = Matrix([384400e3 * cos(2 * pi * t / (27.3 * 24 * 3600)), 384400e3 * sin(2 * pi * t / (27.3 * 24 * 3600)), 0]) # Relative positions (symbolic) sym_rel_moon = r - pos_moon sym_rel_earth = r - Matrix([0, 0, 0]) # Gravitational accelerations (symbolic) sym_accel_earth = -GM_Earth * sym_rel_earth / (sym_rel_earth.norm() ** 3) sym_accel_moon = -GM_Moon * sym_rel_moon / (sym_rel_moon.norm() ** 3) # Drag acceleration (symbolic) sym_drag = -0.5 * rho_sym * v.norm() * CD * S / m * v # J2 Perturbation (symbolic) r_norm = r.norm() factor = 1.5 * J2 * (R_Earth / r_norm) ** 2 sym_j2_accel = factor * Matrix([ (5 * (z / r_norm) ** 2 - 1) * x, (5 * (z / r_norm) ** 2 - 1) * y, (5 * (z / r_norm) ** 2 - 3) * z ]) * GM_Earth / (r_norm ** 3) # Total acceleration (symbolic) sym_accel = sym_accel_earth + sym_accel_moon + sym_j2_accel + sym_drag # Final equations of motion (symbolic) sym_eom = Matrix([vx, vy, vz, sym_accel[0], sym_accel[1], sym_accel[2]]) # Cache to store lambdas lambda_cache = {} #################################################################### # STTS DERIVATIVE EQUATIONS [Ref 1, Eq. 5] #################################################################### def compute_second_order_terms(f_i_j, f_i_jk, x_i_j, x_i_jk): return np.einsum("irs,rj,sk->ijk", f_i_jk, x_i_j, x_i_j) + np.einsum("ir,rjk->ijk", f_i_j, x_i_jk) def compute_third_order_terms(f_i_j, f_i_jk, f_i_jkl, x_i_j, x_i_jk, x_i_jkl): first_term = np.einsum("irsp,rj,sk,pl->ijkl", f_i_jkl, x_i_j, x_i_j, x_i_j) second_term_part1 = np.einsum("irs,rjl,sk->ijkl", f_i_jk, x_i_jk, x_i_j) second_term_part2 = np.einsum("irs,rj,skl->ijkl", f_i_jk, x_i_j, x_i_jk) second_term_part3 = np.einsum("irs,rjk,sl->ijkl", f_i_jk, x_i_jk, x_i_j) second_term = second_term_part1 + second_term_part2 + second_term_part3 third_term = np.einsum("ir,rjkl->ijkl", f_i_j, x_i_jkl) xdot_i_jkl = first_term + second_term + third_term return xdot_i_jkl def compute_fourth_order_terms(f_i_j, f_i_jk, f_i_jkl, f_i_jklm, x_i_j, x_i_jk, x_i_jkl, x_i_jklm): # Rk4 for 4th-order tensor first_term_4th = np.einsum( "irspq,rj,sk,pl,qm->ijklm", f_i_jklm, x_i_j, x_i_j, x_i_j, x_i_j ) second_term_part1_4th = np.einsum( "irsp,rjm,sk,pl->ijklm", f_i_jkl, x_i_jk, x_i_j, x_i_j ) second_term_part2_4th = np.einsum( "irsp,rj,skm,sl->ijklm", f_i_jkl, x_i_j, x_i_jk, x_i_j ) second_term_part3_4th = np.einsum( "irsp,rj,sk,plm->ijklm", f_i_jkl, x_i_j, x_i_j, x_i_jk ) third_term_part1_4th = np.einsum("irsp,rjl,sk,pm->ijklm", f_i_jkl, x_i_jk, x_i_j, x_i_j) third_term_part2_4th = np.einsum("irsp,rj,skl,pm->ijklm", f_i_jkl, x_i_j, x_i_jk, x_i_j) third_term_part3_4th = np.einsum("irsp,rjk,sl,pm->ijklm", f_i_jkl, x_i_jk, x_i_j, x_i_j) fourth_term_part1_4th = np.einsum("irs,rjlm,sk->ijklm", f_i_jk, x_i_jkl, x_i_j) fourth_term_part2_4th = np.einsum("irs,rjl,skm->ijklm", f_i_jk, x_i_jk, x_i_jk) fourth_term_part3_4th = np.einsum("irs,rjm,skl->ijklm", f_i_jk, x_i_jk, x_i_jk) fourth_term_part4_4th = np.einsum("irs,rj,sklm->ijklm", f_i_jk, x_i_j, x_i_jkl) fourth_term_part5_4th = np.einsum("irs,rjkm,sl->ijklm", f_i_jk, x_i_jkl, x_i_j) fourth_term_part6_4th = np.einsum("irs,rjk,slm->ijklm", f_i_jk, x_i_jk, x_i_jk) fourth_term_part7_4th = np.einsum("irs,rjkl,sm->ijklm", f_i_jk, x_i_jkl, x_i_j) fifth_term_4th = np.einsum("ir,rjklm->ijklm", f_i_j, x_i_jklm) return ( first_term_4th + second_term_part1_4th + second_term_part2_4th + second_term_part3_4th + third_term_part1_4th + third_term_part2_4th + third_term_part3_4th + fourth_term_part1_4th + fourth_term_part2_4th + fourth_term_part3_4th + fourth_term_part4_4th + fourth_term_part5_4th + fourth_term_part6_4th + fourth_term_part7_4th + fifth_term_4th ) #################################################################### # UTILITY FOR CACHING SYMBOLIC DIFFERENTIATION FUNCTIONS #################################################################### def get_or_create_lambda(eom: List, wrt_tuple: Tuple, vars: List) -> Callable: """ Get or create lambda functions for given equations of motion (eom) and variables. Parameters ---------- eom : list List of equations of motion. wrt_tuple : tuple Tuple of variables with respect to which differentiation will be done. wrt : list List of all variables. Returns ------- function Lambda function corresponding to the differentiated equation. """ key = (tuple(eom), wrt_tuple) if key not in lambda_cache: diff_expression = eom for var in wrt_tuple: diff_expression = diff(diff_expression, var) f_lambda = lambdify(vars, diff_expression, modules="numpy") lambda_cache[key] = f_lambda return lambda_cache[key] #################################################################### # CONSTRUCTION OF TENSORS #################################################################### def compute_tensors( eom: List, wrt: List, at: Dict[str, Union[float, np.ndarray]], vars: List = None, order: int = 2, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Compute tensors up to a given order for the provided equations of motion (eom). Parameters ---------- eom : list List of equations of motion. wrt : list List of variables with respect to which the tensor will be calculated. at : dict Dictionary containing the points at which to evaluate the tensors. order : int The highest order tensor to compute. Returns ------- tuple Tensors up to the specified order. """ p, q = len(eom), len(wrt) at_str = {str(k): v for k, v in at.items()} first_order = None second_order = None third_order = None fourth_order = None # First order if order >= 1: first_order = np.zeros((p, q)) for j in range(q): f_lambda = get_or_create_lambda(eom, (wrt[j],), vars) # print(f_lambda(**at_str).flatten()) first_order[:, j] = f_lambda(**at_str).flatten() # Second order if order >= 2: second_order = np.zeros((p, q, q)) for j in range(q): for k in range(j + 1): f_lambda = get_or_create_lambda(eom, (wrt[j], wrt[k]), vars) second_order[:, k, j] = second_order[:, j, k] = f_lambda( **at_str ).flatten() # Third order if order >= 3: third_order = np.zeros((p, q, q, q)) for j in range(q): for k in range(j + 1): for l in range(k + 1): f_lambda = get_or_create_lambda(eom, (wrt[j], wrt[k], wrt[l]), vars) base_value = f_lambda(**at_str).flatten() for perm in itertools.permutations([j, k, l]): third_order[:, perm[0], perm[1], perm[2]] = base_value # Fourth order if order >= 4: fourth_order = np.zeros((p, q, q, q, q)) for j in range(q): for k in range(j + 1): for l in range(k + 1): for m in range(l + 1): f_lambda = get_or_create_lambda( eom, (wrt[j], wrt[k], wrt[l], wrt[m]), vars ) base_value = f_lambda(**at_str).flatten() for perm in itertools.permutations([j, k, l, m]): fourth_order[ :, perm[0], perm[1], perm[2], perm[3] ] = base_value return first_order, second_order, third_order, fourth_order #################################################################### # RUNGE-KUTTA 4TH ORDER METHOD WITH TENSOR PROPAGATION #################################################################### def rk4_step( func, sym_func, state, time, x_i_j, x_i_jk, x_i_jkl, x_i_jklm, dt, at, wrt, vars, order=3 ): """Runge-Kutta 4th order method with tensor propagation.""" next_x_i_j = None next_x_i_jk = None next_x_i_jkl = None next_x_i_jklm = None # Pre-allocate memory for xdot terms if order >= 1: xdot_i_j = np.zeros_like(x_i_j) if order >= 2: xdot_i_jk = np.zeros_like(x_i_jk) if order >= 3: xdot_i_jkl = np.zeros_like(x_i_jkl) if order >= 4: xdot_i_jklm = np.zeros_like(x_i_jklm) def get_at(_state, _t): return { **at, **{ x: _state[0], y: _state[1], z: _state[2], vx: _state[3], vy: _state[4], vz: _state[5], t: _t, }, } # Standard RK4 for state k1 = func(*state, time).flatten() k2 = func(*state + k1 / 2 * dt, time + dt / 2).flatten() k3 = func(*state + k2 / 2 * dt, time + dt / 2).flatten() k4 = func(*state + k3 * dt, time + dt).flatten() next_x = state + ((k1 + 2 * k2 + 2 * k3 + k4) / 6) * dt # Compute the Jacobian tensors at the RK4 stages f_i_j_1, f_i_jk_1, f_i_jkl_1, f_i_jklm_1 = compute_tensors( sym_func, wrt, get_at(state, time), vars, order ) f_i_j_2, f_i_jk_2, f_i_jkl_2, f_i_jklm_2 = compute_tensors( sym_func, wrt, get_at(state + k1 / 2 * dt, time + dt / 2), vars, order ) f_i_j_3, f_i_jk_3, f_i_jkl_3, f_i_jklm_3 = compute_tensors( sym_func, wrt, get_at(state + k2 / 2 * dt, time + dt / 2), vars, order ) f_i_j_4, f_i_jk_4, f_i_jkl_4, f_i_jklm_4 = compute_tensors( sym_func, wrt, get_at(state + k3 * dt, time + dt), vars, order ) if order >= 1: # RK4 for 1st-order tensor # First stage x_i_j_1 = x_i_j xdot_i_j_1 = np.einsum("ij,jk->ik", f_i_j_1, x_i_j_1) # Second stage x_i_j_2 = x_i_j + xdot_i_j_1 * dt / 2 xdot_i_j_2 = np.einsum("ij,jk->ik", f_i_j_2, x_i_j_2) # Third stage x_i_j_3 = x_i_j + xdot_i_j_2 * dt / 2 xdot_i_j_3 = np.einsum("ij,jk->ik", f_i_j_3, x_i_j_3) # Fourth stage x_i_j_4 = x_i_j + xdot_i_j_3 * dt xdot_i_j_4 = np.einsum("ij,jk->ik", f_i_j_4, x_i_j_4) # Combine next_x_i_j = x_i_j + ((xdot_i_j_1 + 2 * xdot_i_j_2 + 2 * xdot_i_j_3 + xdot_i_j_4) / 6) * dt if order >= 2: # RK4 for 2nd-order tensor # First stage x_i_jk_1 = x_i_jk xdot_i_jk_1 = compute_second_order_terms(f_i_j_1, f_i_jk_1, x_i_j_1, x_i_jk_1) # Second stage x_i_jk_2 = x_i_jk + xdot_i_jk_1 * dt / 2 xdot_i_jk_2 = compute_second_order_terms(f_i_j_2, f_i_jk_2, x_i_j_2, x_i_jk_2) # Third stage x_i_jk_3 = x_i_jk + xdot_i_jk_2 * dt / 2 xdot_i_jk_3 = compute_second_order_terms(f_i_j_3, f_i_jk_3, x_i_j_3, x_i_jk_3) # Fourth stage x_i_jk_4 = x_i_jk + xdot_i_jk_3 * dt xdot_i_jk_4 = compute_second_order_terms(f_i_j_4, f_i_jk_4, x_i_j_4, x_i_jk_4) # Combine next_x_i_jk = x_i_jk + ((xdot_i_jk_1 + 2 * xdot_i_jk_2 + 2 * xdot_i_jk_3 + xdot_i_jk_4) / 6) * dt if order >= 3: # RK4 for 3rd-order tensor # First stage x_i_jkl_1 = x_i_jkl xdot_i_jkl_1 = compute_third_order_terms(f_i_j_1, f_i_jk_1, f_i_jkl_1, x_i_j_1, x_i_jk_1, x_i_jkl_1) # Second stage x_i_jkl_2 = x_i_jkl + xdot_i_jkl_1 * dt / 2 xdot_i_jkl_2 = compute_third_order_terms(f_i_j_2, f_i_jk_2, f_i_jkl_2, x_i_j_2, x_i_jk_2, x_i_jkl_2) # Third stage x_i_jkl_3 = x_i_jkl + xdot_i_jkl_2 * dt / 2 xdot_i_jkl_3 = compute_third_order_terms(f_i_j_3, f_i_jk_3, f_i_jkl_3, x_i_j_3, x_i_jk_3, x_i_jkl_3) # Fourth stage x_i_jkl_4 = x_i_jkl + xdot_i_jkl_3 * dt xdot_i_jkl_4 = compute_third_order_terms(f_i_j_4, f_i_jk_4, f_i_jkl_4, x_i_j_4, x_i_jk_4, x_i_jkl_4) # Combine all stages next_x_i_jkl = x_i_jkl + ((xdot_i_jkl_1 + 2 * xdot_i_jkl_2 + 2 * xdot_i_jkl_3 + xdot_i_jkl_4) / 6) * dt if order >= 4: # Rk4 for 4th-order tensor # First stage x_i_jklm_1 = x_i_jklm xdot_i_jklm_1 = compute_fourth_order_terms(f_i_j_1, f_i_jk_1, f_i_jkl_1, f_i_jklm_1, x_i_j_1, x_i_jk_1, x_i_jkl_1, x_i_jklm_1) # Second stage x_i_jklm_2 = x_i_jklm + xdot_i_jklm_1 * dt / 2 xdot_i_jklm_2 = compute_fourth_order_terms(f_i_j_2, f_i_jk_2, f_i_jkl_2, f_i_jklm_2, x_i_j_2, x_i_jk_2, x_i_jkl_2, x_i_jklm_2) # Third stage x_i_jklm_3 = x_i_jklm + xdot_i_jklm_2 * dt / 2 xdot_i_jklm_3 = compute_fourth_order_terms(f_i_j_3, f_i_jk_3, f_i_jkl_3, f_i_jklm_3, x_i_j_3, x_i_jk_3, x_i_jkl_3, x_i_jklm_3) # Fourth stage x_i_jklm_4 = x_i_jklm + xdot_i_jklm_3 * dt xdot_i_jklm_4 = compute_fourth_order_terms(f_i_j_4, f_i_jk_4, f_i_jkl_4, f_i_jklm_4, x_i_j_4, x_i_jk_4, x_i_jkl_4, x_i_jklm_4) # Combine all stages next_x_i_jklm = x_i_jklm + ((xdot_i_jklm_1 + 2 * xdot_i_jklm_2 + 2 * xdot_i_jklm_3 + xdot_i_jklm_4) / 6) * dt return next_x, next_x_i_j, next_x_i_jk, next_x_i_jkl, next_x_i_jklm # Variables with respect to which differentiation may be needed wrt = [x, y, z, vx, vy, vz] vars = [x, y, z, vx, vy, vz, t] #################################################################### # CONSTANTS #################################################################### # Constants as a dictionary to replace later const_values = { GM_Earth: 3.986004418e14, GM_Moon: 4.9048695e12, CD: 2.0, S: 2000.0, m: 5000.0, den: 100000, rho0: 3.8 * 10 ** -12, R_Earth: 6378.137e3, r0: 6771e3, J2: 1.08263e-3 } # Substitute the constants when you need numerical evaluation # sym_eom.subs(const_values) sym_eom = sym_eom.subs(const_values) eom = lambdify(vars, sym_eom, modules="numpy") #################################################################### # INITIAL CONDITIONS #################################################################### H = 300e3 r_p = 6371e3 + H r0 = np.array([r_p, 0, 0]) ecc = 0.02 sma = np.linalg.norm(r_p) / (1 - ecc) v_p = np.sqrt(const_values[GM_Earth] * (2 / np.linalg.norm(r_p) - 1 / sma)) v0 = np.array([0, 0, 1]) * v_p state0 = np.concatenate((r0, v0)) #################################################################### # SIMULATION #################################################################### up_to_order = 3 t_orbit = 2 * np.pi * np.sqrt(sma ** 3 / const_values[GM_Earth]) t_now = 0 dt = 100 sst1_j0 = np.zeros((len(wrt), len(wrt))) np.fill_diagonal(sst1_j0, 1) sst1_k0 = np.zeros((len(wrt), len(wrt))) sst2_jk0 = np.zeros((len(wrt), len(wrt), len(wrt))) sst3_jkl0 = np.zeros((len(wrt), len(wrt), len(wrt), len(wrt))) sst4_jklm0 = np.zeros((len(wrt), len(wrt), len(wrt), len(wrt), len(wrt))) # Initialize sst1_j = sst1_j0 sst2_jk = sst2_jk0 sst3_jkl = sst3_jkl0 sst4_jklm = sst4_jklm0 trajectory = [] relevant_keys = ["x", "y", "z", "vx", "vy", "vz"] jd = 86400 t_end = t_orbit * 10.0 # t_end = jd * 10 start = time.time() state = state0 while t_now < t_end: at = { x: state[0], y: state[1], z: state[2], vx: state[3], vy: state[4], vz: state[5], } # Update the next_state, next_sst1_j, next_sst2_jk state, sst1_j, sst2_jk, sst3_jkl, sst4_jklm = rk4_step( eom, sym_eom, state, t_now, sst1_j, sst2_jk, sst3_jkl, sst4_jklm, dt, at, wrt, vars, order=up_to_order ) # Record the state trajectory.append(state) # Update the time t_now += dt end = time.time() print(f"Time elapsed: {end - start} seconds") #################################################################### # SAMPLE PERTURBATIONS #################################################################### v_sigma = 0.1 r_sigma = 10e3 n_perturbations = 200 dx0 = np.zeros((6, n_perturbations)) np.random.seed(0) # reproducibility dx0 = np.random.multivariate_normal( np.zeros(6), np.diag([r_sigma, r_sigma, r_sigma, v_sigma, v_sigma, v_sigma]) ** 2, n_perturbations, ).T #################################################################### # TRANSFORM PERTURBATIONS #################################################################### # First order perturbations dx_1 = np.einsum("ij,js->is", sst1_j, dx0) # Second order; dimensions: # sst2_jk: (6, 6, 6) (i, j, k) # 1st dx0: (6, n_perturbations) (j, s) # 2nd dx0: (6, n_perturbations) (k, s) # dx_1: (6, n_perturbations) (i, s) dx_2 = np.einsum("ijk,js,ks->is", sst2_jk, dx0, dx0) / 2.0 + dx_1 # Third order dx_3 = np.einsum("ijkl,js,ks,ls->is", sst3_jkl, dx0, dx0, dx0) / 6.0 + dx_2 # Forth order if up_to_order >= 4: dx_4 = np.einsum("ijklm,js,ks,ls,ms->is", sst4_jklm, dx0, dx0, dx0, dx0) / 24.0 + dx_3 else: dx_4 = np.zeros((6, n_perturbations)) #################################################################### # NUMERICALLY PROPAGATE PERTURBATIONS FOR BASELINE COMPARISON #################################################################### # List to hold all perturbed trajectories perturbed_trajectories = [] # Iterate through each perturbation for pert_idx in range(n_perturbations): perturbed_state = np.array(state0) + dx0[:, pert_idx] perturbed_trajectory = [perturbed_state] # Initialize the time and state for the perturbed simulation t_pert = 0.0 while t_pert < t_end: at_pert = { x: perturbed_state[0], y: perturbed_state[1], z: perturbed_state[2], vx: perturbed_state[3], vy: perturbed_state[4], vz: perturbed_state[5], } # Update the next_state, for the perturbed trajectory perturbed_state, _, _, _, _ = rk4_step( eom, sym_eom, perturbed_state, t_pert, sst1_j, sst2_jk, sst3_jkl, sst4_jklm, dt, at_pert, wrt, vars, order=0, ) # Record the perturbed state perturbed_trajectory.append(perturbed_state) # Update the time t_pert += dt perturbed_trajectories.append(perturbed_trajectory) final_perturbed_states = np.array( [trajectory[-1] for trajectory in perturbed_trajectories] ) trajectory = np.array(trajectory) #################################################################### # PLOTTING #################################################################### # Assume trajectory is (n, m), n time steps and m dimensions n, m = trajectory.shape # Position Perturbations axis_1 = 0 axis_2 = 2 axis_1_name = relevant_keys[axis_1] axis_2_name = relevant_keys[axis_2] numerical_dx = final_perturbed_states - trajectory[-1, :] plot_with = 'matplotlib' if plot_with == 'matplotlib': import numpy as np import matplotlib.pyplot as plt import matplotlib # Define conversion factor meters_to_km = 1e-3 matplotlib.use("QtAgg") nrows = 3 ncols = 4 if up_to_order == 4 else 3 fig, axes = plt.subplots(nrows, ncols, figsize=(18, 18), constrained_layout=True) marker_config = { "Initial Perturbations": {"c": "black", "s": 10, "alpha": 0.5}, "SST1 Transformed": {"c": "red", "marker": "x", "s": 10, "alpha": 0.5}, "SST2 Transformed": {"c": "blue", "s": 4, "alpha": 0.5}, "SST3 Transformed": {"c": "green", "s": 10, "marker": "+", "alpha": 0.6}, "SST4 Transformed": {"c": "orange", "s": 10, "marker": "o", "alpha": 0.5}, "Numerically Perturbed": {"c": "purple", "marker": "s", "s": 10, "alpha": 0.5}, } data_dicts = [ {"dx": dx0, "name": "Initial Perturbations"}, {"dx": dx_1, "name": "SST1 Transformed"}, {"dx": dx_2, "name": "SST2 Transformed"}, {"dx": dx_3, "name": "SST3 Transformed"}, ] if up_to_order >= 4: data_dicts.append({"dx": dx_4, "name": "SST4 Transformed"}) data_dicts.append({"dx": numerical_dx.T, "name": "Numerically Perturbed"}) for ax_row in axes: for ax in ax_row: ax.grid(True, linestyle='--', linewidth=0.5, alpha=0.7) # Light grid for i, ax_title, xlabel, ylabel in zip( range(ncols), ["Position Perturbations (km)", "Velocity Perturbations (km/s)", "Full Trajectory (km)"], [f"{axis_1_name.upper()} Position (km)", f"{axis_1_name.upper()} Velocity (km/s)", f"{axis_1_name.upper()} Position (km)"], [f"{axis_2_name.upper()} Position (km)", f"{axis_2_name.upper()} Velocity (km/s)", f"{axis_2_name.upper()} Position (km)"] ): ax = axes[0, i] # First row if i == 2: ax.plot(trajectory[:, axis_1] * meters_to_km, trajectory[:, axis_2] * meters_to_km, c='black', label='Nominal Trajectory') for data_dict in data_dicts: dx = data_dict['dx'] * meters_to_km # Conversion to km or km/s name = data_dict['name'] config = marker_config[name] axis_offset = 0 if "Velocity" in ax_title: axis_offset = 3 # For velocity subplot x_data = dx[axis_1 + axis_offset, :] y_data = dx[axis_2 + axis_offset, :] if ax_title == "Full Trajectory (km)": # For the full trajectory subplot x_data += trajectory[-1, axis_1] * meters_to_km y_data += trajectory[-1, axis_2] * meters_to_km ax.scatter(x_data, y_data, **config, label=name) # ax.axis("equal") ax.set_xlabel(xlabel, labelpad=15) ax.set_ylabel(ylabel, labelpad=15) ax.legend() ax.set_title(ax_title, pad=20) # SST-specific error plots for j, error_title in zip(range(1, 3), ["Position Errors (km)", "Velocity Errors (km/s)"]): for i, name in zip(range(1, ncols + 1), ["SST1 Transformed", "SST2 Transformed", "SST3 Transformed", "SST4 Transformed"]): ax = axes[j, i - 1] axis_offset = 0 if "Velocity" in error_title: axis_offset = 3 # For velocity subplot dx = data_dicts[i]['dx'] * meters_to_km # Conversion to km or km/s numerical = data_dicts[-1]['dx'] * meters_to_km # Conversion to km or km/s error_x = dx[axis_1 + axis_offset, :] - numerical[axis_1 + axis_offset, :] error_y = dx[axis_2 + axis_offset, :] - numerical[axis_2 + axis_offset, :] ax.scatter(error_x, error_y, **marker_config[name], label=f"{name.split(' ')[0]} Error") ax.axis("equal") ax.set_xlabel(f"Error in {axis_1_name.upper()} (km)", labelpad=15) ax.set_ylabel(f"Error in {axis_2_name.upper()} (km)", labelpad=15) ax.legend() ax.set_title(name + " " + error_title, pad=20) plt.show() elif plot_with == "mayavi": from mayavi import mlab mlab.figure(size=(400, 300)) # Clear the current Mayavi scene # mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1)) data_norm = np.linalg.norm(trajectory, axis=0)[None, ...] print(data_norm) data_norm = np.ones((1, 6)) * np.max(data_norm) # Create the 3D scatter plot for the position perturbations # mlab.points3d(dx0 / data_norm.T, color=(0, 0, 0), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_1, color=(1, 0, 0), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_2, color=(0, 0, 1), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_3, color=(0, 1, 0), scale_factor=0.1, opacity=0.5) # # Create the 3D scatter plot for the velocity perturbations # mlab.points3d(dx0[3 + axis_1, :], dx0[3 + axis_2, :], color=(0, 0, 0), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_1[3 + axis_1, :], dx_1[3 + axis_2, :], color=(1, 0, 0), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_2[3 + axis_1, :], dx_2[3 + axis_2, :], color=(0, 0, 1), scale_factor=0.1, opacity=0.5) # mlab.points3d(dx_3[3 + axis_1, :], dx_3[3 + axis_2, :], color=(0, 1, 0), scale_factor=0.1, opacity=0.5) # # x = trajectory[:, 0] # y = trajectory[:, 1] # z = trajectory[:, 2] # s = np.linspace(0.5, 1.0, len(trajectory)) # # Create the 3D plot for the full trajectory with position perturbations # mlab.plot3d(x, y, z) # n_mer, n_long = 6, 11 # dphi = np.pi / 1000.0 # phi = np.arange(0.0, 2 * np.pi + 0.5 * dphi, dphi) # mu = phi * n_mer # x = np.cos(mu) * (1 + np.cos(n_long * mu / n_mer) * 0.5) # y = np.sin(mu) * (1 + np.cos(n_long * mu / n_mer) * 0.5) # z = np.sin(n_long * mu / n_mer) * 0.5 # l = mlab.plot3d(x, y, z, np.sin(mu), tube_radius=0.025, colormap='Spectral') # plot test # mlab.points3d([0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]) # every third point # normalize trajectory dimensions to [-1, 1] trajectory = np.array(trajectory) trajectory = trajectory / data_norm # trajectory = trajectory[::3, :] # print(len(trajectory)) x = trajectory[:, 0] y = trajectory[:, 1] z = trajectory[:, 2] s = np.linspace(0.5, 1.0, len(trajectory)) # mlab.points3d(x, y, z, s, colormap="Spectral", scale_factor=0.1, tube_radius=0.025) dx0 /= data_norm.T dx_1 /= data_norm.T dx_1 += trajectory[-1, :][..., None] dx_2 /= data_norm.T dx_2 += trajectory[-1, :][..., None] dx_3 /= data_norm.T dx_3 += trajectory[-1, :][..., None] numerical_dx /= data_norm numerical_dx += trajectory[-1, :][None, ...] # mlab.points3d(*dx0[:3,:], color=(0, 0, 0), scale_factor=0.001, opacity=0.5) mlab.points3d(*dx_1[:3, :], color=(1, 0, 0), scale_factor=0.001, opacity=0.1) mlab.points3d(*dx_2[:3, :], color=(0, 0, 1), scale_factor=0.001, opacity=0.1) mlab.points3d(*dx_3[:3, :], color=(0, 1, 0), scale_factor=0.001, opacity=0.1) mlab.points3d( *numerical_dx[:, :3].T, color=(0.5, 0, 0.5), scale_factor=0.001, opacity=0.1 ) mlab.plot3d(x, y, z, s, colormap="Spectral", tube_radius=0.0005, opacity=0.5) # Add other features like legends, titles, etc., using mlab.text3d(), mlab.title(), etc. # Show the 3D plot # mlab.view(azimuth=0, elevation=90, distance=1000) mlab.show()