Modeling Dynamic Systems with Dymos#

Things you’ll learn through this example

  • How to define a basic Dymos ODE system.

  • How to explicitly propagate the system from some initial state.

Dymos is a library for modeling dynamic systems and performing optimal control with the OpenMDAO framework. Dynamic systems are typically defined by some set of ordinary differential equations (the ODE) which governs their behavior.

Consider a simple damped harmonic oscillator.

Damped harmonic oscillator free-body diagram

(11)#\[\begin{align} \ddot{x} &= -\frac{kx}{m} - \frac{c \dot{x}}{m} \end{align}\]

Converting this to a first order system results in an ODE system with two states:

(12)#\[\begin{align} \dot{x} &= v \\ \dot{v} &= -\frac{kx}{m} - \frac{c \dot{x}}{m} \end{align}\]

The OpenMDAO model of the ODE#

In Dymos, the ODE is an OpenMDAO System (a Component, or a Group of components). The following ExplicitComponent computes the velocity rate for the damped harmonic oscillator.

More detail on the workings of an ExplicitComponent can be found in the OpenMDAO documentation. In summary, an ExplicitComponent used as part of an ODE in Dymos should override the following methods:

  • initialize: Called at setup, and used to define options for the component. ALL Dymos ODE components should have the property num_nodes, which defines the number of points at which the outputs are simultaneously computed.

  • setup: Used to add inputs and outputs to the component, and declare which outputs (and indices of outputs) are dependent on each of the inputs.

  • compute: Used to compute the outputs, given the inputs.

  • compute_partials: Used to compute the derivatives of the outputs w.r.t. each of the inputs analytically. This method may be omitted if finite difference or complex-step approximations are used, though analytic is recommended.

import numpy as np
import openmdao.api as om


class OscillatorODE(om.ExplicitComponent):
    """
    A Dymos ODE for a damped harmonic oscillator.
    """

    def initialize(self):
        self.options.declare('num_nodes', types=int)

    def setup(self):
        nn = self.options['num_nodes']

        # Inputs
        self.add_input('x', shape=(nn,), desc='displacement', units='m')
        self.add_input('v', shape=(nn,), desc='velocity', units='m/s')
        self.add_input('k', shape=(nn,), desc='spring constant', units='N/m')
        self.add_input('c', shape=(nn,), desc='damping coefficient', units='N*s/m')
        self.add_input('m', shape=(nn,), desc='mass', units='kg')

        # self.add_output('x_dot', val=np.zeros(nn), desc='rate of change of displacement', units='m/s')
        self.add_output('v_dot', val=np.zeros(nn), desc='rate of change of velocity', units='m/s**2')

        self.declare_partials(of='*', wrt='*', method='fd')

    def compute(self, inputs, outputs):
        x = inputs['x']
        v = inputs['v']
        k = inputs['k']
        c = inputs['c']
        m = inputs['m']

        f_spring = -k * x
        f_damper = -c * v

        outputs['v_dot'] = (f_spring + f_damper) / m

Things to note about the ODE system

  • In this case, the ODE is a function of both states, but this isn’t always the case. If the dynamics aren’t functions of some states, those states aren’t needed as inputs.

  • This ODE only computes the rate of change of velocity. Since the rate of change of displacement can directly be obtained from another state variable, it doesn’t need to be computed by the ODE. This would also be true if the state’s rate was a control, design parameter, or input parameter value.

  • It’s possible that we might want to use parameters c, k, and m as design variables at some point, so they’re also included as inputs here. Alternatively, if we had no interest in ever treating them as design variables, we could add their values as options to the ODE system in the initialize method.

Hello World: Propagating the dynamics#

One of the first things one might be interested in regarding an ODE is propagating it from some given initial conditions. This is known as solving the initial value problem (IVP), and there are many software packages that can do this. The following is a minimal script that starts the system at some set of initial conditions and propagates them for some fixed duration. Some elements of this code will be explained later, but we’ll hit the highlights now.

import openmdao.api as om
import dymos as dm
import matplotlib.pyplot as plt
# plt.switch_backend('Agg')  # disable plotting to the screen

from dymos.examples.oscillator.oscillator_ode import OscillatorODE

# Instantiate an OpenMDAO Problem instance.
prob = om.Problem()

# Instantiate a Dymos Trajectory and add it to the Problem model.
traj = dm.Trajectory()
prob.model.add_subsystem('traj', traj)

# Instantiate a Phase and add it to the Trajectory.
# Here the transcription is necessary but not particularly relevant.
phase = dm.Phase(ode_class=OscillatorODE, transcription=dm.Radau(num_segments=4))
traj.add_phase('phase0', phase)

# Tell Dymos the states to be propagated using the given ODE.
phase.add_state('v', rate_source='v_dot', targets=['v'], units='m/s')
phase.add_state('x', rate_source='v', targets=['x'], units='m')

# The spring constant, damping coefficient, and mass are inputs to the system
# that are constant throughout the phase.
phase.add_parameter('k', units='N/m', targets=['k'])
phase.add_parameter('c', units='N*s/m', targets=['c'])
phase.add_parameter('m', units='kg', targets=['m'])

# Setup the OpenMDAO problem
prob.setup()

# Assign values to the times and states
prob.set_val('traj.phase0.t_initial', 0.0)
prob.set_val('traj.phase0.t_duration', 15.0)

prob.set_val('traj.phase0.states:x', 10.0)
prob.set_val('traj.phase0.states:v', 0.0)

prob.set_val('traj.phase0.parameters:k', 1.0)
prob.set_val('traj.phase0.parameters:c', 0.5)
prob.set_val('traj.phase0.parameters:m', 1.0)

# Perform a single execution of the model (executing the model is required before simulation).
prob.run_model()

# Perform an explicit simulation of our ODE from the initial conditions.
sim_out = traj.simulate(times_per_seg=50)

# Plot the state values obtained from the phase timeseries objects in the simulation output.
t_sol = prob.get_val('traj.phase0.timeseries.time')
t_sim = sim_out.get_val('traj.phase0.timeseries.time')

states = ['x', 'v']
fig, axes = plt.subplots(len(states), 1)
for i, state in enumerate(states):
    sol = axes[i].plot(t_sol, prob.get_val(f'traj.phase0.timeseries.{state}'), 'o')
    sim = axes[i].plot(t_sim, sim_out.get_val(f'traj.phase0.timeseries.{state}'), '-')
    axes[i].set_ylabel(state)
axes[-1].set_xlabel('time (s)')
fig.legend((sol[0], sim[0]), ('solution', 'simulation'), loc='lower right', ncol=2)
plt.tight_layout()
plt.show()

--- Constraint Report [traj] ---
    --- phase0 ---
        None

Simulating trajectory traj
Done simulating trajectory traj
../../_images/a7b15f369e88dc6d650679bf979b6b238feafdf70baeec066ee555a37ac58919.png

What happened?#

This script consists of the following general steps:

  1. After importing relevant packages, an OpenMDAO Problem is instantiated

  2. A Dymos Trajectory object is instantiated, and a single Phase named 'phase0' is added to it.

    1. That Phase takes the ODE class as one of its arguments. It will instantiate instances of it as needed.

    2. The transcription determines how the implicit integration and optimization are performed. It’s necessary but not particularly relevant in this example.

  3. The states to be integrated are added to the phase.

    1. Each state needs a rate source - an ODE-relative path of the output which provides the time derivative of the state variable. For the rate_source of x, we provide v. Dymos understands this is the name of one of the other states (or time, or controls, or parameters).

    2. Those states which are inputs to the ODE need to provide their targets in the ODE (again, with an ODE-relative path).

  4. The problem is setup (this prepares the model for execution in OpenMDAO).

  5. Default values are assigned to the states and time.

    1. Variables t_initial and t_duration (the initial time and duration of the Phase) are scalars.

    2. State variables are vectors whose values are provided throughout the Phase. Here they’re all being filled with a single value (10.0 for displacement, 0.0 for the velocity)

  6. Problem.run_model is called. This executes the Problem’s model one time. This is a necessary step before using the Trajectory.simulate method.

  7. The trajectory is simulated using the simulate() method and the results returned as an OpenMDAO Problem instance called sim_out.

  8. The states are plotted using values obtained from the timeseries. Phases contain timeseries data that provides contiguous values regardless of the transcription used.

Method simulate exists on both Trajectory and Phase objects. It uses the scipy.integrate.solve_ivp function to propagate the states defined in each phase of the trajectory from their initial values at the initial time to some final value at time = t_initial + t_duration.

In Dymos, the simulate method is useful for testing the functionality of an ODE (making sure it behaves as expected) and for checking the validity of answers after optimization. However, it cannot solve boundary value problems.

Why is the solution different from the simulation results?#

The plots above display both the solution from the implicit transcription (blue dots) and the results of the simulation (orange line). Here they do not match because we only performed a single execution of the model. The purpose of a model execution in Dymos is to calculate the objective and constraints for the optimizer. These constraints include the collocation defect constraints, which (when driven to zero) indicate that the current polynomial representation of the state-time history matches the physically correct trajectory. In this case, no iteration was performed, and thus the solution is not physically valid.

To be clear, the output of Dymos in this case is not a physically valid trajectory. The simulate() call after executing the model is the expected result using the variable step integrator from Scipy.

There are two ways to converge this solution using the implicit transcription techniques in Dymos.

  1. We can run an optimization driver with some “dummy” objective to converge the collocation defect constraints.

  2. We can have Dymos use a nonlinear solver to vary the state time-history until the collocation defect constraints are satisfied.

Traditionally, many collocation optimal control techniques have used an optimizer-based approach because it is extremely efficient. OpenMDAO provides a lot of versatility in adding nonlinear solvers within the optimization problem. In our case, using a solver to converge state-time history means that the optimizer “sees” a physical trajectory at every iteration. In an optimization context, we can use the solver-based convergence of defects to obtain a shooting method with analytic derivatives.

Using a solver to converge the physical trajectory#

We let Dymos know that one or more states should be converged using the solve_segments=forward argument. If passed to the transcription, it applies to all states. Otherwise, we can pass it only to certain states as an argument to add_state or set_state_options.

import openmdao.api as om
import dymos as dm
import matplotlib.pyplot as plt
# plt.switch_backend('Agg')  # disable plotting to the screen

from dymos.examples.oscillator.oscillator_ode import OscillatorODE

# Instantiate an OpenMDAO Problem instance.
prob = om.Problem()

# Instantiate a Dymos Trajectory and add it to the Problem model.
traj = dm.Trajectory()
prob.model.add_subsystem('traj', traj)

# Instantiate a Phase and add it to the Trajectory.
phase = dm.Phase(ode_class=OscillatorODE, transcription=dm.Radau(num_segments=4, solve_segments='forward'))
traj.add_phase('phase0', phase)

# Tell Dymos the states to be propagated using the given ODE.
phase.add_state('x', fix_initial=True, rate_source='v', targets=['x'], units='m')
phase.add_state('v', fix_initial=True, rate_source='v_dot', targets=['v'], units='m/s')

# The spring constant, damping coefficient, and mass are inputs to the system that are
# constant throughout the phase.
phase.add_parameter('k', units='N/m', targets=['k'])
phase.add_parameter('c', units='N*s/m', targets=['c'])
phase.add_parameter('m', units='kg', targets=['m'])

# Setup the OpenMDAO problem
prob.setup()

# Assign values to the times and states
prob.set_val('traj.phase0.t_initial', 0.0)
prob.set_val('traj.phase0.t_duration', 15.0)

prob.set_val('traj.phase0.states:x', 10.0)
prob.set_val('traj.phase0.states:v', 0.0)

prob.set_val('traj.phase0.parameters:k', 1.0)
prob.set_val('traj.phase0.parameters:c', 0.5)
prob.set_val('traj.phase0.parameters:m', 1.0)

# Now we're using the optimization driver to iteratively run the model and vary the
# phase duration until the final y value is 0.
prob.run_model()

# Perform an explicit simulation of our ODE from the initial conditions.
sim_out = traj.simulate(times_per_seg=50)

# Plot the state values obtained from the phase timeseries objects in the simulation output.
t_sol = prob.get_val('traj.phase0.timeseries.time')
t_sim = sim_out.get_val('traj.phase0.timeseries.time')

states = ['x', 'v']
fig, axes = plt.subplots(len(states), 1)
for i, state in enumerate(states):
    sol = axes[i].plot(t_sol, prob.get_val(f'traj.phase0.timeseries.{state}'), 'o')
    sim = axes[i].plot(t_sim, sim_out.get_val(f'traj.phase0.timeseries.{state}'), '-')
    axes[i].set_ylabel(state)
axes[-1].set_xlabel('time (s)')
fig.legend((sol[0], sim[0]), ('solution', 'simulation'), loc='lower right', ncol=2)
plt.tight_layout()
plt.show()

--- Constraint Report [traj] ---
    --- phase0 ---
        None


==================
traj.phases.phase0
==================
NL: Newton 0 ; 64.9519053 1
NL: Newton 1 ; 5.0296006e-08 7.7435767e-10
NL: Newton 2 ; 7.72308087e-15 1.18904609e-16
NL: Newton Converged

Simulating trajectory traj
Done simulating trajectory traj
../../_images/3a67faba3fdfc38b2a0bfc6120d1598c83a8c1dcbb30cf495f46dab6e9f08541.png

Using an optimization driver to converge the physical trajectory#

Alternatively, we can use an optimization driver to converge the state time histories.

In the case of an initial value problem (fixed time duration, fixed initial states, and no controls or parameters as design variables) there are no degrees of freedom to optimize the problem, just single possible trajectory which satisfies the collocation constraints.

In OpenMDAO (and thus Dymos) optimizers require an objective. Even though the initial time and duration of the phase are fixed, we provide the final time as a “dummy” objective here.

import openmdao.api as om
import dymos as dm
import matplotlib.pyplot as plt
# plt.switch_backend('Agg')  # disable plotting to the screen

from dymos.examples.oscillator.oscillator_ode import OscillatorODE

# Instantiate an OpenMDAO Problem instance.
prob = om.Problem()

# We need an optimization driver.  To solve this simple problem ScipyOptimizerDriver will work.
prob.driver = om.ScipyOptimizeDriver()

# Instantiate a Dymos Trajectory and add it to the Problem model.
traj = dm.Trajectory()
prob.model.add_subsystem('traj', traj)

# Instantiate a Phase and add it to the Trajectory.
phase = dm.Phase(ode_class=OscillatorODE, transcription=dm.Radau(num_segments=4))
traj.add_phase('phase0', phase)

# Tell Dymos that the duration of the phase is bounded.
phase.set_time_options(fix_initial=True, fix_duration=True)

# Tell Dymos the states to be propagated using the given ODE.
phase.add_state('x', fix_initial=True, rate_source='v', targets=['x'], units='m')
phase.add_state('v', fix_initial=True, rate_source='v_dot', targets=['v'], units='m/s')

# The spring constant, damping coefficient, and mass are inputs to the system that are
# constant throughout the phase.
phase.add_parameter('k', units='N/m', targets=['k'])
phase.add_parameter('c', units='N*s/m', targets=['c'])
phase.add_parameter('m', units='kg', targets=['m'])

# Since we're using an optimization driver, an objective is required.  We'll minimize
# the final time in this case.
phase.add_objective('time', loc='final')

# Setup the OpenMDAO problem
prob.setup()

# Assign values to the times and states
prob.set_val('traj.phase0.t_initial', 0.0)
prob.set_val('traj.phase0.t_duration', 15.0)

prob.set_val('traj.phase0.states:x', 10.0)
prob.set_val('traj.phase0.states:v', 0.0)

prob.set_val('traj.phase0.parameters:k', 1.0)
prob.set_val('traj.phase0.parameters:c', 0.5)
prob.set_val('traj.phase0.parameters:m', 1.0)

# Now we're using the optimization driver to iteratively run the model and vary the
# phase duration until the final y value is 0.
prob.run_driver()

# Perform an explicit simulation of our ODE from the initial conditions.
sim_out = traj.simulate(times_per_seg=50)

# Plot the state values obtained from the phase timeseries objects in the simulation output.
t_sol = prob.get_val('traj.phase0.timeseries.time')
t_sim = sim_out.get_val('traj.phase0.timeseries.time')

states = ['x', 'v']
fig, axes = plt.subplots(len(states), 1)
for i, state in enumerate(states):
    sol = axes[i].plot(t_sol, prob.get_val(f'traj.phase0.timeseries.{state}'), 'o')
    sim = axes[i].plot(t_sim, sim_out.get_val(f'traj.phase0.timeseries.{state}'), '-')
    axes[i].set_ylabel(state)
axes[-1].set_xlabel('time (s)')
fig.legend((sol[0], sim[0]), ('solution', 'simulation'), loc='lower right', ncol=2)
plt.tight_layout()
plt.show()

--- Constraint Report [traj] ---
    --- phase0 ---
        None

Optimization terminated successfully    (Exit mode 0)
            Current function value: [15.]
            Iterations: 1
            Function evaluations: 2
            Gradient evaluations: 1
Optimization Complete
-----------------------------------

/usr/share/miniconda/envs/test/lib/python3.10/site-packages/openmdao/core/total_jac.py:1782: DerivativesWarning:Constraints or objectives [('traj.phases.phase0.time.t', inds=[15])] cannot be impacted by the design variables of the problem.

Simulating trajectory traj
Done simulating trajectory traj
../../_images/3a67faba3fdfc38b2a0bfc6120d1598c83a8c1dcbb30cf495f46dab6e9f08541.png

But the solution still doesn’t match the simulation#

If you look at the plots from the last two examples, you’ll notice that the state time-history from the solution has some pretty significant deviations from the simulation results. This is an important lesson in using implicit collocation techniques.

Note

A converged trajectory isn’t necessarily correct

As we mentioned before, the simulate() method exists to provide a check on a converged trajectory. In this case, the trajectory found using simulate() doesn’t really interpolate the solution from the collocation technique. In the next section, we’ll explain how to deal with this.