Do linked OpenMDAO phases in a trajectory need to have the same transcription?

Question

I have two phases, one much shorter than the other, but it looks like I have to make the transcription the same for both- same number of nodes and polynomial order for each one, other wise I get a numpy mismatched array size / broadcasting error. Is there a way around this?

Answer 1

You do not need to have the same transcription for both.

Here is an example of setting up a very simple cannonball problem with two phases. The first phase is Radau, the second phase is GaussLabotto. Both use the same ODE, but different orders, numbers of segments, and compression setting.

import numpy as np

from scipy.interpolate import interp1d


import openmdao.api as om
from openmdao.components.interp_util.interp import InterpND

import dymos as dm

from dymos.models.atmosphere.atmos_1976 import USatm1976Data

# CREATE an atmosphere interpolant
english_to_metric_rho = om.unit_conversion('slug/ft**3', 'kg/m**3')[0]
english_to_metric_alt = om.unit_conversion('ft', 'm')[0]


rho_interp = interp1d(np.array(USatm1976Data.alt*english_to_metric_alt, dtype=complex), 
                      np.array(USatm1976Data.rho*english_to_metric_rho, dtype=complex), kind='linear')



GRAVITY = 9.80665


class CannonballODE(om.ExplicitComponent): 

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

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

        # static parameters
        self.add_input('radius', units='m')
        self.add_input('density', units='kg/m**3')
        self.add_input('CD', units=None)

        self.add_input('alt', units='m', shape=nn)
        self.add_input('v', units='m/s', shape=nn)
        self.add_input('gam', units='rad', shape=nn)

        self.add_output('v_dot', shape=(nn,), units='m/s**2')
        self.add_output('gam_dot', shape=(nn,), units='rad/s')
        self.add_output('h_dot', shape=(nn,), units='m/s')
        self.add_output('r_dot', shape=(nn,), units='m/s')
        self.add_output('ke', shape=(nn,), units='J')

        self.add_output('mass', shape=1, units='kg')

        self.declare_partials('*', '*', method='cs')

    def compute(self, inputs, outputs): 

        CD = inputs['CD']
        gam = inputs['gam']
        v = inputs['v']
        alt = inputs['alt']

        radius = inputs['radius']
        dens = inputs['density']

        m = (4/3.)*dens*np.pi*radius**3
        S = np.pi*radius**2

        if np.iscomplexobj(alt): 
            rho = rho_interp(inputs['alt'])
        else: 
            rho = rho_interp(inputs['alt']).real


        q = 0.5*rho*inputs['v']**2

        qS = q * S
        D = qS * CD

        cgam = np.cos(gam)
        sgam = np.sin(gam)

        mv = m*v

        outputs['v_dot'] = - D/m-GRAVITY*sgam
        outputs['gam_dot'] = -(GRAVITY/v)*cgam
        outputs['h_dot'] = v*sgam
        outputs['r_dot'] = v*cgam

        outputs['ke'] = 0.5*m*v**2

if __name__ == "__main__": 

    p = om.Problem()

    p.driver = om.pyOptSparseDriver()
    p.driver.options['optimizer'] = 'SLSQP'
    p.driver.declare_coloring()

    traj = p.model.add_subsystem('traj', dm.Trajectory())
    # constants across the whole trajectory
    traj.add_parameter('radius', units='m', val=0.01, dynamic=False)
    traj.add_parameter('density', units='kg/m**3', dynamic=False)

    p.model.add_design_var('traj.parameters:radius', lower=0.01, upper=0.10, ref0=0.01, ref=0.10)

    transcription = dm.Radau(num_segments=5, order=3, compressed=True)
    ascent = dm.Phase(transcription=transcription, ode_class=CannonballODE)
    traj.add_phase('ascent', ascent)

    ascent.add_state('r', units='m', rate_source='r_dot')
    ascent.add_state('h', units='m', rate_source='h_dot')
    ascent.add_state('gam', units='rad', rate_source='gam_dot')
    ascent.add_state('v', units='m/s', rate_source='v_dot')

    # All initial states except flight path angle are fixed
    # Final flight path angle is fixed (we will set it to zero so that the phase ends at apogee)
    ascent.set_time_options(fix_initial=True, duration_bounds=(1, 100), duration_ref=100, units='s')
    ascent.set_state_options('r', fix_initial=True, fix_final=False)
    ascent.set_state_options('h', fix_initial=True, fix_final=False)
    ascent.set_state_options('gam', fix_initial=False, fix_final=True)
    ascent.set_state_options('v', fix_initial=False, fix_final=False)

    ascent.add_parameter('radius', units='m', dynamic=False)
    ascent.add_parameter('density', units='kg/m**3', dynamic=False)
    ascent.add_parameter('CD', val=0.5, dynamic=False)

    # Limit the muzzle energy
    ascent.add_boundary_constraint('ke', loc='initial', units='J',
                                   upper=400000, lower=0, ref=100000)



    # Second Phase (descent)
    transcription = dm.GaussLobatto(num_segments=2, order=5, compressed=False)
    descent = dm.Phase(transcription=transcription, ode_class=CannonballODE)
    traj.add_phase('descent', descent )


    # All initial states and time are free (they will be linked to the final states of ascent.
    # Final altitude is fixed (we will set it to zero so that the phase ends at ground impact)
    descent.set_time_options(initial_bounds=(.5, 100), duration_bounds=(.5, 100),
                             duration_ref=100, units='s')

    descent.add_state('r', units='m', rate_source='r_dot')
    descent.add_state('h', units='m', rate_source='h_dot')
    descent.add_state('gam', units='rad', rate_source='gam_dot')
    descent.add_state('v', units='m/s', rate_source='v_dot',)
    descent.set_state_options('r', )
    descent.set_state_options('h', fix_initial=False, fix_final=True)
    descent.set_state_options('gam', fix_initial=False, fix_final=False)
    descent.set_state_options('v', fix_initial=False, fix_final=False)

    descent.add_parameter('radius', units='m', dynamic=False)
    descent.add_parameter('density', units='kg/m**3', dynamic=False)
    descent.add_parameter('CD', val=0.5, dynamic=False)

    # Link Phases (link time and all state variables)
    traj.link_phases(phases=['ascent', 'descent'], vars=['*'])


    descent.add_objective('r', loc='final', scaler=-1.0)


    # Finish Problem Setup
    p.setup()

    # Set Initial Guesses
    p.set_val('traj.parameters:radius', 0.05, units='m')
    p.set_val('traj.parameters:density', 7.87, units='g/cm**3')

    # initial guesses
    p.set_val('traj.ascent.t_initial', 0.0)
    p.set_val('traj.ascent.t_duration', 10.0)

    p.set_val('traj.ascent.states:r', ascent.interpolate(ys=[0, 100], nodes='state_input'))
    p.set_val('traj.ascent.states:h', ascent.interpolate(ys=[0, 100], nodes='state_input'))
    p.set_val('traj.ascent.states:v', ascent.interpolate(ys=[200, 150], nodes='state_input'))
    p.set_val('traj.ascent.states:gam', ascent.interpolate(ys=[25, 0], nodes='state_input'),
              units='deg')

    # more intitial guesses
    p.set_val('traj.descent.t_initial', 10.0)
    p.set_val('traj.descent.t_duration', 10.0)

    p.set_val('traj.descent.states:r', descent.interpolate(ys=[100, 200], nodes='state_input'))
    p.set_val('traj.descent.states:h', descent.interpolate(ys=[100, 0], nodes='state_input'))
    p.set_val('traj.descent.states:v', descent.interpolate(ys=[150, 200], nodes='state_input'))
    p.set_val('traj.descent.states:gam', descent.interpolate(ys=[0, -45], nodes='state_input'),
              units='deg')

    dm.run_problem(p, simulate=True, make_plots=True)

    # p.list_problem_vars(print_arrays=True)