Merge pull request #48 from darioizzo/mit

Adding the primer vector to pl2pl

CMakeLists.txt

@@ -153,6 +153,7 @@ set(kep3_SRC_FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/src/core_astro/eq2par2eq.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/core_astro/stm.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/core_astro/propagate_lagrangian.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/core_astro/encodings.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ta/stark.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ta/cr3bp.cpp"
 )

include/kep3/core_astro/encodings.hpp

@@ -0,0 +1,25 @@
+// Copyright 2023, 2024 Dario Izzo ([email protected]), Francesco Biscani
+// ([email protected])
+//
+// This file is part of the kep3 library.
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#ifndef kep3_ENCODING_CONVERSIONS_H
+#define kep3_ENCODING_CONVERSIONS_H
+
+#include <vector>
+
+#include <kep3/detail/visibility.hpp>
+
+namespace kep3
+{
+kep3_DLL_PUBLIC std::vector<double> alpha2direct(const std::vector<double> &alphas, double tof);
+kep3_DLL_PUBLIC std::pair<std::vector<double>, double> direct2alpha(const std::vector<double> &tofs);
+kep3_DLL_PUBLIC std::vector<double> eta2direct(const std::vector<double> &etas, double max_tof);
+kep3_DLL_PUBLIC std::vector<double> direct2eta(const std::vector<double> &tofs, double max_tof);
+} // namespace kep3
+
+#endif

pykep/core.cpp

@@ -18,6 +18,7 @@
 #include <kep3/core_astro/ic2eq2ic.hpp>
 #include <kep3/core_astro/ic2par2ic.hpp>
 #include <kep3/core_astro/propagate_lagrangian.hpp>
+#include <kep3/core_astro/encodings.hpp>
 #include <kep3/epoch.hpp>
 #include <kep3/lambert_problem.hpp>
 #include <kep3/leg/sims_flanagan.hpp>
@@ -109,14 +110,20 @@ PYBIND11_MODULE(core, m) // NOLINT
     m.def("zeta2f_v", py::vectorize(kep3::zeta2f), pk::zeta2f_v_doc().c_str());
     m.def("f2zeta_v", py::vectorize(kep3::f2zeta), pk::f2zeta_v_doc().c_str());
 
-    // Eposing element conversions
+    // Exposing element conversions
     m.def("ic2par", &kep3::ic2par);
     m.def("par2ic", &kep3::par2ic);
     m.def("ic2eq", &kep3::ic2eq);
     m.def("eq2ic", &kep3::eq2ic);
     m.def("par2eq", &kep3::par2eq);
     m.def("eq2par", &kep3::eq2par);
 
+    // Exposing encoding conversions
+    m.def("alpha2direct", &kep3::alpha2direct, py::arg("alphas"), py::arg("tof"), pk::alpha2direct_doc().c_str());
+    m.def("direct2alpha", &kep3::direct2alpha, py::arg("tofs") , pk::direct2alpha_doc().c_str());
+    m.def("eta2direct", &kep3::eta2direct, py::arg("etas"), py::arg("max_tof"), pk::eta2direct_doc().c_str());
+    m.def("direct2eta", &kep3::direct2eta, py::arg("tofs"), py::arg("max_tof"), pk::direct2eta_doc().c_str());
+
     // Class epoch
     py::class_<kep3::epoch> epoch_class(m, "epoch", "Represents a specific point in time.");
 

pykep/docstrings.cpp

@@ -698,6 +698,68 @@ std::string f2zeta_v_doc()
 )";
 }
 
+std::string alpha2direct_doc()
+{
+    return R"(alpha2direct(alphas)
+
+    Converts from alpha encoded to transfer times.
+
+    Args:
+        *alphas* (:class:`list`): a sequence of transfer times encoded using the alpha encoding.
+
+        *T* (:class:`float`): the total transfer time.
+
+    Returns:
+        :class:`list`:: The encoded transfer times
+)";
+}
+
+std::string direct2alpha_doc()
+{
+    return R"(direct2alpha(tofs)
+
+    Converts from transfer times to alpha encoded.
+
+    Args:
+        *tofs* (:class:`list`): a sequence of transfer times.
+
+    Returns:
+        :class:`list`:, :class:`float`: The alpha-encoded transfer times, the total transfer time (for cenvenience)
+)";
+}
+
+std::string eta2direct_doc()
+{
+    return R"(eta2direct(etas, max_tof)
+
+    Converts from eta encoded to transfer times.
+
+    Args:
+        *etas* (:class:`list`): a sequence of transfer times encoded using the eta encoding.
+
+        *max_tof* (:class:`float`): maximum time of flight (might actually be less according to the value of the last eta.)
+
+    Returns:
+        :class:`list`: The encoded transfer times
+)";
+}
+
+std::string direct2eta_doc()
+{
+    return R"(direct2eta(tofs, max_tof)
+
+      Converts from transfer times to eta encoded.
+
+    Args:
+        *tofs* (:class:`list`):  a sequence of transfer times
+
+        *max_tof* (:class:`float`): maximum time of flight (might actually be less according to the value of the last eta.)
+
+    Returns:
+        :class:`list`: The eta-encoded transfer times
+)";
+}
+
 std::string epoch_from_float_doc()
 {
     return R"(__init__(when: float, julian_type = MJD2000)

pykep/docstrings.hpp

@@ -54,6 +54,12 @@ std::string f2h_v_doc();
 std::string zeta2f_v_doc();
 std::string f2zeta_v_doc();
 
+// Encodings
+std::string alpha2direct_doc();
+std::string direct2alpha_doc();
+std::string eta2direct_doc();
+std::string direct2eta_doc();
+
 // Epoch
 std::string epoch_from_float_doc();
 std::string epoch_from_datetime_doc();

pykep/test_utils.py

@@ -23,15 +23,15 @@ def test_alpha_direct_conversion(self):
         import pykep as pk
 
         tofs = [12.34, 232.2, 23.45, 134.3]
-        alphas, T = pk.utils.direct2alpha(tofs)
-        tofs_from_alphas = pk.utils.alpha2direct(alphas, T)
+        alphas, T = pk.direct2alpha(tofs)
+        tofs_from_alphas = pk.alpha2direct(alphas, T)
         err = [a - b for a, b in zip(tofs, tofs_from_alphas)]
         err = np.linalg.norm(err)
         self.assertTrue(err < 1e-13)
 
         tofs = np.random.random((4,)) * 20
-        alphas, T = pk.utils.direct2alpha(tofs)
-        tofs_from_alphas = pk.utils.alpha2direct(alphas, T)
+        alphas, T = pk.direct2alpha(tofs)
+        tofs_from_alphas = pk.alpha2direct(alphas, T)
         err = [a - b for a, b in zip(tofs, tofs_from_alphas)]
         err = np.linalg.norm(err)
         self.assertTrue(err < 1e-13)
@@ -41,16 +41,16 @@ def test_eta_direct_conversion(self):
 
         tofs = [12.34, 232.2, 23.45, 134.3]
         tmax = 300
-        etas = pk.utils.direct2eta(tofs, tmax)
-        tofs_from_etas = pk.utils.eta2direct(etas, tmax)
+        etas = pk.direct2eta(tofs, tmax)
+        tofs_from_etas = pk.eta2direct(etas, tmax)
         err = [a - b for a, b in zip(tofs, tofs_from_etas)]
         err = np.linalg.norm(err)
         self.assertTrue(err < 1e-13)
 
         tofs = np.random.random((4,)) * 100
         tmax = 400
-        etas = pk.utils.direct2eta(tofs, tmax)
-        tofs_from_etas = pk.utils.eta2direct(etas, tmax)
+        etas = pk.direct2eta(tofs, tmax)
+        tofs_from_etas = pk.eta2direct(etas, tmax)
         err = [a - b for a, b in zip(tofs, tofs_from_etas)]
         err = np.linalg.norm(err)
         self.assertTrue(err < 1e-13)

pykep/trajopt/_mga.py

@@ -194,13 +194,13 @@ def get_bounds(self):
     def _decode_tofs(self, x: List[float]) -> List[float]:
         if self.tof_encoding == "alpha":
             # decision vector is  [t0, T, a1, a2, ....]
-            return _pk.utils.alpha2direct(x[2:], x[1])
+            return _pk.alpha2direct(x[2:], x[1])
         elif self.tof_encoding == "direct":
             # decision vector is  [t0, T1, T2, T3, ... ]
             return x[1:]
         elif self.tof_encoding == "eta":
             # decision vector is  [t0, n1, n2, n3, ... ]
-            return _pk.utils.eta2direct(x[1:], self.tof)
+            return _pk.eta2direct(x[1:], self.tof)
 
     @staticmethod
     def alpha2direct(x):
@@ -214,7 +214,7 @@ def alpha2direct(x):
         Returns:
             :class:`numpy.ndarray`: a chromosome encoding the MGA trajectory using the direct encoding
         """
-        retval = _pk.utils.alpha2direct(x[2:], x[1])
+        retval = _pk.alpha2direct(x[2:], x[1])
         retval = _np.insert(retval, 0, x[0])
         return retval
 
@@ -228,7 +228,7 @@ def direct2alpha(x):
         Returns:
             :class:`numpy.ndarray`: a chromosome encoding the MGA trajectory using the alpha encoding
         """
-        alphas, T = _pk.utils.direct2alpha(x[1:])
+        alphas, T = _pk.direct2alpha(x[1:])
         retval = _np.insert(alphas, 0, [x[0], T])
         return retval
 
@@ -269,7 +269,7 @@ def direct2eta(self, x):
         from copy import deepcopy
 
         retval = deepcopy(x)
-        retval[1:] = _pk.utils.direct2eta(x[1:], self.tof)
+        retval[1:] = _pk.direct2eta(x[1:], self.tof)
         return retval
 
     def _compute_dvs(self, x: List[float]) -> Tuple[

pykep/trajopt/_mga_1dsm.py

@@ -220,13 +220,13 @@ def _decode_times_and_vinf(self, x):
         # 1 - we decode the times of flight
         if self._tof_encoding == "alpha":
             # decision vector is  [t0] + [u, v, Vinf, eta1, a1] + [beta, rp/rV, eta2, a2] + ... + [T]
-            retval_T = _pk.utils.alpha2direct(x[5::4], x[-1])
+            retval_T = _pk.alpha2direct(x[5::4], x[-1])
         elif self._tof_encoding == "direct":
             # decision vector is  [t0] + [u, v, Vinf, eta1, T1] + [beta, rp/rV, eta2, T2] + ...
             retval_T = x[5::4]
         elif self._tof_encoding == "eta":
             # decision vector is [t0] + [u, v, Vinf, eta1, n1] + [beta, rp/rV, eta2, n2] + ...
-            retval_T = _pk.utils.eta2direct(x[5::4], self._tof)
+            retval_T = _pk.eta2direct(x[5::4], self._tof)
 
         # 2 - We decode the hyperbolic velocity at departure
         theta = 2 * pi * x[1]
@@ -254,7 +254,7 @@ def alpha2direct(x):
         """
         # decision vector is  [t0] + [u, v, Vinf, eta1, a1] + [beta, rp/rV, eta2, a2] + ... + [T]
         retval = deepcopy(x)
-        retval[5::4] = _pk.utils.alpha2direct(x[5::4], x[-1])
+        retval[5::4] = _pk.alpha2direct(x[5::4], x[-1])
         retval = _np.delete(retval, -1)
         return retval
 
@@ -270,7 +270,7 @@ def direct2alpha(x):
         """
         # decision vector is  [t0] + [u, v, Vinf, eta1, T1] + [beta, rp/rV, eta2, T2] + ...
         retval = deepcopy(x)
-        retval[5::4], T = _pk.utils.direct2alpha(x[5::4])
+        retval[5::4], T = _pk.direct2alpha(x[5::4])
         retval = _np.append(retval,T)
         return retval
 
@@ -286,7 +286,7 @@ def eta2direct(x, max_tof):
         """
         # decision vector is [t0] + [u, v, Vinf, eta1, n1] + [beta, rp/rV, eta2, n2] + ...
         retval = deepcopy(x)
-        retval[5::4] = _pk.utils.eta2direct(x[5::4], max_tof)
+        retval[5::4] = _pk.eta2direct(x[5::4], max_tof)
         return retval
 
     @staticmethod
@@ -300,7 +300,7 @@ def direct2eta(x, max_tof):
             :class:`numpy.ndarray`: a chromosome encoding the MGA trajectory using the eta encoding
         """
         retval = deepcopy(x)
-        retval[5::4] = _pk.utils.direct2eta(x[5::4], max_tof)
+        retval[5::4] = _pk.direct2eta(x[5::4], max_tof)
         return retval
 
     def _compute_dvs(self, x: List[float]) -> Tuple[

pykep/trajopt/_pl2pl_N_impulses.py

@@ -1,5 +1,6 @@
 import pykep as _pk
 import numpy as _np
+import matplotlib.pyplot as _plt
 
 from math import pi, cos, sin, log, acos, sqrt
 
@@ -129,7 +130,7 @@ def decode(self, x):
         """
         retval = []
         # 1 We decode the tofs
-        T = _pk.utils.alpha2direct(x[2::4], x[1])
+        T = _pk.alpha2direct(x[2::4], x[1])
 
         # 2 - We compute the starting and ending position
         r_start, v_start = self.start.eph(_pk.epoch(x[0]))
@@ -268,3 +269,91 @@ def pretty(self, x):
         print(
             "Tofs (days): ", [float(it) for it in T[:-1]]
         )  # last node has a zero TOF by convention
+
+    def plot_primer_vector(self, x, N=200, ax=None):
+        """Plots the primer vector magnitude along the trajectory encoded in *x*.
+        
+        Args:
+            *x* (:class:`list`): The decision vector in the correct tof encoding.
+
+            *N* (:class:`int`, optional): The number of points to use when plotting the primer vector. Defaults to 200.
+
+            *ax* (:class:`matplotlib.axes.Axes`, optional): The axis to plot on. Defaults to None.
+            
+        Returns:
+            :class:`matplotlib.axes.Axes`: The axis where the primer vector was plotted. 
+            :class:`tuple`: A tuple containing the grid and the primer vector magnitude.
+        """
+        # We start by decoding the chromosome into the structure [[r,v], DV, DT]
+        decoded = self.decode(x)
+
+        # We explicitly extract the encoded information
+        dts = [it[2] * _pk.DAY2SEC for it in decoded]
+        DVs = [it[1] for it in decoded]
+        posvels = [it[0] for it in decoded]
+
+        # We create one grid er segment (e.g. part of the trajectory between two impulses)
+        # (this is not guaranteed to have the requested size N, nor has uniform spacing, since all impulses
+        # must belong to the grid points)
+        N = N + len(DVs)  # heuristic to make sure we are close to the requested number of points
+        tgrids = [
+            _np.linspace(
+                sum(dts[:i]),
+                sum(dts[: i + 1]),
+                max(int(dts[i] // (sum(dts) / (N - 1))), 5), # we force a minimum 5 points per segment
+            )
+            for i in range(len(dts) - 1)
+        ]
+        # We assemble all the grids into one single final_grid
+        final_grid = _np.array([0])
+        for i in range(len(dts) - 1):
+            final_grid = _np.concatenate((final_grid, tgrids[i][1:]))
+
+        # These are the indices of the final_grid where the impulses are given.
+        idxs = [0] + [len(tgrids[0]) - 1]
+        for grid in tgrids[1:]:
+            idxs += [idxs[-1] + len(grid) - 1]
+
+        # We now compute the various STMs.
+        retvals = []
+        for posvel, DV, tgrid in zip(posvels, DVs, tgrids):
+            ic = [posvel[0], [a + b for a, b in zip(posvel[1], DV)]]
+            retvals.append(
+                _pk.propagate_lagrangian_grid(ic, tgrid, mu=_pk.MU_SUN, stm=True)
+            )
+
+        # And now assemble them in correspondance to the final_grid and in the Mn0 form.
+        posvels = [item[0] for item in retvals[0]]
+        stms = [item[1] for item in retvals[0]]
+
+        M = stms[-1]
+        for retval in retvals[1:]:
+            posvels = posvels + [item[0] for item in retval[1:]]
+            stms = stms + [item[1] @ M for item in retval[1:]]
+            M = stms[-1]
+
+        res = []
+        # When computing the primer vector we must choose which impulses to use. 
+        # We choose the first and last impulse. But we could choose any pair of impulses,
+        # and if the trajectory is optimal (locally) the primer vector would not change.
+        idx_i = idxs[0]
+        idx_j = idxs[-1]
+        DVi = DVs[0]
+        DVj = DVs[-1]
+        for idx_k, _ in enumerate(final_grid):
+            Mji = stms[idx_j] @ _np.linalg.inv(stms[idx_i])
+            Mjk = stms[idx_j] @ _np.linalg.inv(stms[idx_k])
+            res.append(
+                _np.linalg.norm(_pk.trajopt.primer_vector(DVi, DVj, Mji, Mjk)[0])
+            )
+
+        if ax is None:
+            ax = _plt.figure().add_subplot()
+        ax.plot(res, label="primer vector magnitude")
+        ax.vlines(0, 0.8, 1.2, "k", linestyles="dashed", label="impulse")
+        for idx in idxs:
+            ax.vlines(idx, 0.8, 1.2, "k", linestyles="dashed")
+
+        ax.hlines(1, 0, len(final_grid), "r")
+        ax.legend(loc="lower right")
+        return ax, (final_grid, res)

pykep/utils/__init__.py

@@ -8,4 +8,4 @@
 
 from ._planet_to_keplerian import planet_to_keplerian
 
-from ._encoding_conversions import direct2alpha, alpha2direct, eta2direct, direct2eta, uvV2cartesian, cartesian2uvV
+from ._encoding_conversions import uvV2cartesian, cartesian2uvV