feat: generalize finite difference solver jacobian computation (#268)

* feat: generalize finite difference solver jacobian computation * feat: fix tests
open-space-collective · Nov 7, 2023 · c8fe844 · c8fe844
1 parent f6862fb
commit c8fe844
Show file tree

Hide file tree

Showing 5 changed files with 80 additions and 63 deletions.
diff --git a/bindings/python/src/OpenSpaceToolkitAstrodynamicsPy/Solvers/FiniteDifferenceSolver.cpp b/bindings/python/src/OpenSpaceToolkitAstrodynamicsPy/Solvers/FiniteDifferenceSolver.cpp
@@ -7,6 +7,7 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Solvers_FiniteDifferenceSolver(pybin
     using namespace pybind11;
 
     using ostk::core::ctnr::Array;
+    using ostk::core::ctnr::Size;
 
     using ostk::math::obj::VectorXd;
     using ostk::math::obj::MatrixXd;
@@ -21,7 +22,7 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Solvers_FiniteDifferenceSolver(pybin
         aModule,
         "FiniteDifferenceSolver",
         R"doc(
-            A Finite Difference Solver to compute the gradient, and state transition matrix of a function.
+            A Finite Difference Solver to compute the gradient, and jacobian of a function.
 
             Group:
                 solvers
@@ -97,52 +98,58 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Solvers_FiniteDifferenceSolver(pybin
         )
 
         .def(
-            "compute_state_transition_matrix",
+            "compute_jacobian",
             +[](const ostk::astro::solvers::FiniteDifferenceSolver& solver,
                 const State& aState,
                 const Array<Instant>& anInstantArray,
-                std::function<MatrixXd(const State&, const Array<Instant>&)> generateStateCoordinates) -> MatrixXd
+                const std::function<MatrixXd(const State&, const Array<Instant>&)>& generateStateCoordinates,
+                const Size& aCoordinatesDimension) -> MatrixXd
             {
-                return solver.computeStateTransitionMatrix(aState, anInstantArray, generateStateCoordinates);
+                return solver.computeJacobian(aState, anInstantArray, generateStateCoordinates, aCoordinatesDimension);
             },
             R"doc(
-                Compute the state transition matrix.
+                Compute the jacobian.
 
                 Args:
                     state (State): The state.
                     instants (Array(Instant)): The instants at which to calculate the STM.
                     generate_states_coordinates (function): The function to get the states.
+                    coordinates_dimension (int): The dimension of the coordinates produced by `generate_states_coordinates`.
 
                 Returns:
-                    np.array: The state transition matrix.
+                    np.array: The jacobian.
             )doc",
             arg("state"),
             arg("instants"),
-            arg("generate_states_coordinates")
+            arg("generate_states_coordinates"),
+            arg("coordinates_dimension")
         )
         .def(
-            "compute_state_transition_matrix",
+            "compute_jacobian",
             +[](const ostk::astro::solvers::FiniteDifferenceSolver& solver,
                 const State& aState,
                 const Instant& anInstant,
-                std::function<VectorXd(const State&, const Instant&)> generateStateCoordinates) -> MatrixXd
+                const std::function<VectorXd(const State&, const Instant&)>& generateStateCoordinates,
+                const Size& aCoordinatesDimension) -> MatrixXd
             {
-                return solver.computeStateTransitionMatrix(aState, anInstant, generateStateCoordinates);
+                return solver.computeJacobian(aState, anInstant, generateStateCoordinates, aCoordinatesDimension);
             },
             R"doc(
-                Compute the state transition matrix.
+                Compute the jacobian.
 
                 Args:
                     state (State): The state.
                     instant (Instant): The instant at which to calculate the STM.
                     generate_state_coordinates (function): The function to get the state.
+                    coordinates_dimension (int): The dimension of the coordinates produced by `generate_state_coordinates`
 
                 Returns:
-                    np.array: The state transition matrix.
+                    np.array: The jacobian.
             )doc",
             arg("state"),
             arg("instant"),
-            arg("generate_state_coordinates")
+            arg("generate_state_coordinates"),
+            arg("coordinates_dimension")
         )
         .def(
             "compute_gradient",
@@ -160,7 +167,7 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Solvers_FiniteDifferenceSolver(pybin
                     generate_state_coordinates (function): The function to generate the state coordinates.
 
                 Returns:
-                    np.array: The state transition matrix.
+                    np.array: The jacobian.
             )doc",
             arg("state"),
             arg("generate_state_coordinates")

diff --git a/bindings/python/test/solvers/test_finite_difference_solver.py b/bindings/python/test/solvers/test_finite_difference_solver.py
@@ -126,35 +126,37 @@ def test_string_from_type(self):
             == "Forward"
         )
 
-    def test_compute_state_transition_matrix_array(
+    def test_compute_jacobian_array(
         self,
         finite_difference_solver: FiniteDifferenceSolver,
         state: State,
         instants: list[Instant],
         generate_states_coordinates: callable,
     ):
-        stm = finite_difference_solver.compute_state_transition_matrix(
+        stm = finite_difference_solver.compute_jacobian(
             state=state,
             instants=instants,
             generate_states_coordinates=generate_states_coordinates,
+            coordinates_dimension=2,
         )
         assert isinstance(stm, np.ndarray)
         assert stm.shape == (
-            len(state.get_coordinates()),
             len(state.get_coordinates()) * len(instants),
+            len(state.get_coordinates()),
         )
 
-    def test_compute_state_transition_matrix_single(
+    def test_compute_jacobian_single(
         self,
         finite_difference_solver: FiniteDifferenceSolver,
         state: State,
         generate_state_coordinates: callable,
         instant: Instant,
     ):
-        stm = finite_difference_solver.compute_state_transition_matrix(
+        stm = finite_difference_solver.compute_jacobian(
             state=state,
             instant=instant,
             generate_state_coordinates=generate_state_coordinates,
+            coordinates_dimension=2,
         )
         assert isinstance(stm, np.ndarray)
         assert stm.shape == (

diff --git a/include/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.hpp b/include/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.hpp
@@ -4,6 +4,8 @@
 #define __OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver__
 
 #include <OpenSpaceToolkit/Core/Containers/Array.hpp>
+#include <OpenSpaceToolkit/Core/Types/Real.hpp>
+#include <OpenSpaceToolkit/Core/Types/Size.hpp>
 
 #include <OpenSpaceToolkit/Mathematics/Objects/Vector.hpp>
 
@@ -20,6 +22,7 @@ namespace solvers
 {
 
 using ostk::core::types::Real;
+using ostk::core::types::Size;
 using ostk::core::types::String;
 using ostk::core::ctnr::Array;
 
@@ -86,48 +89,51 @@ class FiniteDifferenceSolver
 
     Duration getStepDuration() const;
 
-    /// @brief                  Compute the State Transition Matrix by perturbing the coordinates
+    /// @brief                  Compute the Jacobian by perturbing the coordinates
     ///
     /// @param                  [in] aState A state.
     /// @param                  [in] anInstantArray An array of instants.
     /// @param                  [in] generateStateCoordinates Callable to generate coordinates of States at the
     /// requested Instants.
-    /// @param                  [in] aStepPercentage The step percentage to use for the perturbation. Defaults to 1e-3.
+    /// @param                  [in] aCoordinatesDimension The dimension of the coordinates produced by
+    /// `generateStateCoordinates`.
     ///
-    /// @return                 The State Transition Matrix
+    /// @return                 The Jacobian
 
-    MatrixXd computeStateTransitionMatrix(
+    MatrixXd computeJacobian(
         const State& aState,
         const Array<Instant>& anInstantArray,
-        std::function<MatrixXd(const State&, const Array<Instant>&)> generateStateCoordinates
+        const std::function<MatrixXd(const State&, const Array<Instant>&)>& generateStateCoordinates,
+        const Size& aCoordinatesDimension
     ) const;
 
-    /// @brief                  Compute the State Transition Matrix by perturbing the coordinates
+    /// @brief                  Compute the Jacobian by perturbing the coordinates
     ///
     /// @param                  [in] aState A state.
     /// @param                  [in] anInstant An instant.
     /// @param                  [in] generateStateCoordinates Callable to generate coordinates of a State at the
     /// requested Instant.
-    /// @param                  [in] aStepPercentage The step percentage to use for the perturbation. Defaults to 1e-3.
-    ///
-    /// @return                 The State Transition Matrix
+    /// @param                  [in] aCoordinatesDimension The dimension of the coordinates produced by
+    /// `generateStateCoordinates`.
+    /// @return                 The Jacobian
 
-    MatrixXd computeStateTransitionMatrix(
+    MatrixXd computeJacobian(
         const State& aState,
         const Instant& anInstant,
-        std::function<VectorXd(const State&, const Instant&)> generateStateCoordinates
+        const std::function<VectorXd(const State&, const Instant&)>& generateStateCoordinates,
+        const Size& aCoordinatesDimension
     ) const;
 
-    /// @brief                  Compute the gradient
+    /// @brief                  Compute the gradient.
     ///
     /// @param                  [in] aState The state to compute the gradient of.
     /// @param                  [in] generateStateCoordinates Callable to generate coordinates of a State at the
     /// requested Instant.
     ///
-    /// @return                 The gradient
+    /// @return                 The gradient.
 
     VectorXd computeGradient(
-        const State& aState, std::function<VectorXd(const State&, const Instant&)> generateStateCoordinates
+        const State& aState, const std::function<VectorXd(const State&, const Instant&)>& generateStateCoordinates
     ) const;
 
     /// @brief                  Print the solver.

diff --git a/src/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.cpp b/src/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.cpp
@@ -50,15 +50,17 @@ Duration FiniteDifferenceSolver::getStepDuration() const
     return stepDuration_;
 }
 
-MatrixXd FiniteDifferenceSolver::computeStateTransitionMatrix(
+MatrixXd FiniteDifferenceSolver::computeJacobian(
     const State& aState,
     const Array<Instant>& anInstantArray,
-    std::function<MatrixXd(const State&, const Array<Instant>&)> generateStateCoordinates
+    const std::function<MatrixXd(const State&, const Array<Instant>&)>& generateStateCoordinates,
+    const Size& aCoordinatesDimension
 ) const
 {
     const StateBuilder stateBuilder = {aState};
 
     const Size stateVectorDimension = aState.getSize();
+    const Size numberOfInstants = anInstantArray.getSize();
 
     std::function<MatrixXd(const Real&, VectorXd, const Index&)> computeBlock;
 
@@ -128,45 +130,48 @@ MatrixXd FiniteDifferenceSolver::computeStateTransitionMatrix(
             throw ostk::core::error::runtime::Wrong("Finite Difference Solver Type.");
     }
 
-    MatrixXd A = MatrixXd::Zero(stateVectorDimension, stateVectorDimension * anInstantArray.size());
+    MatrixXd A = MatrixXd::Zero(aCoordinatesDimension * numberOfInstants, stateVectorDimension);
 
     for (Index i = 0; i < stateVectorDimension; ++i)
     {
         const Real stepSize = (aState.accessCoordinates()(i) * stepPercentage_ != 0.0)
                                 ? aState.accessCoordinates()(i) * stepPercentage_
                                 : stepPercentage_;
 
-        const VectorXd coords = aState.accessCoordinates();
-        const MatrixXd differencedCoordinates = computeBlock(stepSize, aState.getCoordinates(), i);
+        MatrixXd differencedCoordinates = computeBlock(stepSize, aState.getCoordinates(), i);
 
-        for (Index j = 0; j < anInstantArray.size(); ++j)
-        {
-            A.block(0, i + j * stateVectorDimension, stateVectorDimension, 1) = differencedCoordinates.col(j);
-        }
+        const VectorXd columnStackedCoordinates =
+            Eigen::Map<VectorXd>(differencedCoordinates.data(), differencedCoordinates.size());
+
+        A.col(i) = columnStackedCoordinates;
     }
 
     return A;
 }
 
-MatrixXd FiniteDifferenceSolver::computeStateTransitionMatrix(
+MatrixXd FiniteDifferenceSolver::computeJacobian(
     const State& aState,
     const Instant& anInstant,
-    std::function<VectorXd(const State&, const Instant&)> generateStateCoordinates
+    const std::function<VectorXd(const State&, const Instant&)>& generateStateCoordinates,
+    const Size& aCoordinatesDimension
+
 ) const
 {
     const auto generateStatesCoordinates =
         [&generateStateCoordinates](const State& anInputState, const Array<Instant>& anInstantArray) -> MatrixXd
     {
         const VectorXd coordinates = generateStateCoordinates(anInputState, anInstantArray[0]);
 
-        return {coordinates};
+        const MatrixXd coordinatesMatrix = {coordinates};
+
+        return coordinatesMatrix;
     };
 
-    return computeStateTransitionMatrix(aState, {anInstant}, generateStatesCoordinates);
+    return computeJacobian(aState, {anInstant}, generateStatesCoordinates, aCoordinatesDimension);
 }
 
 VectorXd FiniteDifferenceSolver::computeGradient(
-    const State& aState, std::function<VectorXd(const State&, const Instant&)> generateStateCoordinates
+    const State& aState, const std::function<VectorXd(const State&, const Instant&)>& generateStateCoordinates
 ) const
 {
     switch (type_)

diff --git a/test/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.test.cpp b/test/OpenSpaceToolkit/Astrodynamics/Solvers/FiniteDifferenceSolver.test.cpp
@@ -127,26 +127,28 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, Getters)
     }
 }
 
-TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeStateTransitionMatrix)
+TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeJacobian)
 {
     {
         const Array<Instant> instants = {
             Instant::J2000() + Duration::Seconds(100.0),
             Instant::J2000() + Duration::Seconds(200.0),
+            Instant::J2000() + Duration::Seconds(300.0),
         };
 
-        MatrixXd expectedJacobian(2, 4);
-        expectedJacobian << std::cos(100.0), std::sin(100.0), std::cos(200.0), std::sin(200.0), std::sin(-100.0),
-            std::cos(100.0), std::sin(-200.0), std::cos(200.0);
+        MatrixXd expectedJacobian(2 * instants.getSize(), 2);
+
+        expectedJacobian << std::cos(100.0), std::sin(100.0), std::sin(-100.0), std::cos(100.0), std::cos(200.0),
+            std::sin(200.0), std::sin(-200.0), std::cos(200.0), std::cos(300.0), std::sin(300.0), std::sin(-300.0),
+            std::cos(300.0);
 
         {
             const FiniteDifferenceSolver solver = {
                 FiniteDifferenceSolver::Type::Central,
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instants, generateStatesCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instants, generateStatesCoordinates, 2);
 
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }
@@ -157,8 +159,7 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeSta
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instants, generateStatesCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instants, generateStatesCoordinates, 2);
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }
 
@@ -168,8 +169,7 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeSta
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instants, generateStatesCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instants, generateStatesCoordinates, 2);
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }
     }
@@ -186,8 +186,7 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeSta
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instant, generateStateCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instant, generateStateCoordinates, 2);
 
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }
@@ -198,8 +197,7 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeSta
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instant, generateStateCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instant, generateStateCoordinates, 2);
 
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }
@@ -210,8 +208,7 @@ TEST_F(OpenSpaceToolkit_Astrodynamics_Solvers_FiniteDifferenceSolver, ComputeSta
                 defaultStepPercentage_,
                 defaultStepDuration_,
             };
-            const MatrixXd jacobian =
-                solver.computeStateTransitionMatrix(initialState_, instant, generateStateCoordinates);
+            const MatrixXd jacobian = solver.computeJacobian(initialState_, instant, generateStateCoordinates, 2);
 
             EXPECT_TRUE(jacobian.isApprox(expectedJacobian, 1e-12));
         }