Merge pull request #345 from bluescarni/pr/cr3bp

CR3BP

CMakeLists.txt

@@ -259,6 +259,7 @@ set(HEYOKA_SRC_FILES
     "${CMAKE_CURRENT_SOURCE_DIR}/src/model/rotating.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/model/mascon.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/model/vsop2013.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/model/cr3bp.cpp"
     # Math functions.
     "${CMAKE_CURRENT_SOURCE_DIR}/src/math/kepE.cpp"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/math/cos.cpp"
@@ -490,6 +491,11 @@ if(HEYOKA_WITH_MPPP)
         message(FATAL_ERROR "mp++ must be installed with support for Boost.serialization.")
     endif()
 
+    # NOTE: needed for formatting numbers.
+    if(NOT mp++_WITH_FMT)
+        message(FATAL_ERROR "mp++ must be installed with support for fmt.")
+    endif()
+
     target_link_libraries(heyoka PUBLIC mp++::mp++)
 
     # NOTE: _HEYOKA_WITH_REAL128 is used to signal the support

doc/changelog.rst

@@ -7,6 +7,8 @@ Changelog
 New
 ~~~
 
+- Add model for the circular restricted three-body problem
+  (`#345 <https://github.com/bluescarni/heyoka/pull/345>`__).
 - heyoka can now automatically vectorise scalar calls to
   floating-point math functions
   (`#342 <https://github.com/bluescarni/heyoka/pull/342>`__).

doc/install.rst

@@ -37,7 +37,8 @@ Additionally, heyoka has the following **optional** dependencies:
   which provides support for arbitrary-precision integrations on all platforms,
   and for quadruple-precision integrations on platforms
   supporting the non-standard ``__float128`` type. heyoka requires
-  an mp++ installation with support for Boost.serialization
+  an mp++ installation with support for Boost.serialization and for the
+  {fmt} library
   (see the :ref:`mp++ installation instructions <mppp:installation>`).
   The minimum required version of mp++ is 0.27;
 * the `SLEEF <https://sleef.org/>`__ vectorized math library (improves the performance

include/heyoka/kw.hpp

@@ -28,6 +28,7 @@ IGOR_MAKE_NAMED_ARGUMENT(compact_mode);
 IGOR_MAKE_NAMED_ARGUMENT(high_accuracy);
 IGOR_MAKE_NAMED_ARGUMENT(parallel_mode);
 IGOR_MAKE_NAMED_ARGUMENT(prec);
+IGOR_MAKE_NAMED_ARGUMENT(mu);
 
 } // namespace kw
 

include/heyoka/model/cr3bp.hpp

@@ -0,0 +1,66 @@
+// Copyright 2020, 2021, 2022, 2023 Francesco Biscani ([email protected]), Dario Izzo ([email protected])
+//
+// This file is part of the heyoka 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 HEYOKA_MODEL_CR3BP_HPP
+#define HEYOKA_MODEL_CR3BP_HPP
+
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include <heyoka/config.hpp>
+#include <heyoka/detail/igor.hpp>
+#include <heyoka/detail/visibility.hpp>
+#include <heyoka/expression.hpp>
+#include <heyoka/kw.hpp>
+
+HEYOKA_BEGIN_NAMESPACE
+
+namespace model
+{
+
+namespace detail
+{
+
+template <typename... KwArgs>
+auto cr3bp_common_opts(KwArgs &&...kw_args)
+{
+    igor::parser p{kw_args...};
+
+    static_assert(!p.has_unnamed_arguments(), "This function accepts only named arguments");
+
+    expression mu{1e-3};
+    if constexpr (p.has(kw::mu)) {
+        mu = expression{std::forward<decltype(p(kw::mu))>(p(kw::mu))};
+    }
+
+    return std::tuple{std::move(mu)};
+}
+
+HEYOKA_DLL_PUBLIC std::vector<std::pair<expression, expression>> cr3bp_impl(const expression &);
+
+HEYOKA_DLL_PUBLIC expression cr3bp_jacobi_impl(const expression &);
+
+} // namespace detail
+
+// NOTE: non-dimensional c3bp dynamics in the usual rotating (synodic) reference frame.
+// Expressed in terms of canonical state variables.
+inline constexpr auto cr3bp = [](auto &&...kw_args) -> std::vector<std::pair<expression, expression>> {
+    return std::apply(detail::cr3bp_impl, detail::cr3bp_common_opts(std::forward<decltype(kw_args)>(kw_args)...));
+};
+
+inline constexpr auto cr3bp_jacobi = [](auto &&...kw_args) -> expression {
+    return std::apply(detail::cr3bp_jacobi_impl,
+                      detail::cr3bp_common_opts(std::forward<decltype(kw_args)>(kw_args)...));
+};
+
+} // namespace model
+
+HEYOKA_END_NAMESPACE
+
+#endif

include/heyoka/models.hpp

@@ -9,6 +9,7 @@
 #ifndef HEYOKA_MODELS_HPP
 #define HEYOKA_MODELS_HPP
 
+#include <heyoka/model/cr3bp.hpp>
 #include <heyoka/model/fixed_centres.hpp>
 #include <heyoka/model/mascon.hpp>
 #include <heyoka/model/nbody.hpp>

src/model/cr3bp.cpp

@@ -0,0 +1,118 @@
+// Copyright 2020, 2021, 2022, 2023 Francesco Biscani ([email protected]), Dario Izzo ([email protected])
+//
+// This file is part of the heyoka 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/.
+
+#include <cmath>
+#include <stdexcept>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include <fmt/core.h>
+
+#include <heyoka/config.hpp>
+#include <heyoka/expression.hpp>
+#include <heyoka/math/pow.hpp>
+#include <heyoka/math/sqrt.hpp>
+#include <heyoka/model/cr3bp.hpp>
+#include <heyoka/number.hpp>
+
+HEYOKA_BEGIN_NAMESPACE
+
+namespace model::detail
+{
+
+namespace
+{
+
+void cr3bp_check_mu(const expression &mu)
+{
+    if (const auto *nptr = std::get_if<number>(&mu.value())) {
+        std::visit(
+            [](const auto &v) {
+                using std::isfinite;
+
+                if (!isfinite(v) || v <= 0 || v >= .5) {
+                    throw std::invalid_argument(fmt::format("The 'mu' parameter in a CR3BP must be in the range (0, "
+                                                            "0.5), but a value of {} was provided instead",
+                                                            v));
+                }
+            },
+            nptr->value());
+    }
+}
+
+} // namespace
+
+std::vector<std::pair<expression, expression>> cr3bp_impl(const expression &mu)
+{
+    cr3bp_check_mu(mu);
+
+    // Init the state variables,
+    auto [px, py, pz, x, y, z] = make_vars("px", "py", "pz", "x", "y", "z");
+
+    // x - mu.
+    // NOTE: leave it unfixed so that constant folding can take place
+    // in the computation of x - mu + 1 in the common case where mu is a numerical
+    // constant. This allows for better ILP as x - mu and x + (-mu + 1.) can be
+    // computed in parallel.
+    const auto x_m_mu = x - mu;
+    // x - mu + 1.
+    const auto x_m_mu_p1 = x_m_mu + 1.;
+    // y**2 + z**2.
+    const auto y_p_z_2 = fix_nn(y * y + z * z);
+    // rp1**2.
+    const auto rp1_2 = x_m_mu * x_m_mu + y_p_z_2;
+    // rp2**2.
+    const auto rp2_2 = x_m_mu_p1 * x_m_mu_p1 + y_p_z_2;
+    // (1 - mu) / rp1**3.
+    const auto g1 = fix_nn((1. - mu) * pow(rp1_2, -3. / 2));
+    // mu / rp2**3.
+    const auto g2 = fix_nn(mu * pow(rp2_2, -3. / 2));
+    // g1 + g2.
+    const auto g1_g2 = fix_nn(g1 + g2);
+
+    const auto xdot = px + y;
+    const auto ydot = py - x;
+    const auto zdot = pz;
+    const auto pxdot = py - g1 * x_m_mu - g2 * x_m_mu_p1;
+    const auto pydot = -px - g1_g2 * y;
+    const auto pzdot = -g1_g2 * z;
+
+    return {prime(x) = xdot, prime(y) = ydot, prime(z) = zdot, prime(px) = pxdot, prime(py) = pydot, prime(pz) = pzdot};
+}
+
+expression cr3bp_jacobi_impl(const expression &mu)
+{
+    cr3bp_check_mu(mu);
+
+    // Init the state variables,
+    auto [px, py, pz, x, y, z] = make_vars("px", "py", "pz", "x", "y", "z");
+
+    // x - mu.
+    const auto x_m_mu = x - mu;
+    // x - mu + 1.
+    const auto x_m_mu_p1 = x_m_mu + 1.;
+    // y**2 + z**2.
+    const auto y_p_z_2 = fix_nn(y * y + z * z);
+    // rp1**2.
+    const auto rp1_2 = x_m_mu * x_m_mu + y_p_z_2;
+    // rp2**2.
+    const auto rp2_2 = x_m_mu_p1 * x_m_mu_p1 + y_p_z_2;
+    // (1 - mu) / rp1.
+    const auto g1 = fix_nn((1. - mu) / sqrt(rp1_2));
+    // mu / rp2.
+    const auto g2 = fix_nn(mu / sqrt(rp2_2));
+
+    const auto kin = 0.5 * fix_nn(px * px + py * py + pz * pz);
+
+    return kin + y * px - x * py - g1 - g2;
+}
+
+} // namespace model::detail
+
+HEYOKA_END_NAMESPACE

test/CMakeLists.txt

@@ -137,6 +137,7 @@ ADD_HEYOKA_TESTCASE(model_nbody)
 ADD_HEYOKA_TESTCASE(model_fixed_centres)
 ADD_HEYOKA_TESTCASE(model_rotating)
 ADD_HEYOKA_TESTCASE(model_mascon)
+ADD_HEYOKA_TESTCASE(model_cr3bp)
 ADD_HEYOKA_TESTCASE(step_callback)
 ADD_HEYOKA_TESTCASE(llvm_state_mem_cache)
 

test/model_cr3bp.cpp

@@ -0,0 +1,115 @@
+// Copyright 2020, 2021, 2022, 2023 Francesco Biscani ([email protected]), Dario Izzo ([email protected])
+//
+// This file is part of the heyoka 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/.
+
+#include <initializer_list>
+#include <stdexcept>
+#include <vector>
+
+#include <heyoka/expression.hpp>
+#include <heyoka/kw.hpp>
+#include <heyoka/llvm_state.hpp>
+#include <heyoka/model/cr3bp.hpp>
+#include <heyoka/taylor.hpp>
+
+#include "catch.hpp"
+#include "test_utils.hpp"
+
+using namespace heyoka;
+using namespace heyoka_test;
+
+TEST_CASE("basic")
+{
+    using Catch::Matchers::Message;
+
+    {
+        auto dyn = model::cr3bp();
+
+        REQUIRE(dyn.size() == 6u);
+        REQUIRE(dyn[0].first == "x"_var);
+        REQUIRE(dyn[1].first == "y"_var);
+        REQUIRE(dyn[2].first == "z"_var);
+        REQUIRE(dyn[3].first == "px"_var);
+        REQUIRE(dyn[4].first == "py"_var);
+        REQUIRE(dyn[5].first == "pz"_var);
+    }
+
+    // Energy conservation.
+    {
+        auto dyn = model::cr3bp();
+
+        const std::vector init_state = {-0.45, 0.80, 0.00, -0.80, -0.45, 0.58};
+
+        auto ta = taylor_adaptive{dyn, init_state};
+
+        REQUIRE(ta.get_decomposition().size() == 36u);
+
+        ta.propagate_until(20.);
+
+        auto [x, y, z, px, py, pz] = make_vars("x", "y", "z", "px", "py", "pz");
+
+        llvm_state s;
+
+        const auto dc = add_cfunc<double>(s, "jac", {model::cr3bp_jacobi()}, kw::vars = {x, y, z, px, py, pz});
+
+        REQUIRE(dc.size() == 25u);
+
+        s.compile();
+
+        auto *cf
+            = reinterpret_cast<void (*)(double *, const double *, const double *, const double *)>(s.jit_lookup("jac"));
+        double E0 = 0;
+        cf(&E0, init_state.data(), nullptr, nullptr);
+
+        double E = 0;
+        cf(&E, ta.get_state().data(), nullptr, nullptr);
+
+        REQUIRE(E == approximately(E0));
+    }
+
+    // With custom mu.
+    {
+        auto dyn = model::cr3bp(kw::mu = 1e-2);
+
+        const std::vector init_state = {-0.45, 0.80, 0.00, -0.80, -0.45, 0.58};
+
+        auto ta = taylor_adaptive{dyn, init_state};
+
+        REQUIRE(ta.get_decomposition().size() == 36u);
+
+        ta.propagate_until(20.);
+
+        auto [x, y, z, px, py, pz] = make_vars("x", "y", "z", "px", "py", "pz");
+
+        llvm_state s;
+
+        const auto dc
+            = add_cfunc<double>(s, "jac", {model::cr3bp_jacobi(kw::mu = 1e-2)}, kw::vars = {x, y, z, px, py, pz});
+
+        REQUIRE(dc.size() == 25u);
+
+        s.compile();
+
+        auto *cf
+            = reinterpret_cast<void (*)(double *, const double *, const double *, const double *)>(s.jit_lookup("jac"));
+        double E0 = 0;
+        cf(&E0, init_state.data(), nullptr, nullptr);
+
+        double E = 0;
+        cf(&E, ta.get_state().data(), nullptr, nullptr);
+
+        REQUIRE(E == approximately(E0));
+    }
+
+    // Error modes.
+    REQUIRE_THROWS_MATCHES(model::cr3bp(kw::mu = -1.), std::invalid_argument,
+                           Message("The 'mu' parameter in a CR3BP must be in the range (0, "
+                                   "0.5), but a value of -1 was provided instead"));
+    REQUIRE_THROWS_MATCHES(model::cr3bp_jacobi(kw::mu = 1.), std::invalid_argument,
+                           Message("The 'mu' parameter in a CR3BP must be in the range (0, "
+                                   "0.5), but a value of 1 was provided instead"));
+}