Merge pull request #36 from darioizzo/propagate_taylor

Stark problem and heyoka taylor integrators

CMakeLists.txt

@@ -134,6 +134,7 @@ set(kep3_SRC_FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/src/epoch.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/planet.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/lambert_problem.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/stark_problem.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/linalg.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/udpla/keplerian.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/udpla/jpl_lp.cpp"
@@ -144,6 +145,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/ta/lt_kepler.cpp"
 )
 
 # Setup of the kep3 shared library.

include/kep3/core_astro/propagate_lagrangian.hpp

@@ -37,7 +37,7 @@ propagate_lagrangian_v(const std::array<std::array<double, 3>, 2> &pos_vel, std:
                        bool stm = false);
 
 // These are backup functions that use a different algorithm to get the same as propagate_lagrangian.
-// We offer them with an identical interface even if the stm is not implements.
+// We offer them with an identical interface even if the stm is not implemented.
 kep3_DLL_PUBLIC std::pair<std::array<std::array<double, 3>, 2>, std::optional<std::array<double, 36>>>
 propagate_lagrangian_u(const std::array<std::array<double, 3>, 2> &pos_vel, double dt, double mu, bool = false);
 

include/kep3/lambert_problem.hpp

@@ -44,17 +44,13 @@ namespace kep3
  * @author Dario Izzo (dario.izzo _AT_ googlemail.com)
  */
 
-class kep3_DLL_PUBLIC lambert_problem;
-
-// Streaming operator for the class kep_toolbox::lambert_problem.
-kep3_DLL_PUBLIC std::ostream &operator<<(std::ostream &, const lambert_problem &);
-
 class kep3_DLL_PUBLIC lambert_problem
 {
     static const std::array<double, 3> default_r0;
     static const std::array<double, 3> default_r1;
 
 public:
+    // We choose in this case to friend  the streaming operator as to not expose a number of tedious getters (lambda, chord etc...)
     friend kep3_DLL_PUBLIC std::ostream &operator<<(std::ostream &, const lambert_problem &);
     explicit lambert_problem(const std::array<double, 3> &r0  = default_r0, const std::array<double, 3> &r1 = default_r1,
                              double tof = kep3::pi / 2, double mu = 1., bool cw = false, unsigned multi_revs = 1);
@@ -108,9 +104,6 @@ class kep3_DLL_PUBLIC lambert_problem
     unsigned m_multi_revs;
 };
 
-// Streaming operator for the class kep3::lambert_problem.
-kep3_DLL_PUBLIC std::ostream &operator<<(std::ostream &, const lambert_problem &);
-
 } // namespace kep3
 
 template <>

include/kep3/stark_problem.hpp

@@ -0,0 +1,78 @@
+// 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_STARK_PROBLEM_H
+#define kep3_STARK_PROBLEM_H
+
+#include <array>
+#include <utility>
+#include <vector>
+
+#include <fmt/ostream.h>
+
+#include <heyoka/taylor.hpp>
+
+#include <kep3/detail/s11n.hpp>
+#include <kep3/detail/visibility.hpp>
+
+namespace kep3
+{
+
+/// Stark Problem
+/**
+ * This class represent the Stark's problem. When instantiated it gets a
+ * Taylor adaptive integrator from the kep3 ta cache and copies it as to be able to use it
+ * later for numerical propagation. The information on the solution sensitivities can be retreived too
+ * in which case a variational integrator is employed.
+ */
+class kep3_DLL_PUBLIC stark_problem
+{
+public:
+    explicit stark_problem(double mu=1., double veff=1., double tol=1e-16);
+    std::array<double, 7> propagate(const std::array<double, 7> &rvm_state, std::array<double, 3> thrust, double tof);
+    std::tuple<std::array<double, 7>, std::array<double, 49>, std::array<double, 21>>
+    propagate_var(const std::array<double, 7> &rvm_state, std::array<double, 3> thrust, double tof);
+    // Getters
+    [[nodiscard]] double get_mu() const;
+    [[nodiscard]] double get_veff() const;
+    [[nodiscard]] double get_tol() const;
+    [[nodiscard]] heyoka::taylor_adaptive<double> get_ta() const;
+    [[nodiscard]] heyoka::taylor_adaptive<double> get_ta_var() const;
+    // Setters
+    void set_mu(double);
+    void set_veff(double);
+
+private:
+    // Serialization code
+    friend class boost::serialization::access;
+    template <class Archive>
+    void serialize(Archive &ar, const unsigned int)
+    {
+        ar &m_mu;
+        ar &m_veff;
+        ar &m_tol;
+        ar &m_ta;
+        ar &m_ta_var;
+        ar &m_var_ic;
+    }
+    double m_mu;
+    double m_veff;
+    double m_tol;
+    heyoka::taylor_adaptive<double> m_ta;
+    heyoka::taylor_adaptive<double> m_ta_var;
+    std::vector<double> m_var_ic;
+};
+kep3_DLL_PUBLIC std::ostream &operator<<(std::ostream &, const stark_problem &);
+} // namespace kep3
+
+template <>
+struct fmt::formatter<kep3::stark_problem> : fmt::ostream_formatter {
+};
+
+#endif // kep3_STARK_PROBLEM_H

include/kep3/ta/lt_kepler.hpp

@@ -0,0 +1,36 @@
+// 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_TA_LT_KEPLER_H
+#define kep3_TA_LT_KEPLER_H
+
+#include <tuple>
+#include <vector>
+
+#include <kep3/detail/visibility.hpp>
+
+#include <heyoka/expression.hpp>
+#include <heyoka/taylor.hpp>
+
+namespace kep3::ta
+{
+// Returns the low-thrust dynamics (heyoka API) in a Keplerian context and Cartesian throttles. 7 states, 5 parameters: mu, veff, ux, uy, uz.
+std::vector<std::pair<heyoka::expression, heyoka::expression>> lt_kepler_dyn();
+
+// These return const references to function level static variables of type heyoka::taylor_adaptive<double>.
+// NOTE: The object retruned are expected to be copied to then be modified.
+kep3_DLL_PUBLIC const heyoka::taylor_adaptive<double> &get_ta_lt_kepler(double tol);
+kep3_DLL_PUBLIC const heyoka::taylor_adaptive<double> &get_ta_lt_kepler_var(double tol); // variational (x,y,z,vx,vy,vz,ux,uy,uz) first order
+
+// Methods to access the cache dimensions.
+kep3_DLL_PUBLIC size_t get_ta_lt_kepler_cache_dim();
+kep3_DLL_PUBLIC size_t get_ta_lt_kepler_var_cache_dim();
+} // namespace kep3::ta
+
+#endif

kep3_devel.yml

@@ -4,6 +4,7 @@ channels:
 dependencies:
   - c-compiler
   - cxx-compiler
+  - clang-tools
   - ninja
   - cmake >=3.18
   - boost-cpp >=1.73
@@ -19,5 +20,3 @@ dependencies:
   - spiceypy
   - matplotlib
   - pygmo_plugins_nonfree
-
-

pykep/core.cpp

@@ -21,6 +21,7 @@
 #include <kep3/lambert_problem.hpp>
 #include <kep3/leg/sims_flanagan.hpp>
 #include <kep3/planet.hpp>
+#include <kep3/stark_problem.hpp>
 #include <kep3/udpla/keplerian.hpp>
 
 #include <pybind11/chrono.h>
@@ -323,6 +324,78 @@ PYBIND11_MODULE(core, m)
     m.def("propagate_lagrangian_v", &kep3::propagate_lagrangian_v, py::arg("rv") = std::array<std::array<double, 3>, 2>{{{1, 0, 0}, {0, 1, 0}}}, py::arg("tofs") = std::vector<double>{kep3::pi / 2,},
         py::arg("mu") = 1, py::arg("stm") = false, pykep::propagate_lagrangian_v_docstring().c_str());
 
+    // Exposing the Stark problem class
+    py::class_<kep3::stark_problem> stark_problem(m, "stark_problem", pykep::stark_problem_docstring().c_str());
+    stark_problem
+        .def(py::init<double, double, double>(), py::arg("mu") = 1., py::arg("veff") = 1., py::arg("tol") = 1e-16)
+        // repr().
+        .def("__repr__", &pykep::ostream_repr<kep3::stark_problem>)
+        // Copy and deepcopy.
+        .def("__copy__", &pykep::generic_copy_wrapper<kep3::stark_problem>)
+        .def("__deepcopy__", &pykep::generic_deepcopy_wrapper<kep3::stark_problem>)
+        // Pickle support.
+        .def(py::pickle(&pykep::pickle_getstate_wrapper<kep3::stark_problem>,
+                        &pykep::pickle_setstate_wrapper<kep3::stark_problem>))
+        .def_property_readonly("mu", &kep3::stark_problem::get_mu, "The central body gravity parameter.")
+        .def_property_readonly("veff", &kep3::stark_problem::get_veff, "The effective velocity (Isp g0)")
+        .def_property_readonly("tol", &kep3::stark_problem::get_tol, "The Taylor integrator tolerance.")
+        // The actual call to propagators. (we do not here care about copies and allocations as this is 20 times slower
+        // than propagate lagrangian already on c++ side).
+        .def("propagate", &kep3::stark_problem::propagate, py::arg("rvm_state"), py::arg("thrust"), py::arg("tof"),
+             pykep::stark_problem_propagate_docstring().c_str())
+        .def(
+            "propagate_var",
+            [](kep3::stark_problem &sp, const std::array<double, 7> &rvm_state, std::array<double, 3> thrust,
+               double tof) {
+                auto sp_retval = sp.propagate_var(rvm_state, thrust, tof);
+                // Lets transfer ownership of dxdx to python (not sure this is actually needed to
+                // get an efficient return value ... maybe its overkill here). It surely avoid one more copy / allocation of 49+21
+                // values, but in the overall algorithm maybe irrelevant. 
+                std::array<double, 49> &dxdx = std::get<1>(sp_retval);
+
+                // We create a capsule for the py::array_t to manage ownership change.
+                auto vec_ptr = std::make_unique<std::array<double, 49>>(dxdx);
+
+                py::capsule vec_caps(vec_ptr.get(), [](void *ptr) {
+                    std::unique_ptr<std::array<double, 49>> vptr(static_cast<std::array<double, 49> *>(ptr));
+                });
+
+                // NOTE: at this point, the capsule has been created successfully (including
+                // the registration of the destructor). We can thus release ownership from vec_ptr,
+                // as now the capsule is responsible for destroying its contents. If the capsule constructor
+                // throws, the destructor function is not registered/invoked, and the destructor
+                // of vec_ptr will take care of cleaning up.
+                auto *ptr = vec_ptr.release();
+
+                auto computed_dxdx = py::array_t<double>(
+                    py::array::ShapeContainer{static_cast<py::ssize_t>(7), static_cast<py::ssize_t>(7)}, // shape
+                    ptr->data(), std::move(vec_caps));
+
+                // Lets transfer ownership of dxdu to python
+                std::array<double, 21> &dxdu = std::get<2>(sp_retval);
+
+                // We create a capsule for the py::array_t to manage ownership change.
+                auto vec_ptr2 = std::make_unique<std::array<double, 21>>(dxdu);
+
+                py::capsule vec_caps2(vec_ptr2.get(), [](void *ptr) {
+                    std::unique_ptr<std::array<double, 21>> vec_ptr2(static_cast<std::array<double, 21> *>(ptr));
+                });
+
+                // NOTE: at this point, the capsule has been created successfully (including
+                // the registration of the destructor). We can thus release ownership from vec_ptr,
+                // as now the capsule is responsible for destroying its contents. If the capsule constructor
+                // throws, the destructor function is not registered/invoked, and the destructor
+                // of vec_ptr will take care of cleaning up.
+                auto *ptr2 = vec_ptr2.release();
+
+                auto computed_dxdu = py::array_t<double>(
+                    py::array::ShapeContainer{static_cast<py::ssize_t>(7), static_cast<py::ssize_t>(3)}, // shape
+                    ptr->data(), std::move(vec_caps2));
+                return py::make_tuple(std::get<0>(sp_retval), computed_dxdx, computed_dxdu);
+            },
+            py::arg("rvm_state"), py::arg("thrust"), py::arg("tof"),
+            pykep::stark_problem_propagate_docstring().c_str());
+
     // Exposing the sims_flanagan leg
     py::class_<kep3::leg::sims_flanagan> sims_flanagan(m, "_sims_flanagan", pykep::leg_sf_docstring().c_str());
     sims_flanagan
@@ -376,21 +449,19 @@ PYBIND11_MODULE(core, m)
                 });
                 // NOTE: at this point, the capsules have been created successfully (including
                 // the registration of the destructor). We can thus release ownership from vec_ptr_xx,
-                // as now the capsules are responsible for destroying its contents. 
+                // as now the capsules are responsible for destroying its contents.
                 auto *ptr_rs = vec_ptr_rs.release();
                 auto *ptr_rf = vec_ptr_rf.release();
                 auto *ptr_th = vec_ptr_th.release();
                 auto rs_python = py::array_t<double>(
-                    py::array::ShapeContainer{static_cast<py::ssize_t>(7),
-                                              static_cast<py::ssize_t>(7)}, // shape
+                    py::array::ShapeContainer{static_cast<py::ssize_t>(7), static_cast<py::ssize_t>(7)}, // shape
                     ptr_rs->data(), std::move(vec_caps_rs));
                 auto rf_python = py::array_t<double>(
-                    py::array::ShapeContainer{static_cast<py::ssize_t>(7),
-                                              static_cast<py::ssize_t>(7)}, // shape
+                    py::array::ShapeContainer{static_cast<py::ssize_t>(7), static_cast<py::ssize_t>(7)}, // shape
                     ptr_rf->data(), std::move(vec_caps_rf));
                 auto th_python = py::array_t<double>(
                     py::array::ShapeContainer{static_cast<py::ssize_t>(7),
-                                              static_cast<py::ssize_t>(leg.get_nseg()*3+1u)}, // shape
+                                              static_cast<py::ssize_t>(leg.get_nseg() * 3 + 1u)}, // shape
                     ptr_th->data(), std::move(vec_caps_th));
                 return py::make_tuple(rs_python, rf_python, th_python);
             },