Merge pull request #143 from mlooz/feature/mga-refactoring

Feature/mga refactoring

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)

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)

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']

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

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']:

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,25 +244,40 @@ 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(
                 vi, lambert_problem(
                     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