feat: enrich StateBuilder capabilities (#239)

* parent b12aadb
author Kyle Cochran <kyle.cochran@loftorbital.com> 1696962698 +0000
committer Pau Hebrero <pau.hebrero@gmail.com> 1697119993 +0000

chore: rebase

* feat: StateBuilder operators, reduction and expansion

---------

Co-authored-by: Kyle Cochran <kyle.cochran@loftorbital.com>
Co-authored-by: Pau Hebrero <pau.hebrero@gmail.com>

bindings/python/src/OpenSpaceToolkitAstrodynamicsPy/Trajectory/StateBuilder.cpp

@@ -14,6 +14,7 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Trajectory_StateBuilder(pybind11::mo
     using ostk::astro::trajectory::State;
     using ostk::astro::trajectory::StateBuilder;
     using ostk::astro::trajectory::state::CoordinatesBroker;
+    using ostk::astro::trajectory::state::CoordinatesSubset;
 
     class_<StateBuilder>(aModule, "StateBuilder")
 
@@ -27,16 +28,35 @@ inline void OpenSpaceToolkitAstrodynamicsPy_Trajectory_StateBuilder(pybind11::mo
             arg("frame"),
             arg("coordinates_broker")
         )
+        .def(init<const State&>(), arg("state"))
 
         .def(self == self)
         .def(self != self)
+        .def(
+            "__add__",
+            [](const StateBuilder& aStateBuilder, const Shared<const CoordinatesSubset>& aCoordinatesSubsetSPtr)
+            {
+                return aStateBuilder + aCoordinatesSubsetSPtr;
+            },
+            is_operator()
+        )
+        .def(
+            "__sub__",
+            [](const StateBuilder& aStateBuilder, const Shared<const CoordinatesSubset>& aCoordinatesSubsetSPtr)
+            {
+                return aStateBuilder - aCoordinatesSubsetSPtr;
+            },
+            is_operator()
+        )
 
         .def("__str__", &(shiftToString<StateBuilder>))
         .def("__repr__", &(shiftToString<StateBuilder>))
 
         .def("is_defined", &StateBuilder::isDefined)
 
-        .def("build_state", &StateBuilder::buildState)
+        .def("build", &StateBuilder::build)
+        .def("reduce", &StateBuilder::reduce)
+        .def("expand", &StateBuilder::expand)
 
         .def("get_coordinates_subsets", &StateBuilder::getCoordinatesSubsets)
         .def("get_frame", &StateBuilder::getFrame)

bindings/python/test/trajectory/test_state_builder.py

@@ -38,17 +38,32 @@ def coordinates_subsets() -> list[CoordinatesSubset]:
     return [CartesianPosition.default(), CartesianVelocity.default()]
 
 
+@pytest.fixture
+def coordinates() -> list[float]:
+    return [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
+
+
 @pytest.fixture
 def coordinates_broker(coordinates_subsets: list[CoordinatesSubset]) -> CoordinatesBroker:
     return CoordinatesBroker(coordinates_subsets)
 
 
+@pytest.fixture
+def state(
+    instant: Instant,
+    coordinates: list[float],
+    frame: Frame,
+    coordinates_broker: CoordinatesBroker,
+) -> State:
+    return State(instant, coordinates, frame, coordinates_broker)
+
+
 @pytest.fixture
 def state_builder(frame: Frame, coordinates_broker: CoordinatesBroker) -> State:
     return StateBuilder(frame, coordinates_broker)
 
 
-class TestState:
+class TestStateBuilder:
     def test_broker_constructor(
         self,
         frame: Frame,
@@ -69,17 +84,38 @@ def test_subsets_constructor(
         assert isinstance(builder, StateBuilder)
         assert builder.is_defined()
 
+    def test_state_constructor(
+        self,
+        state: State,
+    ):
+        builder = StateBuilder(state)
+        assert builder is not None
+        assert isinstance(builder, StateBuilder)
+        assert builder.is_defined()
+
     def test_comparators(self, state_builder: StateBuilder):
         assert (state_builder == state_builder) is True
         assert (state_builder != state_builder) is False
 
-    def test_build_state(
+    def test_operators(
+        self,
+        state_builder: StateBuilder,
+    ):
+        added_builder: StateBuilder = state_builder + CoordinatesSubset.mass()
+        assert isinstance(added_builder, StateBuilder)
+        assert state_builder != added_builder
+
+        subtracted_builder: StateBuilder = state_builder - CartesianPosition.default()
+        assert isinstance(subtracted_builder, StateBuilder)
+        assert state_builder != subtracted_builder
+
+    def test_build(
         self,
         instant: Instant,
         state_builder: StateBuilder,
     ):
         coordinates = [1, 2, 3, 1, 2, 3]
-        state: State = state_builder.build_state(instant, coordinates)
+        state: State = state_builder.build(instant, coordinates)
 
         assert state is not None
         assert isinstance(state, State)
@@ -89,6 +125,40 @@ def test_build_state(
         assert state.get_frame() == state_builder.get_frame()
         assert state.get_coordinates_subsets() == state_builder.get_coordinates_subsets()
 
+    def test_reduce(
+        self,
+        state: State,
+    ):
+        builder = StateBuilder(state.get_frame(), [CartesianPosition.default()])
+        reduced_state: State = builder.reduce(state)
+
+        assert isinstance(reduced_state, State)
+        assert state != reduced_state
+
+    def test_expand(
+        self,
+        state: State,
+    ):
+        builder = StateBuilder(
+            state.get_frame(),
+            [
+                CartesianPosition.default(),
+                CartesianVelocity.default(),
+                CoordinatesSubset.mass(),
+            ],
+        )
+        default_state: State = State(
+            state.get_instant(),
+            [100],
+            state.get_frame(),
+            CoordinatesBroker([CoordinatesSubset.mass()]),
+        )
+        expanded_state: State = builder.expand(state, default_state)
+
+        assert isinstance(expanded_state, State)
+        assert state != expanded_state
+        assert default_state != expanded_state
+
     def test_getters(
         self,
         state_builder: StateBuilder,

include/OpenSpaceToolkit/Astrodynamics/Trajectory/State/CoordinatesSubset.hpp

@@ -41,7 +41,7 @@ class CoordinatesSubset
     /// @brief              Constructor
     ///
     ///                     The default CoordinatesSubset instance is frame-invariant and implements element-wise
-    ///                     addition/substraction.
+    ///                     addition/subtraction.
     /// @code
     ///                     CoordinateSubset coordinateSubset = {aName, aSize};
     /// @endcode

include/OpenSpaceToolkit/Astrodynamics/Trajectory/StateBuilder.hpp

@@ -4,7 +4,9 @@
 #define __OpenSpaceToolkit_Astrodynamics_Trajectory_StateBuilder__
 
 #include <OpenSpaceToolkit/Core/Containers/Array.hpp>
+#include <OpenSpaceToolkit/Core/Types/Integer.hpp>
 #include <OpenSpaceToolkit/Core/Types/Shared.hpp>
+#include <OpenSpaceToolkit/Core/Types/Size.hpp>
 
 #include <OpenSpaceToolkit/Mathematics/Objects/Vector.hpp>
 
@@ -22,7 +24,9 @@ namespace astro
 namespace trajectory
 {
 
+using ostk::core::types::Integer;
 using ostk::core::types::Shared;
+using ostk::core::types::Size;
 using ostk::core::ctnr::Array;
 
 using ostk::math::obj::VectorXd;
@@ -59,6 +63,12 @@ class StateBuilder
 
     StateBuilder(const Shared<const Frame>& aFrameSPtr, const Shared<const CoordinatesBroker>& aCoordinatesBrokerSPtr);
 
+    /// @brief                  Constructor.
+    ///
+    /// @param                  [in] aState The state to be used as a template
+
+    StateBuilder(const State& aState);
+
     /// @brief                  Equality operator.
     ///
     /// @param                  [in] aStateBuilder The StateBuilder to compare to
@@ -73,6 +83,22 @@ class StateBuilder
 
     bool operator!=(const StateBuilder& aStateBuilder) const;
 
+    /// @brief                  Return a new StateBuilder with the additional CoordinatesSubset.
+    ///
+    /// @param                  [in] aCoordinatesSubsetSPtr The CoordinatesSubset to append
+    ///
+    /// @return                 A new StateBuilder
+
+    const StateBuilder operator+(const Shared<const CoordinatesSubset>& aCoordinatesSubsetSPtr) const;
+
+    /// @brief                  Return a new StateBuilder without the given CoordinatesSubset.
+    ///
+    /// @param                  [in] aCoordinatesSubsetSPtr The CoordinatesSubset to remove
+    ///
+    /// @return                 A new StateBuilder
+
+    const StateBuilder operator-(const Shared<const CoordinatesSubset>& aCoordinatesSubsetSPtr) const;
+
     /// @brief                  Stream insertion operator.
     ///
     /// @param                  [in] anOutputStream The output stream to insert into
@@ -91,7 +117,22 @@ class StateBuilder
     ///
     /// @return                 A State linked to the Frame and Coordinates Broker of the StateBuilder
 
-    const State buildState(const Instant& anInstant, const VectorXd& aCoordinates) const;
+    const State build(const Instant& anInstant, const VectorXd& aCoordinates) const;
+
+    /// @brief                  Produce a State with the CoordinatesSubsets specified by the StateBuilder.
+    ///
+    /// @param                  [in] aState the state from which the coordinates will be taken.
+    /// @return                 A State with the CoordinatesSubsets of the StateBuilder.
+
+    const State reduce(const State& aState) const;
+
+    /// @brief                  Produce a State with the CoordinatesSubsets specified by the StateBuilder.
+    ///
+    /// @param                  [in] aState the state from which the coordinates will be taken.
+    /// @param                  [in] defaultState the state from which missing coordinates will be taken.
+    /// @return                 A State with the CoordinatesSubsets of the StateBuilder.
+
+    const State expand(const State& aState, const State& defaultState) const;
 
     /// @brief                  Accessor for the reference frame.
     ///

src/OpenSpaceToolkit/Astrodynamics/Trajectory/State.cpp

@@ -101,6 +101,11 @@ bool State::operator==(const State& aState) const
         return false;
     }
 
+    if (this->getSize() != aState.getSize())
+    {
+        return false;
+    }
+
     for (const Shared<const CoordinatesSubset>& subset : this->coordinatesBrokerSPtr_->accessSubsets())
     {
         if (!aState.coordinatesBrokerSPtr_->hasSubset(subset))