The Brachistochrone with Externally-Sourced initial state values#

Things you’ll learn through this example

- How to link phase state boundary values to an externally provided value.

This is another modification of the brachistochrone in which the target initial value of a state is provided by an external source (an IndepVarComp in this case).

Rather than the external value being directly connected to the phase, the values are “linked” via constraint. This is exactly how phase linkages in trajectories work as well, but the trajectory hides some of the implementation.

The following script fully defines the brachistochrone problem with Dymos and solves it. A new IndepVarComp is added before the trajectory which provides x0. An ExecComp then computes the error between x0_target (taken from the IndepVarComp) and x0_actual (taken from the phase timeseries output). The result of this calculation (x0_error) is then constrained as a normal OpenMDAO constraint.

Hide code cell outputs
import numpy as np
import openmdao.api as om


class BrachistochroneODE(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('num_nodes', types=int)
        self.options.declare('g', default=9.80665, desc='gravitational acceleration in m/s**2')

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

        # Inputs
        self.add_input('v', val=np.zeros(nn), desc='velocity', units='m/s')

        self.add_input('theta', val=np.ones(nn), desc='angle of wire', units='rad')

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

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

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

        self.add_output('check', val=np.zeros(nn), desc='check solution: v/sin(theta) = constant',
                        units='m/s')

        # Setup partials
        ar = np.arange(self.options['num_nodes'], dtype=int)

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

        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)

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

    def compute(self, inputs, outputs):
        theta = inputs['theta']
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        g = self.options['g']
        v = inputs['v']

        outputs['vdot'] = g * cos_theta
        outputs['xdot'] = v * sin_theta
        outputs['ydot'] = -v * cos_theta
        outputs['check'] = v / sin_theta

    def compute_partials(self, inputs, partials):
        theta = inputs['theta']
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        g = self.options['g']
        v = inputs['v']

        partials['vdot', 'theta'] = -g * sin_theta

        partials['xdot', 'v'] = sin_theta
        partials['xdot', 'theta'] = v * cos_theta

        partials['ydot', 'v'] = -cos_theta
        partials['ydot', 'theta'] = v * sin_theta

        partials['check', 'v'] = 1 / sin_theta
        partials['check', 'theta'] = -v * cos_theta / sin_theta**2
import openmdao.api as om
import dymos as dm

import matplotlib.pyplot as plt
plt.switch_backend('Agg')

from dymos.examples.brachistochrone.brachistochrone_ode import BrachistochroneODE

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

# Instantiate the transcription so we can get the number of nodes from it while
# building the problem.
tx = dm.GaussLobatto(num_segments=10, order=3)

# Add an indep var comp to provide the external control values
ivc = p.model.add_subsystem('states_ivc', om.IndepVarComp(), promotes_outputs=['*'])

# Add the output to provide the values of theta at the control input nodes of the transcription.
ivc.add_output('x0', shape=(1,), units='m')

# Connect x0 to the state error component so we can constrain the given value of x0
# to be equal to the value chosen in the phase.
p.model.connect('x0', 'state_error_comp.x0_target')
p.model.connect('traj.phase0.timeseries.x', 'state_error_comp.x0_actual', src_indices=[0])

#
# Define a Trajectory object
#
traj = dm.Trajectory()

p.model.add_subsystem('traj', subsys=traj)

p.model.add_subsystem('state_error_comp',
                      om.ExecComp('x0_error = x0_target - x0_actual',
                                  x0_error={'units': 'm'},
                                  x0_target={'units': 'm'},
                                  x0_actual={'units': 'm'}))

p.model.add_constraint('state_error_comp.x0_error', equals=0.0)

#
# Define a Dymos Phase object with GaussLobatto Transcription
#
phase = dm.Phase(ode_class=BrachistochroneODE,
                 transcription=tx)

traj.add_phase(name='phase0', phase=phase)

#
# Set the time options
# Time has no targets in our ODE.
# We fix the initial time so that the it is not a design variable in the optimization.
# The duration of the phase is allowed to be optimized, but is bounded on [0.5, 10].
#
phase.set_time_options(fix_initial=True, duration_bounds=(0.5, 10.0), units='s')

#
# Set the time options
# Initial values of positions and velocity are all fixed.
# The final value of position are fixed, but the final velocity is a free variable.
# The equations of motion are not functions of position, so 'x' and 'y' have no targets.
# The rate source points to the output in the ODE which provides the time derivative of the
# given state.
phase.add_state('x', fix_initial=False, fix_final=True, units='m', rate_source='xdot')
phase.add_state('y', fix_initial=True, fix_final=True, units='m', rate_source='ydot')
phase.add_state('v', fix_initial=True, fix_final=False, units='m/s',
                rate_source='vdot', targets=['v'])

# Define theta as a control.
# Use opt=False to allow it to be connected to an external source.
# Arguments lower and upper are no longer valid for an input control.
phase.add_control(name='theta', units='rad', targets=['theta'])

# 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()

# Setup the problem
p.setup(check=True)

# Now that the OpenMDAO problem is setup, we can set the values of the states.
p.set_val('x0', 0.0, units='m')

# Here we're intentially setting the intiial x value to something other than zero, just
# to demonstrate that the optimizer brings it back in line with the value of x0 set above.
phase.set_state_val('x', [1, 10],
                   units='m')

phase.set_state_val('y', [10, 5],
                   units='m')

phase.set_state_val('v', [0, 5],
                   units='m/s')

phase.set_control_val('theta',
                      [90, 90],
                     units='deg')

# Run the driver to solve the problem
dm.run_problem(p, make_plots=True)

print(p.get_val('traj.phase0.timeseries.time'))

# Check the validity of our results by using scipy.integrate.solve_ivp to
# integrate the solution.
sim_out = traj.simulate()

# Plot the results
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 4.5))

axes[0].plot(p.get_val('traj.phase0.timeseries.x'),
             p.get_val('traj.phase0.timeseries.y'),
             'ro', label='solution')

axes[0].plot(sim_out.get_val('traj.phase0.timeseries.x'),
             sim_out.get_val('traj.phase0.timeseries.y'),
             'b-', label='simulation')

axes[0].set_xlabel('x (m)')
axes[0].set_ylabel('y (m/s)')
axes[0].legend()
axes[0].grid()

axes[1].plot(p.get_val('traj.phase0.timeseries.time'),
             p.get_val('traj.phase0.timeseries.theta', units='deg'),
             'ro', label='solution')

axes[1].plot(sim_out.get_val('traj.phase0.timeseries.time'),
             sim_out.get_val('traj.phase0.timeseries.theta', units='deg'),
             'b-', label='simulation')

axes[1].set_xlabel('time (s)')
axes[1].set_ylabel(r'$\theta$ (deg)')
axes[1].legend()
axes[1].grid()

plt.show()
--- Constraint Report [traj] ---
    --- phase0 ---
        None

INFO: checking out_of_order
INFO: checking system
INFO: checking solvers
INFO: checking dup_inputs
INFO: checking missing_recorders
INFO: checking unserializable_options
INFO: checking comp_has_no_outputs
INFO: checking auto_ivc_warnings
Model viewer data has already been recorded for Driver.
INFO: checking out_of_order
INFO: checking system
INFO: checking solvers
INFO: checking dup_inputs
INFO: checking missing_recorders
INFO: checking unserializable_options
INFO: checking comp_has_no_outputs
INFO: checking auto_ivc_warnings
/usr/share/miniconda/envs/test/lib/python3.11/site-packages/openmdao/recorders/sqlite_recorder.py:231: UserWarning:The existing case recorder file, /home/runner/work/dymos/dymos/docs/dymos_book/examples/brachistochrone/problem_out/dymos_solution.db, is being overwritten.
Full total jacobian for problem 'problem' was computed 3 times, taking 0.07994873500001631 seconds.
Total jacobian shape: (41, 51) 


Jacobian shape: (41, 51)  (12.91% nonzero)
FWD solves: 8   REV solves: 0
Total colors vs. total size: 8 vs 51  (84.31% improvement)

Sparsity computed using tolerance: 1e-25
Time to compute sparsity:   0.0799 sec
Time to compute coloring:   0.0220 sec
Memory to compute coloring:   0.1250 MB
Coloring created on: 2024-11-23 14:34:50
Optimization terminated successfully    (Exit mode 0)
            Current function value: 1.8016061502983167
            Iterations: 63
            Function evaluations: 70
            Gradient evaluations: 63
Optimization Complete
-----------------------------------
[[0.        ]
 [0.09008031]
 [0.18016062]
 [0.18016062]
 [0.27024092]
 [0.36032123]
 [0.36032123]
 [0.45040154]
 [0.54048185]
 [0.54048185]
 [0.63056215]
 [0.72064246]
 [0.72064246]
 [0.81072277]
 [0.90080308]
 [0.90080308]
 [0.99088338]
 [1.08096369]
 [1.08096369]
 [1.171044  ]
 [1.26112431]
 [1.26112431]
 [1.35120461]
 [1.44128492]
 [1.44128492]
 [1.53136523]
 [1.62144554]
 [1.62144554]
 [1.71152584]
 [1.80160615]]

Simulating trajectory traj
Done simulating trajectory traj
/usr/share/miniconda/envs/test/lib/python3.11/site-packages/openmdao/core/group.py:1137: DerivativesWarning:Constraints or objectives [ode_eval.control_interp.control_rates:theta_rate, ode_eval.control_interp.control_rates:theta_rate2, ode_eval.control_interp.control_values:theta] cannot be impacted by the design variables of the problem because no partials were defined for them in their parent component(s).