Very slow (factor 10000) propagation for spacecraft with mass and custom state

[briha]

Hi! When propagating a spacecraft with custom_state and mass propagators, it takes forever (~200 seconds for rkdp_89 with $abs=rel=10^{-10}$, but comparable for other options, too) when doing it through tudatpy directly (using a SingleArcSimulator). When using the exact same integrator manually, i.e.

std::shared_ptr<tni::NumericalIntegrator<double, Eigen::Vector14d>> integrator
        = tni::createIntegrator(stateDerivativeFunction, initialState, integratorSettings->initialTimeDeprecated_, integratorSettings);

std::vector<Eigen::Vector15d> stateHistory;

double currentTime = integratorSettings->initialTimeDeprecated_;
double timeStep = integratorSettings->initialTimeStep_;
Eigen::Vector15d currentState;
currentState << currentTime, initialState;
stateHistory.push_back(currentState);

unsigned int counter = 0;
while (currentTime < finalTime && counter++ < maxSteps) {
    currentState.tail<14>() = integrator->performIntegrationStep(timeStep);
    currentTime = currentState[0] = integrator->getCurrentIndependentVariable();
    timeStep = integrator->getNextStepSize();

    stateHistory.push_back(currentState);
}

where the state-derivative represents the fullstate vector (i.e., state, mass, custom state) it takes about 0.02s (a factor 10,000 difference).

Here's an example of what my code looks like. This is not a minimal working example, but should give an idea of how I invoke the integration. Please let me know if anything is unclear (DynamicsSimulator, Spacecraft, constants, and a number of other classes, constants, and function are locally imported which is not shown here):

from typing import Optional, Type

import numpy as np
from tudatpy.kernel.astro.element_conversion import mee_to_cartesian
from tudatpy.kernel.numerical_simulation import environment_setup as env_setup
from tudatpy.kernel.numerical_simulation import propagation_setup as prop_setup
from tudatpy.kernel import numerical_simulation

from fast_trajectory_optimization import FastOptimalSundmanControlTrajectory, FastOptimalControlTrajectoryController, FastOptimalControlTrajectoryBase

class TudatPropagatorV2(ManualPropagatorV2):
    """
    Propagation routine using the fullstate derivatives as supplied by the :py:module:`tudat` library. This is the only implementation of
    :py:class:`trajectory_generation.propagation2.BasePropagatorV2` that is limited to the tudat integration routines.
    """
    def __init__(self, spacecraft: Optional[Spacecraft] = None,
                 continuation_parameter: float = 1e-6,
                 gravitational_parameter: float = constants.mu_sun,
                 g0: float = constants.g_0,
                 trajectory_class: Type[FastOptimalControlTrajectoryBase] = FastOptimalSundmanControlTrajectory):

        super().__init__(spacecraft,
                         continuation_parameter=continuation_parameter,
                         gravitational_parameter=gravitational_parameter,
                         g0=g0,
                         trajectory_class=trajectory_class)

        bodies_to_create = ["Sun"]
        body_settings = env_setup.get_default_body_settings(bodies_to_create, "Sun", "ECLIPJ2000")
        body_settings.get("Sun").gravity_field_settings = env_setup.gravity_field.central(self.gravitational_parameter_nd)

        bodies = env_setup.create_system_of_bodies(body_settings)
        bodies.create_empty_body("Spacecraft")
        self.controller = FastOptimalControlTrajectoryController(bodies, "Spacecraft", self.trajectory)

        self.bodies = bodies

    # overriding the tudat integrate function to use the actual tudat propagation routine
    def _tudat_integrate_from_settings(self, integration_settings, initial_time: float, final_time: float, initial_fullstate: np.ndarray):
        if initial_fullstate.size == 15:
            initial_fullstate = initial_fullstate[1:]
        bodies = self.bodies
        central_bodies = ["Sun"]
        bodies_to_propagate = ["Spacecraft"]

        rotation_model_settings = \
            env_setup.rotation_model.custom_inertial_direction_based(
                inertial_body_axis_direction=self.controller.get_thrust_direction,
                base_frame="ECLIPJ2000",
                target_frame="SpacecraftFixed"
            )

        env_setup.add_rotation_model(bodies, 'Spacecraft', rotation_model_settings)
        thrust_magnitude_settings = (
            prop_setup.thrust.custom_thrust_magnitude_fixed_isp(self.controller.get_thrust_force, self.specific_impulse_nd / self.unit.acceleration)
        )

        env_setup.add_engine_model(body_name='Spacecraft',
                                   engine_name='MainEngine',
                                   thrust_magnitude_settings=thrust_magnitude_settings,
                                   bodies=bodies,
                                   body_fixed_thrust_direction=np.array([-1, 0, 0]))

        # setup accelerations
        acceleration_settings_spacecraft = dict(
            Sun=[prop_setup.acceleration.point_mass_gravity()],
            Spacecraft=[prop_setup.acceleration.thrust_from_engine('MainEngine')]
        )

        acceleration_models = prop_setup.create_acceleration_models(
            bodies, {"Spacecraft": acceleration_settings_spacecraft}, bodies_to_propagate, central_bodies
        )

        # mass propagator
        mass_rate_models = prop_setup.create_mass_rate_models(
            body_system=bodies,
            selected_mass_rates_per_body={"Spacecraft": [prop_setup.mass_rate.from_thrust()]},
            acceleration_models=acceleration_models
        )

        termination_condition = prop_setup.propagator.time_termination(final_time, terminate_exactly_on_final_condition=True)

        initial_state_cartesian = mee_to_cartesian(initial_fullstate[:6], self.gravitational_parameter_nd, singularity_at_zero_inclination=False)
        self.controller.trajectory.set_fullstate(initial_fullstate)

        # regular equations of motions propagator
        translational_propagator_settings = prop_setup.propagator.translational(
            central_bodies=central_bodies,
            acceleration_models=acceleration_models,
            bodies_to_integrate=["Spacecraft"],
            initial_states=initial_state_cartesian,
            termination_settings=termination_condition,
            propagator=prop_setup.propagator.cowell
        )

        mass_propagator_settings = prop_setup.propagator.mass(
            bodies_with_mass_to_propagate=bodies_to_propagate,
            mass_rate_models=mass_rate_models,
            initial_body_masses=[initial_fullstate[6]],
            termination_settings=termination_condition)

        # co-state propagator
        custom_propagator_settings = prop_setup.propagator.custom_state(
            state_derivative_function=self.controller.get_costate_derivative,
            initial_state=initial_fullstate[7:],
            initial_time=initial_time,
            integrator_settings=integration_settings,
            termination_settings=termination_condition
        )

        output_variables = []

        # Combine the propagators
        propagator_settings = prop_setup.propagator.multitype(
            [translational_propagator_settings,
             mass_propagator_settings,
             custom_propagator_settings],
            termination_condition,
            output_variables=output_variables,
        )

        # run simulation
        result = DynamicsSimulator(numerical_simulation.SingleArcSimulator(
            bodies=bodies,
            integrator_settings=integration_settings,
            propagator_settings=propagator_settings,
            print_dependent_variable_data=False
        ))

        return result.state_history.joint_array

The integration itself is invoked like this, and identical settings are used for when the integration routine is 'manually' called and also when tudat's SingleArcSimulator is used:

def tudat_integrate(initial_time: float,
                    final_time: float,
                    initial_fullstate: np.ndarray,
                    rk_set: tudat_integ.CoefficientSets = rk_set):
    integration_settings = prop_setup.integrator.runge_kutta_variable_step_size(
        initial_time,  # initial_time
        1e-6,  # initial_time_step
        rk_set,  # coefficient_set
        np.finfo(float).eps,  # minimum_step_size
        (final_time - initial_time) / self.min_steps if self.min_steps > 0 else np.inf,  # maximum_step_size
        self.rtol,  # relative_error_tolerance
        self.atol,  # absolute_error_tolerance,
        safety_factor=0.9,
        maximum_factor_increase=4.0,
        minimum_factor_increase=0.5
    )

    return self._tudat_integrate_from_settings(integration_settings, initial_time, final_time, initial_fullstate)

The directives for the costates (i.e. the custom propagator) and the fullstate are both implemented in C++ and exposed to Python, through the FastOptimalControlTrajectoryBase class or one of its children.

[DominicDirkx]

Hi, this factor 10,000 is not something that should be there, so thanks for posting.

There are two possible (broadly) reasons why this could be happening:

• The different setups are taking different steps, and as a result a different number of function evaluations. Have you checked whether the results are at the same epochs, and whether the same number of function evaluations are performed in the two cases?
• A single function evaluation takes much longer in the Python implementation. To assess this, we'd need to see the full implementation of a single state derivative.

One question for you, before we dive into it further:

• Have you tried simplifying (for testing purposes) or removing altogether, the custom models to check the influence on the run time and the number of function evaluations?

[briha]

The different setups are taking different steps, and as a result a different number of function evaluations. Have you checked whether the results are at the same epochs, and whether the same number of function evaluations are performed in the two cases?

This is a good point. I have not checked the function evaluations, as they are integrating the same trajectory with the same integration settings. That is clearly where the difference comes from, although I do not know why that is (all using the RKDP_87 from tudat):		steps	runtime
Tudat	127836	300.14 s
"manually"	234	66.87 ms
pykep	185	134.57 ms

I included pykep to show the Python overhead, as it is implementing the same dynamics in Python (in Cartesian instead of MEE).

A single function evaluation takes much longer in the Python implementation. To assess this, we'd need to see the full implementation of a single state derivative.

Has can be seen by just looking at the difference in steps, this is not the issue at hand. To clarify nonetheless: Both functions use the same derivative function, implemented in C++. It is exposed to python and passed to tudat through tudatpy, and thus has to go from C++ –> Python –> C++. I do not know whether the overhead added this way is significant, but, again, if it is, it should affect for both implementations equally.

The fullstate derivative is computed in the following way (within the context of a class):

Eigen::Vector14d computeFullstateDerivative() {
    Eigen::Vector14d fullstateDerivative;
    fullstateDerivative << computeStateDerivative(), computeCostateDerivative();
    return fullstateDerivative;
}

which should make evident how the costate derivative is computed.

Have you tried simplifying (for testing purposes) or removing altogether, the custom models to check the influence on the run time and the number of function evaluations?

Only propagating kinematic state and mass, the problem seems to indeed mostly vanish (it runs in 690.14ms, taking 143 steps). Letting go of the costate-propagation means a constant thrust direction and magnitude, of course, which may even be zero, if that is the value at t=0.