Implicit Duration

In some cases, using dymos without the need for an optimizer is desirable. Sometimes in the absense of a control, you may just want to propagate the problem. Solve segments allows us to do this for a specified amount of time, but if the duration of the phase is itself an implicit output, then we need dymos to understand the residual to associate with the phase duration.

The requirement for this procedure to work is that the dynamics have to be obeyed at each iteration. A time-stepping transcription (ExplicitShooting) or a pseudospectral transcription with solve_segments=True is required.

One important note: In the pseudospectral transcriptions, using solve_segments with option compressed=False, results in a “multiple shooting” approach, whereby the state discontinuitites between segments are treated as constraints to be nulled out by the optimizer. This approach should therefore only be used with “single shooting” approaches (compressed=True in the case of the pseudospectral transcriptions)

Solving the brachistochrone without an optimizer.

Consider the brachistochrone problem. Below we implement the brachistochrone equations of motion using OpenMDAO’s JaxExplicitComponent so that we can use automatic differentiation to provide the partials.

The implicit solution to the brachistochrone is very robust, but what if we wanted to simplify it? For instance, we leave the duration of the bead’s trajectory as a design variable that we want to minimize, but we also require the enforcement of the final position of the brachistochrone to some (x, y) coordinate pair.

import jax.numpy as jnp
import openmdao.api as om


class BrachistochroneODE(om.JaxExplicitComponent):
    """
    The brachistochrone EOM assuming 
    """

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

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

        # static parameters
        self.add_input('v', desc='speed of the bead on the wire', shape=(nn,), units='m/s')
        self.add_input('theta', desc='angle between wire tangent and the nadir', shape=(nn,), units='rad')
        self.add_input('g', desc='gravitational acceleration', tags=['dymos.static_target'], units='m/s**2')

        self.add_output('xdot', desc='velocity component in x', shape=(nn,), units='m/s',
                        tags=['dymos.state_rate_source:x', 'dymos.state_units:m'])

        self.add_output('ydot', desc='velocity component in y', shape=(nn,), units='m/s',
                        tags=['dymos.state_rate_source:y', 'dymos.state_units:m'])

        self.add_output('vdot', desc='acceleration magnitude', shape=(nn,), units='m/s**2',
                        tags=['dymos.state_rate_source:v', 'dymos.state_units:m/s'])

        # Setup partials
        ar = jnp.arange(nn, dtype=int)

        self.declare_partials(of='vdot', wrt='v', rows=ar, cols=ar)
        self.declare_partials(of='vdot', wrt='theta', rows=ar, cols=ar)
        self.declare_partials(of='vdot', wrt='g')

        self.declare_partials(of='xdot', wrt='v', rows=ar, cols=ar)
        self.declare_partials(of='xdot', wrt='theta', rows=ar, cols=ar)

        self.declare_partials(of='ydot', wrt='v', rows=ar, cols=ar)
        self.declare_partials(of='ydot', wrt='theta', rows=ar, cols=ar)

    def compute(self, v, theta, g):
        sin_theta = jnp.sin(theta)
        cos_theta = jnp.cos(theta)

        vdot = g * cos_theta
        xdot = v * sin_theta
        ydot = -v * cos_theta

        return xdot, ydot, vdot

Propagating the equations of motion using solve_segments

In the following example, we set solve_segments='forward' such that all of the collocation defects associate with the Radau transcription are satisfied by a Newton solver rather than the optimizer. This is the first step to solving the problem without an optimizer. This ensures that the equations of motion are satisfied to the extent possible assuming the transcription grid.

# Define the OpenMDAO problem
p = om.Problem(model=om.Group())

# Define a Trajectory object
traj = dm.Trajectory()
p.model.add_subsystem('traj', subsys=traj)

# Define a Dymos Phase object with Radau Transcription
phase = dm.Phase(ode_class=BrachistochroneODE,
                 transcription=dm.Radau(num_segments=5, order=3, solve_segments='forward'))
traj.add_phase(name='phase0', phase=phase)

# Set the time options
phase.set_time_options(fix_initial=True,
                       duration_bounds=(0.5, 10.0))

# Set the state options
phase.add_state('x', rate_source='xdot',
                fix_initial=True)
phase.add_state('y', rate_source='ydot',
                fix_initial=True)
phase.add_state('v', rate_source='vdot',
                fix_initial=True)

phase.add_control('theta', lower=0.0, upper=1.5, units='rad')
phase.add_parameter('g', opt=False, val=9.80665, units='m/s**2')

# Minimize final time.
phase.add_objective('time', loc='final')

# Set the driver.
p.driver = om.ScipyOptimizeDriver()

# Allow OpenMDAO to automatically determine our sparsity pattern.
# Doing so can significant speed up the execution of dymos.
p.driver.declare_coloring()

phase.nonlinear_solver = om.NewtonSolver(maxiter=1000, solve_subsystems=True, stall_limit=3, iprint=2)
phase.linear_solver = om.DirectSolver()

# Setup the problem
p.setup()

phase.set_time_val(initial=0.0, duration=2.0)
phase.set_state_val('x', vals=[0, 10])
phase.set_state_val('y', vals=[10, 5])
phase.set_state_val('v', vals=[0.1, 100])
phase.set_control_val('theta', vals=[0.0, 90.0], units='deg')
phase.set_parameter_val('g', val=9.80665)

# Run the driver to solve the problem
p.run_model()

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


traj.phases.phase0.rhs_all:
def compute_primal(self, v, theta, g):
    sin_theta = jnp.sin(theta)
    cos_theta = jnp.cos(theta)
    vdot = g * cos_theta
    xdot = v * sin_theta
    ydot = -v * cos_theta
    return (xdot, ydot, vdot)
    return (xdot, ydot, vdot)


==================
traj.phases.phase0
==================
NL: Newton 0 ; 52.4703333 1

NL: Newton 1 ; 2.81042974e-13 5.35622621e-15
NL: Newton Converged

If we plot the resulting trajectory of the bead, we notice that our guesses for time and the control history didn’t bring the bead to the desired target at (10, 5):

%matplotlib inline
import matplotlib.pyplot as plt

x = p.get_val('traj.phase0.timeseries.x')
y = p.get_val('traj.phase0.timeseries.y')

plt.plot(0.0, 10.0, 'ko')
plt.plot(10.0, 5.0, 'ko')
plt.plot(x, y)
plt.show()

Stopping the propagation at the desired time.

We can utilize the set_implicit_duration method on Phase to turn t_duration into an implicit output and provide it with a residual.

In our case, we set phase.set_duration_balance('x', 10.0)

This specifies that we want t_duration to be associated with the following residual:

(96)\[\begin{align} \mathcal{R}(t_d) = x - 10 = 0 \end{align}\]

Note that we limit the angle \(\theta\) to be between 0 and 180, the bead must move to the right and we can just terminate the propagation when x=10.

# Define the OpenMDAO problem
p = om.Problem(model=om.Group())

# Define a Trajectory object
traj = dm.Trajectory()
p.model.add_subsystem('traj', subsys=traj)

# Define a Dymos Phase object with Radau Transcription
phase = dm.Phase(ode_class=BrachistochroneODE,
                 transcription=dm.Radau(num_segments=5, order=3, solve_segments='forward', compressed=True))
traj.add_phase(name='phase0', phase=phase)

# Set the time options
phase.set_time_options(fix_initial=True,
                       duration_bounds=(0.5, 10.0))
phase.set_duration_balance('x', 10.0)

# Set the state options
phase.add_state('x', rate_source='xdot',
                fix_initial=True)
phase.add_state('y', rate_source='ydot',
                fix_initial=True)
phase.add_state('v', rate_source='vdot',
                fix_initial=True)

phase.add_control('theta', lower=0.0, upper=1.5, units='rad')
phase.add_parameter('g', opt=False, val=9.80665, units='m/s**2')

# Minimize final time.
phase.add_objective('time', loc='final')

# Set the driver.
p.driver = om.ScipyOptimizeDriver()

# Allow OpenMDAO to automatically determine our sparsity pattern.
# Doing so can significant speed up the execution of dymos.
p.driver.declare_coloring()

phase.nonlinear_solver = om.NewtonSolver(maxiter=1000, solve_subsystems=True, stall_limit=3)
phase.nonlinear_solver.linesearch = None
phase.linear_solver = om.DirectSolver()

# Setup the problem
p.setup()

phase.set_time_val(initial=0.0, duration=1.8)
phase.set_state_val('x', vals=[0, 10])
phase.set_state_val('y', vals=[10, 5])
phase.set_state_val('v', vals=[0.1, 100])
phase.set_control_val('theta', vals=[0.0, 90.0], units='deg')
phase.set_parameter_val('g', val=9.80665)

# Run the driver to solve the problem
p.run_model()

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


traj.phases.phase0.rhs_all:
def compute_primal(self, v, theta, g):
    sin_theta = jnp.sin(theta)
    cos_theta = jnp.cos(theta)
    vdot = g * cos_theta
    xdot = v * sin_theta
    ydot = -v * cos_theta
    return (xdot, ydot, vdot)
    return (xdot, ydot, vdot)

==================
traj.phases.phase0
==================
NL: Newton Converged in 4 iterations

%matplotlib inline
import matplotlib.pyplot as plt

x = p.get_val('traj.phase0.timeseries.x')
y = p.get_val('traj.phase0.timeseries.y')
v = p.get_val('traj.phase0.timeseries.y')
theta = p.get_val('traj.phase0.timeseries.theta')

plt.plot(0.0, 10.0, 'ko')
plt.plot(10.0, 5.0, 'ko')
plt.plot(x, y)
plt.show()

print('final time', p.get_val('traj.phase0.timeseries.time')[-1])
print('final y', y[-1])
print('final theta', jnp.degrees(theta[-1]))

final time [1.77967375]
final y [3.59263414]
final theta [90.]

Solving the brachistochrone without an optimizer

The brachistochrone has an analytic solution such that rate of change of the tangent to the wire (angle $\theta) is constant. We can use this information to pose the brachistochrone as a boundary value problem and solve as a shooting problem.

We add another implicit output that provides the rate of change as a constant value throughout the trajectory. That means that theta_rate is a parameter as far as dymos is concerned. We then make theta a state variable whose rate is provided by the theta_rate parameter. We’ll set the initial value of theta to zero degrees, since we know that is approximately correct for this problem.

(97)\[\begin{align} \mathcal{R}(\dot{\theta}) &= y_f - 5 \end{align}\]

Dymos currently doesn’t support making parameters implicit outputs, so we’re going to do this the manual OpenMDAO way.

# Define the OpenMDAO problem
p = om.Problem(model=om.Group())

# Define a Trajectory object
traj = dm.Trajectory()
p.model.add_subsystem('traj', subsys=traj)

# Define a Dymos Phase object with Radau Transcription
phase = dm.Phase(ode_class=BrachistochroneODE,
                 transcription=dm.Radau(num_segments=10, order=3, solve_segments='forward', compressed=True))
traj.add_phase(name='phase0', phase=phase)

# Set the time options
phase.set_time_options(fix_initial=True)
phase.set_duration_balance('x', 10.0)

# Set the state options
phase.add_state('x', rate_source='xdot',
                fix_initial=True)
phase.add_state('y', rate_source='ydot',
                fix_initial=True)
phase.add_state('v', rate_source='vdot',
                fix_initial=True)
phase.add_state('theta', rate_source='theta_rate',
                fix_initial=True, units='rad')

phase.add_parameter('theta_rate', opt=False, units='rad/s')
phase.add_parameter('g', opt=False, val=9.80665, units='m/s**2')

# Minimize final time.
phase.add_objective('time', loc='final')

# Set the driver.
p.driver = om.ScipyOptimizeDriver()

# Allow OpenMDAO to automatically determine our sparsity pattern.
# Doing so can significant speed up the execution of dymos.
p.driver.declare_coloring()

phase.nonlinear_solver = om.NewtonSolver(maxiter=1000, solve_subsystems=True, stall_limit=3,
                                         iprint=0, atol=1.0E-8, rtol=1.0E-8, debug_print=True)
phase.nonlinear_solver.linesearch = None
phase.linear_solver = om.DirectSolver()

theta_rate_balance = p.model.add_subsystem('theta_rate_balance', om.BalanceComp())

theta_rate_balance.add_balance('theta_rate', units='rad/s', eq_units='m', lhs_name='y_final', rhs_val=5.0)
p.model.connect('theta_rate_balance.theta_rate', 'traj.phase0.parameters:theta_rate')
p.model.connect('traj.phase0.timeseries.y', 'theta_rate_balance.y_final', src_indices=[-1])

# Now we need a solver to converge the loop around the entire problem.
p.model.nonlinear_solver = om.NewtonSolver(maxiter=100, solve_subsystems=True, stall_limit=3,
                                           iprint=0, atol=1.0E-8, rtol=1.0E-8, debug_print=True)
p.model.nonlinear_solver.linesearch = None
p.model.linear_solver = om.DirectSolver()

# Setup the problem
p.setup()

phase.set_time_val(initial=0.0, duration=1.8)
phase.set_state_val('x', vals=[0, 10])
phase.set_state_val('y', vals=[10, 5])
phase.set_state_val('v', vals=[0., 10])
phase.set_state_val('theta', vals=[0.0, 45], units='deg')
phase.set_parameter_val('g', val=9.80665)

# Note that we set theta_rate at the balance comp, which is then passed into the phase as a parameter.
p.model.set_val('theta_rate_balance.theta_rate', val=0.5, units='rad/s')

# Run the driver to solve the problem
p.run_model()

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


traj.phases.phase0.rhs_all:
def compute_primal(self, v, theta, g):
    sin_theta = jnp.sin(theta)
    cos_theta = jnp.cos(theta)
    vdot = g * cos_theta
    xdot = v * sin_theta
    ydot = -v * cos_theta
    return (xdot, ydot, vdot)
    return (xdot, ydot, vdot)

%matplotlib inline
import matplotlib.pyplot as plt

x = p.get_val('traj.phase0.timeseries.x')
y = p.get_val('traj.phase0.timeseries.y')
v = p.get_val('traj.phase0.timeseries.y')
theta = p.get_val('traj.phase0.timeseries.theta')

plt.plot(0.0, 10.0, 'ko')
plt.plot(10.0, 5.0, 'ko')
plt.plot(x, y)
plt.show()

print('final time', p.get_val('traj.phase0.timeseries.time')[-1])
print('final y', y[-1])
print('final theta', jnp.degrees(theta[-1]))

final time [1.80160313]
final y [5.]
final theta [100.50736083]

Using a combination of solve_segments, implicit phase duration, and an outer loop to make a phase parameter an implicit output, we’ve managed to solve a two point boundary value problem without the need of an optimizer. Since we’re solving the segments in an inner loop here, each converged phase provides a physically valid traejctory, though it may not satisfy our desired final conditions unless the solvers are converged at all levels.

Note that solving problems in this way can be filled with pitfalls. Here we have a formulation of the brachistochrone that is robust to a variety of values for the parameter theta_rate. An alternative formulation is to keep the ratio \(\frac{\sin{\theta}}{v}\) constant. However, using the inverse sine in that formulation can easily result in domain issues which makes convergence difficult.

Even in this case, setting the solve_subsystems option on the outermost Newton solver to False seems to be too challenging and the solver fails to converge.