diff --git a/doc/sphinx/documentation/orbitplots.rst b/doc/sphinx/documentation/orbitplots.rst index 5999786a..1fa8a865 100644 --- a/doc/sphinx/documentation/orbitplots.rst +++ b/doc/sphinx/documentation/orbitplots.rst @@ -12,6 +12,7 @@ Name Type D :func:`pykep.orbit_plots.plot_kepler` function Plots a keplerian propagated arc :func:`pykep.orbit_plots.plot_taylor` function Plots a constant thrust propagated arc :func:`pykep.orbit_plots.plot_sf_leg` function Plots a Sims_Flanagan leg +:func:`pykep.orbit_plots.plot_distance_and_flybys function Plots distance from sun and flybys over time ======================================================== ========= ================================================ Detailed Documentation @@ -34,3 +35,7 @@ Detailed Documentation ------------ .. autofunction:: pykep.orbit_plots.plot_sf_leg(*args) + +------------ + +.. autofunction:: pykep.orbit_plots.plot_distance_and_flybys(*args) diff --git a/doc/sphinx/documentation/trajopt.rst b/doc/sphinx/documentation/trajopt.rst index f463b7f5..f3840c66 100644 --- a/doc/sphinx/documentation/trajopt.rst +++ b/doc/sphinx/documentation/trajopt.rst @@ -73,6 +73,10 @@ Detailed Documentation .. automethod:: pykep.trajopt.mga.plot(*args) + .. automethod:: pykep.trajopt.mga.get_eph_function(*args) + + .. automethod:: pykep.trajopt.mga.plot_distance_and_flybys(*args) + ------------ .. autoclass:: pykep.trajopt.mga_1dsm(*args) @@ -83,6 +87,10 @@ Detailed Documentation .. automethod:: pykep.trajopt.mga_1dsm.plot(*args) + .. automethod:: pykep.trajopt.mga.get_eph_function(*args) + + .. automethod:: pykep.trajopt.mga_1dsm.plot_distance_and_flybys(*args) + ------------ .. autoclass:: pykep.trajopt.pl2pl_N_impulses(*args) diff --git a/pykep/__init__.py b/pykep/__init__.py index 4a987aa0..fdb1c2a1 100644 --- a/pykep/__init__.py +++ b/pykep/__init__.py @@ -59,7 +59,7 @@ ########################################################### # We define pykep module ########################################################### -version = '2.6' +version = '2.6.1' __doc__ = 'pykep is the answer ... but what was the question?' __all__ = ['core', 'sims_flanagan', 'pontryagin', 'orbit_plots', 'examples', 'trajopt', 'phasing', 'util', 'planet', 'test'] diff --git a/pykep/orbit_plots/_plots.py b/pykep/orbit_plots/_plots.py index 78aaa213..5b09a7a8 100644 --- a/pykep/orbit_plots/_plots.py +++ b/pykep/orbit_plots/_plots.py @@ -595,3 +595,52 @@ def plot_sf_leg(leg, N=5, units=1, color='b', legend=False, plot_line=True, plot if axes is None: # show only if axis is not set plt.show() return ax + +def plot_flybys(seq, ep, eph_function, probename="Probe", axes=None, N: int=200, extension: float=300): + import matplotlib.pyplot as plt + import numpy as np + from pykep.core import epoch, AU + + if len(seq) != len(ep): + raise ValueError("Got sequence of length " + len(seq) + ", but " + len(ep) + " time epochs.") + + timeframe = np.linspace(0, ep[-1] - ep[0] + extension, N) + + planets = set(seq) + + distances = [] + pl_distances = {pl : [] for pl in planets} + + xeph = eph_function + + for day in timeframe: + # position of spacecraft + t = ep[0] + day + pos, vel = xeph(t) + distances.append(np.linalg.norm(pos) / AU) + + # positions of planets + for pl in planets: + ppos, _ = pl.eph(t) + pl_distances[pl].append(np.linalg.norm(ppos) / AU) + + # flyby markers + fl_times = list() + fl_distances = list() + for pl, t in zip(seq, ep): + fl_times.append(t - ep[0]) + pos, _ = pl.eph(t) + fl_distances.append(np.linalg.norm(pos) / AU) + + if axes is None: + fig, axes = plt.subplots() + axes.plot(list(timeframe), distances, label=probename) + for pl in planets: + axes.plot(list(timeframe), pl_distances[pl], label=pl.name.capitalize()) + plt.scatter(fl_times, fl_distances, marker="o", color="r") + axes.set_xlabel("Days") + axes.set_ylabel("AU") + axes.set_title("Distance to Sun") + axes.legend() + + return axes \ No newline at end of file diff --git a/pykep/test.py b/pykep/test.py index 6ed39361..e4a0a7e0 100644 --- a/pykep/test.py +++ b/pykep/test.py @@ -69,6 +69,7 @@ class mga_1dsm_test_case(_ut.TestCase): def runTest(self): self.run_construction_test() self.run_decode_times_and_vinf_test() + self.run_eph_function_test() def run_construction_test(self): from .trajopt import mga_1dsm @@ -131,6 +132,54 @@ def run_decode_times_and_vinf_test(self): self.assertAlmostEqual(retval[2], 0.) self.assertAlmostEqual(retval[3], 1.) + def run_eph_function_test(self): + from .trajopt import mga_1dsm + udp = mga_1dsm(tof_encoding='direct', tof=[[10, 400], [10, 400]]) + x = [10] + [0., 0., 1., 0.9, 123] + [0.4, 1.3, 0.5, 321] + retval = udp._decode_tofs(x) + eph = udp.get_eph_function(x) + # check closeness at first flyby + ep = x[0] + retval[0] + pos = eph(ep)[0] + pl_pos = udp._seq[1].eph(ep)[0] + self.assertLess(np.linalg.norm(np.array(pos) - np.array(pl_pos)), 0.01) + # check closeness at second flyby + ep = x[0] + sum(retval) + pos = eph(ep)[0] + pl_pos = udp._seq[2].eph(ep)[0] + self.assertLess(np.linalg.norm(np.array(pos) - np.array(pl_pos)), 0.01) + # check error for too early epoch + with self.assertRaises(ValueError): + eph(9) + +class mga_test_case(_ut.TestCase): + """Test case for the mga1_dsm class + + """ + + def runTest(self): + self.run_eph_function_test() + + def run_eph_function_test(self): + from .trajopt import mga + udp = mga(tof_encoding='direct', tof=[[10, 400], [10, 400]]) + x = [10] + [100, 250] + retval = udp._decode_tofs(x) + eph = udp.get_eph_function(x) + # check closeness at first flyby + ep = x[0] + retval[0] + pos = eph(ep)[0] + pl_pos = udp.seq[1].eph(ep)[0] + self.assertLess(np.linalg.norm(np.array(pos) - np.array(pl_pos)), 0.01) + # check closeness at second flyby + ep = x[0] + sum(retval) + pos = eph(ep)[0] + pl_pos = udp.seq[2].eph(ep)[0] + self.assertLess(np.linalg.norm(np.array(pos) - np.array(pl_pos)), 0.01) + # check error for too early epoch + with self.assertRaises(ValueError): + eph(9) + class gym_test_case(_ut.TestCase): """Test case for the gym @@ -419,6 +468,7 @@ def run_test_suite(level=0): suite = _ut.TestLoader().loadTestsFromTestCase(core_functions_test_case) suite.addTest(lambert_test_case()) suite.addTest(mga_1dsm_test_case()) + suite.addTest(mga_test_case()) suite.addTest(gym_test_case()) if __extensions__['numba']: diff --git a/pykep/trajopt/_mga.py b/pykep/trajopt/_mga.py index b5532171..28f20645 100644 --- a/pykep/trajopt/_mga.py +++ b/pykep/trajopt/_mga.py @@ -1,8 +1,10 @@ -from pykep.core import epoch, lambert_problem, DAY2SEC, fb_vel, AU +from pykep.core import epoch, lambert_problem, propagate_lagrangian, DAY2SEC, fb_vel, AU from pykep.planet import jpl_lp from pykep.trajopt._lambert import lambert_problem_multirev import numpy as np +from typing import Any, Dict, List, Tuple +from bisect import bisect_left class mga: @@ -149,7 +151,7 @@ def get_bounds(self): ub = [t0[1].mjd2000] + [1.0 - 1e-3] * (n_legs) return (lb, ub) - def _decode_tofs(self, x): + def _decode_tofs(self, x: List[float]) -> List[float]: if self.tof_encoding == 'alpha': # decision vector is [t0, T, a1, a2, ....] T = np.log(x[2:]) @@ -242,18 +244,31 @@ def direct2eta(self, x): retval[i] = x[i] / (self.tof - sum(x[1:i])) return retval - def _compute_dvs(self, x): + def _compute_dvs(self, x: List[float]) -> Tuple[ + float, # DVlaunch + List[float], # DVs + float, # DVarrival, + List[Any], # Lambert legs + float, #DVlaunch_tot + List[float], # T + List[Tuple[List[float], List[float]]], # ballistic legs + List[float], # epochs of ballistic legs + ]: # 1 - we 'decode' the times of flights and compute epochs (mjd2000) - T = self._decode_tofs(x) # [T1, T2 ...] + T: List[float] = self._decode_tofs(x) # [T1, T2 ...] ep = np.insert(T, 0, x[0]) # [t0, T1, T2 ...] ep = np.cumsum(ep) # [t0, t1, t2, ...] # 2 - we compute the ephemerides r = [0] * len(self.seq) v = [0] * len(self.seq) for i in range(len(self.seq)): - r[i], v[i] = self.seq[i].eph(ep[i]) - # 3 - we solve the lambert problems + r[i], v[i] = self.seq[i].eph(float(ep[i])) + l = list() + ballistic_legs: List[Tuple[List[float],List[float]]] = [] + ballistic_ep: List[float] = [] + + # 3 - we solve the lambert problems vi = v[0] for i in range(self._n_legs): lp = lambert_problem_multirev( @@ -261,6 +276,8 @@ def _compute_dvs(self, x): r[i], r[i + 1], T[i] * DAY2SEC, self._common_mu, False, self.max_revs)) l.append(lp) vi = lp.get_v2()[0] + ballistic_legs.append((r[i], lp.get_v1()[0])) + ballistic_ep.append(ep[i]) # 4 - we compute the various dVs needed at fly-bys to match incoming # and outcoming DVfb = list() @@ -282,11 +299,11 @@ def _compute_dvs(self, x): DVper2 = np.sqrt(2 * self.seq[-1].mu_self / self.rp_target - self.seq[-1].mu_self / self.rp_target * (1. - self.e_target)) DVarrival = np.abs(DVper - DVper2) - return (DVlaunch, DVfb, DVarrival, l, DVlaunch_tot) + return (DVlaunch, DVfb, DVarrival, l, DVlaunch_tot, T, ballistic_legs, ballistic_ep) # Objective function def fitness(self, x): - DVlaunch, DVfb, DVarrival, _, _ = self._compute_dvs(x) + DVlaunch, DVfb, DVarrival, _, _, _, _, _ = self._compute_dvs(x) if self.tof_encoding == 'direct': T = sum(x[1:]) elif self.tof_encoding == 'alpha': @@ -309,7 +326,7 @@ def pretty(self, x): T = self._decode_tofs(x) ep = np.insert(T, 0, x[0]) # [t0, T1, T2 ...] ep = np.cumsum(ep) # [t0, t1, t2, ...] - DVlaunch, DVfb, DVarrival, l, DVlaunch_tot = self._compute_dvs(x) + DVlaunch, DVfb, DVarrival, l, DVlaunch_tot, _, _, _ = self._compute_dvs(x) print("Multiple Gravity Assist (MGA) problem: ") print("Planet sequence: ", [pl.name for pl in self.seq]) @@ -354,10 +371,10 @@ def plot(self, x, axes=None, units=AU, N=60): fig = plt.figure() axes = fig.gca(projection='3d') - T = self._decode_tofs(x) + _, _, _, l, _, T, _, _ = self._compute_dvs(x) ep = np.insert(T, 0, x[0]) # [t0, T1, T2 ...] ep = np.cumsum(ep) # [t0, t1, t2, ...] - _, _, _, l, _ = self._compute_dvs(x) + for pl, e in zip(self.seq, ep): plot_planet(pl, epoch(e), units=units, legend=True, color=(0.7, 0.7, 1), axes=axes) @@ -366,6 +383,78 @@ def plot(self, x, axes=None, units=AU, N=60): legend=False, axes=axes, alpha=0.8) return axes + def get_eph_function(self, x: List[float]): + """ + For a chromosome x, returns a function object eph to compute the ephemerides of the spacecraft + + Args: + - x (``list``, ``tuple``, ``numpy.ndarray``): Decision chromosome, e.g. (``pygmo.population.champion_x``). + + Example: + + eph = prob.get_eph_function(population.champion_x) + pos, vel = eph(pykep.epoch(7000)) + + """ + if len(x) != len(self.get_bounds()[0]): + raise ValueError( + "Expected chromosome of length " + + str(len(self.get_bounds()[0])) + + " but got length " + + str(len(x)) + ) + + _, _, _, _, _, _, b_legs, b_ep = self._compute_dvs(x) + + def eph( + t: float + ) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]: + + if t < b_ep[0]: + raise ValueError( + "Given epoch " + str(t) + " is before launch date " + str(b_ep[0]) + ) + + if t == b_ep[0]: + # exactly at launch + return self.seq[0].eph(t) + + i = bisect_left(b_ep, t) # ballistic leg i goes from planet i to planet i+1 + + assert i >= 1 and i <= len(b_ep) + if i < len(b_ep): + assert t <= b_ep[i] + + # get start of ballistic leg + r_b, v_b = b_legs[i - 1] + + elapsed_seconds = (t - b_ep[i - 1]) * DAY2SEC + assert elapsed_seconds >= 0 + + # propagate the lagrangian + r, v = propagate_lagrangian(r_b, v_b, elapsed_seconds, self.seq[0].mu_central_body) + + return r, v + + return eph + + def plot_distance_and_flybys(self, x: List[float], **kwargs): + + from pykep.orbit_plots import plot_flybys + eph_function = self.get_eph_function(x) + T = self._decode_tofs(x) + ep = np.insert(T, 0, x[0]) # [t0, T1, T2 ...] + ep = np.cumsum(ep) # [t0, t1, t2, ...] + + return plot_flybys(self.seq, ep, eph_function, probename=self.get_name(), **kwargs) + + def get_name(self): + return "MGA Trajectory" + + def __repr__(self): + return self.get_name() + + if __name__ == "__main__": import pygmo as pg diff --git a/pykep/trajopt/_mga_1dsm.py b/pykep/trajopt/_mga_1dsm.py index 38203cba..fa87032e 100644 --- a/pykep/trajopt/_mga_1dsm.py +++ b/pykep/trajopt/_mga_1dsm.py @@ -3,6 +3,9 @@ from pykep.trajopt._lambert import lambert_problem_multirev from math import pi, cos, sin, acos, log, sqrt import numpy as np +from typing import Any, Dict, List, Tuple +from bisect import bisect_left + # Avoiding scipy dependency def norm(x): @@ -198,8 +201,17 @@ def _decode_times_and_vinf(self, x): return (retval_T, Vinfx, Vinfy, Vinfz) - # Objective function - def fitness(self, x): + def _decode_tofs(self, x): + T, _, _, _ = self._decode_times_and_vinf(x) + return T + + def _compute_dvs(self, x: List[float]) -> Tuple[ + List[float], # DVs + List[Any], # Lambert legs + List[float], # T + List[Tuple[List[float], List[float]]], # ballistic legs + List[float], # epochs of ballistic legs + ]: # 1 - we 'decode' the chromosome recording the various times of flight # (days) in the list T and the cartesian components of vinf T, Vinfx, Vinfy, Vinfz = self._decode_times_and_vinf(x) @@ -212,9 +224,14 @@ def fitness(self, x): for i in range(len(self._seq)): t_P[i] = epoch(x[0] + sum(T[0:i])) r_P[i], v_P[i] = self._seq[i].eph(t_P[i]) + ballistic_legs: List[Tuple[List[float],List[float]]] = [] + ballistic_ep: List[float] = [] + lamberts = [] # 3 - We start with the first leg v0 = [a + b for a, b in zip(v_P[0], [Vinfx, Vinfy, Vinfz])] + ballistic_legs.append((r_P[0], v0)) + ballistic_ep.append(t_P[0].mjd2000) r, v = propagate_lagrangian( r_P[0], v0, x[4] * T[0] * DAY2SEC, self.common_mu) @@ -224,6 +241,10 @@ def fitness(self, x): r, r_P[1], dt, self.common_mu, cw=False, max_revs=self.max_revs)) v_end_l = l.get_v2()[0] v_beg_l = l.get_v1()[0] + lamberts.append(l) + + ballistic_legs.append((r, v_beg_l)) + ballistic_ep.append(t_P[0].mjd2000 + x[4] * T[0]) # First DSM occuring at time nu1*T1 DV[0] = norm([a - b for a, b in zip(v_beg_l, v)]) @@ -233,6 +254,8 @@ def fitness(self, x): # Fly-by v_out = fb_prop(v_end_l, v_P[i], x[ 7 + (i - 1) * 4] * self._seq[i].radius, x[6 + (i - 1) * 4], self._seq[i].mu_self) + ballistic_legs.append((r_P[i], v_out)) + ballistic_ep.append(t_P[i].mjd2000) # s/c propagation before the DSM r, v = propagate_lagrangian( r_P[i], v_out, x[8 + (i - 1) * 4] * T[i] * DAY2SEC, self.common_mu) @@ -242,9 +265,13 @@ def fitness(self, x): self.common_mu, cw=False, max_revs=self.max_revs)) v_end_l = l.get_v2()[0] v_beg_l = l.get_v1()[0] + lamberts.append(l) # DSM occuring at time nu2*T2 DV[i] = norm([a - b for a, b in zip(v_beg_l, v)]) + ballistic_legs.append((r, v_beg_l)) + ballistic_ep.append(t_P[i].mjd2000 + x[8 + (i - 1) * 4] * T[i]) + # Last Delta-v if self._add_vinf_arr: DV[-1] = norm([a - b for a, b in zip(v_end_l, v_P[-1])]) @@ -260,14 +287,19 @@ def fitness(self, x): if self._add_vinf_dep: DV[0] += x[3] + return (DV, lamberts, T, ballistic_legs, ballistic_ep) + + # Objective function + def fitness(self, x): + DV, _, T, _, _ = self._compute_dvs(x) if not self._multi_objective: - return (sum(DV),) + return [sum(DV),] else: return (sum(DV), sum(T)) def pretty(self, x): """ - prob.plot(x) + prob.pretty(x) - x: encoded trajectory @@ -454,3 +486,84 @@ def get_extra_info(self): return ("\n\t Sequence: " + [pl.name for pl in self._seq].__repr__() + "\n\t Add launcher vinf to the objective?: " + self._add_vinf_dep.__repr__() + "\n\t Add final vinf to the objective?: " + self._add_vinf_arr.__repr__()) + + def plot_distance_and_flybys(self, x: List[float], **kwargs): + """ + Plot solar distance and flybys of the trajectory defined by x. + + Args: + - x (``list``, ``tuple``, ``numpy.ndarray``): Decision chromosome, e.g. (``pygmo.population.champion_x``). + - axes: Matplotlib plotting axes + - N: number of sample points + - extension: number of days to propagate the trajectory after the last flyby + + """ + + from pykep.orbit_plots import plot_flybys + eph_function = self.get_eph_function(x) + T = self._decode_tofs(x) + ep = np.insert(T, 0, x[0]) # [t0, T1, T2 ...] + ep = np.cumsum(ep) # [t0, t1, t2, ...] + + + return plot_flybys(self._seq, ep, eph_function, probename=self.get_name(), **kwargs) + + def get_name(self): + return "MGA_1DSM Trajectory" + + def __repr__(self): + return self.get_name() + + def get_eph_function(self, x): + """ + For a chromosome x, returns a function object eph to compute the ephemerides of the spacecraft + + Args: + - x (``list``, ``tuple``, ``numpy.ndarray``): Decision chromosome, e.g. (``pygmo.population.champion_x``). + + Example: + + eph = prob.get_eph_function(population.champion_x) + pos, vel = eph(pykep.epoch(7000)) + + """ + if len(x) != len(self.get_bounds()[0]): + raise ValueError( + "Expected chromosome of length " + + str(len(self.get_bounds()[0])) + + " but got length " + + str(len(x)) + ) + _, _, _, b_legs, b_ep = self._compute_dvs(x) + + def eph( + t: float + ) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]: + + if t < b_ep[0]: + raise ValueError( + "Given epoch " + str(t) + " is before launch date " + str(b_ep[0]) + ) + + if t == b_ep[0]: + # exactly at launch + return self._seq[0].eph(t) + + i = bisect_left(b_ep, t) # ballistic leg i goes from planet i to planet i+1 + + assert i >= 1 and i <= len(b_ep) + if i < len(b_ep): + assert t <= b_ep[i] + + # get start of ballistic leg + r_b, v_b = b_legs[i - 1] + + elapsed_seconds = (t - b_ep[i - 1]) * DAY2SEC + assert elapsed_seconds >= 0 + + # propagate the lagrangian + r, v = propagate_lagrangian(r_b, v_b, elapsed_seconds, self._seq[0].mu_central_body) + + return r, v + + return eph